From f2bae31763a0be4719701e16d5effe6ea648a3bf Mon Sep 17 00:00:00 2001 From: Ruben Fiszel Date: Tue, 18 Nov 2025 09:40:57 +0000 Subject: [PATCH] buffer cloud hosted usage --- backend/Cargo.lock | 2 + backend/Cargo.toml | 1 + backend/windmill-api/src/lib.rs | 9 +- backend/windmill-queue/Cargo.toml | 2 + backend/windmill-queue/src/cloud_usage.rs | 203 ++++++++++++++++++++++ backend/windmill-queue/src/jobs.rs | 45 +---- backend/windmill-queue/src/lib.rs | 5 + 7 files changed, 221 insertions(+), 46 deletions(-) create mode 100644 backend/windmill-queue/src/cloud_usage.rs diff --git a/backend/Cargo.lock b/backend/Cargo.lock index f535756953..2ea6c7fb7b 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -15767,12 +15767,14 @@ dependencies = [ "chrono", "chrono-tz", "cron", + "dashmap 6.1.0", "futures", "futures-core", "hex", "hmac", "itertools 0.14.0", "lazy_static", + "once_cell", "prometheus", "quick_cache", "regex", diff --git a/backend/Cargo.toml b/backend/Cargo.toml index 39f6d688f7..6a5197d215 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -332,6 +332,7 @@ dyn-iter = "0.2.0" rsa = "^0" async_zip = { version = "0.0.17", features = ["tokio", "tokio-fs", "deflate", "chrono"] } once_cell = "1.17.1" +dashmap = "6.1.0" gosyn = "0.2.6" bytes = "1.4.0" gethostname = "0.4.3" diff --git a/backend/windmill-api/src/lib.rs b/backend/windmill-api/src/lib.rs index 3a83b11ce8..96bf0ed085 100644 --- a/backend/windmill-api/src/lib.rs +++ b/backend/windmill-api/src/lib.rs @@ -75,12 +75,12 @@ pub mod agent_workers_ee; #[cfg(feature = "agent_worker_server")] mod agent_workers_oss; mod ai; -mod bedrock; mod apps; pub mod args; mod assets; mod audit; pub mod auth; +mod bedrock; mod capture; mod concurrency_groups; mod configs; @@ -155,6 +155,7 @@ mod smtp_server_oss; pub mod teams_approvals_ee; mod teams_approvals_oss; +mod public_app_layer; mod static_assets; #[cfg(all(feature = "stripe", feature = "enterprise", feature = "private"))] pub mod stripe_ee; @@ -184,7 +185,6 @@ pub mod workspaces_ee; mod workspaces_export; mod workspaces_extra; mod workspaces_oss; -mod public_app_layer; #[cfg(feature = "mcp")] mod mcp; @@ -325,6 +325,11 @@ pub async fn run_server( #[cfg(feature = "embedding")] load_embeddings_db(&db); + #[cfg(feature = "cloud")] + if *CLOUD_HOSTED { + windmill_queue::init_usage_buffer(db.clone()); + } + let mut start_smtp_server = false; if let Some(smtp_settings) = load_value_from_global_settings(&db, EMAIL_DOMAIN_SETTING).await? diff --git a/backend/windmill-queue/Cargo.toml b/backend/windmill-queue/Cargo.toml index 6349aee84e..16178502fe 100644 --- a/backend/windmill-queue/Cargo.toml +++ b/backend/windmill-queue/Cargo.toml @@ -47,3 +47,5 @@ regex.workspace = true backon.workspace = true quick_cache.workspace = true thiserror.workspace = true +dashmap.workspace = true +once_cell.workspace = true diff --git a/backend/windmill-queue/src/cloud_usage.rs b/backend/windmill-queue/src/cloud_usage.rs new file mode 100644 index 0000000000..7ded6977ef --- /dev/null +++ b/backend/windmill-queue/src/cloud_usage.rs @@ -0,0 +1,203 @@ +/* + * Author: Ruben Fiszel + * Copyright: Windmill Labs, Inc 2022 + * This file and its contents are licensed under the AGPLv3 License. + * Please see the included NOTICE for copyright information and + * LICENSE-AGPL for a copy of the license. + */ + +use chrono::Datelike; +use dashmap::DashMap; +use sqlx::{Pool, Postgres}; +use std::sync::Arc; +use tokio::sync::Notify; + +#[derive(Debug, Clone, Hash, Eq, PartialEq)] +struct UsageKey { + id: String, + is_workspace: bool, + month: i32, +} + +pub struct UsageBuffer { + buffer: Arc>, + db: Pool, + shutdown_notify: Arc, +} + +impl UsageBuffer { + pub fn new(db: Pool) -> Arc { + let buffer = Arc::new(Self { + buffer: Arc::new(DashMap::new()), + db, + shutdown_notify: Arc::new(Notify::new()), + }); + + // Spawn the periodic flush task + let buffer_clone = buffer.clone(); + tokio::spawn(async move { + buffer_clone.flush_loop().await; + }); + + buffer + } + + pub fn increment(&self, workspace_id: String, email: Option) { + let month = Self::current_month(); + + // Increment workspace usage + self.buffer + .entry(UsageKey { id: workspace_id, is_workspace: true, month }) + .and_modify(|counter| *counter += 1) + .or_insert(1); + + // Increment user usage if email is provided + if let Some(email) = email { + self.buffer + .entry(UsageKey { id: email, is_workspace: false, month }) + .and_modify(|counter| *counter += 1) + .or_insert(1); + } + } + + async fn flush_loop(&self) { + let mut interval = tokio::time::interval(std::time::Duration::from_secs(60)); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + loop { + tokio::select! { + _ = interval.tick() => { + self.flush().await; + } + _ = self.shutdown_notify.notified() => { + // Final flush on shutdown + self.flush().await; + break; + } + } + } + } + + async fn flush(&self) { + if self.buffer.is_empty() { + return; + } + + // Drain all buffered usage counts + let mut to_flush = Vec::new(); + self.buffer.retain(|key, value| { + to_flush.push((key.clone(), *value)); + false + }); + + if to_flush.is_empty() { + return; + } + + tracing::debug!( + "Flushing {} buffered usage entries to database", + to_flush.len() + ); + + // Batch update to database + for (key, count) in to_flush { + let result = tokio::time::timeout( + std::time::Duration::from_secs(10), + sqlx::query!( + "INSERT INTO usage (id, is_workspace, month_, usage) + VALUES ($1, $2, $3, $4) + ON CONFLICT (id, is_workspace, month_) DO UPDATE SET usage = usage.usage + $4", + &key.id, + key.is_workspace, + key.month, + count + ) + .execute(&self.db), + ) + .await; + + match result { + Ok(Ok(_)) => {} + Ok(Err(e)) => { + tracing::error!( + "Failed to flush usage for {} (is_workspace: {}): {:#}", + key.id, + key.is_workspace, + e + ); + } + Err(_) => { + tracing::error!( + "Usage flush timed out for {} (is_workspace: {})", + key.id, + key.is_workspace + ); + } + } + } + } + + fn current_month() -> i32 { + let now = chrono::Utc::now(); + (now.year() * 12 + now.month() as i32) as i32 + } +} + +lazy_static::lazy_static! { + static ref USAGE_BUFFER: once_cell::sync::OnceCell> = once_cell::sync::OnceCell::new(); +} + +pub fn init_usage_buffer(db: Pool) { + USAGE_BUFFER.get_or_init(|| UsageBuffer::new(db)); +} + +pub fn increment_usage_async(db: Pool, workspace_id: String, email: Option) { + if let Some(buffer) = USAGE_BUFFER.get() { + buffer.increment(workspace_id, email); + } else { + tracing::warn!("Usage buffer not initialized, falling back to direct database update"); + // Fallback to old implementation if buffer not initialized + 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 + ); + } + }); + } +} diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs index 7783751944..c050d56659 100644 --- a/backend/windmill-queue/src/jobs.rs +++ b/backend/windmill-queue/src/jobs.rs @@ -3809,50 +3809,7 @@ async fn check_usage_limits( } #[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 - ); - } - }); -} +use crate::cloud_usage::increment_usage_async; // #[instrument(level = "trace", skip_all)] pub async fn push<'c, 'd>( diff --git a/backend/windmill-queue/src/lib.rs b/backend/windmill-queue/src/lib.rs index 9b6025e515..3ea2e2c1aa 100644 --- a/backend/windmill-queue/src/lib.rs +++ b/backend/windmill-queue/src/lib.rs @@ -14,3 +14,8 @@ pub mod schedule; pub use jobs::*; pub mod flow_status; pub mod tags; + +#[cfg(feature = "cloud")] +pub mod cloud_usage; +#[cfg(feature = "cloud")] +pub use cloud_usage::init_usage_buffer;