From 083646eda18391a9b5c0d4c7f7e1245052076ac0 Mon Sep 17 00:00:00 2001 From: Durable Workflow Date: Mon, 19 Oct 2026 22:16:27 +0000 Subject: [PATCH 1/3] Stop replaying live query leases to fresh polls from the same worker Native returned a worker's live query lease to any later poll from that worker. Published PHP replays a live lease only to a poll carrying the matching poll_request_id, and native does not advertise query poll-request idempotency, so its SDK workers send no id. A late handler cancelled at its view of the query deadline polled again a fraction of a millisecond before the stored expiry and received the identical claim a second time. That was the intermittent "Compare published query deadlines and installed Worker authority" failure, reproduced locally with full logs. A live lease held by a current registration now stays with the poll that claimed it, whatever its owner. A lease held by a stale registration can still be reclaimed. Three native controls asserted the old replay; they now assert that the later poll receives no task while the original claim still completes. With the previous runtime, two of them fail; the third polls through a reconfigured runtime where the old replay never applied. Co-Authored-By: Claude Opus 5.5 --- docs/rust-server-port.md | 18 ++++-------------- rust/src/runtime/queries.rs | 8 +++----- rust/tests/execution.rs | 17 +++-------------- rust/tests/execution/namespaces.rs | 2 +- rust/tests/execution/query_authority.rs | 5 +---- 5 files changed, 12 insertions(+), 38 deletions(-) diff --git a/docs/rust-server-port.md b/docs/rust-server-port.md index 349f23089..ac964afa9 100644 --- a/docs/rust-server-port.md +++ b/docs/rust-server-port.md @@ -26,13 +26,13 @@ Updated October 19, 2026, after merging #460. - **Shared-suite reliability.** Each 13-job run still tends to show one intermittent failure in a step its change does not touch, and qualification records classify each one. Diagnosed and corrected: published PHP SQLite - reference lock pressure, and CPU observer client timing. Diagnosed, decision - pending: the query-authority deadline-boundary replay (see known gaps). + reference lock pressure, CPU observer client timing, and native replay of a + live query lease to a fresh poll, which published PHP never does without a + matching poll request id. Not yet diagnosed: one native-test failure in `execution-and-codec` step 8 on main's post-#478 run, which four local repetitions under the CI limits did not reproduce. -- **Next action.** Decide the query deadline-boundary replay contract and add - its fixture. Then correct native timer dispatch under CPU load with a +- **Next action.** Correct native timer dispatch under CPU load with a deadline-health control. Then continue the remaining parity families, the backup-first in-place takeover and the performance baseline on disposable hosts. @@ -74,16 +74,6 @@ contract fixture before its correction. and native then correctly fences the attempt. A heartbeat timeout comfortably above the request timeout, still inside the 9-second load wave, needs its own reviewed profile change. -- **Query replay at the deadline boundary.** Both implementations replay a live - query lease to its owner on a fresh poll without a request id. The installed - Rust query client cancels a late handler at its view of the query deadline and - polls again. Native can answer that poll a fraction of a millisecond before - the stored microsecond expiry and replay the task, so the client records a - second claim. Published PHP's request overhead keeps it past the deadline in - practice. One of six local SQLite cohort runs failed this way. Decide whether - native stops replaying a lease with less than a minimum delivery budget left, - or whether the fixture counts a boundary replay as the same claim, then add the - fixture before changing either side. ## Contract and verification inventory diff --git a/rust/src/runtime/queries.rs b/rust/src/runtime/queries.rs index f91c4e04e..ef3502096 100644 --- a/rust/src/runtime/queries.rs +++ b/rust/src/runtime/queries.rs @@ -495,15 +495,13 @@ where config, ) .await?; + // A live lease belongs to the poll that claimed it. Published PHP replays + // it only to a matching poll_request_id, which native does not advertise, + // so a later poll from the same worker is an independent slot. if existing_proof .as_deref() .is_some_and(|proof| task["registration_incarnation"] == proof) { - if existing_owner == owner { - let delivered = Self::query_delivery(&task); - tx.commit().await?; - return Ok(QueryClaim::Task(delivered)); - } continue; } } diff --git a/rust/tests/execution.rs b/rust/tests/execution.rs index 06e0bfe2e..b9770cf43 100644 --- a/rust/tests/execution.rs +++ b/rust/tests/execution.rs @@ -2248,19 +2248,6 @@ async fn await_query(app: &Router, worker: &str) -> Value { .expect("query must be admitted without holding the database pool") } -fn assert_query_replay(first: &Value, repeated: Value) { - let initial = - chrono::DateTime::parse_from_rfc3339(first["server_time"].as_str().unwrap()).unwrap(); - let current = - chrono::DateTime::parse_from_rfc3339(repeated["server_time"].as_str().unwrap()).unwrap(); - assert!(current >= initial); - let mut replay = repeated; - // Only this observation clock is nondeterministic. Original query timing, - // lease, identities, values and every business-history field remain exact. - replay["server_time"] = first["server_time"].clone(); - assert_eq!(replay, *first); -} - #[tokio::test] async fn queries_fence_recovered_leases_across_pools_without_mutating_the_run() { let database = TestDatabase::new().await; @@ -2357,7 +2344,9 @@ async fn queries_fence_recovered_leases_across_pools_without_mutating_the_run() ); assert_eq!(projected, before["events"]); assert_eq!(claim["query_arguments"], input["input"]); - assert_query_replay(&claim, poll(&a, "query-a", "query").await); + // A later poll from the claiming worker is an independent slot. The live + // lease stays with the original claim, as in published PHP. + assert!(poll(&a, "query-a", "query").await.is_null()); assert!(poll(&b, "query-b", "query").await.is_null()); let finish_path = format!( "/api/worker/query-tasks/{}/complete", diff --git a/rust/tests/execution/namespaces.rs b/rust/tests/execution/namespaces.rs index 17dc71cd0..2fc43d238 100644 --- a/rust/tests/execution/namespaces.rs +++ b/rust/tests/execution/namespaces.rs @@ -281,7 +281,7 @@ async fn named_idle_query_poll_preserves_same_id_default_namespace_claim() { idle.1["task"].is_null(), "named poll must not replay another namespace's same-ID lease" ); - assert_query_replay(&original, poll(&app, "shared-query-worker", "query").await); + assert!(poll(&app, "shared-query-worker", "query").await.is_null()); assert_eq!(request(&app,"POST", &format!("/api/worker/query-tasks/{}/complete",original["query_task_id"].as_str().unwrap()), json!({"lease_owner":"shared-query-worker","query_task_attempt":original["query_task_attempt"],"result_envelope":result})).await.0, diff --git a/rust/tests/execution/query_authority.rs b/rust/tests/execution/query_authority.rs index 21737681f..e41c3e39e 100644 --- a/rust/tests/execution/query_authority.rs +++ b/rust/tests/execution/query_authority.rs @@ -875,10 +875,7 @@ async fn zero_is_immediate_and_later_policy_does_not_rebase_an_existing_query() .await; assert_eq!(expired.0, StatusCode::GATEWAY_TIMEOUT, "{}", expired.1); assert_eq!(expired.1["reason"], "query_task_not_claimed"); - assert_query_replay( - &original.claim, - poll(&changed, "authority-worker", "query").await, - ); + assert!(poll(&changed, "authority-worker", "query").await.is_null()); let complete_path = format!( "/api/worker/query-tasks/{}/complete", original.claim["query_task_id"].as_str().unwrap() From 6663ab64aea43b6fda8ed27fb1eab46ac1b64314 Mon Sep 17 00:00:00 2001 From: Durable Workflow Date: Tue, 20 Oct 2026 03:49:55 +0000 Subject: [PATCH 2/3] Replay a live query lease to the poll request that claimed it The previous commit stopped replaying a live lease to any later poll from the same worker. The installed Python SDK retries a timed-out query poll once with the same poll_request_id, so that retry would no longer recover its own claim and the task would stay leased until its lease expired. As in published PHP, the claim now records its poll_request_id, a retry of that same request receives the lease again, and every other poll stays an independent slot. The ID is not delivered to the worker. A non-string ID is refused like other invalid poll requests. A raw-protocol probe against published PHP Server 2.8.0 and this build leases, replays and withholds the same tasks for repeated, newer, missing, foreign-worker and concurrent duplicate request IDs. Native still lacks PHP's cached empty results, supersession of an older waiting poll and its leased/empty poll statuses, so it keeps advertising query_task_poll_request_idempotency as false. Part of #325. Co-Authored-By: Claude Opus 5.5 --- rust/src/runtime/queries.rs | 20 ++++- rust/tests/execution.rs | 3 + rust/tests/execution/query_poll_requests.rs | 93 +++++++++++++++++++++ 3 files changed, 113 insertions(+), 3 deletions(-) create mode 100644 rust/tests/execution/query_poll_requests.rs diff --git a/rust/src/runtime/queries.rs b/rust/src/runtime/queries.rs index ef3502096..093bf9846 100644 --- a/rust/src/runtime/queries.rs +++ b/rust/src/runtime/queries.rs @@ -421,6 +421,7 @@ where "terminal_receipt", "registration_incarnation", "legacy_refusal_at", + "poll_request_id", ] { delivered .as_object_mut() @@ -439,6 +440,11 @@ where ) -> Result { let owner = text(body, "worker_id")?; let queue = text(body, "task_queue")?; + let request_id = body + .get("poll_request_id") + .filter(|value| !value.is_null()) + .map(|_| text(body, "poll_request_id")) + .transpose()?; let mut tx = self.begin().await?; let worker = Self::query("SELECT * FROM workflow_worker_registrations WHERE namespace=$1 AND worker_id=$2 AND status='active'") .bind(self.namespace.as_ref()).bind(owner) @@ -495,13 +501,20 @@ where config, ) .await?; - // A live lease belongs to the poll that claimed it. Published PHP replays - // it only to a matching poll_request_id, which native does not advertise, - // so a later poll from the same worker is an independent slot. + // A live lease belongs to the poll request that claimed it. As in + // published PHP, only a retry of that same request receives it + // again; any other poll from the same worker is an independent slot. if existing_proof .as_deref() .is_some_and(|proof| task["registration_incarnation"] == proof) { + if existing_owner == owner + && request_id.is_some_and(|id| task["poll_request_id"] == id) + { + let delivered = Self::query_delivery(&task); + tx.commit().await?; + return Ok(QueryClaim::Task(delivered)); + } continue; } } @@ -520,6 +533,7 @@ where } task["lease_expires_at"] = json!(authority::wire(expires)); task["registration_incarnation"] = json!(proof); + task["poll_request_id"] = request_id.map_or(Value::Null, |id| json!(id)); Self::put_query(&mut tx, &key, &task, authority::expiration(&task)?).await?; tx.commit().await?; return Ok(QueryClaim::Task(Self::query_delivery(&task))); diff --git a/rust/tests/execution.rs b/rust/tests/execution.rs index b9770cf43..437f95640 100644 --- a/rust/tests/execution.rs +++ b/rust/tests/execution.rs @@ -129,6 +129,9 @@ mod query_arguments; #[path = "execution/query_authority.rs"] mod query_authority; +#[path = "execution/query_poll_requests.rs"] +mod query_poll_requests; + #[path = "execution/cross_kind_poll.rs"] mod cross_kind_poll; diff --git a/rust/tests/execution/query_poll_requests.rs b/rust/tests/execution/query_poll_requests.rs new file mode 100644 index 000000000..6a960acea --- /dev/null +++ b/rust/tests/execution/query_poll_requests.rs @@ -0,0 +1,93 @@ +//! A live query lease belongs to the poll request that claimed it. As in +//! published PHP, a retry of that same request receives it again; any other +//! poll, from the same worker or another, is an independent slot. +use super::*; + +async fn poll_with(app: &Router, worker: &str, request_id: Option<&str>) -> Value { + let mut body = json!({"worker_id":worker,"task_queue":"test","timeout_seconds":0}); + if let Some(id) = request_id { + body["poll_request_id"] = json!(id); + } + let (status, response) = request(app, "POST", "/api/worker/query-tasks/poll", body).await; + assert_eq!(status, StatusCode::OK, "{response}"); + response["task"].clone() +} + +fn assert_replay(first: &Value, repeated: &Value) { + let initial = + chrono::DateTime::parse_from_rfc3339(first["server_time"].as_str().unwrap()).unwrap(); + let current = + chrono::DateTime::parse_from_rfc3339(repeated["server_time"].as_str().unwrap()).unwrap(); + assert!(current >= initial); + let mut replay = repeated.clone(); + // Only the observation clock may differ. The claim, its attempt, lease, + // timing and every payload field stay exactly as first delivered. + replay["server_time"] = first["server_time"].clone(); + assert_eq!(replay, *first); +} + +#[tokio::test] +async fn query_lease_replays_only_to_the_poll_request_that_claimed_it() { + let database = TestDatabase::new().await; + let runtime = database.open().await.unwrap(); + let app = router(runtime.clone()); + query_worker(&app, "request-worker").await; + query_worker(&app, "other-worker").await; + let started = start(&app, "poll-request-query").await; + let workflow_task = poll(&app, "request-worker", "workflow").await; + let completed = finish_task( + &app, + &workflow_task, + json!([{"type":"complete_workflow","result":envelope(Payload::Long(7))}]), + ) + .await; + assert_eq!(completed.0, StatusCode::OK, "{}", completed.1); + let query_path = format!( + "/api/workflows/poll-request-query/runs/{}/query/state", + started["run_id"].as_str().unwrap() + ); + let waiting = { + let app = app.clone(); + tokio::spawn(async move { + request(&app, "POST", &query_path, + json!({"input":envelope(Payload::Array(vec![Payload::Long(1)]))})).await + }) + }; + let claim = tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + let task = poll_with(&app, "request-worker", Some("first")).await; + if !task.is_null() { + break task; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await + .expect("query is admitted"); + assert!(claim.get("poll_request_id").is_none(), "{claim}"); + + assert_replay(&claim, &poll_with(&app, "request-worker", Some("first")).await); + assert!(poll_with(&app, "request-worker", Some("second")).await.is_null()); + assert!(poll_with(&app, "request-worker", None).await.is_null()); + assert!(poll_with(&app, "other-worker", Some("first")).await.is_null()); + // A newer request does not take the lease away from the original one. + assert_replay(&claim, &poll_with(&app, "request-worker", Some("first")).await); + + let invalid = request(&app, "POST", "/api/worker/query-tasks/poll", + json!({"worker_id":"request-worker","task_queue":"test","timeout_seconds":0,"poll_request_id":7})).await; + assert_eq!(invalid.0, StatusCode::UNPROCESSABLE_ENTITY, "{}", invalid.1); + + let finish = request( + &app, + "POST", + &format!("/api/worker/query-tasks/{}/complete", claim["query_task_id"].as_str().unwrap()), + json!({"lease_owner":"request-worker","query_task_attempt":claim["query_task_attempt"], + "result_envelope":envelope(Payload::Long(7))}), + ) + .await; + assert_eq!(finish.0, StatusCode::OK, "{}", finish.1); + assert_eq!(waiting.await.unwrap().0, StatusCode::OK); + assert!(poll_with(&app, "request-worker", Some("first")).await.is_null()); + runtime.close().await; + database.remove().await; +} From 3c11e7b435347b081cda80d11abb5fd8112fe07c Mon Sep 17 00:00:00 2001 From: Durable Workflow Date: Tue, 20 Oct 2026 04:41:15 +0000 Subject: [PATCH 3/3] Remember query poll answers per request, as published PHP does Published PHP keeps a short-lived record for every query poll request and advertises query_task_poll_request_idempotency. Native now does the same: - a retry of a request gets the answer the original got, including an empty one, and a remembered task is delivered again only while this worker's registration still holds its live lease; - a duplicate of a request still waiting shares that request's answer; - a new request becomes the worker's current one, and a waiting poll whose request is no longer current ends at once with an empty answer. Records live in dw_query_cache under their own prefix, with PHP's lifetime rules measured from the native query poll timeout. Each new request retires up to 32 expired poll records. Polls without a request ID are unchanged. A raw-protocol probe now matches published PHP Server 2.8.0 on every step except the leased/empty status names. New tests cover remembered empty answers, supersession and shared duplicates on SQLite, PostgreSQL 17 and MySQL 8.0. Part of #325. Co-Authored-By: Claude Opus 5.5 --- rust/src/runtime/backend.rs | 3 + rust/src/runtime/queries.rs | 46 ++- rust/src/runtime/query_poll_requests.rs | 314 ++++++++++++++++++++ rust/src/runtime/store.rs | 2 +- rust/tests/execution/query_poll_requests.rs | 133 ++++++++- 5 files changed, 484 insertions(+), 14 deletions(-) create mode 100644 rust/src/runtime/query_poll_requests.rs diff --git a/rust/src/runtime/backend.rs b/rust/src/runtime/backend.rs index 7aaa69405..82fa1dc12 100644 --- a/rust/src/runtime/backend.rs +++ b/rust/src/runtime/backend.rs @@ -44,6 +44,8 @@ pub(super) trait Backend: Database { const QUERY_CACHE_UPDATE: &'static str = "UPDATE dw_query_cache SET value=$1,expiration=$2 WHERE \"key\"=$3"; const QUERY_CACHE_DELETE: &'static str = "DELETE FROM dw_query_cache WHERE \"key\"=$1"; + // A bounded batch of expired per-request query poll rows. + const QUERY_POLL_EXPIRED: &'static str = "SELECT \"key\" FROM dw_query_cache WHERE \"key\" LIKE 'dw-native-query-poll:%' AND expiration<=$1 ORDER BY \"key\" LIMIT 32"; const MEMO_INITIAL_INSERT: &'static str = "INSERT INTO workflow_memos(workflow_run_id,workflow_instance_id,\"key\",value,upserted_at_sequence,inherited_from_parent,created_at,updated_at,portable_value,portable_value_sequence) VALUES ($1,$2,$3,$4,0,false,$5,$6,$7,0)"; const MEMO_INHERIT_INSERT: &'static str = "INSERT INTO workflow_memos(workflow_run_id,workflow_instance_id,\"key\",value,upserted_at_sequence,inherited_from_parent,created_at,updated_at,portable_value,portable_value_sequence) VALUES ($1,$2,$3,$4,1,true,$5,$6,$7,1)"; const MEMO_READ: &'static str = "SELECT \"key\",value,portable_value,portable_value_sequence,upserted_at_sequence FROM workflow_memos WHERE workflow_run_id=$1 ORDER BY \"key\",id LIMIT 101"; @@ -388,6 +390,7 @@ impl Backend for MySql { const QUERY_CACHE_UPDATE: &'static str = "UPDATE dw_query_cache SET value=$1,expiration=$2 WHERE `key`=$3"; const QUERY_CACHE_DELETE: &'static str = "DELETE FROM dw_query_cache WHERE `key`=$1"; + const QUERY_POLL_EXPIRED: &'static str = "SELECT `key` FROM dw_query_cache WHERE `key` LIKE 'dw-native-query-poll:%' AND expiration<=$1 ORDER BY `key` LIMIT 32"; const TIMER_TASK_SQL: &'static str = "SELECT id,payload FROM workflow_tasks WHERE workflow_run_id=? AND task_type='timer' AND status='ready' AND JSON_UNQUOTE(JSON_EXTRACT(payload,'$.timer_id'))=? AND BINARY namespace=BINARY ?"; const TIMER_SCHEDULE_SQL: &str = "SELECT payload FROM workflow_history_events WHERE workflow_run_id=? AND event_type='TimerScheduled' AND JSON_UNQUOTE(JSON_EXTRACT(payload,'$.timer_id'))=? ORDER BY sequence DESC LIMIT 1"; // First use the preserved run/sequence index to materialize only IDs and diff --git a/rust/src/runtime/queries.rs b/rust/src/runtime/queries.rs index 093bf9846..53e9b5f3c 100644 --- a/rust/src/runtime/queries.rs +++ b/rust/src/runtime/queries.rs @@ -23,10 +23,15 @@ const SNAPSHOT_BYTES: usize = 1024 * 1024; pub(super) const ARGUMENT_BYTES: usize = 64 * 1024; pub(super) const RESULT_BYTES: usize = 256 * 1024; +#[path = "query_poll_requests.rs"] +mod poll_requests; + enum QueryClaim { Task(Value), Empty, StaleWorker, + // A newer poll request from the same worker replaced this one. + Superseded, } fn poll_result(task: Option, status: &str) -> Value { @@ -344,23 +349,53 @@ where let deadline = tokio::time::Instant::now() + Duration::from_secs(timeout); // One request belongs to the registration observed before its wait. // A later same-ID registration cannot silently upgrade this poll. + let request_id = body + .get("poll_request_id") + .filter(|value| !value.is_null()) + .map(|_| text(&body, "poll_request_id")) + .transpose()?; let Some(proof) = Self::query_worker_proof(&self.pool, self.namespace.as_ref(), owner, config).await? else { return Ok(poll_result(None, "stale_worker_registration")); }; + match request_id { + None => { + self.wait_for_query(&body, config, &proof, timeout, deadline, None) + .await + } + Some(request_id) => { + let keys = self.poll_keys(&body, request_id).await?; + self.coordinated_query_poll(&body, config, &proof, timeout, deadline, &keys) + .await + } + } + } + + // One poll's long wait. With request keys, a newer request from the same + // worker ends this wait, as in published PHP. + async fn wait_for_query( + &self, + body: &Value, + config: &QueryTaskConfig, + proof: &str, + timeout: u64, + deadline: tokio::time::Instant, + keys: Option<&poll_requests::PollKeys>, + ) -> Result { let mut cross_kind = if timeout > 0 { - self.poll_wake_snapshot(&body, "query").await? + self.poll_wake_snapshot(body, "query").await? } else { None }; loop { let wake = self.wake.notified(); - match self.claim_query(&body, config, &proof).await? { + match self.claim_query(body, config, proof, keys).await? { QueryClaim::Task(task) => return Ok(poll_result(Some(task), "task")), QueryClaim::StaleWorker => { return Ok(poll_result(None, "stale_worker_registration")); } + QueryClaim::Superseded => return Ok(poll_result(None, "timeout")), QueryClaim::Empty => {} } if tokio::time::Instant::now() >= deadline { @@ -437,6 +472,7 @@ where body: &Value, config: &QueryTaskConfig, expected_proof: &str, + keys: Option<&poll_requests::PollKeys>, ) -> Result { let owner = text(body, "worker_id")?; let queue = text(body, "task_queue")?; @@ -457,6 +493,12 @@ where let Some(proof) = proof.filter(|proof| proof == expected_proof) else { return Ok(QueryClaim::StaleWorker); }; + if let Some(keys) = keys + && !Self::poll_request_is_current(&mut tx, keys).await? + { + tx.commit().await?; + return Ok(QueryClaim::Superseded); + } if DB::string(&worker, "task_queue")? != queue { return Err(refuse(StatusCode::CONFLICT, "task_queue_mismatch")); } diff --git a/rust/src/runtime/query_poll_requests.rs b/rust/src/runtime/query_poll_requests.rs new file mode 100644 index 000000000..1f996d57a --- /dev/null +++ b/rust/src/runtime/query_poll_requests.rs @@ -0,0 +1,314 @@ +//! Per-request query poll results, kept as published PHP keeps them +//! (`QueryTaskPollRequestStore`). A retry of a poll request gets the answer the +//! original got, a duplicate still in flight waits for that answer, and a newer +//! request from the same worker ends an older one that is still waiting. +use super::*; +use sha2::{Digest, Sha256}; + +// Distinct from the query task prefix, so query scans never see these rows. +const POLL_PREFIX: &str = "dw-native-query-poll:"; + +pub(super) struct PollKeys { + request_id: String, + pending: String, + result: String, + current: String, +} + +fn pending_ttl(config: &QueryTaskConfig) -> i64 { + (config.poll_timeout_seconds + 5).max(5) +} + +fn empty_ttl(config: &QueryTaskConfig) -> i64 { + pending_ttl(config).min(60).max(5) +} + +impl Store +where + for<'c> &'c mut DB::Connection: Executor<'c, Database = DB>, + DB::Arguments: IntoArguments, + for<'q> &'q str: Encode<'q, DB> + Type, + for<'q> Option<&'q str>: Encode<'q, DB>, + for<'q> Option: Encode<'q, DB>, + for<'q> String: Encode<'q, DB> + Type, + for<'q> i64: Encode<'q, DB> + Type, + for<'r> i64: Decode<'r, DB>, + usize: ColumnIndex, +{ + // Keys cover namespace, queue, the registration's build ID, the worker and + // the request ID, as PHP's cache keys do. + pub(super) async fn poll_keys(&self, body: &Value, request_id: &str) -> Result { + let owner = text(body, "worker_id")?; + let queue = text(body, "task_queue")?; + let build = Self::query( + "SELECT build_id FROM workflow_worker_registrations WHERE namespace=$1 AND worker_id=$2", + ) + .bind(self.namespace.as_ref()) + .bind(owner) + .fetch_optional(&self.pool) + .await? + .map(|row| DB::optional_string(&row, "build_id")) + .transpose()? + .flatten(); + // MySQL and MariaDB keys are varchar(64): keep 128 bits of the digest. + let key = |kind: &str, id: Option<&str>| { + let material = json!([kind, self.namespace.as_ref(), queue, build, owner, id]); + let digest: String = Sha256::digest(material.to_string().as_bytes())[..16] + .iter() + .map(|byte| format!("{byte:02x}")) + .collect(); + format!("{POLL_PREFIX}{digest}") + }; + Ok(PollKeys { + request_id: request_id.to_owned(), + pending: key("pending", Some(request_id)), + result: key("result", Some(request_id)), + current: key("current", None), + }) + } + + async fn poll_row(tx: &mut Transaction<'_, DB>, key: &str) -> Result> { + let Some(row) = Self::query(DB::QUERY_CACHE_GET) + .bind(key) + .fetch_optional(&mut **tx) + .await? + else { + return Ok(None); + }; + if DB::number(&row, "expiration")? <= now() { + Self::query(DB::QUERY_CACHE_DELETE) + .bind(key) + .execute(&mut **tx) + .await?; + return Ok(None); + } + Ok(Some(serde_json::from_str(&DB::string(&row, "value")?)?)) + } + + async fn put_poll_row( + tx: &mut Transaction<'_, DB>, + key: &str, + value: &Value, + ttl: i64, + ) -> Result<()> { + let exists = Self::query(DB::QUERY_CACHE_GET) + .bind(key) + .fetch_optional(&mut **tx) + .await? + .is_some(); + let expiration = now() + ttl; + if exists { + Self::query(DB::QUERY_CACHE_UPDATE) + .bind(value.to_string()) + .bind(expiration) + .bind(key) + .execute(&mut **tx) + .await?; + } else { + Self::query(DB::QUERY_CACHE_INSERT) + .bind(key) + .bind(value.to_string()) + .bind(expiration) + .execute(&mut **tx) + .await?; + } + Ok(()) + } + + // Workers send a fresh request ID with every poll, so retire a bounded + // batch of expired poll rows whenever a new request arrives. + async fn sweep_poll_rows(tx: &mut Transaction<'_, DB>) -> Result<()> { + let expired = Self::query(DB::QUERY_POLL_EXPIRED) + .bind(now()) + .fetch_all(&mut **tx) + .await?; + for row in expired { + Self::query(DB::QUERY_CACHE_DELETE) + .bind(DB::string(&row, "key")?) + .execute(&mut **tx) + .await?; + } + Ok(()) + } + + pub(super) async fn poll_request_is_current( + tx: &mut Transaction<'_, DB>, + keys: &PollKeys, + ) -> Result { + Ok(Self::poll_row(tx, &keys.current) + .await? + .is_some_and(|current| current["poll_request_id"] == keys.request_id.as_str())) + } + + // A request becomes the worker's current one only when it is new. A retry + // of an older request never takes that place back. + async fn mark_current_if_fresh(&self, keys: &PollKeys, config: &QueryTaskConfig) -> Result<()> { + let mut tx = self.begin().await?; + Self::sweep_poll_rows(&mut tx).await?; + if Self::poll_row(&mut tx, &keys.pending).await?.is_none() + && Self::poll_row(&mut tx, &keys.result).await?.is_none() + { + Self::put_poll_row( + &mut tx, + &keys.current, + &json!({"poll_request_id": keys.request_id}), + pending_ttl(config), + ) + .await?; + } + tx.commit().await?; + Ok(()) + } + + // The remembered answer for this request. A remembered task is delivered + // again only while this worker's registration still holds its live lease. + async fn cached_query_poll( + &self, + keys: &PollKeys, + body: &Value, + proof: &str, + config: &QueryTaskConfig, + ) -> Result> { + let owner = text(body, "worker_id")?; + let queue = text(body, "task_queue")?; + let mut tx = self.begin().await?; + let Some(result) = Self::poll_row(&mut tx, &keys.result).await? else { + tx.commit().await?; + return Ok(None); + }; + let Some(task_id) = result["query_task_id"].as_str() else { + tx.commit().await?; + let status = result["poll_status"].as_str().unwrap_or("timeout"); + return Ok(Some(poll_result(None, status))); + }; + let at = authority::now(); + let task = Self::poll_row(&mut tx, &key(task_id)?).await?; + let deliverable = task.as_ref().filter(|task| { + task["status"] == "leased" + && task["namespace"] == self.namespace.as_ref() + && task["task_queue"] == queue + && task["lease_owner"] == owner + && task["registration_incarnation"] == proof + && authority::timestamp(&task["lease_expires_at"]).is_some_and(|lease| lease > at) + && authority::deadline(task).is_some_and(|deadline| deadline > at) + }); + let response = match deliverable { + Some(task) => poll_result(Some(Self::query_delivery(task)), "task"), + None => { + Self::put_poll_row( + &mut tx, + &keys.result, + &json!({"query_task_id": null, "poll_status": "timeout"}), + empty_ttl(config), + ) + .await?; + poll_result(None, "timeout") + } + }; + tx.commit().await?; + Ok(Some(response)) + } + + async fn try_start_query_poll(&self, keys: &PollKeys, config: &QueryTaskConfig) -> Result { + let mut tx = self.begin().await?; + if Self::poll_row(&mut tx, &keys.pending).await?.is_some() { + tx.commit().await?; + return Ok(false); + } + Self::put_poll_row(&mut tx, &keys.pending, &json!({"pending": true}), pending_ttl(config)) + .await?; + tx.commit().await?; + Ok(true) + } + + async fn remember_query_poll( + &self, + keys: &PollKeys, + response: &Value, + config: &QueryTaskConfig, + ) -> Result<()> { + let task_id = response["task"]["query_task_id"].as_str(); + let ttl = match authority::timestamp(&response["task"]["lease_expires_at"]) { + Some(lease) if task_id.is_some() => { + ((lease - authority::now()).num_seconds() + 5).clamp(5, 3600) + } + _ => empty_ttl(config), + }; + let mut tx = self.begin().await?; + Self::put_poll_row( + &mut tx, + &keys.result, + &json!({"query_task_id": task_id, "poll_status": response["poll_status"]}), + ttl, + ) + .await?; + Self::query(DB::QUERY_CACHE_DELETE) + .bind(&keys.pending) + .execute(&mut *tx) + .await?; + tx.commit().await?; + Ok(()) + } + + async fn forget_pending_query_poll(&self, keys: &PollKeys) -> Result<()> { + let mut tx = self.begin().await?; + Self::query(DB::QUERY_CACHE_DELETE) + .bind(&keys.pending) + .execute(&mut *tx) + .await?; + tx.commit().await?; + Ok(()) + } + + pub(super) async fn coordinated_query_poll( + &self, + body: &Value, + config: &QueryTaskConfig, + proof: &str, + timeout: u64, + deadline: tokio::time::Instant, + keys: &PollKeys, + ) -> Result { + self.mark_current_if_fresh(keys, config).await?; + let owner = text(body, "worker_id")?; + for _ in 0..3 { + let current = + Self::query_worker_proof(&self.pool, self.namespace.as_ref(), owner, config).await?; + if current.as_deref() != Some(proof) { + return Ok(poll_result(None, "stale_worker_registration")); + } + if let Some(cached) = self.cached_query_poll(keys, body, proof, config).await? { + return Ok(cached); + } + if self.try_start_query_poll(keys, config).await? { + let outcome = self + .wait_for_query(body, config, proof, timeout, deadline, Some(keys)) + .await; + match &outcome { + Ok(response) => self.remember_query_poll(keys, response, config).await?, + Err(_) => self.forget_pending_query_poll(keys).await?, + } + return outcome; + } + // Another poll with this request ID is waiting. Share its answer. + let follower_deadline = tokio::time::Instant::now() + + Duration::from_millis(config.poll_timeout_seconds.max(1) as u64 * 1000 + 250); + loop { + if let Some(cached) = self.cached_query_poll(keys, body, proof, config).await? { + return Ok(cached); + } + let mut tx = self.begin().await?; + let pending = Self::poll_row(&mut tx, &keys.pending).await?.is_some(); + tx.commit().await?; + if !pending || tokio::time::Instant::now() >= follower_deadline { + break; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + } + Ok(self + .cached_query_poll(keys, body, proof, config) + .await? + .unwrap_or_else(|| poll_result(None, "timeout"))) + } +} diff --git a/rust/src/runtime/store.rs b/rust/src/runtime/store.rs index 7b9362e9c..c3bb6c9d2 100644 --- a/rust/src/runtime/store.rs +++ b/rust/src/runtime/store.rs @@ -2911,7 +2911,7 @@ pub(crate) fn capabilities() -> Value { "worker_deregistration_fencing":{"schema":"durable-workflow.v2.worker-deregistration.v1","supported":true,"receipt_retention_seconds":600,"endpoint":"/worker/registrations/{workerId}/deregister"}, "workflow_memo_updates": false, "cooperative_cancellation": true, "prepared_local_activities": false, "worker_sessions": false, "sticky_execution": false, "local_activities": false, "message_streams": false, - "workflow_updates": true, "query_tasks": true, "query_task_authority":true, "query_task_poll_request_idempotency": false}); + "workflow_updates": true, "query_tasks": true, "query_task_authority":true, "query_task_poll_request_idempotency": true}); if super::memos::UPDATES_SUPPORTED { value["supported_workflow_task_commands"] .as_array_mut() diff --git a/rust/tests/execution/query_poll_requests.rs b/rust/tests/execution/query_poll_requests.rs index 6a960acea..1180755eb 100644 --- a/rust/tests/execution/query_poll_requests.rs +++ b/rust/tests/execution/query_poll_requests.rs @@ -53,17 +53,14 @@ async fn query_lease_replays_only_to_the_poll_request_that_claimed_it() { json!({"input":envelope(Payload::Array(vec![Payload::Long(1)]))})).await }) }; - let claim = tokio::time::timeout(std::time::Duration::from_secs(5), async { - loop { - let task = poll_with(&app, "request-worker", Some("first")).await; - if !task.is_null() { - break task; - } - tokio::time::sleep(std::time::Duration::from_millis(10)).await; - } - }) - .await - .expect("query is admitted"); + // One waiting poll claims it. Repeating an ID that already got an empty + // answer would only get that answer again. + let (status, claimed) = request(&app, "POST", "/api/worker/query-tasks/poll", + json!({"worker_id":"request-worker","task_queue":"test","timeout_seconds":5, + "poll_request_id":"first"})).await; + assert_eq!(status, StatusCode::OK, "{claimed}"); + let claim = claimed["task"].clone(); + assert!(claim.is_object(), "query is admitted: {claimed}"); assert!(claim.get("poll_request_id").is_none(), "{claim}"); assert_replay(&claim, &poll_with(&app, "request-worker", Some("first")).await); @@ -91,3 +88,117 @@ async fn query_lease_replays_only_to_the_poll_request_that_claimed_it() { runtime.close().await; database.remove().await; } + +async fn poll_waiting(app: &Router, request_id: &str, timeout: u64) -> (Value, std::time::Duration) { + let started = std::time::Instant::now(); + let (status, response) = request(app, "POST", "/api/worker/query-tasks/poll", + json!({"worker_id":"request-worker","task_queue":"test","timeout_seconds":timeout, + "poll_request_id":request_id})).await; + assert_eq!(status, StatusCode::OK, "{response}"); + (response, started.elapsed()) +} + +// One workflow that a "request-worker" can be queried about, plus a helper +// that admits one query and returns its pending caller. +async fn queryable(name: &str) -> (TestDatabase, Runtime, Router, String) { + let database = TestDatabase::new().await; + let runtime = database.open().await.unwrap(); + let app = router(runtime.clone()); + query_worker(&app, "request-worker").await; + let started = start(&app, name).await; + let workflow_task = poll(&app, "request-worker", "workflow").await; + let completed = finish_task( + &app, + &workflow_task, + json!([{"type":"complete_workflow","result":envelope(Payload::Long(7))}]), + ) + .await; + assert_eq!(completed.0, StatusCode::OK, "{}", completed.1); + let path = format!("/api/workflows/{name}/runs/{}/query/state", started["run_id"].as_str().unwrap()); + (database, runtime, app, path) +} + +fn admit_query(app: &Router, path: &str) -> tokio::task::JoinHandle<(StatusCode, Value)> { + let app = app.clone(); + let path = path.to_owned(); + tokio::spawn(async move { + request(&app, "POST", &path, json!({"input":envelope(Payload::Array(vec![Payload::Long(1)]))})).await + }) +} + +async fn finish_query(app: &Router, task: &Value) { + let finish = request( + app, + "POST", + &format!("/api/worker/query-tasks/{}/complete", task["query_task_id"].as_str().unwrap()), + json!({"lease_owner":"request-worker","query_task_attempt":task["query_task_attempt"], + "result_envelope":envelope(Payload::Long(7))}), + ) + .await; + assert_eq!(finish.0, StatusCode::OK, "{}", finish.1); +} + +#[tokio::test] +async fn a_retried_empty_poll_request_returns_its_remembered_answer() { + let (database, runtime, app, path) = queryable("remembered-empty-query").await; + let (first, _) = poll_waiting(&app, "empty-first", 0).await; + assert!(first["task"].is_null(), "{first}"); + let waiting = admit_query(&app, &path); + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + // The retry gets the original's answer, even though a query now waits. + let (retry, _) = poll_waiting(&app, "empty-first", 0).await; + assert!(retry["task"].is_null(), "{retry}"); + assert_eq!(retry["poll_status"], first["poll_status"]); + let (fresh, _) = poll_waiting(&app, "empty-fresh", 5).await; + assert!(fresh["task"].is_object(), "{fresh}"); + finish_query(&app, &fresh["task"]).await; + assert_eq!(waiting.await.unwrap().0, StatusCode::OK); + runtime.close().await; + database.remove().await; +} + +#[tokio::test] +async fn a_newer_poll_request_ends_an_older_waiting_poll() { + let (database, runtime, app, _) = queryable("superseded-query-poll").await; + let older = { + let app = app.clone(); + tokio::spawn(async move { poll_waiting(&app, "older", 10).await }) + }; + tokio::time::sleep(std::time::Duration::from_millis(500)).await; + let (newer, _) = poll_waiting(&app, "newer", 0).await; + assert!(newer["task"].is_null(), "{newer}"); + let (ended, elapsed) = older.await.unwrap(); + assert!(ended["task"].is_null(), "{ended}"); + assert!(elapsed < std::time::Duration::from_secs(3), "the older poll ended early: {elapsed:?}"); + // Its retry keeps the same answer instead of waiting again. + let (retry, elapsed) = poll_waiting(&app, "older", 10).await; + assert!(retry["task"].is_null(), "{retry}"); + assert!(elapsed < std::time::Duration::from_secs(2), "{elapsed:?}"); + runtime.close().await; + database.remove().await; +} + +#[tokio::test] +async fn duplicate_polls_with_one_request_id_share_one_claim() { + let (database, runtime, app, path) = queryable("shared-query-claim").await; + let polls: Vec<_> = (0..2) + .map(|_| { + let app = app.clone(); + tokio::spawn(async move { poll_waiting(&app, "shared", 5).await }) + }) + .collect(); + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + let waiting = admit_query(&app, &path); + let mut claims = vec![]; + for poll in polls { + claims.push(poll.await.unwrap().0["task"].clone()); + } + assert!(claims[0].is_object(), "{}", claims[0]); + assert_eq!(claims[0]["query_task_id"], claims[1]["query_task_id"]); + assert_eq!(claims[0]["query_task_attempt"], claims[1]["query_task_attempt"]); + assert_eq!(claims[0]["query_task_attempt"], 1); + finish_query(&app, &claims[0]).await; + assert_eq!(waiting.await.unwrap().0, StatusCode::OK); + runtime.close().await; + database.remove().await; +}