buffer cloud hosted usage

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",

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"

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?

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

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<DashMap<UsageKey, i32>>,
+    db: Pool<Postgres>,
+    shutdown_notify: Arc<Notify>,
+}
+
+impl UsageBuffer {
+    pub fn new(db: Pool<Postgres>) -> Arc<Self> {
+        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<String>) {
+        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<Arc<UsageBuffer>> = once_cell::sync::OnceCell::new();
+}
+
+pub fn init_usage_buffer(db: Pool<Postgres>) {
+    USAGE_BUFFER.get_or_init(|| UsageBuffer::new(db));
+}
+
+pub fn increment_usage_async(db: Pool<Postgres>, workspace_id: String, email: Option<String>) {
+    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
+                );
+            }
+        });
+    }
+}

backend/windmill-queue/src/jobs.rs

@@ -3809,50 +3809,7 @@ async fn check_usage_limits(
 }
 
 #[cfg(feature = "cloud")]
-fn increment_usage_async(db: Pool<Postgres>, workspace_id: String, email: Option<String>) {
-    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>(

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;