Reduce Native poll and completion database work

rust/src/runtime/active_summary.rs

@@ -13,6 +13,11 @@ use super::{
 use chrono::{DateTime, Utc};
 use sqlx::{ColumnIndex, Decode, Encode, Executor, IntoArguments, Type};
 
+// One aggregate scan for each source table, including the empty-table case.
+const PROJECTION_CENSUS_SQL: &str = "SELECT waits.*,tasks.*,(SELECT COUNT(*) FROM workflow_signal_records WHERE workflow_run_id=$1 AND status='received') AS received_signals,(SELECT COUNT(*) FROM workflow_failures WHERE workflow_run_id=$2) AS failures FROM (SELECT COUNT(CASE WHEN status='open' AND kind NOT IN ('condition','signal','timer','activity','child','update') THEN 1 END) AS other_waits,COUNT(CASE WHEN (status='open' OR (status='resolved' AND source_status='timed_out')) AND kind IN ('condition','signal') THEN 1 END) AS external_waits,COUNT(CASE WHEN kind='timer' AND status='open' THEN 1 END) AS timers,COUNT(CASE WHEN kind='activity' AND status='open' THEN 1 END) AS activities,COUNT(CASE WHEN kind='child' THEN 1 END) AS children,COUNT(CASE WHEN kind='update' THEN 1 END) AS updates FROM workflow_run_waits WHERE workflow_run_id=$3) waits CROSS JOIN (SELECT COUNT(CASE WHEN task_type='workflow' AND repair_count>0 THEN 1 END) AS repairs,COUNT(CASE WHEN task_type='workflow' AND status='failed' THEN 1 END) AS failed_tasks FROM workflow_tasks WHERE workflow_run_id=$4) tasks";
+#[cfg(test)]
+const PREVIOUS_CENSUS_SQL: &str = "SELECT (SELECT COUNT(*) FROM workflow_run_waits WHERE workflow_run_id=$1 AND status='open' AND kind NOT IN ('condition','signal','timer','activity','child','update')) AS other_waits,(SELECT COUNT(*) FROM workflow_tasks WHERE workflow_run_id=$2 AND task_type='workflow' AND repair_count>0) AS repairs,(SELECT COUNT(*) FROM workflow_tasks WHERE workflow_run_id=$3 AND task_type='workflow' AND status='failed') AS failed_tasks,(SELECT COUNT(*) FROM workflow_signal_records WHERE workflow_run_id=$4 AND status='received') AS received_signals,(SELECT COUNT(*) FROM workflow_run_waits w WHERE w.workflow_run_id=$5 AND (w.status='open' OR (w.status='resolved' AND w.source_status='timed_out')) AND w.kind IN ('condition','signal')) AS external_waits,(SELECT COUNT(*) FROM workflow_run_waits w WHERE w.workflow_run_id=$6 AND w.kind='timer' AND w.status='open') AS timers,(SELECT COUNT(*) FROM workflow_run_waits w WHERE w.workflow_run_id=$7 AND w.kind='activity' AND w.status='open') AS activities,(SELECT COUNT(*) FROM workflow_run_waits w WHERE w.workflow_run_id=$8 AND w.kind='child') AS children,(SELECT COUNT(*) FROM workflow_run_waits w WHERE w.workflow_run_id=$9 AND w.kind='update') AS updates,(SELECT COUNT(*) FROM workflow_failures WHERE workflow_run_id=$10) AS failures";
+
 pub(super) struct ActiveSummary {
     pub(super) state: String,
     pub(super) reason: String,
@@ -27,6 +32,7 @@ pub(super) struct ActiveSummary {
 }
 
 /// Each zero count skips only the corresponding read, never its validation.
+#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
 pub(super) struct ProjectionCensus {
     other_waits: i64,
     repairs: i64,
@@ -63,9 +69,8 @@ where
     usize: ColumnIndex<DB::Row>,
 {
     pub(super) async fn projection_census(tx: &mut Transaction<'_, DB>, run_id: &str) -> Result<ProjectionCensus> {
-        let row = Self::query("SELECT (SELECT COUNT(*) FROM workflow_run_waits WHERE workflow_run_id=$1 AND status='open' AND kind NOT IN ('condition','signal','timer','activity','child','update')) AS other_waits,(SELECT COUNT(*) FROM workflow_tasks WHERE workflow_run_id=$2 AND task_type='workflow' AND repair_count>0) AS repairs,(SELECT COUNT(*) FROM workflow_tasks WHERE workflow_run_id=$3 AND task_type='workflow' AND status='failed') AS failed_tasks,(SELECT COUNT(*) FROM workflow_signal_records WHERE workflow_run_id=$4 AND status='received') AS received_signals,(SELECT COUNT(*) FROM workflow_run_waits w WHERE w.workflow_run_id=$5 AND (w.status='open' OR (w.status='resolved' AND w.source_status='timed_out')) AND w.kind IN ('condition','signal')) AS external_waits,(SELECT COUNT(*) FROM workflow_run_waits w WHERE w.workflow_run_id=$6 AND w.kind='timer' AND w.status='open') AS timers,(SELECT COUNT(*) FROM workflow_run_waits w WHERE w.workflow_run_id=$7 AND w.kind='activity' AND w.status='open') AS activities,(SELECT COUNT(*) FROM workflow_run_waits w WHERE w.workflow_run_id=$8 AND w.kind='child') AS children,(SELECT COUNT(*) FROM workflow_run_waits w WHERE w.workflow_run_id=$9 AND w.kind='update') AS updates,(SELECT COUNT(*) FROM workflow_failures WHERE workflow_run_id=$10) AS failures")
-            .bind(run_id).bind(run_id).bind(run_id).bind(run_id).bind(run_id)
-            .bind(run_id).bind(run_id).bind(run_id).bind(run_id).bind(run_id)
+        let row = Self::query(PROJECTION_CENSUS_SQL)
+            .bind(run_id).bind(run_id).bind(run_id).bind(run_id)
             .fetch_one(&mut **tx).await?;
         Ok(ProjectionCensus {
             other_waits: DB::number(&row,"other_waits")?, repairs: DB::number(&row,"repairs")?,
@@ -76,6 +81,90 @@ where
         })
     }
 
+    #[cfg(test)]
+    async fn census_values(tx: &mut Transaction<'_,DB>, run_id: &str, sql: &'static str, binds: usize) -> [i64;10] {
+        let mut query = Self::query(sql);
+        for _ in 0..binds { query = query.bind(run_id); }
+        let row = query.fetch_one(&mut **tx).await.unwrap();
+        ["other_waits","repairs","failed_tasks","received_signals","external_waits","timers","activities","children","updates","failures"]
+            .map(|name| DB::number(&row,name).unwrap())
+    }
+
+    #[cfg(test)]
+    async fn census_scans(tx: &mut Transaction<'_,DB>, run_id: &str, sql: &'static str, binds: usize) -> (usize,usize) {
+        let prefix = match DB::NAME {"PostgreSQL"=>"EXPLAIN (FORMAT JSON) ","SQLite"=>"EXPLAIN QUERY PLAN ",_=>"EXPLAIN "};
+        let mut query = sqlx::query(sqlx::AssertSqlSafe(format!("{prefix}{}",DB::statement(sql))));
+        for _ in 0..binds { query = query.bind(run_id); }
+        let rows = query.fetch_all(&mut **tx).await.unwrap();
+        if DB::NAME == "PostgreSQL" {
+            fn visit(value: &serde_json::Value, counts: &mut (usize,usize)) {
+                if value["Relation Name"] == "workflow_run_waits" { counts.0 += 1; }
+                if value["Relation Name"] == "workflow_tasks" { counts.1 += 1; }
+                match value {
+                    serde_json::Value::Object(map) => { for child in map.values() { visit(child,counts); } }
+                    serde_json::Value::Array(items) => { for child in items { visit(child,counts); } }
+                    _ => {}
+                }
+            }
+            let mut counts = (0,0);
+            visit(&DB::document_row(&rows[0],"QUERY PLAN").unwrap(),&mut counts);
+            return counts;
+        }
+        let mut counts = (0,0);
+        for row in rows {
+            let table = if DB::NAME == "SQLite" {
+                let detail = DB::string(&row,"detail").unwrap();
+                let mut parts = detail.split_whitespace();
+                match parts.next() {
+                    Some("SCAN" | "SEARCH") => parts.next().map(str::to_owned),
+                    _ => None,
+                }
+            } else { DB::optional_string(&row,"table").unwrap() };
+            match table.as_deref() {
+                Some("workflow_run_waits" | "w") => counts.0 += 1,
+                Some("workflow_tasks") => counts.1 += 1,
+                _ => {}
+            }
+        }
+        counts
+    }
+
+    #[cfg(test)]
+    pub(super) async fn qualify_projection_census_passes(&self, run_id: &str, instance_id: &str) {
+        let mut tx = self.begin().await.unwrap();
+        for absent in [run_id,"missing-census-run"] {
+            assert_eq!(Self::census_values(&mut tx,absent,PROJECTION_CENSUS_SQL,4).await,
+                       Self::census_values(&mut tx,absent,PREVIOUS_CENSUS_SQL,10).await);
+        }
+        let mut position = 0_i64;
+        for kind in ["other","condition","signal","timer","activity","child","update"] {
+            for status in ["open","resolved","cancelled"] {
+                for source in [None,Some("timed_out"),Some("completed")] {
+                    Self::query("INSERT INTO workflow_run_waits(id,workflow_run_id,workflow_instance_id,wait_id,position,kind,status,source_status) VALUES ($1,$2,$3,$4,$5,$6,$7,$8)")
+                        .bind(super::store::id()).bind(run_id).bind(instance_id).bind(format!("census-{position}"))
+                        .bind(position).bind(kind).bind(status).bind(source).execute(&mut *tx).await.unwrap();
+                    position += 1;
+                }
+            }
+        }
+        assert_eq!(Self::census_values(&mut tx,run_id,PROJECTION_CENSUS_SQL,4).await,
+                   Self::census_values(&mut tx,run_id,PREVIOUS_CENSUS_SQL,10).await);
+        for kind in ["workflow","activity"] {
+            for status in ["ready","leased","failed","completed"] {
+                for repairs in [0_i64,1,2] {
+                    Self::query("INSERT INTO workflow_tasks(id,workflow_run_id,namespace,queue,task_type,status,repair_count) VALUES ($1,$2,'default','census',$3,$4,$5)")
+                        .bind(super::store::id()).bind(run_id).bind(kind).bind(status).bind(repairs)
+                        .execute(&mut *tx).await.unwrap();
+                }
+            }
+        }
+        assert_eq!(Self::census_values(&mut tx,run_id,PROJECTION_CENSUS_SQL,4).await,
+                   Self::census_values(&mut tx,run_id,PREVIOUS_CENSUS_SQL,10).await);
+        assert_eq!(Self::census_scans(&mut tx,run_id,PREVIOUS_CENSUS_SQL,10).await,(6,2));
+        assert_eq!(Self::census_scans(&mut tx,run_id,PROJECTION_CENSUS_SQL,4).await,(1,1));
+        tx.rollback().await.unwrap();
+    }
+
     #[cfg(test)]
     pub(super) async fn projection_census_reference(tx: &mut Transaction<'_, DB>, run_id: &str) -> Result<ProjectionCensus> {
         // The old projector always performed these reads. Keep that path

rust/src/runtime/backend.rs

@@ -92,6 +92,7 @@ pub(super) trait Backend: Database {
     // Waits whose stored position differs from their chronological ordinal
     // over every retained wait of the run, as Workflow #941's projector stores.
     const WAIT_POSITION_MOVES: &str = "WITH clocks AS (SELECT id,wait_id,position,opened_at,EXTRACT(EPOCH FROM opened_at) AS seconds FROM workflow_run_waits WHERE workflow_run_id=$1), ordinals AS (SELECT id,position,ROW_NUMBER() OVER (ORDER BY CASE WHEN opened_at IS NULL THEN 1 ELSE 0 END,ROUND((CASE WHEN seconds<0 THEN 2*FLOOR(seconds)-seconds ELSE seconds END)*1000),wait_id)-1 AS ordinal FROM clocks) SELECT id,ordinal FROM ordinals WHERE position IS DISTINCT FROM ordinal ORDER BY ordinal LIMIT 100";
+    const SUMMARY_UPSERT_SQL: &str = super::terminal_summary::SUMMARY_UPSERT_SQL;
     const SUMMARY_HISTORY_SQL: &str = "SELECT id,sequence,event_type,CAST(payload AS TEXT) AS payload_json FROM workflow_history_events WHERE workflow_run_id=$1 AND sequence>$2 ORDER BY sequence LIMIT 16";
     const OBSERVATION_FAILURE_EVENT: &str = "SELECT id,sequence,event_type,recorded_at FROM workflow_history_events WHERE workflow_run_id=$1 AND payload->>'failure_id'=$2 AND event_type IN ('ActivityFailed','ActivityTimedOut','ChildRunFailed','ChildRunCancelled','ChildRunTerminated','WorkflowFailed','WorkflowTimedOut','WorkflowCancelled','WorkflowTerminated','UpdateCompleted') ORDER BY sequence DESC LIMIT 1";
     const OBSERVATION_ATTRIBUTES: &str = "SELECT \"key\" AS attribute_key,type,value_string,value_keyword,value_int,value_float,CASE WHEN value_bool IS NULL THEN NULL WHEN value_bool THEN 1 ELSE 0 END AS value_bool,value_datetime,value_keyword_list FROM workflow_search_attributes WHERE workflow_run_id=$1 AND workflow_instance_id=$2 AND \"key\" IN ($3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21,$22) ORDER BY \"key\" LIMIT 20";
@@ -126,6 +127,7 @@ pub(super) trait Backend: Database {
     const QUERY_CACHE_SCAN: &str = "SELECT \"key\",value,expiration FROM dw_query_cache WHERE \"key\" LIKE 'dw-native-query:%' AND value::jsonb->>'namespace'=$1 ORDER BY \"key\" LIMIT 65";
     const QUERY_CACHE_GET: &str =
         "SELECT value,expiration FROM dw_query_cache WHERE \"key\"=$1";
+    const QUERY_CACHE_GET_PAIR: &str = "SELECT 0 AS cache_part,value,expiration FROM dw_query_cache WHERE \"key\"=$1 UNION ALL SELECT 1 AS cache_part,value,expiration FROM dw_query_cache WHERE \"key\"=$2";
     const QUERY_CACHE_INSERT: &str =
         "INSERT INTO dw_query_cache(\"key\",value,expiration) VALUES ($1,$2,$3)";
     const QUERY_CACHE_UPDATE: &str =
@@ -623,6 +625,7 @@ impl Backend for MySql {
     // The run-and-sequence index keeps MySQL from sorting whole rows: a
     // filesort of one large payload exceeds the default sort buffer and
     // fails the transaction ("Out of sort memory").
+    const SUMMARY_UPSERT_SQL: &str = super::terminal_summary::summary_projection_sql!("REPLACE", "");
     const SUMMARY_HISTORY_SQL: &str = "SELECT id,sequence,event_type,CAST(payload AS CHAR) AS payload_json FROM workflow_history_events FORCE INDEX (workflow_history_events_workflow_run_id_sequence_unique) WHERE workflow_run_id=$1 AND sequence>$2 ORDER BY sequence LIMIT 16";
     const OBSERVATION_FAILURE_EVENT: &str = "SELECT id,sequence,event_type,recorded_at FROM workflow_history_events WHERE workflow_run_id=? AND CAST(JSON_UNQUOTE(JSON_EXTRACT(payload,'$.failure_id')) AS BINARY)=CAST(? AS BINARY) AND event_type IN ('ActivityFailed','ActivityTimedOut','ChildRunFailed','ChildRunCancelled','ChildRunTerminated','WorkflowFailed','WorkflowTimedOut','WorkflowCancelled','WorkflowTerminated','UpdateCompleted') ORDER BY sequence DESC LIMIT 1";
     const OBSERVATION_ATTRIBUTES: &str = "SELECT `key` AS attribute_key,type,value_string,value_keyword,value_int,value_float,CASE WHEN value_bool IS NULL THEN NULL WHEN value_bool THEN 1 ELSE 0 END AS value_bool,value_datetime,value_keyword_list FROM workflow_search_attributes WHERE workflow_run_id=$1 AND workflow_instance_id=$2 AND `key` IN ($3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21,$22) ORDER BY `key` LIMIT 20";
@@ -650,6 +653,7 @@ impl Backend for MySql {
     type EncodedDocument = Json<Value>;
     const QUERY_CACHE_GET: &str =
         "SELECT value,expiration FROM dw_query_cache WHERE `key`=$1";
+    const QUERY_CACHE_GET_PAIR: &str = "SELECT 0 AS cache_part,value,expiration FROM dw_query_cache WHERE `key`=$1 UNION ALL SELECT 1 AS cache_part,value,expiration FROM dw_query_cache WHERE `key`=$2";
     const QUERY_CACHE_INSERT: &str =
         "INSERT INTO dw_query_cache(`key`,value,expiration) VALUES ($1,$2,$3)";
     const QUERY_CACHE_UPDATE: &str =

rust/src/runtime/carried_polls.rs

@@ -165,9 +165,10 @@ store_methods! {
         &self, tx: &mut super::transaction::Transaction<'_, DB>, poll: &ActivityPoll, completion_owner: &str,
     ) -> Result<Outcome> {
         let owner = super::text(&poll.body, "worker_id")?;
-        if owner != completion_owner || !self.worker_poll_fence_holds_in(tx, &poll.fence, owner).await? {
+        if owner != completion_owner { return Ok(Outcome::Declined); }
+        let Some(registration) = self.fenced_claim_registration(tx,&poll.fence,owner).await? else {
             return Ok(Outcome::Declined);
-        }
+        };
         // Validation work uses the workflow slot too. An unverifiable lease
         // deadline remains protected, rather than being treated as free.
         let at = super::store::now();
@@ -179,7 +180,7 @@ store_methods! {
             return Ok(Outcome::Declined);
         }
         let saved = tx.poll_savepoint().await?;
-        match self.claim_in(tx, &poll.body, "activity", &poll.protocol, None, None).await {
+        match self.claim_with_registration(tx, &poll.body, "activity", &poll.protocol, None, None, Some(registration)).await {
             Ok(task) => {
                 tx.release_poll_claim().await?;
                 Ok(task.map_or(Outcome::Empty, Outcome::Leased))
@@ -213,9 +214,10 @@ store_methods! {
         &self, tx: &mut super::transaction::Transaction<'_, DB>, poll: &WorkflowPoll, completion_owner: &str,
     ) -> Result<Outcome> {
         let owner = super::text(&poll.body, "worker_id")?;
-        if owner != completion_owner || !self.worker_poll_fence_holds_in(tx, &poll.fence, owner).await? {
+        if owner != completion_owner { return Ok(Outcome::Declined); }
+        let Some(registration) = self.fenced_claim_registration(tx,&poll.fence,owner).await? else {
             return Ok(Outcome::Declined);
-        }
+        };
         let occupied = Self::scalar("SELECT COUNT(*) FROM workflow_tasks WHERE namespace=$1 AND lease_owner=$2 AND task_type='activity' AND status='leased' AND (lease_expires_at IS NULL OR lease_expires_at>$3)")
             .bind(self.namespace.as_ref()).bind(owner).bind(DB::bind_time(super::store::now())).fetch_one(&mut **tx).await?;
         if occupied > 0 { return Ok(Outcome::Declined); }
@@ -226,7 +228,7 @@ store_methods! {
             if let Some(answer) = self.replay_workflow_poll_in(tx, &poll.request, owner, super::text(&poll.body, "task_queue")?).await? {
                 return Ok(if answer["task"].is_object() { Outcome::Leased(answer["task"].clone()) } else { Outcome::Empty });
             }
-            let task = self.claim_in(tx, &poll.body, "workflow", &poll.protocol, poll.page, Some(&poll.request)).await?;
+            let task = self.claim_with_registration(tx, &poll.body, "workflow", &poll.protocol, poll.page, Some(&poll.request), Some(registration)).await?;
             if let Some(task) = task {
                 Self::remember_workflow_poll(tx, &poll.request, &json!({"task":task,"poll_status":"leased"})).await?;
                 Ok(Outcome::Leased(task))

rust/src/runtime/commit_handoff_tests.rs

@@ -159,7 +159,7 @@ store_methods! {
             Default::default(), None, None,
         )).await;
         assert_eq!(result.unwrap()["recorded"], true);
-        assert_eq!(statements.count("SELECT sequence FROM activity_executions WHERE workflow_run_id=$1 ORDER BY sequence DESC LIMIT 1"), 1,
+        assert_eq!(statements.count(super::store::AUTHORED_RUN_SQL), 1,
             "scope validation and command application share one authored-sequence scan");
         let rows = Self::query("SELECT sequence,event_type FROM workflow_history_events WHERE workflow_run_id=$1 ORDER BY sequence")
             .bind(started["run_id"].as_str().unwrap()).fetch_all(&self.pool).await.unwrap();

rust/src/runtime/mysql.rs

@@ -1197,6 +1197,9 @@ mod tests {
     async fn qualification_poll_cache_single_statement_writes() {
         let db = Database::new().await;
         let store = super::super::store::Store::new(MySqlStorage::open(db.options.clone()).await.unwrap().into_pool());
+        store.qualify_completion_waiting_status().await;
+        store.qualify_authored_run_sequences().await;
+        store.qualify_transaction_run_tags().await;
         store.qualify_query_poll_cache_writes().await;
         store.qualify_workflow_poll_cache_writes().await;
         store.qualify_summary_projection_bytes().await;
@@ -1226,6 +1229,7 @@ mod tests {
         store.qualify_workflow_poll_transactions().await;
         store.qualify_first_candidate_batch().await;
         store.qualify_claim_history_sequence().await;
+        store.qualify_fenced_registration_reuse().await;
         store.qualify_query_poll_transactions().await;
         store.close().await;
         db.remove().await;