Send sticky workers only the history after their cached boundary (#1017 part 2)

rust/src/runtime/completion_rules.rs

@@ -61,6 +61,9 @@ const TOP: &[(&str, &[Check])] = &[
     ("sticky_cache.run_id", &[STICKY, STR, max(255)]),
     ("sticky_cache.build_id", &[STICKY, STR, max(255)]),
     ("sticky_cache.ttl_seconds", &[STICKY, INT, min(1), max(3600)]),
+    // Native only (Server #1017): the cached boundary of a worker that
+    // advertises `sticky_history_suffix`.
+    ("sticky_cache.history_last_sequence", &[NUL, INT, min(1)]),
     ("sticky_cache.metrics", &[NUL, ARR]),
     ("sticky_cache.metrics.hit", &[NUL, INT, min(0)]),
     ("sticky_cache.metrics.miss", &[NUL, INT, min(0)]),

rust/src/runtime/local_activities.rs

@@ -997,14 +997,18 @@ store_methods! {
         let Some(worker) = worker else {
             return Err(rejected("worker_capability_not_supported"));
         };
-        let listed = DB::optional_document_row(&worker, "capabilities")?.as_ref().and_then(Value::as_array)
-            .is_some_and(|capabilities| capabilities.iter().any(|capability| capability.as_str() == Some(CAPABILITY)));
-        let supported = DB::optional_document_row(&worker, "capability_manifest")?.as_ref().and_then(|manifest| manifest.get(CAPABILITY))
-            .is_some_and(|entry| (entry.is_object() || entry.is_array()) && entry["supported"] == Value::Bool(true));
-        if !listed || !supported {
+        let capabilities = DB::optional_document_row(&worker, "capabilities")?;
+        let manifest = DB::optional_document_row(&worker, "capability_manifest")?;
+        if !super::sticky_execution::registration_supports(capabilities.as_ref(), manifest.as_ref(), CAPABILITY) {
             return Err(rejected("worker_capability_not_supported"));
         }
         let claim = &body["sticky_cache"];
+        // Only a worker that takes history suffixes reports its cached boundary.
+        if !claim["history_last_sequence"].is_null()
+            && !super::sticky_execution::takes_suffixes(capabilities.as_ref())
+        {
+            return Err(refuse(StatusCode::UNPROCESSABLE_ENTITY, "sticky_history_suffix_not_supported"));
+        }
         let build = DB::optional_string(&worker, "build_id")?.or(DB::optional_string(&worker, "sdk_version")?);
         let text = |field: &str| claim[field].as_str().unwrap_or_default();
         if text("worker_id") != owner || text("run_id") != DB::string(task, "workflow_run_id")? || text("workflow_id") != workflow_id
@@ -1019,8 +1023,15 @@ store_methods! {
     /// transaction, the run keeps the worker's affinity until the claim's
     /// expiry, or loses it when it closed. The worker's reported cache
     /// metrics are kept with its registration.
-    pub(super) async fn record_sticky_cache(tx: &mut Transaction<'_, DB>, namespace: &str, run_id: &str, body: &Value) -> Result<()> {
+    pub(super) async fn record_sticky_cache(tx: &mut Transaction<'_, DB>, namespace: &str, run_id: &str, task_id: &str, body: &Value) -> Result<()> {
         let claim = &body["sticky_cache"];
+        if let Some(through) = claim["history_last_sequence"].as_i64() {
+            let row = Self::query("SELECT payload FROM workflow_tasks WHERE id=$1").bind(task_id).fetch_one(&mut **tx).await?;
+            let mut payload = DB::optional_document_row(&row, "payload")?.filter(Value::is_object).unwrap_or_else(|| json!({}));
+            payload[super::sticky_execution::CACHED_THROUGH_KEY] = json!(through);
+            Self::query("UPDATE workflow_tasks SET payload=$1 WHERE id=$2").bind(DB::document(&payload)).bind(task_id)
+                .execute(&mut **tx).await?;
+        }
         let owner = super::text(body, "lease_owner")?;
         let status = Self::query("SELECT status FROM workflow_runs WHERE id=$1").bind(run_id).fetch_one(&mut **tx).await?;
         let retain = !super::sticky_execution::terminal(&DB::string(&status, "status")?);

rust/src/runtime/sticky_execution.rs

@@ -56,6 +56,30 @@ pub(super) fn claim_mode(worker: Option<&str>, until: Option<DateTime<Utc>>, lea
     if worker == lease_owner { STICKY_HIT_EXPECTED } else { FORCED_COLD_REPLAY }
 }
 
+/// The worker capability and server capability for history suffixes.
+pub(super) const SUFFIX_CAPABILITY: &str = "sticky_history_suffix";
+/// The task payload key holding a completed task's reported cached boundary.
+pub(super) const CACHED_THROUGH_KEY: &str = "sticky_history_last_sequence";
+
+/// `server_capabilities.sticky_history_suffix`.
+pub(super) fn suffix_capability() -> Value {
+    json!({"supported": true, "version": 1})
+}
+
+/// Whether a registration takes history suffixes: it lists the capability.
+/// The stored manifest keeps only PHP's three affinity entries, so the
+/// listing is the opt-in; sticky execution itself must still be supported.
+pub(super) fn takes_suffixes(capabilities: Option<&Value>) -> bool {
+    capabilities.and_then(Value::as_array).is_some_and(|listed| listed.iter().any(|item| item.as_str() == Some(SUFFIX_CAPABILITY)))
+}
+
+/// Whether a registration lists and supports `capability`.
+pub(super) fn registration_supports(capabilities: Option<&Value>, manifest: Option<&Value>, capability: &str) -> bool {
+    capabilities.and_then(Value::as_array).is_some_and(|listed| listed.iter().any(|item| item.as_str() == Some(capability)))
+        && manifest.and_then(|manifest| manifest.get(capability))
+            .is_some_and(|entry| (entry.is_object() || entry.is_array()) && entry["supported"] == Value::Bool(true))
+}
+
 /// A run in one of these states keeps no affinity.
 pub(super) fn terminal(status: &str) -> bool {
     matches!(status, "completed" | "failed" | "cancelled" | "terminated" | "continued")

rust/src/runtime/store.rs

@@ -2844,10 +2844,32 @@ where
                     }
                 }
                 let last = DB::number(&run, "last_history_sequence")?;
+                // Server #1017: a worker that takes history suffixes and is
+                // expected to hit its cache gets only the events after the
+                // boundary it reported with its previous completion of the run.
+                let capabilities = DB::optional_document_row(&worker, "capabilities")?;
+                let suffix_worker = super::sticky_execution::takes_suffixes(capabilities.as_ref())
+                    && super::sticky_execution::registration_supports(capabilities.as_ref(),
+                        DB::optional_document_row(&worker, "capability_manifest")?.as_ref(), "sticky_execution");
+                let after = if suffix_worker && sticky_mode == Some(super::sticky_execution::STICKY_HIT_EXPECTED) {
+                    Self::query("SELECT payload FROM workflow_tasks WHERE workflow_run_id=$1 AND task_type='workflow' AND status='completed' AND lease_owner=$2 AND id<>$3 ORDER BY id DESC LIMIT 1")
+                        .bind(&run_id).bind(worker_id).bind(&task_id).fetch_optional(&mut **tx).await?
+                        .map(|row| DB::optional_document_row(&row, "payload")).transpose()?.flatten()
+                        .and_then(|payload| payload[super::sticky_execution::CACHED_THROUGH_KEY].as_i64())
+                        .filter(|through| (1..=last).contains(through))
+                        .unwrap_or(0)
+                } else {
+                    0
+                };
                 let (mut events, next) = match page {
-                    Some(size) => Self::history_page_connection_with_total(&mut **tx, &run_id, 0, size, Some(last)).await?,
+                    Some(size) => Self::history_page_connection_with_total(&mut **tx, &run_id, after, size, Some(last)).await?,
+                    None if after > 0 => Self::history_page_connection_with_total(&mut **tx, &run_id, after, last - after, Some(last)).await?,
                     None => (Self::whole_history(&mut **tx, &run_id, last).await?, None),
                 };
+                if suffix_worker {
+                    claim["history_after_sequence"] = json!(after);
+                    claim["history_last_sequence"] = json!(last);
+                }
                 Self::with_signal_sequences(&mut **tx, &run_id, &mut events).await?;
                 let budget = Self::workflow_history_budget(
                     tx,
@@ -3717,7 +3739,7 @@ where
             }
         };
         if sticky {
-            Self::record_sticky_cache(&mut tx, self.namespace.as_ref(), &run_id, &body).await?;
+            Self::record_sticky_cache(&mut tx, self.namespace.as_ref(), &run_id, task_id, &body).await?;
         }
         let transaction = if deferred { Some(tx) } else {
             tx.commit().await?;
@@ -5813,7 +5835,7 @@ fn capabilities_template() -> Value {
         "workflow_memo_updates": false, "cooperative_cancellation": true, "prepared_local_activities": false,
         "prepared_local_activity_groups": false, "prepared_local_activity_cancellation_policies": [],
         "worker_sessions": super::worker_sessions::capability(), "worker_session_verbs": ["create", "heartbeat", "close"],
-        "sticky_execution": super::sticky_execution::describe(), "local_activities": super::local_activities::manifest(), "message_streams": super::message_streams::worker_capability(),
+        "sticky_execution": super::sticky_execution::describe(), "sticky_history_suffix": super::sticky_execution::suffix_capability(), "local_activities": super::local_activities::manifest(), "message_streams": super::message_streams::worker_capability(),
         "update_validation_tasks": true, "synchronous_update_validation": super::update_validation::capability(),
         "workflow_updates": true, "query_tasks": true, "query_task_authority":true, "query_task_poll_request_idempotency": true,
         "claimed_task_ready_age": true, "activity_lease_renewal": true, "activity_cancellation_acknowledgement": true,

rust/tests/execution/sticky_execution.rs

@@ -98,3 +98,65 @@ async fn an_expired_claim_forces_a_cold_replay() {
     let second = poll(&app, "sticky-a", "workflow").await;
     assert_eq!(second["sticky_replay_mode"], "forced_cold_replay", "{second}");
 }
+
+async fn suffix_worker(app: &Router, worker: &str) {
+    let entry = json!({"supported":true,"minimum_protocol_version":"1.18"});
+    let (status, body) = request(app, "POST", "/api/worker/register", json!({"worker_id":worker,"task_queue":"test",
+        "runtime":"php","build_id":"build-1","supported_workflow_types":["echo"],"supported_activity_types":[],
+        "capabilities":["sticky_execution","sticky_history_suffix"],
+        "capability_manifest":{"local_activities":{"supported":false,"minimum_protocol_version":"1.18"},
+            "worker_sessions":{"supported":false,"minimum_protocol_version":"1.18"},"sticky_execution":entry,
+            "sticky_history_suffix":{"supported":true,"minimum_protocol_version":"1.18","version":1}}})).await;
+    assert_eq!(status, StatusCode::CREATED, "{body}");
+}
+
+fn sequences(task: &Value) -> Vec<i64> {
+    task["history_events"].as_array().unwrap().iter().map(|event| event["sequence"].as_i64().unwrap()).collect()
+}
+
+// Server #1017 part 2: a worker that takes suffixes gets only the events
+// after the boundary it reported with its previous completion.
+#[tokio::test]
+async fn a_suffix_worker_gets_only_the_events_after_its_cached_boundary() {
+    let database = TestDatabase::new().await;
+    let app = router(database.open().await.unwrap());
+    let (_, cluster) = request(&app, "GET", "/api/cluster/info", Value::Null).await;
+    assert_eq!(cluster["worker_protocol"]["server_capabilities"]["sticky_history_suffix"], json!({"supported":true,"version":1}));
+    suffix_worker(&app, "suffix-a").await;
+    start(&app, "suffix").await;
+
+    let first = poll(&app, "suffix-a", "workflow").await;
+    assert_eq!(first["history_after_sequence"], 0, "{first}");
+    let delivered = sequences(&first);
+    assert_eq!(delivered.first(), Some(&1), "{first}");
+    assert_eq!(first["history_last_sequence"], json!(delivered.last().unwrap()), "{first}");
+    let mut sticky = claim(&first, "suffix-a", 300);
+    sticky["history_last_sequence"] = json!(delivered.last().unwrap());
+    let (status, body) = complete_sticky(&app, &first, sticky).await;
+    assert_eq!(status, StatusCode::OK, "{body}");
+
+    signal(&app, "suffix").await;
+    let second = poll(&app, "suffix-a", "workflow").await;
+    assert_eq!(second["sticky_replay_mode"], "sticky_hit_expected", "{second}");
+    let after = *delivered.last().unwrap();
+    assert_eq!(second["history_after_sequence"], after, "{second}");
+    let suffix = sequences(&second);
+    assert_eq!(suffix.first(), Some(&(after + 1)), "{second}");
+    assert!(suffix.windows(2).all(|pair| pair[1] == pair[0] + 1), "{suffix:?}");
+    assert_eq!(second["history_last_sequence"], json!(suffix.last().unwrap()), "{second}");
+}
+
+#[tokio::test]
+async fn only_suffix_workers_report_or_receive_a_boundary() {
+    let database = TestDatabase::new().await;
+    let app = router(database.open().await.unwrap());
+    sticky_worker(&app, "plain").await;
+    start(&app, "plain-suffix").await;
+    let first = poll(&app, "plain", "workflow").await;
+    assert!(first.get("history_after_sequence").is_none(), "{first}");
+    let mut sticky = claim(&first, "plain", 300);
+    sticky["history_last_sequence"] = json!(3);
+    let (status, body) = complete_sticky(&app, &first, sticky).await;
+    assert_eq!(status, StatusCode::UNPROCESSABLE_ENTITY, "{body}");
+    assert_eq!(body["reason"], "sticky_history_suffix_not_supported", "{body}");
+}