From dfd08d8a4bcf270a0579e68906f207e15ad52da0 Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Sat, 18 Oct 2025 07:47:31 +0000 Subject: [PATCH] refactor: improve usage table behavior to eliminate synchronous row locks (#6863) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor: improve usage table behavior to eliminate synchronous row locks Replace synchronous INSERT...RETURNING with SELECT + async UPDATE pattern: - Add check_usage_limits() to read current usage without row locks - Add increment_usage_async() to update usage in background task - Refactor job push logic to use optimistic validation - Simplify job completion tracking with better error handling This eliminates blocking row locks on the usage table during job creation, significantly improving throughput and reducing contention. Note: Requires running 'cargo sqlx prepare' with database access to update the query cache in .sqlx/ directory. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-authored-by: Ruben Fiszel * Update SQLx metadata * refactor: optimize cloud usage checks with caching and conditional queries - Add 60s cache for superadmin status checks to reduce DB load - Skip unnecessary user usage query for premium workspaces - Use existing team plan status cache (already implemented in windmill-common) - Update check_usage_limits to accept check_user_usage parameter - Add sqlx query cache for conditional user usage query This optimization eliminates redundant database queries during job creation, particularly for premium workspaces where user usage tracking is not needed. Co-authored-by: Ruben Fiszel --------- Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com> Co-authored-by: Ruben Fiszel Co-authored-by: windmill-internal-app[bot] --- ...984ff5ce7086ac511e8647e2024d9dbe0af56.json | 14 + ...ee42efb674f61bd2f3507278199773d5d6479.json | 22 ++ ...35a6a41a13f39159e106834c62f3b46c44227.json | 22 ++ ...311afe67d339ca022ba61c2767c004b038ef0.json | 14 + ...8566101dd27a0f3363884781071d417f8b7e7.json | 22 -- ...edb07863c362cf2a8c4624a1fc6ae3136a568.json | 15 ++ ...be69959987cbac1d83b64225fd1bf9ab047c9.json | 22 -- ...8db2b1061c22698393c2e42d94da11e697b8b.json | 15 -- ...65d93abfbd68b6d21c1382895238b23dd78dd.json | 22 ++ ...78c45e2a2615a8fdda3cc70778f8d037da2cc.json | 15 ++ ...a13481db39eea0b52cc513b3ee6571cf44ed4.json | 15 -- backend/windmill-queue/src/jobs.rs | 245 ++++++++++++++---- 12 files changed, 312 insertions(+), 131 deletions(-) create mode 100644 backend/.sqlx/query-045b26db0cefe6eaac0e572661d984ff5ce7086ac511e8647e2024d9dbe0af56.json create mode 100644 backend/.sqlx/query-08643ecaac35008b37bfd71a136ee42efb674f61bd2f3507278199773d5d6479.json create mode 100644 backend/.sqlx/query-0b43d1f0c0d205d978cdb41d30835a6a41a13f39159e106834c62f3b46c44227.json create mode 100644 backend/.sqlx/query-42baae3c69367bbb21771a9711c311afe67d339ca022ba61c2767c004b038ef0.json delete mode 100644 backend/.sqlx/query-56b2326015fde12b1a4efa226518566101dd27a0f3363884781071d417f8b7e7.json create mode 100644 backend/.sqlx/query-7580917b8c791207556e2ed6734edb07863c362cf2a8c4624a1fc6ae3136a568.json delete mode 100644 backend/.sqlx/query-83f64dd93b1ddc03b84681d65d9be69959987cbac1d83b64225fd1bf9ab047c9.json delete mode 100644 backend/.sqlx/query-9aebf706529889dc044e0ef41da8db2b1061c22698393c2e42d94da11e697b8b.json create mode 100644 backend/.sqlx/query-c13d40a1c8db041137e0990937a65d93abfbd68b6d21c1382895238b23dd78dd.json create mode 100644 backend/.sqlx/query-c41787a3efbb2c520a8fee93b2078c45e2a2615a8fdda3cc70778f8d037da2cc.json delete mode 100644 backend/.sqlx/query-ceb10ca124d9425b24d81d8a279a13481db39eea0b52cc513b3ee6571cf44ed4.json diff --git a/backend/.sqlx/query-045b26db0cefe6eaac0e572661d984ff5ce7086ac511e8647e2024d9dbe0af56.json b/backend/.sqlx/query-045b26db0cefe6eaac0e572661d984ff5ce7086ac511e8647e2024d9dbe0af56.json new file mode 100644 index 0000000000..9affa78033 --- /dev/null +++ b/backend/.sqlx/query-045b26db0cefe6eaac0e572661d984ff5ce7086ac511e8647e2024d9dbe0af56.json @@ -0,0 +1,14 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO usage (id, is_workspace, month_, usage)\n VALUES ($1, FALSE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), 1)\n ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + 1", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Varchar" + ] + }, + "nullable": [] + }, + "hash": "045b26db0cefe6eaac0e572661d984ff5ce7086ac511e8647e2024d9dbe0af56" +} diff --git a/backend/.sqlx/query-08643ecaac35008b37bfd71a136ee42efb674f61bd2f3507278199773d5d6479.json b/backend/.sqlx/query-08643ecaac35008b37bfd71a136ee42efb674f61bd2f3507278199773d5d6479.json new file mode 100644 index 0000000000..3b1b8e3825 --- /dev/null +++ b/backend/.sqlx/query-08643ecaac35008b37bfd71a136ee42efb674f61bd2f3507278199773d5d6479.json @@ -0,0 +1,22 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT usage FROM usage\n WHERE id = $1\n AND is_workspace = FALSE\n AND month_ = EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date)", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "usage", + "type_info": "Int4" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + false + ] + }, + "hash": "08643ecaac35008b37bfd71a136ee42efb674f61bd2f3507278199773d5d6479" +} diff --git a/backend/.sqlx/query-0b43d1f0c0d205d978cdb41d30835a6a41a13f39159e106834c62f3b46c44227.json b/backend/.sqlx/query-0b43d1f0c0d205d978cdb41d30835a6a41a13f39159e106834c62f3b46c44227.json new file mode 100644 index 0000000000..7091c464ad --- /dev/null +++ b/backend/.sqlx/query-0b43d1f0c0d205d978cdb41d30835a6a41a13f39159e106834c62f3b46c44227.json @@ -0,0 +1,22 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT usage FROM usage\n WHERE id = $1\n AND is_workspace = FALSE\n AND month_ = EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date)", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "usage", + "type_info": "Int4" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + false + ] + }, + "hash": "0b43d1f0c0d205d978cdb41d30835a6a41a13f39159e106834c62f3b46c44227" +} diff --git a/backend/.sqlx/query-42baae3c69367bbb21771a9711c311afe67d339ca022ba61c2767c004b038ef0.json b/backend/.sqlx/query-42baae3c69367bbb21771a9711c311afe67d339ca022ba61c2767c004b038ef0.json new file mode 100644 index 0000000000..de5b99791f --- /dev/null +++ b/backend/.sqlx/query-42baae3c69367bbb21771a9711c311afe67d339ca022ba61c2767c004b038ef0.json @@ -0,0 +1,14 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO usage (id, is_workspace, month_, usage)\n VALUES ($1, TRUE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), 1)\n ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + 1", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Varchar" + ] + }, + "nullable": [] + }, + "hash": "42baae3c69367bbb21771a9711c311afe67d339ca022ba61c2767c004b038ef0" +} diff --git a/backend/.sqlx/query-56b2326015fde12b1a4efa226518566101dd27a0f3363884781071d417f8b7e7.json b/backend/.sqlx/query-56b2326015fde12b1a4efa226518566101dd27a0f3363884781071d417f8b7e7.json deleted file mode 100644 index 9da0554266..0000000000 --- a/backend/.sqlx/query-56b2326015fde12b1a4efa226518566101dd27a0f3363884781071d417f8b7e7.json +++ /dev/null @@ -1,22 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "INSERT INTO usage (id, is_workspace, month_, usage)\n VALUES ($1, FALSE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), 1)\n ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + 1 \n RETURNING usage.usage", - "describe": { - "columns": [ - { - "ordinal": 0, - "name": "usage", - "type_info": "Int4" - } - ], - "parameters": { - "Left": [ - "Varchar" - ] - }, - "nullable": [ - false - ] - }, - "hash": "56b2326015fde12b1a4efa226518566101dd27a0f3363884781071d417f8b7e7" -} diff --git a/backend/.sqlx/query-7580917b8c791207556e2ed6734edb07863c362cf2a8c4624a1fc6ae3136a568.json b/backend/.sqlx/query-7580917b8c791207556e2ed6734edb07863c362cf2a8c4624a1fc6ae3136a568.json new file mode 100644 index 0000000000..5ea14fd63f --- /dev/null +++ b/backend/.sqlx/query-7580917b8c791207556e2ed6734edb07863c362cf2a8c4624a1fc6ae3136a568.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO usage (id, is_workspace, month_, usage)\n VALUES ($1, TRUE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), $2)\n ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + EXCLUDED.usage", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Varchar", + "Int4" + ] + }, + "nullable": [] + }, + "hash": "7580917b8c791207556e2ed6734edb07863c362cf2a8c4624a1fc6ae3136a568" +} diff --git a/backend/.sqlx/query-83f64dd93b1ddc03b84681d65d9be69959987cbac1d83b64225fd1bf9ab047c9.json b/backend/.sqlx/query-83f64dd93b1ddc03b84681d65d9be69959987cbac1d83b64225fd1bf9ab047c9.json deleted file mode 100644 index 3ca69e27dd..0000000000 --- a/backend/.sqlx/query-83f64dd93b1ddc03b84681d65d9be69959987cbac1d83b64225fd1bf9ab047c9.json +++ /dev/null @@ -1,22 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "INSERT INTO usage (id, is_workspace, month_, usage)\n VALUES ($1, TRUE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), 1)\n ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + 1 \n RETURNING usage.usage", - "describe": { - "columns": [ - { - "ordinal": 0, - "name": "usage", - "type_info": "Int4" - } - ], - "parameters": { - "Left": [ - "Varchar" - ] - }, - "nullable": [ - false - ] - }, - "hash": "83f64dd93b1ddc03b84681d65d9be69959987cbac1d83b64225fd1bf9ab047c9" -} diff --git a/backend/.sqlx/query-9aebf706529889dc044e0ef41da8db2b1061c22698393c2e42d94da11e697b8b.json b/backend/.sqlx/query-9aebf706529889dc044e0ef41da8db2b1061c22698393c2e42d94da11e697b8b.json deleted file mode 100644 index 1cdf7354e2..0000000000 --- a/backend/.sqlx/query-9aebf706529889dc044e0ef41da8db2b1061c22698393c2e42d94da11e697b8b.json +++ /dev/null @@ -1,15 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "INSERT INTO usage (id, is_workspace, month_, usage) \n VALUES ($1, TRUE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), $2) \n ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + EXCLUDED.usage", - "describe": { - "columns": [], - "parameters": { - "Left": [ - "Varchar", - "Int4" - ] - }, - "nullable": [] - }, - "hash": "9aebf706529889dc044e0ef41da8db2b1061c22698393c2e42d94da11e697b8b" -} diff --git a/backend/.sqlx/query-c13d40a1c8db041137e0990937a65d93abfbd68b6d21c1382895238b23dd78dd.json b/backend/.sqlx/query-c13d40a1c8db041137e0990937a65d93abfbd68b6d21c1382895238b23dd78dd.json new file mode 100644 index 0000000000..32a06ab680 --- /dev/null +++ b/backend/.sqlx/query-c13d40a1c8db041137e0990937a65d93abfbd68b6d21c1382895238b23dd78dd.json @@ -0,0 +1,22 @@ +{ + "db_name": "PostgreSQL", + "query": "SELECT usage FROM usage\n WHERE id = $1\n AND is_workspace = TRUE\n AND month_ = EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date)", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "usage", + "type_info": "Int4" + } + ], + "parameters": { + "Left": [ + "Text" + ] + }, + "nullable": [ + false + ] + }, + "hash": "c13d40a1c8db041137e0990937a65d93abfbd68b6d21c1382895238b23dd78dd" +} diff --git a/backend/.sqlx/query-c41787a3efbb2c520a8fee93b2078c45e2a2615a8fdda3cc70778f8d037da2cc.json b/backend/.sqlx/query-c41787a3efbb2c520a8fee93b2078c45e2a2615a8fdda3cc70778f8d037da2cc.json new file mode 100644 index 0000000000..88a6899759 --- /dev/null +++ b/backend/.sqlx/query-c41787a3efbb2c520a8fee93b2078c45e2a2615a8fdda3cc70778f8d037da2cc.json @@ -0,0 +1,15 @@ +{ + "db_name": "PostgreSQL", + "query": "INSERT INTO usage (id, is_workspace, month_, usage)\n VALUES ($1, FALSE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), $2)\n ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + EXCLUDED.usage", + "describe": { + "columns": [], + "parameters": { + "Left": [ + "Varchar", + "Int4" + ] + }, + "nullable": [] + }, + "hash": "c41787a3efbb2c520a8fee93b2078c45e2a2615a8fdda3cc70778f8d037da2cc" +} diff --git a/backend/.sqlx/query-ceb10ca124d9425b24d81d8a279a13481db39eea0b52cc513b3ee6571cf44ed4.json b/backend/.sqlx/query-ceb10ca124d9425b24d81d8a279a13481db39eea0b52cc513b3ee6571cf44ed4.json deleted file mode 100644 index 0b5137067a..0000000000 --- a/backend/.sqlx/query-ceb10ca124d9425b24d81d8a279a13481db39eea0b52cc513b3ee6571cf44ed4.json +++ /dev/null @@ -1,15 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "INSERT INTO usage (id, is_workspace, month_, usage) \n VALUES ($1, FALSE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), $2) \n ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + EXCLUDED.usage", - "describe": { - "columns": [], - "parameters": { - "Left": [ - "Varchar", - "Int4" - ] - }, - "nullable": [] - }, - "hash": "ceb10ca124d9425b24d81d8a279a13481db39eea0b52cc513b3ee6571cf44ed4" -} diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index 3d23caee58..22c29e2eb1 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -1431,35 +1431,47 @@ fn apply_completed_job_cloud_usage( let premium_workspace = windmill_common::workspaces::get_team_plan_status(&db, &w_id) .await .premium; - tokio::time::timeout(std::time::Duration::from_secs(10), async move { - let _ = sqlx::query!( - "INSERT INTO usage (id, is_workspace, month_, usage) - VALUES ($1, TRUE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), $2) + let result = tokio::time::timeout(std::time::Duration::from_secs(10), async move { + // Update workspace usage + let workspace_result = sqlx::query!( + "INSERT INTO usage (id, is_workspace, month_, usage) + VALUES ($1, TRUE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), $2) ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + EXCLUDED.usage", - w_id, + &w_id, additional_usage as i32 ) .execute(&db) - .await - .map_err(|e| { - Error::internal_err(format!("updating usage: {e:#}")) - }); + .await; + if let Err(e) = workspace_result { + tracing::error!("Failed to update workspace usage for {}: {:#}", w_id, e); + } + + // Update user usage for non-premium workspaces if !premium_workspace { - let _ = sqlx::query!( - "INSERT INTO usage (id, is_workspace, month_, usage) - VALUES ($1, FALSE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), $2) + let user_result = sqlx::query!( + "INSERT INTO usage (id, is_workspace, month_, usage) + VALUES ($1, FALSE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), $2) ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + EXCLUDED.usage", - email, + &email, additional_usage as i32 ) .execute(&db) - .await - .map_err(|e| Error::internal_err(format!("updating usage: {e:#}"))); + .await; + + if let Err(e) = user_result { + tracing::error!("Failed to update user usage for {}: {:#}", email, e); + } } - }).await.unwrap_or_else(|_| { - tracing::error!("Could not update usage for workspace {w_id2} and permissioned as {email2}, stopped after 10s"); - }); + }).await; + + if let Err(_) = result { + tracing::error!( + "Could not update usage for workspace {} and permissioned as {}, stopped after 10s", + w_id2, + email2 + ); + } }); } } @@ -3243,6 +3255,140 @@ lazy_static::lazy_static! { pub static ref RE_ARG_TAG: Regex = Regex::new(r#"\$args\[((?:\w+\.)*\w+)\]"#).unwrap(); } +#[cfg(feature = "cloud")] +lazy_static::lazy_static! { + // Cache for superadmin status: email -> (is_super_admin, expiry_timestamp) + static ref SUPERADMIN_CACHE: Arc>> = + Arc::new(RwLock::new(HashMap::new())); +} + +#[cfg(feature = "cloud")] +const SUPERADMIN_CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(60); + +#[cfg(feature = "cloud")] +async fn is_superadmin_cached( + db: &Pool, + email: &str, +) -> Result { + let now = std::time::Instant::now(); + + // Try to get from cache first + { + let cache = SUPERADMIN_CACHE.read().await; + if let Some((is_super_admin, expiry)) = cache.get(email) { + if *expiry > now { + return Ok(*is_super_admin); + } + } + } + + // Cache miss or expired, fetch from database + let is_super_admin = sqlx::query_scalar!( + "SELECT super_admin FROM password WHERE email = $1", + email + ) + .fetch_optional(db) + .await? + .unwrap_or(false); + + // Update cache + { + let mut cache = SUPERADMIN_CACHE.write().await; + cache.insert(email.to_string(), (is_super_admin, now + SUPERADMIN_CACHE_TTL)); + } + + Ok(is_super_admin) +} + +#[cfg(feature = "cloud")] +async fn check_usage_limits( + db: &Pool, + workspace_id: &str, + email: &str, + check_user_usage: bool, +) -> Result<(i32, Option), Error> { + // Get current workspace usage with a simple SELECT (no row lock) + let workspace_usage = sqlx::query_scalar!( + "SELECT usage FROM usage + WHERE id = $1 + AND is_workspace = TRUE + AND month_ = EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date)", + workspace_id + ) + .fetch_optional(db) + .await + .map_err(|e| Error::internal_err(format!("fetching workspace usage: {e:#}")))? + .unwrap_or(0); + + // Get current user usage (only for non-premium workspaces) + let user_usage = if check_user_usage { + sqlx::query_scalar!( + "SELECT usage FROM usage + WHERE id = $1 + AND is_workspace = FALSE + AND month_ = EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date)", + email + ) + .fetch_optional(db) + .await + .map_err(|e| Error::internal_err(format!("fetching user usage: {e:#}")))? + } else { + None + }; + + Ok((workspace_usage, user_usage)) +} + +#[cfg(feature = "cloud")] +fn increment_usage_async( + db: Pool, + workspace_id: String, + email: Option, +) { + tokio::task::spawn(async move { + let result = tokio::time::timeout(std::time::Duration::from_secs(10), async { + // Update workspace usage + let workspace_result = sqlx::query!( + "INSERT INTO usage (id, is_workspace, month_, usage) + VALUES ($1, TRUE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), 1) + ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + 1", + &workspace_id + ) + .execute(&db) + .await; + + if let Err(e) = workspace_result { + tracing::error!("Failed to update workspace usage for {}: {:#}", workspace_id, e); + } + + // Update user usage if email is provided (non-premium workspaces only) + if let Some(ref email) = email { + let user_result = sqlx::query!( + "INSERT INTO usage (id, is_workspace, month_, usage) + VALUES ($1, FALSE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), 1) + ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + 1", + email + ) + .execute(&db) + .await; + + if let Err(e) = user_result { + tracing::error!("Failed to update user usage for {}: {:#}", email, e); + } + } + }) + .await; + + if let Err(_) = result { + tracing::error!( + "Usage update timed out after 10s for workspace {} and email {:?}", + workspace_id, + email + ); + } + }); +} + // #[instrument(level = "trace", skip_all)] pub async fn push<'c, 'd>( _db: &Pool, @@ -3276,54 +3422,39 @@ pub async fn push<'c, 'd>( ) -> Result<(Uuid, Transaction<'c, Postgres>), Error> { #[cfg(feature = "cloud")] if *CLOUD_HOSTED { - let team_plan_status = - windmill_common::workspaces::get_team_plan_status(_db, workspace_id).await; + let team_plan_status = windmill_common::workspaces::get_team_plan_status(_db, workspace_id).await; // we track only non flow steps let (workspace_usage, user_usage) = if !matches!( job_payload, JobPayload::Flow { .. } | JobPayload::RawFlow { .. } ) { - tokio::time::timeout(std::time::Duration::from_secs(10), async move { - let workspace_usage = sqlx::query_scalar!( - "INSERT INTO usage (id, is_workspace, month_, usage) - VALUES ($1, TRUE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), 1) - ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + 1 - RETURNING usage.usage", - workspace_id - ) - .fetch_one(_db) - .await - .map_err(|e| Error::internal_err(format!("updating usage: {e:#}")))?; + // Check current usage with SELECT (fast, no row locks) + // Only check user usage for non-premium workspaces + let (current_workspace_usage, current_user_usage) = + check_usage_limits(_db, workspace_id, email, !team_plan_status.premium).await?; - let user_usage = if !team_plan_status.premium { - Some(sqlx::query_scalar!( - "INSERT INTO usage (id, is_workspace, month_, usage) - VALUES ($1, FALSE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), 1) - ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + 1 - RETURNING usage.usage", - email - ) - .fetch_one(_db) - .await - .map_err(|e| Error::internal_err(format!("updating usage: {e:#}")))?) - } else { - None - }; - Ok((Some(workspace_usage), user_usage)) - }).await.unwrap_or_else(|e| { - tracing::error!("Could not update usage for workspace {workspace_id} and permissioned as {email}, stopped after 10s: {e:#}"); - Err(Error::internal_err(format!("Could not update usage for workspace {workspace_id} and permissioned as {email}, stopped after 10s: {e:#}"))) - }) + // Spawn async task to update usage counters in the background + increment_usage_async( + _db.clone(), + workspace_id.to_string(), + if !team_plan_status.premium { Some(email.to_string()) } else { None }, + ); + + // Return the current usage + 1 to account for this job + let workspace_usage_with_new_job = current_workspace_usage + 1; + let user_usage_with_new_job = if !team_plan_status.premium { + Some(current_user_usage.unwrap_or(0) + 1) + } else { + None + }; + + (Some(workspace_usage_with_new_job), user_usage_with_new_job) } else { - Ok((None, None)) - }?; + (None, None) + }; if !team_plan_status.premium || team_plan_status.is_past_due { - let is_super_admin = - sqlx::query_scalar!("SELECT super_admin FROM password WHERE email = $1", email) - .fetch_optional(_db) - .await? - .unwrap_or(false); + let is_super_admin = is_superadmin_cached(_db, email).await?; #[cfg(feature = "private")] let recovery_email = crate::jobs_ee::SCHEDULE_RECOVERY_HANDLER_USER_EMAIL;