diff --git a/rust/src/runtime/backend.rs b/rust/src/runtime/backend.rs index a2faead85..7161edcdc 100644 --- a/rust/src/runtime/backend.rs +++ b/rust/src/runtime/backend.rs @@ -6,6 +6,29 @@ use std::borrow::Cow; use super::{Result, mysql, postgres, schema}; +macro_rules! sqlite_task_candidates_sql { + ($limit:literal) => { + concat!(r#" +WITH utc_tasks AS ( + SELECT t.*, + replace(replace(replace(t.available_at,'T',' '),'Z',''),'+00:00','') AS _dw_available_utc, + replace(replace(replace(t.lease_expires_at,'T',' '),'Z',''),'+00:00','') AS _dw_lease_utc + FROM workflow_tasks t WHERE t.namespace=$1 AND t.queue=$2 AND t.task_type=$3 AND t.status IN ('ready','leased') +), ordered_tasks AS ( + SELECT *, + rtrim(rtrim(CASE WHEN length(_dw_available_utc)=19 THEN _dw_available_utc||'.0' ELSE _dw_available_utc END,'0'),'.') AS _dw_available_order, + rtrim(rtrim(CASE WHEN length(_dw_lease_utc)=19 THEN _dw_lease_utc||'.0' ELSE _dw_lease_utc END,'0'),'.') AS _dw_lease_order + FROM utc_tasks +) +SELECT t.*,r.workflow_type FROM ordered_tasks t JOIN workflow_runs r ON r.id=t.workflow_run_id +WHERE r.status IN ('pending','running','waiting') + AND (t.status='ready' OR (t.status='leased' AND t.task_type='workflow' AND t._dw_lease_order<=rtrim(rtrim($4,'0'),'.'))) + AND (t.available_at IS NULL OR t._dw_available_order<=rtrim(rtrim($5,'0'),'.')) + AND (COALESCE(NULLIF(t.compatibility,''),NULLIF(r.compatibility,'')) IS NULL OR COALESCE(NULLIF(t.compatibility,''),NULLIF(r.compatibility,''))=$6) +ORDER BY t.priority,t._dw_available_order,t.id LIMIT "#, $limit, "\n") + }; +} + /// Counts every durable stream item, durable stream and inbound message that /// holds an external payload reference. The arguments name the JSON payload /// and metadata columns as the backend reads them as text. @@ -74,6 +97,7 @@ pub(super) trait Backend: Database { const WORKER_REGISTRATION_SQL: &str = super::store::WORKER_REGISTRATION_SQL; const TASK_CANDIDATES_SQL: &str = super::store::TASK_CANDIDATES_SQL; const TASK_CANDIDATE_BIND_SETS: usize = 1; + const FIRST_TASK_CANDIDATES_SQL: &str = super::store::FIRST_TASK_CANDIDATES_SQL; const POLL_WAKE_TASKS_SQL: &str = super::cross_kind_poll::TASKS_SQL; const POLL_RECEIPT_CLEANUP_SQL: &str = super::store::POLL_RECEIPT_CLEANUP_SQL; const TIMER_TASK_SQL: &str; @@ -319,25 +343,8 @@ ORDER BY t.available_order,t.id LIMIT 100 // Normalize only the two supported UTC storage shapes. Removing trailing // fractional zeros preserves exact decimal ordering, including earlier // native precision; SQLite's floating-point date functions are avoided. - const TASK_CANDIDATES_SQL: &str = r#" -WITH utc_tasks AS ( - SELECT t.*, - replace(replace(replace(t.available_at,'T',' '),'Z',''),'+00:00','') AS _dw_available_utc, - replace(replace(replace(t.lease_expires_at,'T',' '),'Z',''),'+00:00','') AS _dw_lease_utc - FROM workflow_tasks t WHERE t.namespace=$1 AND t.queue=$2 AND t.task_type=$3 AND t.status IN ('ready','leased') -), ordered_tasks AS ( - SELECT *, - rtrim(rtrim(CASE WHEN length(_dw_available_utc)=19 THEN _dw_available_utc||'.0' ELSE _dw_available_utc END,'0'),'.') AS _dw_available_order, - rtrim(rtrim(CASE WHEN length(_dw_lease_utc)=19 THEN _dw_lease_utc||'.0' ELSE _dw_lease_utc END,'0'),'.') AS _dw_lease_order - FROM utc_tasks -) -SELECT t.*,r.workflow_type FROM ordered_tasks t JOIN workflow_runs r ON r.id=t.workflow_run_id -WHERE r.status IN ('pending','running','waiting') - AND (t.status='ready' OR (t.status='leased' AND t.task_type='workflow' AND t._dw_lease_order<=rtrim(rtrim($4,'0'),'.'))) - AND (t.available_at IS NULL OR t._dw_available_order<=rtrim(rtrim($5,'0'),'.')) - AND (COALESCE(NULLIF(t.compatibility,''),NULLIF(r.compatibility,'')) IS NULL OR COALESCE(NULLIF(t.compatibility,''),NULLIF(r.compatibility,''))=$6) -ORDER BY t.priority,t._dw_available_order,t.id LIMIT 100 -"#; + const TASK_CANDIDATES_SQL: &str = sqlite_task_candidates_sql!("100"); + const FIRST_TASK_CANDIDATES_SQL: &str = sqlite_task_candidates_sql!("10"); const POLL_RECEIPT_CLEANUP_SQL: &str = r#" DELETE FROM dw_poll_receipts WHERE (namespace,worker_id,kind,request_id) IN ( SELECT p.namespace,p.worker_id,p.kind,p.request_id @@ -430,9 +437,10 @@ impl Encode<'_, Postgres> for PostgresDocument { } impl Backend for Postgres { - const TASK_CANDIDATES_SQL: &str = super::postgres::TASK_CANDIDATES_SQL; const COMMIT_HANDOFF: bool = true; const HISTORY_SEQUENCE_MAX: i64 = i32::MAX as i64; + const TASK_CANDIDATES_SQL: &str = super::postgres::TASK_CANDIDATES_SQL; + const FIRST_TASK_CANDIDATES_SQL: &str = super::postgres::FIRST_TASK_CANDIDATES_SQL; type EncodedInteger = Option; type EncodedFloat = Option; type EncodedFlag = Option; @@ -526,14 +534,22 @@ impl Backend for Postgres { // Equality on the dispatch index's status prefix preserves each range's // order. The outer sort sees at most two bounded candidate pages. macro_rules! mysql_status_candidates { - ($status:literal, $namespace:literal, $queue:literal, $kind:literal, $lease_at:literal, $ready_at:literal, $build:literal) => { + ($status:literal, $namespace:literal, $queue:literal, $kind:literal, $lease_at:literal, $ready_at:literal, $build:literal, $limit:literal) => { concat!("(SELECT t.*,r.workflow_type FROM workflow_tasks t FORCE INDEX (dw_workflow_tasks_scoped_dispatch) STRAIGHT_JOIN workflow_runs r FORCE INDEX (PRIMARY) ON r.id=t.workflow_run_id ", "WHERE t.namespace=", $namespace, " AND t.queue=", $queue, " AND t.task_type=", $kind, " AND t.status='", $status, "' ", "AND r.status IN ('pending','running','waiting') ", "AND (t.status='ready' OR (t.task_type='workflow' AND t.lease_expires_at<=", $lease_at, ")) ", "AND (t.available_at IS NULL OR t.available_at<=", $ready_at, ") ", "AND (COALESCE(NULLIF(t.compatibility,''),NULLIF(r.compatibility,'')) IS NULL OR COALESCE(NULLIF(t.compatibility,''),NULLIF(r.compatibility,''))=", $build, ") ", - "ORDER BY t.priority,t.available_at,t.id LIMIT 100)") + "ORDER BY t.priority,t.available_at,t.id LIMIT ", $limit, ")") + }; +} + +macro_rules! mysql_task_candidates_sql { + ($limit:literal) => { + concat!("SELECT * FROM (", mysql_status_candidates!("ready", "$1", "$2", "$3", "$4", "$5", "$6", $limit), + " UNION ALL ", mysql_status_candidates!("leased", "$7", "$8", "$9", "$10", "$11", "$12", $limit), + ") AS candidates ORDER BY priority,available_at,id LIMIT ", $limit) }; } @@ -545,8 +561,8 @@ impl Backend for MySql { fn advanced_history_sequence(result: &Self::QueryResult) -> Option { i64::try_from(result.last_insert_id()).ok() } - const TASK_CANDIDATES_SQL: &str = concat!("SELECT * FROM (", mysql_status_candidates!("ready", "$1", "$2", "$3", "$4", "$5", "$6"), - " UNION ALL ", mysql_status_candidates!("leased", "$7", "$8", "$9", "$10", "$11", "$12"), ") AS candidates ORDER BY priority,available_at,id LIMIT 100"); + const TASK_CANDIDATES_SQL: &str = mysql_task_candidates_sql!("100"); + const FIRST_TASK_CANDIDATES_SQL: &str = mysql_task_candidates_sql!("10"); const TASK_CANDIDATE_BIND_SETS: usize = 2; // The wake query can otherwise start from a run and range-check every // task in its queue, even with the explicit live-status predicate. Keep @@ -797,8 +813,10 @@ mod timestamp_tests { fn mysql_status_candidate_template_keeps_positions_unique_across_both_ranges() { // Convert the actual production template. Reusing $1 in the second // range must fail the existing adapter assertion before any query. - let statement = MySql::statement(MySql::TASK_CANDIDATES_SQL); - assert_eq!(statement.matches('?').count(), 12); + for sql in [MySql::TASK_CANDIDATES_SQL, MySql::FIRST_TASK_CANDIDATES_SQL] { + let statement = MySql::statement(sql); + assert_eq!(statement.matches('?').count(), 12); + } assert_eq!(MySql::TASK_CANDIDATE_BIND_SETS, 2); } diff --git a/rust/src/runtime/mysql.rs b/rust/src/runtime/mysql.rs index dfd5cda61..4a387e4ac 100644 --- a/rust/src/runtime/mysql.rs +++ b/rust/src/runtime/mysql.rs @@ -1224,6 +1224,7 @@ mod tests { let db = Database::new().await; let store = super::super::store::Store::new(MySqlStorage::open(db.options.clone()).await.unwrap().into_pool()); store.qualify_workflow_poll_transactions().await; + store.qualify_first_candidate_batch().await; store.qualify_query_poll_transactions().await; store.close().await; db.remove().await; diff --git a/rust/src/runtime/mysql_candidate_tests.rs b/rust/src/runtime/mysql_candidate_tests.rs index 64ee0623f..5defc03cb 100644 --- a/rust/src/runtime/mysql_candidate_tests.rs +++ b/rust/src/runtime/mysql_candidate_tests.rs @@ -124,7 +124,7 @@ store_methods! { async fn mysql_candidate_scope_work(&self, queue: &str, at: chrono::DateTime) { let mut connection = self.pool.acquire().await.unwrap(); let before = Self::mysql_candidate_read_next(&mut connection).await; - let expected = self.task_candidate_query(queue,"workflow",None,at,at) + let expected = self.task_candidate_query(queue,"workflow",None,at,at,false) .fetch_all(&mut *connection).await.unwrap(); let bounded = Self::mysql_candidate_read_next(&mut connection).await - before; // Reproduce the rejected cross-tenant index choice with the same @@ -203,7 +203,7 @@ store_methods! { .execute(&mut *transaction).await.unwrap(); } transaction.commit().await.unwrap(); - let predecessor = super::store::task_candidates_sql!("FROM workflow_tasks t FORCE INDEX (dw_workflow_tasks_poll) STRAIGHT_JOIN workflow_runs r"); + let predecessor = super::store::task_candidates_sql!("FROM workflow_tasks t FORCE INDEX (dw_workflow_tasks_poll) STRAIGHT_JOIN workflow_runs r", "100"); for statistics in ["before", "after"] { if statistics == "after" { Self::query("ANALYZE TABLE workflow_tasks,workflow_runs").execute(&self.pool).await.unwrap(); @@ -213,26 +213,31 @@ store_methods! { for build in [None,Some("build"),Some("other")] { let old = Self::query(predecessor).bind(self.namespace.as_ref()).bind(&queue).bind(kind) .bind(DB::bind_time(at)).bind(DB::bind_time(at)).bind(build).fetch_all(&self.pool).await.unwrap(); - let new = self.task_candidate_query(&queue, kind, build, at, at).fetch_all(&self.pool).await.unwrap(); + let new = self.task_candidate_query(&queue, kind, build, at, at, false).fetch_all(&self.pool).await.unwrap(); let identities = |rows: &[DB::Row]| rows.iter().map(|row| (DB::string(row,"id").unwrap(),DB::string(row,"workflow_type").unwrap())).collect::>(); assert!(!old.is_empty()); assert_eq!(identities(&new), identities(&old), "{statistics} / {kind} / {build:?}"); + let first = self.task_candidate_query(&queue, kind, build, at, at, true) + .fetch_all(&self.pool).await.unwrap(); + assert_eq!(identities(&first), identities(&old[..old.len().min(super::priority_fairness::FAIRNESS_BATCH)])); assert!(new.len() <= 100); if kind == "activity" { assert!(new.iter().all(|row| DB::string(row,"status").unwrap() == "ready")); } else if build.is_none() { assert!(new.iter().any(|row| DB::string(row,"status").unwrap() == "leased")); } - let mut explain = sqlx::query::(AssertSqlSafe(format!("EXPLAIN FORMAT=JSON {}",DB::statement(DB::TASK_CANDIDATES_SQL)))); - for _ in 0..2 { - explain = explain.bind(self.namespace.as_ref()).bind(&queue).bind(kind) - .bind(DB::bind_time(at)).bind(DB::bind_time(at)).bind(build); + for (first, sql) in [(false, DB::TASK_CANDIDATES_SQL), (true, DB::FIRST_TASK_CANDIDATES_SQL)] { + let mut explain = sqlx::query::(AssertSqlSafe(format!("EXPLAIN FORMAT=JSON {}",DB::statement(sql)))); + for _ in 0..2 { + explain = explain.bind(self.namespace.as_ref()).bind(&queue).bind(kind) + .bind(DB::bind_time(at)).bind(DB::bind_time(at)).bind(build); + } + let row = explain.fetch_one(&self.pool).await.unwrap(); + let plan: Value = serde_json::from_str(&DB::string(&row,"EXPLAIN").unwrap()).unwrap(); + println!("mysql_status_candidate_plan statistics={statistics} kind={kind} build={build:?} first={first}: {plan}"); + assert!(ordered_ranges(&plan, kind == "activity"), "each status range must arrive in dispatch order: {plan}"); } - let row = explain.fetch_one(&self.pool).await.unwrap(); - let plan: Value = serde_json::from_str(&DB::string(&row,"EXPLAIN").unwrap()).unwrap(); - println!("mysql_status_candidate_plan statistics={statistics} kind={kind} build={build:?}: {plan}"); - assert!(ordered_ranges(&plan, kind == "activity"), "each status range must arrive in dispatch order: {plan}"); } } } diff --git a/rust/src/runtime/poll_sweeps_tests.rs b/rust/src/runtime/poll_sweeps_tests.rs index b428bb895..0d2be2581 100644 --- a/rust/src/runtime/poll_sweeps_tests.rs +++ b/rust/src/runtime/poll_sweeps_tests.rs @@ -201,7 +201,7 @@ store_methods! { } for (sql, wake) in [(DB::TASK_CANDIDATES_SQL, false),(DB::POLL_WAKE_TASKS_SQL,true)] { let compiled = DB::statement(sql); - let predecessor = if DB::NAME == "MySQL" && !wake { + let predecessor = if matches!(DB::NAME, "MySQL" | "PostgreSQL") && !wake { DB::statement(super::store::TASK_CANDIDATES_SQL) } else { compiled.clone() }; let predecessor = predecessor.replace(" AND t.status IN ('ready','leased')", "") diff --git a/rust/src/runtime/poll_transaction_tests.rs b/rust/src/runtime/poll_transaction_tests.rs index 7e0ea1388..6f0214ade 100644 --- a/rust/src/runtime/poll_transaction_tests.rs +++ b/rust/src/runtime/poll_transaction_tests.rs @@ -134,6 +134,56 @@ store_methods! { } } + pub(super) async fn qualify_first_candidate_batch(&self) { + self.count_worker().await; + let mut runs = Vec::new(); + for index in 0..12 { + let started = self.count_start(&format!("first-batch-{index}")).await; + runs.push(started["run_id"].as_str().unwrap().to_owned()); + } + let at = super::store::now(); + let first = self.task_candidate_query("counted","workflow",None,at,at,true) + .fetch_all(&self.pool).await.unwrap(); + let full = self.task_candidate_query("counted","workflow",None,at,at,false) + .fetch_all(&self.pool).await.unwrap(); + assert_eq!(first.len(),super::priority_fairness::FAIRNESS_BATCH); + assert_eq!(full.len(),12); + let ids = |rows: &[DB::Row]| rows.iter().map(|row| DB::string(row,"id").unwrap()).collect::>(); + assert_eq!(ids(&first),ids(&full[..first.len()])); + let (answer, statements) = super::sql_counts::measure(self.count_poll("first-batch-poll",false)).await; + assert_eq!(answer["task"]["task_id"],DB::string(&first[0],"id").unwrap()); + assert_eq!(statements.count(DB::FIRST_TASK_CANDIDATES_SQL),1); + assert_eq!(statements.count(DB::TASK_CANDIDATES_SQL),0,"a claimable first batch must avoid the full page"); + for run in runs { + Self::query("UPDATE workflow_tasks SET status='completed' WHERE workflow_run_id=$1") + .bind(run).execute(&self.pool).await.unwrap(); + } + // The raw first batch contains ten unsupported workflow types. It + // must not hide the supported candidate at the end of the full page. + let mut runs = Vec::new(); + let mut expected = Value::Null; + for index in 0..=10 { + let started = self.start(json!({"workflow_id":format!("unsupported-first-batch-{index}"), + "workflow_type":if index < 10 { "unsupported" } else { "counted" }, + "task_queue":"counted","input":input()}), BTreeMap::new(), None, None).await.unwrap(); + let run = started["run_id"].as_str().unwrap(); + // Fix the order independently of timestamp precision or ULID ties. + Self::query("UPDATE workflow_tasks SET available_at=$1 WHERE workflow_run_id=$2") + .bind(DB::bind_time(at - chrono::Duration::seconds(30) + chrono::Duration::milliseconds(index))) + .bind(run).execute(&self.pool).await.unwrap(); + if index == 10 { expected = started["run_id"].clone(); } + runs.push(run.to_owned()); + } + let (answer, statements) = super::sql_counts::measure(self.count_poll("unsupported-first-batch",false)).await; + assert_eq!(answer["task"]["run_id"],expected); + assert_eq!(statements.count(DB::FIRST_TASK_CANDIDATES_SQL),1); + assert_eq!(statements.count(DB::TASK_CANDIDATES_SQL),1,"an incomplete eligible prefix must read the full page"); + for run in runs { + Self::query("UPDATE workflow_tasks SET status='completed' WHERE workflow_run_id=$1") + .bind(run).execute(&self.pool).await.unwrap(); + } + } + pub(super) async fn qualify_query_poll_transactions(&self) { self.count_worker().await; let workflow_id = "counted-query-workflow"; @@ -199,6 +249,7 @@ async fn sqlite_workflow_and_query_poll_transaction_counts() { let dir = tempfile::tempdir().unwrap(); let store = Store::::new(super::sqlite::open(dir.path().join("counts.sqlite").to_str().unwrap()).await.unwrap()); store.qualify_workflow_poll_transactions().await; + store.qualify_first_candidate_batch().await; store.qualify_query_poll_transactions().await; store.close().await; } diff --git a/rust/src/runtime/postgres.rs b/rust/src/runtime/postgres.rs index ffecee93b..ff601d923 100644 --- a/rust/src/runtime/postgres.rs +++ b/rust/src/runtime/postgres.rs @@ -46,6 +46,7 @@ macro_rules! postgres_task_candidates_sql { // default collation, through the dedicated C-collated index. Their collation // does not change the priority/availability/id ordering. pub(super) const TASK_CANDIDATES_SQL: &str = postgres_task_candidates_sql!("100"); +pub(super) const FIRST_TASK_CANDIDATES_SQL: &str = postgres_task_candidates_sql!("10"); fn migrator() -> &'static Migrator { static MIGRATOR: OnceLock = OnceLock::new(); @@ -743,6 +744,7 @@ mod tests { let db = Database::new().await; let store = super::super::store::Store::new(PostgresStorage::open(db.options.clone()).await.unwrap().into_pool()); store.qualify_workflow_poll_transactions().await; + store.qualify_first_candidate_batch().await; store.qualify_query_poll_transactions().await; store.close().await; db.remove().await; diff --git a/rust/src/runtime/postgres_candidate_tests.rs b/rust/src/runtime/postgres_candidate_tests.rs index 53c0d947c..ad7dd94f4 100644 --- a/rust/src/runtime/postgres_candidate_tests.rs +++ b/rust/src/runtime/postgres_candidate_tests.rs @@ -114,6 +114,8 @@ impl Store { let expected = candidates(transaction,old,at,kind,build).await; let actual = candidates(transaction,new,at,kind,build).await; assert_eq!(actual,expected,"analyzed={analyzed} / kind={kind} / build={build:?}"); + let first = candidates(transaction,Postgres::FIRST_TASK_CANDIDATES_SQL,at,kind,build).await; + assert_eq!(first,expected[..expected.len().min(super::priority_fairness::FAIRNESS_BATCH)]); if kind=="workflow" && build.is_none() { assert_eq!(actual.len(),100); } } } @@ -136,11 +138,13 @@ impl Store { if analyzed { Self::query("ANALYZE workflow_tasks,workflow_runs").execute(&mut *transaction).await.unwrap(); } Self::require_candidate_parity(&mut transaction,at,old,new,analyzed).await; let before = explain(&mut transaction,old,at).await; - let after = explain(&mut transaction,new,at).await; - println!("postgres_status_candidate_plan analyzed={analyzed} old={before} new={after}"); - let (mut tasks,mut run_lookups) = (0,0); - require_index_ranges(&after[0]["Plan"],&mut tasks,&mut run_lookups,false); - assert_eq!((tasks,run_lookups),(2,2),"{after}"); + for (first, sql) in [(false,new),(true,Postgres::FIRST_TASK_CANDIDATES_SQL)] { + let after = explain(&mut transaction,sql,at).await; + println!("postgres_status_candidate_plan analyzed={analyzed} first={first} old={before} new={after}"); + let (mut tasks,mut run_lookups) = (0,0); + require_index_ranges(&after[0]["Plan"],&mut tasks,&mut run_lookups,false); + assert_eq!((tasks,run_lookups),(2,2),"{after}"); + } } transaction.rollback().await.unwrap(); } diff --git a/rust/src/runtime/store.rs b/rust/src/runtime/store.rs index 27e33c230..f4a7963f9 100644 --- a/rust/src/runtime/store.rs +++ b/rust/src/runtime/store.rs @@ -94,13 +94,14 @@ pub(super) struct SignalWrite<'a> { /// The runtime-owned message-stream delivery signal (`WorkflowStub::MESSAGE_STREAM_RUNTIME_SIGNAL`). const MESSAGE_STREAM_SIGNAL: &str = super::message_streams::INTERNAL_SIGNAL; macro_rules! task_candidates_sql { - ($tables:literal) => { - concat!("SELECT t.*,r.workflow_type ", $tables, " ON r.id=t.workflow_run_id WHERE t.namespace=$1 AND t.queue=$2 AND t.task_type=$3 AND t.status IN ('ready','leased') AND r.status IN ('pending','running','waiting') AND (t.status='ready' OR (t.status='leased' AND t.task_type='workflow' AND t.lease_expires_at<=$4)) AND (t.available_at IS NULL OR t.available_at<=$5) AND (COALESCE(NULLIF(t.compatibility,''),NULLIF(r.compatibility,'')) IS NULL OR COALESCE(NULLIF(t.compatibility,''),NULLIF(r.compatibility,''))=$6) ORDER BY t.priority,t.available_at,t.id LIMIT 100") + ($tables:literal, $limit:literal) => { + concat!("SELECT t.*,r.workflow_type ", $tables, " ON r.id=t.workflow_run_id WHERE t.namespace=$1 AND t.queue=$2 AND t.task_type=$3 AND t.status IN ('ready','leased') AND r.status IN ('pending','running','waiting') AND (t.status='ready' OR (t.status='leased' AND t.task_type='workflow' AND t.lease_expires_at<=$4)) AND (t.available_at IS NULL OR t.available_at<=$5) AND (COALESCE(NULLIF(t.compatibility,''),NULLIF(r.compatibility,'')) IS NULL OR COALESCE(NULLIF(t.compatibility,''),NULLIF(r.compatibility,''))=$6) ORDER BY t.priority,t.available_at,t.id LIMIT ", $limit) }; } #[cfg(test)] pub(super) use task_candidates_sql; -pub(super) const TASK_CANDIDATES_SQL: &str = task_candidates_sql!("FROM workflow_tasks t JOIN workflow_runs r"); +pub(super) const TASK_CANDIDATES_SQL: &str = task_candidates_sql!("FROM workflow_tasks t JOIN workflow_runs r", "100"); +pub(super) const FIRST_TASK_CANDIDATES_SQL: &str = task_candidates_sql!("FROM workflow_tasks t JOIN workflow_runs r", "10"); /// A signal record's place in `PendingMessageTask::signalFollows`: its /// command sequence (unsequenced last), receipt and creation times, id. struct SignalOrder { @@ -2145,16 +2146,16 @@ where return Ok(true); } let at = now(); - Ok(self.task_candidate_query(queue, kind, build.as_deref(), at, at) + Ok(self.task_candidate_query(queue, kind, build.as_deref(), at, at, true) .fetch_optional(&self.pool) .await? .is_some()) } pub(super) fn task_candidate_query( - &self, queue: &str, kind: &str, build: Option<&str>, lease_at: DateTime, ready_at: DateTime, + &self, queue: &str, kind: &str, build: Option<&str>, lease_at: DateTime, ready_at: DateTime, first: bool, ) -> sqlx::query::Query<'static, DB, DB::Arguments> { - let mut query = Self::query(DB::TASK_CANDIDATES_SQL); + let mut query = Self::query(if first { DB::FIRST_TASK_CANDIDATES_SQL } else { DB::TASK_CANDIDATES_SQL }); // MySQL's positional protocol needs the same values for both ranges. for _ in 0..DB::TASK_CANDIDATE_BIND_SETS { query = query.bind(self.namespace.as_ref()).bind(queue).bind(kind) @@ -2334,11 +2335,41 @@ where }, )?; let build = worker_build::(&worker)?; - let candidates = self.task_candidate_query(queue, kind, build.as_deref(), now(), now()) - .fetch_all(&mut **tx) - .await?; - let candidates = Self::fairness_batches(candidates, kind, &types)?; - for task in candidates { + let lease_at = now(); + let ready_at = now(); + let first = self.task_candidate_query(queue, kind, build.as_deref(), lease_at, ready_at, true) + .fetch_all(&mut **tx).await?; + let first_is_full = first.len() == super::priority_fairness::FAIRNESS_BATCH; + let mut candidates = Self::fairness_batches(first, kind, &types)?; + // Unsupported workflow types are removed before fairness batching. + // A short eligible prefix cannot decide the first complete batch. + let mut full_read = first_is_full && candidates.len() < super::priority_fairness::FAIRNESS_BATCH; + if full_read { + let full = self.task_candidate_query(queue, kind, build.as_deref(), lease_at, ready_at, false) + .fetch_all(&mut **tx).await?; + candidates = Self::fairness_batches(full, kind, &types)?; + } + let mut candidates = candidates.into_iter(); + let mut tried = std::collections::HashSet::new(); + loop { + let task = match candidates.next() { + Some(task) => task, + None if first_is_full && !full_read => { + full_read = true; + let full = self.task_candidate_query(queue, kind, build.as_deref(), lease_at, ready_at, false) + .fetch_all(&mut **tx).await?; + // Reorder the original full page before dropping tried IDs, + // so later fairness batches retain their original boundaries. + let mut remaining = Vec::new(); + for task in Self::fairness_batches(full, kind, &types)? { + if !tried.contains(&DB::string(&task, "id")?) { remaining.push(task); } + } + candidates = remaining.into_iter(); + continue; + } + None => break, + }; + tried.insert(DB::string(&task, "id")?); // Unsupported stored timestamp shapes fail before leasing work. // Takeover preflight must validate all records, including rows // that are not eligible for this poll's queue/type/deadline. diff --git a/rust/tests/execution/task_fairness.rs b/rust/tests/execution/task_fairness.rs index 2eaed918e..27709d2c1 100644 --- a/rust/tests/execution/task_fairness.rs +++ b/rust/tests/execution/task_fairness.rs @@ -49,3 +49,74 @@ async fn a_skipped_candidate_is_followed_by_another_fairness_class() { runtime.close().await; database.remove().await; } + +async fn register_batch_workers(app: &Router) { + let (status, body) = request(app, "POST", "/api/search-attributes", json!({"name":"Tenant","type":"keyword"})).await; + assert_eq!(status, StatusCode::CREATED, "{body}"); + for (worker, types, capabilities) in [ + ("capable", json!(["echo","other"]), json!(["typed_search_attributes"])), + ("plain", json!(["echo"]), json!([])), + ] { + let (status, body) = request(app, "POST", "/api/worker/register", json!({"worker_id":worker,"task_queue":"test", + "runtime":"php","supported_workflow_types":types,"supported_activity_types":[],"capabilities":capabilities})).await; + assert_eq!(status, StatusCode::CREATED, "{body}"); + } +} + +async fn park_typed_candidate(app: &Router, workflow_id: &str) { + start_keyed(app,workflow_id,"a").await; + let task = poll(app,"capable","workflow").await; + assert_eq!(task["workflow_id"],workflow_id,"{task}"); + let (status, body) = finish_task(app,&task,json!([ + {"type":"upsert_search_attributes","attributes":{"Tenant":"a"},"attribute_types":{"Tenant":"keyword"}}, + {"type":"open_signal_wait","signal_name":"go"} + ])).await; + assert_eq!(status,StatusCode::OK,"{body}"); +} + +async fn ready_typed_candidate(app: &Router, workflow_id: &str) { + let (status, body) = request(app,"POST",&format!("/api/workflows/{workflow_id}/signal/go"), + json!({"input":envelope(Payload::Array(vec![]))})).await; + assert_eq!(status,StatusCode::ACCEPTED,"{body}"); + tokio::time::sleep(std::time::Duration::from_millis(20)).await; +} + +#[tokio::test] +async fn a_poll_claims_past_a_whole_unclaimable_fairness_batch() { + let database = TestDatabase::new().await; + let runtime = database.open().await.unwrap(); + let app = router(runtime.clone()); + register_batch_workers(&app).await; + for index in 0..10 { park_typed_candidate(&app,&format!("blocked-{index}")).await; } + for index in 0..10 { ready_typed_candidate(&app,&format!("blocked-{index}")).await; } + start_keyed(&app,"after-blocked","b").await; + let task = poll(&app,"plain","workflow").await; + assert_eq!(task["workflow_id"],"after-blocked","{task}"); + assert!(poll(&app,"plain","workflow").await.is_null()); + runtime.close().await; + database.remove().await; +} + +#[tokio::test] +async fn unsupported_types_are_removed_before_deciding_the_first_fairness_batch() { + let database = TestDatabase::new().await; + let runtime = database.open().await.unwrap(); + let app = router(runtime.clone()); + register_batch_workers(&app).await; + park_typed_candidate(&app,"a1").await; + ready_typed_candidate(&app,"a1").await; + let (status, body) = request(&app,"POST","/api/workflows",json!({"workflow_id":"unsupported", + "workflow_type":"other","task_queue":"test","input":envelope(Payload::Array(vec![])),"fairness_key":"ignored"})).await; + assert_eq!(status,StatusCode::CREATED,"{body}"); + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + for index in 2..=9 { start_keyed(&app,&format!("a{index}"),"a").await; } + start_keyed(&app,"b1","b").await; + // The first ten raw rows contain only nine supported types. Reordering + // that incomplete prefix would incorrectly pick a2 after blocked a1. + // The full eligible batch includes b1, which fairness tries before a2. + let task = poll(&app,"plain","workflow").await; + assert_eq!(task["workflow_id"],"b1","{task}"); + assert_eq!(poll(&app,"plain","workflow").await["workflow_id"],"a2"); + runtime.close().await; + database.remove().await; +}