diff --git a/backend/ee-repo-ref.txt b/backend/ee-repo-ref.txt index e61c7a7c19..20ac124f54 100644 --- a/backend/ee-repo-ref.txt +++ b/backend/ee-repo-ref.txt @@ -1 +1 @@ -ac53164fd27fe962a020f3ee944949df1f760603 +c465c042eeae3c5d9a0180f62a56e8f8d3e9e162 diff --git a/backend/windmill-common/src/lib.rs b/backend/windmill-common/src/lib.rs index d02ceddd1d..fb37721114 100644 --- a/backend/windmill-common/src/lib.rs +++ b/backend/windmill-common/src/lib.rs @@ -2338,13 +2338,13 @@ pub async fn legacy_on_behalf_of_email( /// Shared by [`ScriptHashInfo::on_behalf_of`] and [`FlowVersionInfo::on_behalf_of`]. /// -/// Reads identity data through the non-RLS pool and enforces nothing itself: it answers who a -/// row already says it runs as. Callers must have authorized `w_id` — and the row they read it +/// Reads identity tables, which carry no RLS (a user transaction reads them as the pool does), +/// and enforces nothing itself: it answers who a row already says it runs as. Callers must have authorized `w_id` — and the row they read it /// from — before dispatching a job with what it returns. -pub async fn on_behalf_of_from_permissioned_as( +pub async fn on_behalf_of_from_permissioned_as<'c>( permissioned_as: Option<&str>, w_id: &str, - db: &DB, + db: impl sqlx::PgExecutor<'c>, ) -> error::Result> { let Some(permissioned_as) = permissioned_as else { return Ok(None); @@ -3034,9 +3034,8 @@ async fn get_latest_flow_version_for_path<'e, E: sqlx::PgExecutor<'e>>( Ok(version) } -pub async fn get_latest_hash_for_path<'c, E: sqlx::PgExecutor<'c>>( - db: E, - db2: &DB, +pub async fn get_latest_hash_for_path( + db: &mut sqlx::PgConnection, w_id: &str, script_path: &str, require_locked: bool, @@ -3067,13 +3066,13 @@ pub async fn get_latest_hash_for_path<'c, E: sqlx::PgExecutor<'c>>( w_id, require_locked ) - .fetch_optional(db) + .fetch_optional(&mut *db) .await?; let script = utils::not_found_if_none(r_o, "script", script_path)?; let on_behalf_of = - on_behalf_of_from_permissioned_as(script.on_behalf_of.as_deref(), w_id, db2).await?; + on_behalf_of_from_permissioned_as(script.on_behalf_of.as_deref(), w_id, db).await?; Ok(( scripts::ScriptHash(script.hash), diff --git a/backend/windmill-common/src/variables.rs b/backend/windmill-common/src/variables.rs index f3ad826136..65eb912ced 100644 --- a/backend/windmill-common/src/variables.rs +++ b/backend/windmill-common/src/variables.rs @@ -297,9 +297,14 @@ pub async fn generate_approval_token( job_id: uuid::Uuid, db: &DB, ) -> crate::error::Result { + let key = get_workspace_key(w_id, db).await?; + approval_token_from_key(&key, job_id) +} + +/// The token [`generate_approval_token`] mints, from a workspace key already in hand. +pub fn approval_token_from_key(key: &str, job_id: uuid::Uuid) -> crate::error::Result { use hmac::{Hmac, Mac}; use sha2::Sha256; - let key = get_workspace_key(w_id, db).await?; let mut mac = Hmac::::new_from_slice(key.as_bytes()) .map_err(|e| crate::Error::internal_err(format!("HMAC key error: {e}")))?; mac.update(job_id.as_bytes()); diff --git a/backend/windmill-queue/src/schedule.rs b/backend/windmill-queue/src/schedule.rs index 0041dceea0..e404e07f05 100644 --- a/backend/windmill-queue/src/schedule.rs +++ b/backend/windmill-queue/src/schedule.rs @@ -42,7 +42,6 @@ use windmill_common::{ /// Helper to fetch metadata for a schedule's script or flow async fn get_schedule_metadata<'c>( tx: &mut sqlx::Transaction<'c, sqlx::Postgres>, - db: &DB, schedule: &Schedule, ) -> Result<( Option, // tag @@ -78,7 +77,12 @@ async fn get_schedule_metadata<'c>( Ok(( flow_info.tag.clone(), None, - flow_info.on_behalf_of(&schedule.workspace_id, db).await?, + windmill_common::on_behalf_of_from_permissioned_as( + flow_info.on_behalf_of.as_deref(), + &schedule.workspace_id, + &mut **tx, + ) + .await?, None, Some(version), parsed_retry, @@ -106,7 +110,6 @@ async fn get_schedule_metadata<'c>( _job_token_scopes, ) = windmill_common::get_latest_hash_for_path( &mut **tx, - db, &schedule.workspace_id, &schedule.script_path, false, @@ -117,6 +120,10 @@ async fn get_schedule_metadata<'c>( } } +/// Read identity and runnable settings through `tx` rather than `db`: every caller holds `tx` +/// open on the same pool, so a second checkout per tick starves it when a burst of schedules +/// fires at once. `tx` may run as `windmill_user` (schedule and flow API); those tables carry +/// no RLS, so they resolve the same as on the pool. pub async fn push_scheduled_job<'c>( db: &DB, mut tx: Transaction<'c, Postgres>, @@ -319,7 +326,7 @@ pub async fn push_scheduled_job<'c>( // Get metadata from the scheduled script/flow for tag, timeout, etc. let (tag, timeout, on_behalf_of, hash, flow_version, retry) = - get_schedule_metadata(&mut tx, db, schedule).await?; + get_schedule_metadata(&mut tx, schedule).await?; ( JobPayload::SingleStepFlow { @@ -373,7 +380,12 @@ pub async fn push_scheduled_job<'c>( ) .warn_after_seconds_with_sql(1, "get_flow_version_info_from_version".to_string()) .await?; - let on_behalf_of = flow_info.on_behalf_of(&schedule.workspace_id, db).await?; + let on_behalf_of = windmill_common::on_behalf_of_from_permissioned_as( + flow_info.on_behalf_of.as_deref(), + &schedule.workspace_id, + &mut *tx, + ) + .await?; let FlowVersionInfo { version, tag, dedicated_worker, labels, job_token_scopes, .. } = flow_info; @@ -453,7 +465,6 @@ pub async fn push_scheduled_job<'c>( job_token_scopes, ) = windmill_common::get_latest_hash_for_path( &mut *tx, - db, &schedule.workspace_id, &schedule.script_path, false, @@ -461,17 +472,11 @@ pub async fn push_scheduled_job<'c>( .warn_after_seconds_with_sql(1, "get_latest_hash_for_path".to_string()) .await?; - // NB: read on the non-RLS pool (`db`), not `tx`. push_scheduled_job is - // also invoked with an RLS user_db transaction (api-schedule/api-flows), - // under which these lookups would resolve against the caller's row - // visibility rather than the full table. The dual-connection here is - // intentional and required for correctness. + let rs = + windmill_common::runnable_settings::from_handle(runnable_settings_handle, &mut *tx) + .await?; let (debouncing_settings, concurrency_settings) = - windmill_common::runnable_settings::prefetch_cached_from_handle( - runnable_settings_handle, - db, - ) - .await?; + windmill_common::runnable_settings::prefetch_cached_tx(&rs, &mut tx).await?; if schedule.retry.is_some() { let parsed_retry = serde_json::from_value::(schedule.retry.clone().unwrap()) @@ -598,7 +603,7 @@ pub async fn push_scheduled_job<'c>( let resolved_email = windmill_common::users::get_email_from_permissioned_as( &permissioned_as, &schedule.workspace_id, - db, + &mut *tx, ) .await?; (resolved_email, permissioned_as, authed, false) @@ -621,7 +626,7 @@ pub async fn push_scheduled_job<'c>( }; if let Some(tag) = tag.as_deref().filter(|t| !t.is_empty()) { - let is_super_admin = windmill_common::auth::is_super_admin_email(db, &email).await?; + let is_super_admin = windmill_common::auth::is_super_admin_email(&mut *tx, &email).await?; crate::check_tag_available_for_push( db, &schedule.workspace_id, diff --git a/backend/windmill-queue/tests/schedule_push.rs b/backend/windmill-queue/tests/schedule_push.rs index ed50215473..e79d4ff976 100644 --- a/backend/windmill-queue/tests/schedule_push.rs +++ b/backend/windmill-queue/tests/schedule_push.rs @@ -222,6 +222,43 @@ mod schedule_push { Ok(()) } + // Callers hold `tx` on the pool they pass in, so a burst of schedules the size of the pool + // starves it if a push checks out a second connection. On a one-connection pool any such + // checkout hangs until the acquire timeout. + #[sqlx::test(migrations = "../migrations", fixtures("base", "schedule_push"))] + async fn test_push_holds_one_connection(db: Pool) -> anyhow::Result<()> { + let one = sqlx::postgres::PgPoolOptions::new() + .max_connections(1) + .acquire_timeout(std::time::Duration::from_secs(5)) + .connect_with((*db.connect_options()).clone()) + .await?; + + let schedules = [ + make_schedule(|s| { + s.script_path = "f/system/tagged_script".to_string(); + s.path = "f/system/tagged_schedule".to_string(); + }), + make_schedule(|s| { + s.script_path = "f/system/obo_script".to_string(); + s.path = "f/system/obo_schedule".to_string(); + }), + make_schedule(|s| { + s.is_flow = true; + s.script_path = "f/system/obo_flow".to_string(); + s.path = "f/system/obo_flow_schedule".to_string(); + }), + ]; + for schedule in &schedules { + let tx = one.begin().await?; + let push = push_scheduled_job(&one, tx, schedule, None, None); + let tx = tokio::time::timeout(std::time::Duration::from_secs(15), push).await??; + tx.commit().await?; + } + + assert_eq!(count_queued_jobs(&db).await, 3); + Ok(()) + } + // ----------------------------------------------------------------------- // push_scheduled_job: with retry wraps in SingleStepFlow // ----------------------------------------------------------------------- diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index 01ac2becdd..6456a85bdf 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/backend/windmill-worker/src/ai_executor.rs @@ -1041,8 +1041,7 @@ pub async fn handle_ai_agent_job( Ok(Some(hub_script.schema)) } else { let hash = get_latest_hash_for_path( - db, - db, + &mut *db.acquire().await?, &job.workspace_id, path.as_str(), true, diff --git a/backend/windmill-worker/src/bun_executor.rs b/backend/windmill-worker/src/bun_executor.rs index 6da968e422..7d78958775 100644 --- a/backend/windmill-worker/src/bun_executor.rs +++ b/backend/windmill-worker/src/bun_executor.rs @@ -4142,6 +4142,43 @@ pub async fn handle_wac_v2_output( job_ids: serde_json::Map::new(), }); + // Generate resume URLs for the inline approval buttons. The resume_id + // is derived from the step key so each waitForApproval() in the same + // workflow gets a unique resume_job record, and so URLs the workflow + // minted for this step ahead of time (getApprovalUrls) address the + // same one. + let resume_id: u32 = windmill_common::wac::approval_resume_id(&key); + // Read before the tx: it holds the queue row lock, and a pool checkout under it + // waits on the connections every other suspending parent is holding. + let wkey = + windmill_common::variables::get_workspace_key(&job.workspace_id, db).await?; + let approval_token = + windmill_common::variables::approval_token_from_key(&wkey, job.id)?; + + let (resume_url, cancel_url, approval_page_url) = { + use hmac::{Hmac, Mac}; + use sha2::Sha256; + let mut mac = Hmac::::new_from_slice(wkey.as_bytes()) + .map_err(|e| error::Error::internal_err(format!("HMAC key error: {e}")))?; + mac.update(job.id.as_bytes()); + mac.update(resume_id.to_be_bytes().as_ref()); + let signature = hex::encode(mac.finalize().into_bytes()); + + let base_url = (**windmill_common::BASE_URL.load()).clone(); + let w_id = &job.workspace_id; + let job_id = &job.id; + + let resume = format!( + "{base_url}/api/w/{w_id}/jobs_u/resume/{job_id}/{resume_id}/{signature}" + ); + let cancel = format!( + "{base_url}/api/w/{w_id}/jobs_u/cancel/{job_id}/{resume_id}/{signature}" + ); + let approval_page = + format!("{base_url}/approve/{w_id}/{job_id}?token={approval_token}"); + (resume, cancel, approval_page) + }; + let mut tx = db.begin().await?; // Park first: the queue row is locked before the status row, the order every @@ -4216,44 +4253,6 @@ pub async fn handle_wac_v2_output( } } - // Generate resume URLs for the inline approval buttons. The resume_id - // is derived from the step key so each waitForApproval() in the same - // workflow gets a unique resume_job record, and so URLs the workflow - // minted for this step ahead of time (getApprovalUrls) address the - // same one. - let resume_id: u32 = windmill_common::wac::approval_resume_id(&key); - // Generate stateless approval token using shared utility - let approval_token = - windmill_common::variables::generate_approval_token(&job.workspace_id, job.id, db) - .await?; - - let (resume_url, cancel_url, approval_page_url) = { - use hmac::{Hmac, Mac}; - use sha2::Sha256; - use windmill_common::variables::get_workspace_key; - - let wkey = get_workspace_key(&job.workspace_id, db).await?; - let mut mac = Hmac::::new_from_slice(wkey.as_bytes()) - .map_err(|e| error::Error::internal_err(format!("HMAC key error: {e}")))?; - mac.update(job.id.as_bytes()); - mac.update(resume_id.to_be_bytes().as_ref()); - let signature = hex::encode(mac.finalize().into_bytes()); - - let base_url = (**windmill_common::BASE_URL.load()).clone(); - let w_id = &job.workspace_id; - let job_id = &job.id; - - let resume = format!( - "{base_url}/api/w/{w_id}/jobs_u/resume/{job_id}/{resume_id}/{signature}" - ); - let cancel = format!( - "{base_url}/api/w/{w_id}/jobs_u/cancel/{job_id}/{resume_id}/{signature}" - ); - let approval_page = - format!("{base_url}/approve/{w_id}/{job_id}?token={approval_token}"); - (resume, cancel, approval_page) - }; - // Store approval form metadata for the approval page endpoint let mut approval_meta = serde_json::json!({ "key": key, diff --git a/backend/windmill-worker/src/worker_flow.rs b/backend/windmill-worker/src/worker_flow.rs index e4be72086e..e0c4026804 100644 --- a/backend/windmill-worker/src/worker_flow.rs +++ b/backend/windmill-worker/src/worker_flow.rs @@ -328,6 +328,25 @@ fn get_stop_after_if_data(stop_after_if: Option<&StopAfterIf>) -> (bool, Option< return (false, None, false); } +/// Takes the caller's executor: under the status-update tx a pool read would wait on a second +/// connection, which a burst of completions holding every one never frees. +async fn fetch_flow_args<'c>( + db: impl sqlx::PgExecutor<'c>, + flow: Uuid, +) -> Result>, sqlx::Error> { + let args = sqlx::query_scalar!( + r#" + SELECT args AS "args: Json>>" + FROM v2_job + WHERE id = $1 + "#, + flow + ) + .fetch_one(db) + .await; + args.map(|flow_args| flow_args.map(|flow_args| flow_args.0).unwrap_or_default()) +} + async fn get_id_ctx_for_expr<'c>( expr: &str, flow: uuid::Uuid, @@ -633,22 +652,7 @@ pub async fn update_flow_status_after_job_completion_internal( } }; - let args = Arc::pin(Lazy::new(async move { - let args = sqlx::query_scalar!( - r#" - SELECT args AS "args: Json>>" - FROM v2_job - WHERE id = $1 - "#, - flow - ) - .fetch_one(db) - .await; - let args = - args.map(|flow_args| flow_args.map(|flow_args| flow_args.0).unwrap_or_default()); - - args - })); + let args = Arc::pin(Lazy::new(fetch_flow_args(db, flow))); let from_result_to_args = |args: &Result>, sqlx::Error>| { let args = args @@ -992,7 +996,7 @@ pub async fn update_flow_status_after_job_completion_internal( .as_ref() .and_then(|x| x.stop_after_all_iters_if.as_ref()) { - let args = from_result_to_args(args.as_ref().await.get_ref())?; + let args = from_result_to_args(&fetch_flow_args(&mut *tx, flow).await)?; // Isolate the reads in a savepoint on the same connection: the // caller below swallows our error and keeps using `tx`, so a DB // read failure must not leave the outer transaction aborted (that @@ -1227,7 +1231,7 @@ pub async fn update_flow_status_after_job_completion_internal( .as_ref() .and_then(|x| x.stop_after_all_iters_if.as_ref()) { - let args = from_result_to_args(args.as_ref().await.get_ref())?; + let args = from_result_to_args(&fetch_flow_args(&mut *tx, flow).await)?; // See the matching savepoint comment above: isolate the reads so a // DB read failure (whose error the caller swallows) cannot abort @@ -1329,7 +1333,7 @@ pub async fn update_flow_status_after_job_completion_internal( .unwrap_or_default(); tracing::info!("update flow status on retry: {retry:#?} "); - let args = from_result_to_args(args.as_ref().await.get_ref())?; + let args = from_result_to_args(&fetch_flow_args(&mut *tx, flow).await)?; evaluate_retry( &retry,