Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 44 additions & 26 deletions rust/src/runtime/backend.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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<i64>;
type EncodedFloat = Option<f64>;
type EncodedFlag = Option<bool>;
Expand Down Expand Up @@ -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)
};
}

Expand All @@ -545,8 +561,8 @@ impl Backend for MySql {
fn advanced_history_sequence(result: &Self::QueryResult) -> Option<i64> {
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
Expand Down Expand Up @@ -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);
}

Expand Down
1 change: 1 addition & 0 deletions rust/src/runtime/mysql.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
27 changes: 16 additions & 11 deletions rust/src/runtime/mysql_candidate_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,7 @@ store_methods! {
async fn mysql_candidate_scope_work(&self, queue: &str, at: chrono::DateTime<chrono::Utc>) {
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
Expand Down Expand Up @@ -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();
Expand All @@ -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::<Vec<_>>();
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::<DB>(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::<DB>(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}");
}
}
}
Expand Down
2 changes: 1 addition & 1 deletion rust/src/runtime/poll_sweeps_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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')", "")
Expand Down
51 changes: 51 additions & 0 deletions rust/src/runtime/poll_transaction_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::<Vec<_>>();
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";
Expand Down Expand Up @@ -199,6 +249,7 @@ async fn sqlite_workflow_and_query_poll_transaction_counts() {
let dir = tempfile::tempdir().unwrap();
let store = Store::<sqlx::Sqlite>::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;
}
2 changes: 2 additions & 0 deletions rust/src/runtime/postgres.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<Migrator> = OnceLock::new();
Expand Down Expand Up @@ -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;
Expand Down
Loading
Loading