diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 380d6d200f..8699332e3a 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -15135,6 +15135,7 @@ dependencies = [ "gethostname", "git-version", "globset", + "indexmap 2.11.1", "k8s-openapi", "kube", "lazy_static", diff --git a/backend/Cargo.toml b/backend/Cargo.toml index c5c8edfb9e..e0803196ae 100644 --- a/backend/Cargo.toml +++ b/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 diff --git a/backend/src/main.rs b/backend/src/main.rs index d202adfbc7..55cf1c7c9d 100644 --- a/backend/src/main.rs +++ b/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); } diff --git a/backend/windmill-api/openapi.yaml b/backend/windmill-api/openapi.yaml index 77d760dc4a..8878054c98 100644 --- a/backend/windmill-api/openapi.yaml +++ b/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 diff --git a/backend/windmill-api/src/jobs.rs b/backend/windmill-api/src/jobs.rs index 67d2056f86..dce5158369 100644 --- a/backend/windmill-api/src/jobs.rs +++ b/backend/windmill-api/src/jobs.rs @@ -2259,6 +2259,7 @@ struct QueueStats { pub struct CountQueueJobsQuery { all_workspaces: Option, tags: Option, + num_shards: Option, } 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, +} + async fn count_completed_jobs( Extension(db): Extension, Path(w_id): Path, + Query(cq): Query, ) -> error::JsonResult { 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?; diff --git a/backend/windmill-common/src/lib.rs b/backend/windmill-common/src/lib.rs index 38209eb11e..035a7928e3 100644 --- a/backend/windmill-common/src/lib.rs +++ b/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> = OnceCell::new(); -pub static SHARD_ID_TO_DB_INSTANCE: OnceCell>> = OnceCell::new(); +pub static SHARD_ID_TO_DB_INSTANCE: OnceCell>> = OnceCell::new(); const LATEST_VERSION_ID_CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(60); diff --git a/benchmarks/benchmark_sharding.ts b/benchmarks/benchmark_sharding.ts index b553929136..2dba838fbb 100644 --- a/benchmarks/benchmark_sharding.ts +++ b/benchmarks/benchmark_sharding.ts @@ -51,10 +51,11 @@ interface BenchmarkContext { }; nStepsFlow: number; bodyTemplate: any; - getQueueCount: (tags?: string[]) => Promise; + getQueueCount: (tags?: string[], num_shards?: number) => Promise; getCompletedJobsCount: ( tags?: string[], - baseline?: number + baseline?: number, + num_shards?: number ) => Promise; } @@ -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 { + 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)); }