update

backend/Cargo.lock

@@ -15135,6 +15135,7 @@ dependencies = [
  "gethostname",
  "git-version",
  "globset",
+ "indexmap 2.11.1",
  "k8s-openapi",
  "kube",
  "lazy_static",

backend/Cargo.toml

@@ -135,6 +135,7 @@ lazy_static.workspace = true
 once_cell.workspace = true
 prometheus = { workspace = true, optional = true }
 uuid.workspace = true
+indexmap.workspace = true
 gethostname.workspace = true
 serde_json.workspace = true
 serde_yml.workspace = true

backend/src/main.rs

@@ -27,7 +27,7 @@ use strum::IntoEnumIterator;
 use tokio::{fs::File, io::AsyncReadExt, task::JoinHandle};
 use uuid::Uuid;
 use windmill_api::HTTP_CLIENT;
-
+use indexmap::map::IndexMap;
 #[cfg(feature = "enterprise")]
 use windmill_common::ee_oss::{
     maybe_renew_license_key_on_start, LICENSE_KEY_ID, LICENSE_KEY_VALID,
@@ -320,10 +320,10 @@ async fn initialize_server_shard_instances() -> anyhow::Result<()> {
         .as_ref()
         .ok_or_else(|| anyhow!("SHARD_URLS environment variable is required for server shard mode. Please set it as: SHARD_URLS=dburl1,dburl2,..."))?;
 
-    let mut shard_to_db = HashMap::new();
+    let mut shard_to_db = IndexMap::new();
     for (i, shard_url) in shard_urls.iter().enumerate() {
-        println!("Url: {}", &shard_url);
-        let shard = connect_db(Some(&shard_url), true, false, false).await?;
+        tracing::info!("Connecting to shard {}: {}", i, &shard_url);
+        let shard = connect_db(Some(&shard_url), true, false, false).await.with_context(|| format!("Failed to connect to shard {}", i))?;
         shard_to_db.insert(i, shard);
     }
 

backend/windmill-api/openapi.yaml

@@ -7581,6 +7581,16 @@ paths:
           in: query
           schema:
             type: boolean
+        - name: tags
+          description: filter by tags (comma-separated)
+          in: query
+          schema:
+            type: string
+        - name: num_shards
+          description: number of shards to query (only applicable in shard mode)
+          in: query
+          schema:
+            type: integer
       responses:
         "200":
           description: queue count
@@ -7604,6 +7614,11 @@ paths:
         - job
       parameters:
         - $ref: "#/components/parameters/WorkspaceId"
+        - name: num_shards
+          description: number of shards to query (only applicable in shard mode)
+          in: query
+          schema:
+            type: integer
       responses:
         "200":
           description: completed count

backend/windmill-api/src/jobs.rs

@@ -2259,6 +2259,7 @@ struct QueueStats {
 pub struct CountQueueJobsQuery {
     all_workspaces: Option<bool>,
     tags: Option<String>,
+    num_shards: Option<usize>,
 }
 
 async fn count_queue_jobs(
@@ -2274,13 +2275,13 @@ async fn count_queue_jobs(
         sqlx::query_as!(
             QueueStats,
             "
-            SELECT 
-                coalesce(COUNT(*) FILTER(WHERE suspend = 0 AND running = false), 0) as \"database_length!\", 
-                coalesce(COUNT(*) FILTER(WHERE suspend > 0), 0) as \"suspended!\" 
-            FROM 
-                v2_as_queue 
-            WHERE 
-                (workspace_id = $1 OR $2) AND 
+            SELECT
+                coalesce(COUNT(*) FILTER(WHERE suspend = 0 AND running = false), 0) as \"database_length!\",
+                coalesce(COUNT(*) FILTER(WHERE suspend > 0), 0) as \"suspended!\"
+            FROM
+                v2_as_queue
+            WHERE
+                (workspace_id = $1 OR $2) AND
                 scheduled_for <= now() AND ($3::text[] IS NULL OR tag = ANY($3))
             ",
             w_id,
@@ -2292,11 +2293,13 @@ async fn count_queue_jobs(
     };
 
     let count_queue_jobs = if *SHARD_MODE {
-        let count_futures = SHARD_ID_TO_DB_INSTANCE
-            .get()
-            .unwrap()
-            .iter()
-            .map(|(_, db)| count_queue_jobs(db))
+        let shard_db_store = SHARD_ID_TO_DB_INSTANCE.get().unwrap();
+        println!("{:#?}", &cq.num_shards);
+        let num_shards_to_query = cq.num_shards.unwrap_or(shard_db_store.len());
+
+        let count_futures = (0..num_shards_to_query)
+            .filter_map(|shard_id| shard_db_store.get(&shard_id))
+            .map(|db| count_queue_jobs(db))
             .collect_vec();
 
         let completed_futures = futures::future::try_join_all(count_futures).await?;
@@ -2373,20 +2376,26 @@ async fn count_completed_jobs_detail(
     Ok(Json(stats))
 }
 
+#[derive(Deserialize)]
+pub struct CountCompletedJobsSimpleQuery {
+    num_shards: Option<usize>,
+}
+
 async fn count_completed_jobs(
     Extension(db): Extension<DB>,
     Path(w_id): Path<String>,
+    Query(cq): Query<CountCompletedJobsSimpleQuery>,
 ) -> error::JsonResult<QueueStats> {
     let count_completed_jobs = async |db: &DB| {
         sqlx::query_as!(
             QueueStats,
             "
-            SELECT 
-                coalesce(COUNT(*), 0) as \"database_length!\", 
-                null::bigint as suspended 
-            FROM 
-                v2_job_completed 
-            WHERE 
+            SELECT
+                coalesce(COUNT(*), 0) as \"database_length!\",
+                null::bigint as suspended
+            FROM
+                v2_job_completed
+            WHERE
                 workspace_id = $1
             ",
             w_id
@@ -2396,11 +2405,15 @@ async fn count_completed_jobs(
     };
 
     let count_completed_jobs = if *SHARD_MODE {
-        let count_futures = SHARD_ID_TO_DB_INSTANCE
-            .get()
-            .unwrap()
-            .iter()
-            .map(|(_, db)| count_completed_jobs(db))
+        let shard_db_store = SHARD_ID_TO_DB_INSTANCE.get().unwrap();
+
+        println!("{:#?}", &cq.num_shards);
+
+        let num_shards_to_query = cq.num_shards.unwrap_or(shard_db_store.len());
+
+        let count_futures = (0..num_shards_to_query)
+            .filter_map(|shard_id| shard_db_store.get(&shard_id))
+            .map(|db| count_completed_jobs(db))
             .collect_vec();
 
         let completed_futures = futures::future::try_join_all(count_futures).await?;

backend/windmill-common/src/lib.rs

@@ -6,6 +6,7 @@
  * LICENSE-AGPL for a copy of the license.
  */
 
+use indexmap::IndexMap;
 use itertools::Itertools as _;
 use once_cell::sync::OnceCell;
 use quick_cache::sync::Cache;
@@ -172,7 +173,7 @@ lazy_static::lazy_static! {
 }
 
 pub static SHARD_DB_INSTANCE: OnceCell<Pool<Postgres>> = OnceCell::new();
-pub static SHARD_ID_TO_DB_INSTANCE: OnceCell<HashMap<usize, Pool<Postgres>>> = OnceCell::new();
+pub static SHARD_ID_TO_DB_INSTANCE: OnceCell<IndexMap<usize, Pool<Postgres>>> = OnceCell::new();
 
 const LATEST_VERSION_ID_CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(60);
 

benchmarks/benchmark_sharding.ts

@@ -51,10 +51,11 @@ interface BenchmarkContext {
   };
   nStepsFlow: number;
   bodyTemplate: any;
-  getQueueCount: (tags?: string[]) => Promise<number>;
+  getQueueCount: (tags?: string[], num_shards?: number) => Promise<number>;
   getCompletedJobsCount: (
     tags?: string[],
-    baseline?: number
+    baseline?: number,
+    num_shards?: number
   ) => Promise<number>;
 }
 
@@ -112,7 +113,15 @@ async function initializeBenchmarkContext(
   config.token = final_token;
   windmill.setClient(final_token, host);
 
-  async function getQueueCount(tags?: string[]) {
+  async function getQueueCount(tags?: string[], numShards?: number) {
+    const params = new URLSearchParams();
+    if (tags && tags.length > 0) {
+      params.set("tags", tags.join(","));
+    }
+    if (numShards !== undefined) {
+      params.set("num_shards", numShards.toString());
+    }
+    const queryString = params.toString();
     return (
       await (
         await fetch(
@@ -120,7 +129,7 @@ async function initializeBenchmarkContext(
             "/api/w/" +
             config.workspace_id +
             "/jobs/queue/count" +
-            (tags && tags.length > 0 ? "?tags=" + tags.join(",") : ""),
+            (queryString ? "?" + queryString : ""),
           { headers: { ["Authorization"]: "Bearer " + config.token } }
         )
       ).json()
@@ -153,8 +162,17 @@ async function initializeBenchmarkContext(
 
   async function getCompletedJobsCount(
     tags?: string[],
-    baseline: number = 0
+    baseline: number = 0,
+    numShards?: number
   ): Promise<number> {
+    const params = new URLSearchParams();
+    if (tags && tags.length > 0) {
+      params.set("tags", tags.join(","));
+    }
+    if (numShards !== undefined) {
+      params.set("num_shards", numShards.toString());
+    }
+    const queryString = params.toString();
     const completedJobs = (
       await (
         await fetch(
@@ -162,7 +180,7 @@ async function initializeBenchmarkContext(
             "/api/w/" +
             config.workspace_id +
             "/jobs/completed/count" +
-            (tags && tags.length > 0 ? "?tags=" + tags.join(",") : ""),
+            (queryString ? "?" + queryString : ""),
           { headers: { ["Authorization"]: "Bearer " + config.token } }
         )
       ).json()
@@ -255,7 +273,7 @@ async function runShardingBenchmark(
   } = context;
   console.log(`\n=== Running benchmark with ${numShards} shard(s) ===`);
 
-  const pastJobs = await getCompletedJobsCount(NON_TEST_TAGS, 0);
+  const pastJobs = await getCompletedJobsCount(NON_TEST_TAGS, 0, numShards);
 
   const enc = (s: string) => new TextEncoder().encode(s);
   const jobsSent = jobs;
@@ -312,7 +330,7 @@ async function runShardingBenchmark(
   while (completedJobs < jobsSent) {
     if (!didStart) {
       const queueCheckStart = Date.now();
-      const actual_queue = await getQueueCount(NON_TEST_TAGS);
+      const actual_queue = await getQueueCount(NON_TEST_TAGS, numShards);
       totalMonitoringOverhead += Date.now() - queueCheckStart;
 
       if (actual_queue < jobsSent) {
@@ -323,10 +341,16 @@ async function runShardingBenchmark(
     } else {
       await sleep(1);
       const monitoringStart = Date.now();
-      completedJobs = await getCompletedJobsCount(NON_TEST_TAGS, pastJobs);
+      completedJobs = await getCompletedJobsCount(
+        NON_TEST_TAGS,
+        pastJobs,
+        numShards
+      );
       totalMonitoringOverhead += Date.now() - monitoringStart;
 
-      const elapsed = start ? Math.max(0, Date.now() - start - totalMonitoringOverhead) : 0;
+      const elapsed = start
+        ? Math.max(0, Date.now() - start - totalMonitoringOverhead)
+        : 0;
       if (nStepsFlow > 0) {
         completedJobs = Math.floor(completedJobs / (nStepsFlow + 1));
       }