Batch local activity writes within row and byte limits

rust/src/runtime/backend.rs

@@ -88,6 +88,7 @@ pub(super) trait Backend: Database {
     const HISTORY_SEQUENCE_MAX: i64 = i64::MAX;
     const HISTORY_ADVANCE_RETURNS_ROW: bool = true;
     const HISTORY_ADVANCE_SQL: &str = "UPDATE workflow_runs SET last_history_sequence=last_history_sequence+1 WHERE id=$1 AND last_history_sequence>=0 AND last_history_sequence<$2 RETURNING last_history_sequence";
+    const HISTORY_BATCH_ADVANCE_SQL: &str = "UPDATE workflow_runs SET last_history_sequence=last_history_sequence+$1 WHERE id=$2 AND last_history_sequence>=0 AND last_history_sequence<=$3 RETURNING last_history_sequence";
     const ACTIVITY_METADATA_HISTORY: &str = "SELECT sequence,event_type,CASE WHEN projection_bytes<=8388608 THEN payload ELSE NULL END AS payload,workflow_task_id,recorded_at,projection_count FROM (SELECT sequence,event_type,payload,workflow_task_id,recorded_at,COUNT(*) OVER () AS projection_count,SUM(octet_length(CAST(payload AS TEXT))) OVER (ORDER BY sequence) AS projection_bytes FROM workflow_history_events WHERE workflow_run_id=$1 AND event_type IN ('ActivityScheduled','ActivityStarted','ActivityHeartbeatRecorded','ActivityRetryScheduled','ActivityCompleted','ActivityFailed','ActivityTimedOut','ActivityCancelled') ORDER BY sequence LIMIT 10001) AS page ORDER BY sequence";
     // Waits whose stored position differs from their chronological ordinal
     // over every retained wait of the run, as Workflow #941's projector stores.
@@ -253,6 +254,7 @@ fn sqlite_instant(value: String) -> Result<DateTime<Utc>> {
 
 impl Backend for Sqlite {
     const HISTORY_ADVANCE_SQL: &str = "UPDATE workflow_runs SET last_history_sequence=last_history_sequence+1 WHERE id=$1 AND typeof(last_history_sequence)='integer' AND last_history_sequence>=0 AND last_history_sequence<$2 RETURNING last_history_sequence";
+    const HISTORY_BATCH_ADVANCE_SQL: &str = "UPDATE workflow_runs SET last_history_sequence=last_history_sequence+$1 WHERE id=$2 AND typeof(last_history_sequence)='integer' AND last_history_sequence>=0 AND last_history_sequence<=$3 RETURNING last_history_sequence";
     type EncodedInteger = Option<i64>;
     type EncodedFloat = Option<f64>;
     type EncodedFlag = Option<bool>;
@@ -564,6 +566,7 @@ impl Backend for MySql {
     const HISTORY_SEQUENCE_MAX: i64 = u32::MAX as i64;
     const HISTORY_ADVANCE_RETURNS_ROW: bool = false;
     const HISTORY_ADVANCE_SQL: &str = "UPDATE workflow_runs SET last_history_sequence=LAST_INSERT_ID(last_history_sequence+1) WHERE id=$1 AND last_history_sequence>=0 AND last_history_sequence<$2";
+    const HISTORY_BATCH_ADVANCE_SQL: &str = "UPDATE workflow_runs SET last_history_sequence=LAST_INSERT_ID(last_history_sequence+$1) WHERE id=$2 AND last_history_sequence>=0 AND last_history_sequence<=$3";
     fn advanced_history_sequence(result: &Self::QueryResult) -> Option<i64> {
         i64::try_from(result.last_insert_id()).ok()
     }

rust/src/runtime/cancellation_scopes.rs

@@ -1160,6 +1160,7 @@ store_methods! {
     /// after each descendant or fence fetches only the rows written since,
     /// instead of every row again while it holds the transition lock.
     pub(super) async fn scope_events(tx: &mut Transaction<'_, DB>, run_id: &str) -> Result<Arc<Vec<Event>>> {
+        Self::write_local_activities(tx).await?;
         let mut events = match tx.scope_history.take() {
             Some(cache) if cache.run_id == run_id => cache.events,
             _ => Arc::new(Vec::new()),

rust/src/runtime/commit_handoff_tests.rs

@@ -14,6 +14,10 @@ fn allocation_outcome(result: super::Result<i64>) -> std::result::Result<i64, St
     })
 }
 
+fn batch_allocation_outcome(result: super::Result<Vec<i64>>) -> std::result::Result<Vec<i64>, String> {
+    result.map_err(|error| allocation_outcome(Err(error)).unwrap_err())
+}
+
 #[derive(Default)]
 pub(super) struct CommitGate {
     pub(super) entered: Notify,
@@ -137,6 +141,73 @@ store_methods! {
         }
         Self::query("UPDATE workflow_runs SET last_history_sequence=$1 WHERE id=$2")
             .bind(initial_sequence).bind(run_id).execute(&self.pool).await.unwrap();
+        self.qualify_history_batch_allocation().await;
+    }
+
+    async fn qualify_history_batch_allocation(&self) {
+        use crate::codec::{Value as Payload, ValueCodec};
+        use serde_json::json;
+        let blob = ValueCodec::new().unwrap().encode(&Payload::Array(vec![])).unwrap();
+        let run = self.start(json!({"workflow_id":"history-batch-allocation","workflow_type":"probe","task_queue":"batch-allocation",
+            "input":{"codec":"avro","blob":blob}}),Default::default(),None,None).await.unwrap();
+        let run_id = run["run_id"].as_str().unwrap();
+        let original_sequence = DB::number(&Self::query("SELECT last_history_sequence FROM workflow_runs WHERE id=$1")
+            .bind(run_id).fetch_one(&self.pool).await.unwrap(),"last_history_sequence").unwrap();
+        let maximum = DB::HISTORY_SEQUENCE_MAX;
+        for value in [0,42,maximum-3,maximum-2,maximum-1,maximum] {
+            Self::query("UPDATE workflow_runs SET last_history_sequence=$1 WHERE id=$2")
+                .bind(value).bind(run_id).execute(&self.pool).await.unwrap();
+            let mut original = self.begin().await.unwrap();
+            let expected: super::Result<Vec<i64>> = async {
+                let mut values = Vec::new();
+                for _ in 0..3 { values.push(Self::allocate_history_sequence_fallback(&mut original,run_id).await?); }
+                Ok(values)
+            }.await;
+            let expected = batch_allocation_outcome(expected);
+            original.rollback().await.unwrap();
+            let mut candidate = self.begin().await.unwrap();
+            let (answer,statements) = super::sql_counts::measure(Self::allocate_history_sequences(&mut candidate,run_id,3)).await;
+            assert_eq!(batch_allocation_outcome(answer),expected,"batch at {value}");
+            assert_eq!(statements.count(DB::HISTORY_BATCH_ADVANCE_SQL),1);
+            if value <= maximum-3 {
+                assert_eq!(statements.total(),1,"canonical range allocation is one statement");
+            }
+            candidate.rollback().await.unwrap();
+            let row = Self::query("SELECT last_history_sequence FROM workflow_runs WHERE id=$1")
+                .bind(run_id).fetch_one(&self.pool).await.unwrap();
+            assert_eq!(DB::number(&row,"last_history_sequence").unwrap(),value);
+        }
+        let missing = "00000000000000000000000000";
+        let mut original = self.begin().await.unwrap();
+        let expected = allocation_outcome(Self::allocate_history_sequence_fallback(&mut original,missing).await).unwrap_err();
+        original.rollback().await.unwrap();
+        let mut candidate = self.begin().await.unwrap();
+        assert_eq!(batch_allocation_outcome(Self::allocate_history_sequences(&mut candidate,missing,3).await).unwrap_err(),expected);
+        let (empty,statements) = super::sql_counts::measure(Self::allocate_history_sequences(&mut candidate,missing,0)).await;
+        assert!(empty.unwrap().is_empty());
+        assert_eq!(statements.total(),0);
+        candidate.rollback().await.unwrap();
+        if DB::NAME != "MySQL" {
+            Self::query("UPDATE workflow_runs SET last_history_sequence=-1 WHERE id=$1")
+                .bind(run_id).execute(&self.pool).await.unwrap();
+            let mut candidate = self.begin().await.unwrap();
+            assert_eq!(Self::allocate_history_sequences(&mut candidate,run_id,3).await.unwrap(),vec![0,1,2]);
+            candidate.rollback().await.unwrap();
+        }
+        if DB::NAME == "SQLite" {
+            for sql in ["UPDATE workflow_runs SET last_history_sequence=1.5 WHERE id=$1",
+                "UPDATE workflow_runs SET last_history_sequence='not-an-integer' WHERE id=$1"] {
+                Self::query(sql).bind(run_id).execute(&self.pool).await.unwrap();
+                let mut original = self.begin().await.unwrap();
+                let expected = allocation_outcome(Self::allocate_history_sequence_fallback(&mut original,run_id).await).unwrap_err();
+                original.rollback().await.unwrap();
+                let mut candidate = self.begin().await.unwrap();
+                assert_eq!(batch_allocation_outcome(Self::allocate_history_sequences(&mut candidate,run_id,3).await).unwrap_err(),expected);
+                candidate.rollback().await.unwrap();
+            }
+        }
+        Self::query("UPDATE workflow_runs SET last_history_sequence=$1 WHERE id=$2")
+            .bind(original_sequence).bind(run_id).execute(&self.pool).await.unwrap();
     }
 
     pub(super) async fn qualify_authored_sequence_reads(&self) {