From d5d880fd754b044f1c13ba48cee4fcfd87f054df Mon Sep 17 00:00:00 2001 From: Durable Workflow Date: Tue, 20 Oct 2026 07:41:11 +0000 Subject: [PATCH 1/3] Let held long polls wait without the transition lock Each held native poll used to repeat its full claim every 50 ms: a transaction under the global transition lock that read the registration, deleted expired poll receipts and selected candidates. With no work, idle cost grew by about 0.035 cores per waiting poll, two-thirds of it in MySQL, and every probe queued behind, and delayed, real transitions. A held poll now claims under the lock on its first probe. Afterwards it claims only when a lock-free read finds something it could claim or must refuse: an eligible task or query, a receipt for its own poll request, or a removed, moved or replaced registration. The read runs at once when a commit on the same node wakes the poll, and otherwise every 50 ms, which is how a poll still observes other server processes. A false positive costs one empty claim, and a commit racing a read is seen on the next read. New tests hold the transition lock from outside the runtime and check that idle workflow, activity and query polls still end at their deadline (each fails with the previous probing), and that a held poll still claims a task or query, and refuses a removed registration, committed by another node. Co-Authored-By: Claude Opus 5.5 --- docs/rust-server-transition-locking.md | 9 ++ rust/src/runtime/backend.rs | 4 +- rust/src/runtime/mod.rs | 1 + rust/src/runtime/poll_readiness.rs | 97 +++++++++++++++++ rust/src/runtime/queries.rs | 22 ++-- rust/src/runtime/store.rs | 17 ++- rust/tests/execution.rs | 2 + rust/tests/execution/lock_free_waiting.rs | 123 ++++++++++++++++++++++ rust/tests/execution/read_snapshots.rs | 6 +- 9 files changed, 265 insertions(+), 16 deletions(-) create mode 100644 rust/src/runtime/poll_readiness.rs create mode 100644 rust/tests/execution/lock_free_waiting.rs diff --git a/docs/rust-server-transition-locking.md b/docs/rust-server-transition-locking.md index e593eab1a..4e6b6bad5 100644 --- a/docs/rust-server-transition-locking.md +++ b/docs/rust-server-transition-locking.md @@ -26,6 +26,15 @@ read-only snapshot (`Store::begin_read()`, `READ_SNAPSHOT_BEGIN` in `backend.rs`). Control-plane reads such as describe and run history query the pool without a transaction. +A held long poll takes the lock for its first probe, and afterwards only when +a lock-free read finds something to claim or refuse: an eligible task (or +query), a receipt for the poll's own request ID, or a removed, moved or +replaced registration (`poll_readiness.rs`). The read runs at once when a +commit on the same node wakes the poll, and otherwise every 50 ms, which is +how a poll observes another server process. An idle poll therefore neither +waits for the lock nor delays transitions, and an unrelated commit does not +send every waiting poll to the lock. + ## Measured cost Under the CPU deadline-health profile on a local MySQL 8.0 cell, per-step diff --git a/rust/src/runtime/backend.rs b/rust/src/runtime/backend.rs index 7aaa69405..7dc404387 100644 --- a/rust/src/runtime/backend.rs +++ b/rust/src/runtime/backend.rs @@ -337,8 +337,8 @@ impl Backend for Postgres { async fn begin(pool: &Pool) -> Result> { let mut tx = pool.begin().await?; // The first single-namespace slice serializes its durable transitions - // across independent pools/processes. Empty long polls release this - // transaction before waiting. Throughput/idle cost remain unqualified. + // across independent pools/processes. Held long polls wait without it + // (poll_readiness.rs). Throughput/idle cost remain unqualified. sqlx::query("SELECT pg_advisory_xact_lock(4924501473315676163)") .execute(&mut *tx) .await?; diff --git a/rust/src/runtime/mod.rs b/rust/src/runtime/mod.rs index 0d051c048..e83d15aa2 100644 --- a/rust/src/runtime/mod.rs +++ b/rust/src/runtime/mod.rs @@ -35,6 +35,7 @@ mod memos; pub mod mysql; mod namespaces; pub mod postgres; +mod poll_readiness; mod previews; mod queries; mod query_authority; diff --git a/rust/src/runtime/poll_readiness.rs b/rust/src/runtime/poll_readiness.rs new file mode 100644 index 000000000..f0727bf34 --- /dev/null +++ b/rust/src/runtime/poll_readiness.rs @@ -0,0 +1,97 @@ +//! Lock-free waiting for held long polls. +//! +//! A held poll takes the transition lock only to claim or refuse. After its +//! first claim it reads, without the lock, whether anything now exists that it +//! could claim or must refuse: an eligible task, a receipt for its own request, +//! or a changed registration. A same-node commit wakes it to read at once; +//! timed reads observe other server processes. A false positive costs one locked +//! claim that finds nothing, and a commit racing a read is seen on the next one. +use super::{ + Result, backend::Backend, query_authority::{self, QueryTaskConfig}, store::Store, text, +}; +use serde_json::Value; +use sqlx::{ColumnIndex, Decode, Encode, Executor, IntoArguments, Type}; +use std::time::Duration; + +/// How often a held poll reads for work committed by another server process. +pub(super) const READ_INTERVAL: Duration = Duration::from_millis(50); + +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, +{ + /// Whether a held workflow or activity poll should claim again. + pub(super) async fn task_poll_ready(&self, body: &Value, kind: &str) -> Result { + let worker_id = text(body, "worker_id")?; + let queue = text(body, "task_queue")?; + // The claim refuses a removed or moved registration. + let registration = Self::query( + "SELECT task_queue FROM workflow_worker_registrations WHERE worker_id=$1 AND namespace=$2", + ) + .bind(worker_id) + .bind(self.namespace.as_ref()) + .fetch_optional(&self.pool) + .await?; + match registration { + Some(row) if DB::string(&row, "task_queue")? == queue => {} + _ => return Ok(true), + } + // A concurrent poll with the same request ID may have claimed already; + // the claim replays its receipt. + if let Some(request_id) = body.get("poll_request_id").and_then(Value::as_str) + && Self::query("SELECT task_id FROM dw_poll_receipts WHERE worker_id=$1 AND kind=$2 AND request_id=$3 AND namespace=$4") + .bind(worker_id).bind(kind).bind(request_id).bind(self.namespace.as_ref()) + .fetch_optional(&self.pool).await?.is_some() + { + return Ok(true); + } + let at = super::store::now(); + Ok(Self::query(DB::TASK_CANDIDATES_SQL) + .bind(self.namespace.as_ref()) + .bind(queue) + .bind(kind) + .bind(DB::bind_time(at)) + .bind(DB::bind_time(at)) + .fetch_optional(&self.pool) + .await? + .is_some()) + } + + /// Whether a held query poll should claim again. `proof` is the + /// registration this poll belongs to. + pub(super) async fn query_poll_ready( + &self, + body: &Value, + config: &QueryTaskConfig, + proof: &str, + ) -> Result { + let owner = text(body, "worker_id")?; + // The claim refuses a replaced, removed or stale registration. + if Self::query_worker_proof(&self.pool, self.namespace.as_ref(), owner, config) + .await? + .as_deref() + != Some(proof) + { + return Ok(true); + } + let at = super::store::now(); + let query_now = query_authority::wire(at); + Ok(Self::query(DB::POLL_WAKE_QUERIES_SQL) + .bind(at.timestamp()) + .bind(self.namespace.as_ref()) + .bind(text(body, "task_queue")?) + .bind(query_now.clone()) + .bind(query_now) + .fetch_optional(&self.pool) + .await? + .is_some()) + } +} diff --git a/rust/src/runtime/queries.rs b/rust/src/runtime/queries.rs index f91c4e04e..ed9a5588e 100644 --- a/rust/src/runtime/queries.rs +++ b/rust/src/runtime/queries.rs @@ -354,14 +354,19 @@ where } else { None }; + // Only the first probe claims under the lock unconditionally; later + // probes claim after a lock-free read finds something to claim or refuse. + let mut first = true; loop { let wake = self.wake.notified(); - match self.claim_query(&body, config, &proof).await? { - QueryClaim::Task(task) => return Ok(poll_result(Some(task), "task")), - QueryClaim::StaleWorker => { - return Ok(poll_result(None, "stale_worker_registration")); + if std::mem::take(&mut first) || self.query_poll_ready(&body, config, &proof).await? { + match self.claim_query(&body, config, &proof).await? { + QueryClaim::Task(task) => return Ok(poll_result(Some(task), "task")), + QueryClaim::StaleWorker => { + return Ok(poll_result(None, "stale_worker_registration")); + } + QueryClaim::Empty => {} } - QueryClaim::Empty => {} } if tokio::time::Instant::now() >= deadline { return Ok(poll_result(None, "timeout")); @@ -375,11 +380,14 @@ where json!({"task":null,"poll_status":"task_queue_changed","server_capabilities":super::store::capabilities()}), ); } - tokio::select! { _ = wake => {}, _ = tokio::time::sleep(Duration::from_millis(50)) => {} } + tokio::select! { + _ = wake => {}, + _ = tokio::time::sleep(super::poll_readiness::READ_INTERVAL) => {}, + } } } - async fn query_worker_proof<'e, E>( + pub(super) async fn query_worker_proof<'e, E>( executor: E, namespace: &str, owner: &str, diff --git a/rust/src/runtime/store.rs b/rust/src/runtime/store.rs index 6490d0f95..1dddc5dab 100644 --- a/rust/src/runtime/store.rs +++ b/rust/src/runtime/store.rs @@ -933,11 +933,17 @@ where } else { None }; + // Only the first probe claims under the lock unconditionally. Later + // probes, on a same-node wake or a timed read that also observes + // other server processes, claim only after a lock-free read finds + // something to claim or refuse. + let mut first = true; loop { - // Register the notification before probing; periodic probing also - // observes work committed by an independent process. + // Register the notification before probing. let wake = self.wake.notified(); - if let Some(task) = self.claim(&body, kind, protocol_version).await? { + if (std::mem::take(&mut first) || self.task_poll_ready(&body, kind).await?) + && let Some(task) = self.claim(&body, kind, protocol_version).await? + { return Ok( json!({"task": task, "poll_status": "task", "server_capabilities": capabilities()}), ); @@ -956,7 +962,10 @@ where json!({"task":null,"poll_status":"task_queue_changed","server_capabilities":capabilities()}), ); } - tokio::select! { _ = wake => {}, _ = tokio::time::sleep(Duration::from_millis(50)) => {} } + tokio::select! { + _ = wake => {}, + _ = tokio::time::sleep(super::poll_readiness::READ_INTERVAL) => {}, + } } } diff --git a/rust/tests/execution.rs b/rust/tests/execution.rs index 8c2b7ba54..7eeceb171 100644 --- a/rust/tests/execution.rs +++ b/rust/tests/execution.rs @@ -61,6 +61,8 @@ mod version_markers; #[path = "execution/activity_deadlines.rs"] mod activity_deadlines; +#[path = "execution/lock_free_waiting.rs"] +mod lock_free_waiting; #[path = "execution/read_snapshots.rs"] mod read_snapshots; diff --git a/rust/tests/execution/lock_free_waiting.rs b/rust/tests/execution/lock_free_waiting.rs new file mode 100644 index 000000000..bd22759dd --- /dev/null +++ b/rust/tests/execution/lock_free_waiting.rs @@ -0,0 +1,123 @@ +//! A held long poll waits without the transition lock, so idle polls neither +//! wait behind transitions nor delay them. It still claims work, and refuses +//! a removed registration, committed by another server process. +use super::read_snapshots::hold_transition_lock; +use super::*; +use std::time::Duration; + +fn held_poll( + app: &Router, + worker: &str, + kind: &str, + seconds: u64, +) -> tokio::task::JoinHandle<(StatusCode, Value)> { + let app = app.clone(); + let path = format!("/api/worker/{kind}-tasks/poll"); + let body = json!({"worker_id":worker,"task_queue":"test","timeout_seconds":seconds}); + tokio::spawn(async move { request(&app, "POST", &path, body).await }) +} + +#[tokio::test] +async fn held_polls_reach_their_deadline_while_a_transition_holds_the_lock() { + let database = TestDatabase::new().await; + let runtime = database.open().await.unwrap(); + let app = router(runtime.clone()); + register(&app, "waiting-worker", json!(["echo"]), json!(["echo-activity"])).await; + query_worker(&app, "waiting-query-worker").await; + let polls = [ + held_poll(&app, "waiting-worker", "workflow", 2), + held_poll(&app, "waiting-worker", "activity", 2), + held_poll(&app, "waiting-query-worker", "query", 2), + ]; + // Each poll's first claim takes the lock; let them finish and wait. + tokio::time::sleep(Duration::from_millis(500)).await; + let writer = hold_transition_lock(&database).await; + for poll in polls { + // The polls end about 1.5 s from now. A probe that took the lock + // would wait for it, and the lock is held past this bound. + let (status, body) = tokio::time::timeout(Duration::from_secs(3), poll) + .await + .expect("an idle poll must reach its deadline without the transition lock") + .unwrap(); + assert_eq!(status, StatusCode::OK, "{body}"); + assert!(body["task"].is_null(), "{body}"); + } + writer.release().await; + runtime.close().await; + database.remove().await; +} + +#[tokio::test] +async fn held_polls_observe_work_and_removed_registrations_from_another_node() { + let database = TestDatabase::new().await; + let node_a = database.open().await.unwrap(); + let node_b = database.open().await.unwrap(); + let app_a = router(node_a.clone()); + let app_b = router(node_b.clone()); + register(&app_a, "observing-worker", json!(["echo"]), json!([])).await; + query_worker(&app_a, "observing-query-worker").await; + + // Node B's commits cannot wake node A's poll in process. + let waiting = held_poll(&app_a, "observing-worker", "workflow", 20); + tokio::time::sleep(Duration::from_millis(300)).await; + let started = start(&app_b, "observed-elsewhere").await; + let (status, claimed) = tokio::time::timeout(Duration::from_secs(5), waiting) + .await + .expect("a held poll claims work another node committed") + .unwrap(); + assert_eq!(status, StatusCode::OK, "{claimed}"); + assert_eq!(claimed["task"]["run_id"], started["run_id"], "{claimed}"); + let completed = request( + &app_a, + "POST", + &format!( + "/api/worker/workflow-tasks/{}/complete", + claimed["task"]["task_id"].as_str().unwrap() + ), + completion( + &claimed["task"], + json!([{"type":"complete_workflow","result":envelope(Payload::Null)}]), + ), + ) + .await; + assert_eq!(completed.0, StatusCode::OK, "{}", completed.1); + + let waiting = held_poll(&app_a, "observing-query-worker", "query", 5); + tokio::time::sleep(Duration::from_millis(300)).await; + let query = { + let app = app_b.clone(); + let path = format!( + "/api/workflows/observed-elsewhere/runs/{}/query/state", + started["run_id"].as_str().unwrap() + ); + let input = json!({"input":envelope(Payload::Array(vec![Payload::Long(1)]))}); + tokio::spawn(async move { request(&app, "POST", &path, input).await }) + }; + let (status, claimed) = tokio::time::timeout(Duration::from_secs(3), waiting) + .await + .expect("a held query poll claims a query another node admitted") + .unwrap(); + assert_eq!(status, StatusCode::OK, "{claimed}"); + assert_eq!(claimed["task"]["query_name"], "state", "{claimed}"); + query.abort(); + + let waiting = held_poll(&app_a, "observing-worker", "workflow", 20); + tokio::time::sleep(Duration::from_millis(300)).await; + let removed = request( + &app_b, + "DELETE", + "/api/worker/registrations/observing-worker", + Value::Null, + ) + .await; + assert_eq!(removed.0, StatusCode::OK, "{}", removed.1); + let (status, refused) = tokio::time::timeout(Duration::from_secs(5), waiting) + .await + .expect("a held poll refuses a registration another node removed") + .unwrap(); + assert_eq!(status, StatusCode::CONFLICT, "{refused}"); + assert_eq!(refused["reason"], "worker_not_registered", "{refused}"); + node_a.close().await; + node_b.close().await; + database.remove().await; +} diff --git a/rust/tests/execution/read_snapshots.rs b/rust/tests/execution/read_snapshots.rs index 8dcb0d814..2c0fdcfb9 100644 --- a/rust/tests/execution/read_snapshots.rs +++ b/rust/tests/execution/read_snapshots.rs @@ -3,7 +3,7 @@ use super::*; use sqlx::Connection; -enum Writer { +pub(super) enum Writer { Sqlite(sqlx::SqliteConnection), Postgres(sqlx::PgConnection), MySql(sqlx::MySqlConnection), @@ -11,7 +11,7 @@ enum Writer { // Takes the lock a native transition takes, on a connection outside the // runtime's pool, as a transition from another server process would. -async fn hold_transition_lock(database: &TestDatabase) -> Writer { +pub(super) async fn hold_transition_lock(database: &TestDatabase) -> Writer { match database { TestDatabase::Sqlite(dir) => { let mut connection = sqlx::SqliteConnection::connect(&format!( @@ -56,7 +56,7 @@ async fn hold_transition_lock(database: &TestDatabase) -> Writer { } impl Writer { - async fn release(self) { + pub(super) async fn release(self) { match self { Self::Sqlite(mut connection) => { sqlx::raw_sql("ROLLBACK") From 98ef1af7af18bb13c12f15c55452a68b37071399 Mon Sep 17 00:00:00 2001 From: Durable Workflow Date: Tue, 20 Oct 2026 08:09:56 +0000 Subject: [PATCH 2/3] Back off a held poll's timed reads up to 250 ms On the October 20 profile host, lock-free reads every 50 ms cut MySQL's idle cost at 24 waiting polls only from about 0.53 to 0.46 cores, and the native process stayed near 0.28 cores. The cost is the read rate, 20 per second per poll, not the lock. A held poll now reads 50 ms after a probe and then at doubling intervals up to 250 ms. A same-node wake still reads at once and starts again from 50 ms, so same-node work is unaffected, and work committed by another server process is observed within 250 ms. No wait runs past the poll's deadline. Co-Authored-By: Claude Opus 5.5 --- docs/rust-server-transition-locking.md | 9 +++++---- rust/src/runtime/poll_readiness.rs | 28 ++++++++++++++++++++++++-- rust/src/runtime/queries.rs | 6 ++---- rust/src/runtime/store.rs | 6 ++---- 4 files changed, 35 insertions(+), 14 deletions(-) diff --git a/docs/rust-server-transition-locking.md b/docs/rust-server-transition-locking.md index 4e6b6bad5..9bd74a9c8 100644 --- a/docs/rust-server-transition-locking.md +++ b/docs/rust-server-transition-locking.md @@ -30,10 +30,11 @@ A held long poll takes the lock for its first probe, and afterwards only when a lock-free read finds something to claim or refuse: an eligible task (or query), a receipt for the poll's own request ID, or a removed, moved or replaced registration (`poll_readiness.rs`). The read runs at once when a -commit on the same node wakes the poll, and otherwise every 50 ms, which is -how a poll observes another server process. An idle poll therefore neither -waits for the lock nor delays transitions, and an unrelated commit does not -send every waiting poll to the lock. +commit on the same node wakes the poll. Otherwise it runs 50 ms after a probe +and then at doubling intervals up to 250 ms, so a poll observes work committed +by another server process within 250 ms. An idle poll therefore neither waits +for the lock nor delays transitions, and an unrelated commit does not send +every waiting poll to the lock. ## Measured cost diff --git a/rust/src/runtime/poll_readiness.rs b/rust/src/runtime/poll_readiness.rs index f0727bf34..459db2fe1 100644 --- a/rust/src/runtime/poll_readiness.rs +++ b/rust/src/runtime/poll_readiness.rs @@ -12,9 +12,33 @@ use super::{ use serde_json::Value; use sqlx::{ColumnIndex, Decode, Encode, Executor, IntoArguments, Type}; use std::time::Duration; +use tokio::{sync::futures::Notified, time::Instant}; -/// How often a held poll reads for work committed by another server process. -pub(super) const READ_INTERVAL: Duration = Duration::from_millis(50); +/// Paces a held poll's timed reads, which are how it observes another server +/// process. The first read comes 50 ms after a probe, and each later one waits +/// twice as long, up to 250 ms. A same-node wake reads at once and starts again +/// from 50 ms. No wait runs past the poll's deadline. +pub(super) struct ReadPace { + next: Duration, +} + +impl ReadPace { + const FIRST: Duration = Duration::from_millis(50); + const LONGEST: Duration = Duration::from_millis(250); + + pub(super) fn new() -> Self { + Self { next: Self::FIRST } + } + + /// Returns at a same-node wake or when the next timed read is due. + pub(super) async fn wait(&mut self, wake: Notified<'_>, deadline: Instant) { + let pause = self.next.min(deadline.saturating_duration_since(Instant::now())); + tokio::select! { + _ = wake => self.next = Self::FIRST, + _ = tokio::time::sleep(pause) => self.next = (self.next * 2).min(Self::LONGEST), + } + } +} impl Store where diff --git a/rust/src/runtime/queries.rs b/rust/src/runtime/queries.rs index ed9a5588e..cdca5859c 100644 --- a/rust/src/runtime/queries.rs +++ b/rust/src/runtime/queries.rs @@ -357,6 +357,7 @@ where // Only the first probe claims under the lock unconditionally; later // probes claim after a lock-free read finds something to claim or refuse. let mut first = true; + let mut pace = super::poll_readiness::ReadPace::new(); loop { let wake = self.wake.notified(); if std::mem::take(&mut first) || self.query_poll_ready(&body, config, &proof).await? { @@ -380,10 +381,7 @@ where json!({"task":null,"poll_status":"task_queue_changed","server_capabilities":super::store::capabilities()}), ); } - tokio::select! { - _ = wake => {}, - _ = tokio::time::sleep(super::poll_readiness::READ_INTERVAL) => {}, - } + pace.wait(wake, deadline).await; } } diff --git a/rust/src/runtime/store.rs b/rust/src/runtime/store.rs index 1dddc5dab..9b045bee8 100644 --- a/rust/src/runtime/store.rs +++ b/rust/src/runtime/store.rs @@ -938,6 +938,7 @@ where // other server processes, claim only after a lock-free read finds // something to claim or refuse. let mut first = true; + let mut pace = super::poll_readiness::ReadPace::new(); loop { // Register the notification before probing. let wake = self.wake.notified(); @@ -962,10 +963,7 @@ where json!({"task":null,"poll_status":"task_queue_changed","server_capabilities":capabilities()}), ); } - tokio::select! { - _ = wake => {}, - _ = tokio::time::sleep(super::poll_readiness::READ_INTERVAL) => {}, - } + pace.wait(wake, deadline).await; } } From eb272676d352c75fc79c08b0d929170ff3e19402 Mon Sep 17 00:00:00 2001 From: Durable Workflow Date: Tue, 20 Oct 2026 08:46:36 +0000 Subject: [PATCH 3/3] Keep the readiness reads beside the claims they mirror The task read now sits next to the task claim in store.rs and the query read next to the query claim in queries.rs, inside their existing Store impls. poll_readiness.rs keeps only the read pacing. No behavior changes. Co-Authored-By: Claude Opus 5.5 --- docs/rust-server-transition-locking.md | 2 +- rust/src/runtime/backend.rs | 2 +- rust/src/runtime/poll_readiness.rs | 92 ++------------------------ rust/src/runtime/queries.rs | 32 ++++++++- rust/src/runtime/store.rs | 37 +++++++++++ 5 files changed, 74 insertions(+), 91 deletions(-) diff --git a/docs/rust-server-transition-locking.md b/docs/rust-server-transition-locking.md index 9bd74a9c8..392f7743c 100644 --- a/docs/rust-server-transition-locking.md +++ b/docs/rust-server-transition-locking.md @@ -29,7 +29,7 @@ pool without a transaction. A held long poll takes the lock for its first probe, and afterwards only when a lock-free read finds something to claim or refuse: an eligible task (or query), a receipt for the poll's own request ID, or a removed, moved or -replaced registration (`poll_readiness.rs`). The read runs at once when a +replaced registration (`Store::task_poll_ready` and `Store::query_poll_ready`). The read runs at once when a commit on the same node wakes the poll. Otherwise it runs 50 ms after a probe and then at doubling intervals up to 250 ms, so a poll observes work committed by another server process within 250 ms. An idle poll therefore neither waits diff --git a/rust/src/runtime/backend.rs b/rust/src/runtime/backend.rs index 7dc404387..b31a8349d 100644 --- a/rust/src/runtime/backend.rs +++ b/rust/src/runtime/backend.rs @@ -338,7 +338,7 @@ impl Backend for Postgres { let mut tx = pool.begin().await?; // The first single-namespace slice serializes its durable transitions // across independent pools/processes. Held long polls wait without it - // (poll_readiness.rs). Throughput/idle cost remain unqualified. + // (Store::task_poll_ready). Throughput/idle cost remain unqualified. sqlx::query("SELECT pg_advisory_xact_lock(4924501473315676163)") .execute(&mut *tx) .await?; diff --git a/rust/src/runtime/poll_readiness.rs b/rust/src/runtime/poll_readiness.rs index 459db2fe1..6c0265829 100644 --- a/rust/src/runtime/poll_readiness.rs +++ b/rust/src/runtime/poll_readiness.rs @@ -3,14 +3,10 @@ //! A held poll takes the transition lock only to claim or refuse. After its //! first claim it reads, without the lock, whether anything now exists that it //! could claim or must refuse: an eligible task, a receipt for its own request, -//! or a changed registration. A same-node commit wakes it to read at once; -//! timed reads observe other server processes. A false positive costs one locked -//! claim that finds nothing, and a commit racing a read is seen on the next one. -use super::{ - Result, backend::Backend, query_authority::{self, QueryTaskConfig}, store::Store, text, -}; -use serde_json::Value; -use sqlx::{ColumnIndex, Decode, Encode, Executor, IntoArguments, Type}; +//! or a changed registration (`Store::task_poll_ready`, `Store::query_poll_ready`). +//! A same-node commit wakes it to read at once; timed reads observe other +//! server processes. A false positive costs one locked claim that finds +//! nothing, and a commit racing a read is seen on the next one. use std::time::Duration; use tokio::{sync::futures::Notified, time::Instant}; @@ -39,83 +35,3 @@ impl ReadPace { } } } - -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, -{ - /// Whether a held workflow or activity poll should claim again. - pub(super) async fn task_poll_ready(&self, body: &Value, kind: &str) -> Result { - let worker_id = text(body, "worker_id")?; - let queue = text(body, "task_queue")?; - // The claim refuses a removed or moved registration. - let registration = Self::query( - "SELECT task_queue FROM workflow_worker_registrations WHERE worker_id=$1 AND namespace=$2", - ) - .bind(worker_id) - .bind(self.namespace.as_ref()) - .fetch_optional(&self.pool) - .await?; - match registration { - Some(row) if DB::string(&row, "task_queue")? == queue => {} - _ => return Ok(true), - } - // A concurrent poll with the same request ID may have claimed already; - // the claim replays its receipt. - if let Some(request_id) = body.get("poll_request_id").and_then(Value::as_str) - && Self::query("SELECT task_id FROM dw_poll_receipts WHERE worker_id=$1 AND kind=$2 AND request_id=$3 AND namespace=$4") - .bind(worker_id).bind(kind).bind(request_id).bind(self.namespace.as_ref()) - .fetch_optional(&self.pool).await?.is_some() - { - return Ok(true); - } - let at = super::store::now(); - Ok(Self::query(DB::TASK_CANDIDATES_SQL) - .bind(self.namespace.as_ref()) - .bind(queue) - .bind(kind) - .bind(DB::bind_time(at)) - .bind(DB::bind_time(at)) - .fetch_optional(&self.pool) - .await? - .is_some()) - } - - /// Whether a held query poll should claim again. `proof` is the - /// registration this poll belongs to. - pub(super) async fn query_poll_ready( - &self, - body: &Value, - config: &QueryTaskConfig, - proof: &str, - ) -> Result { - let owner = text(body, "worker_id")?; - // The claim refuses a replaced, removed or stale registration. - if Self::query_worker_proof(&self.pool, self.namespace.as_ref(), owner, config) - .await? - .as_deref() - != Some(proof) - { - return Ok(true); - } - let at = super::store::now(); - let query_now = query_authority::wire(at); - Ok(Self::query(DB::POLL_WAKE_QUERIES_SQL) - .bind(at.timestamp()) - .bind(self.namespace.as_ref()) - .bind(text(body, "task_queue")?) - .bind(query_now.clone()) - .bind(query_now) - .fetch_optional(&self.pool) - .await? - .is_some()) - } -} diff --git a/rust/src/runtime/queries.rs b/rust/src/runtime/queries.rs index cdca5859c..27b78179c 100644 --- a/rust/src/runtime/queries.rs +++ b/rust/src/runtime/queries.rs @@ -385,7 +385,7 @@ where } } - pub(super) async fn query_worker_proof<'e, E>( + async fn query_worker_proof<'e, E>( executor: E, namespace: &str, owner: &str, @@ -437,6 +437,36 @@ where delivered } + /// Whether a held query poll should claim again. `proof` is the + /// registration this poll belongs to. + async fn query_poll_ready( + &self, + body: &Value, + config: &QueryTaskConfig, + proof: &str, + ) -> Result { + let owner = text(body, "worker_id")?; + // The claim refuses a replaced, removed or stale registration. + if Self::query_worker_proof(&self.pool, self.namespace.as_ref(), owner, config) + .await? + .as_deref() + != Some(proof) + { + return Ok(true); + } + let at = super::store::now(); + let query_now = authority::wire(at); + Ok(Self::query(DB::POLL_WAKE_QUERIES_SQL) + .bind(at.timestamp()) + .bind(self.namespace.as_ref()) + .bind(text(body, "task_queue")?) + .bind(query_now.clone()) + .bind(query_now) + .fetch_optional(&self.pool) + .await? + .is_some()) + } + async fn claim_query( &self, body: &Value, diff --git a/rust/src/runtime/store.rs b/rust/src/runtime/store.rs index 9b045bee8..c593973b6 100644 --- a/rust/src/runtime/store.rs +++ b/rust/src/runtime/store.rs @@ -967,6 +967,43 @@ where } } + /// Whether a held workflow or activity poll should claim again. + async fn task_poll_ready(&self, body: &Value, kind: &str) -> Result { + let worker_id = text(body, "worker_id")?; + let queue = text(body, "task_queue")?; + // The claim refuses a removed or moved registration. + let registration = Self::query( + "SELECT task_queue FROM workflow_worker_registrations WHERE worker_id=$1 AND namespace=$2", + ) + .bind(worker_id) + .bind(self.namespace.as_ref()) + .fetch_optional(&self.pool) + .await?; + match registration { + Some(row) if DB::string(&row, "task_queue")? == queue => {} + _ => return Ok(true), + } + // A concurrent poll with the same request ID may have claimed already; + // the claim replays its receipt. + if let Some(request_id) = body.get("poll_request_id").and_then(Value::as_str) + && Self::query("SELECT task_id FROM dw_poll_receipts WHERE worker_id=$1 AND kind=$2 AND request_id=$3 AND namespace=$4") + .bind(worker_id).bind(kind).bind(request_id).bind(self.namespace.as_ref()) + .fetch_optional(&self.pool).await?.is_some() + { + return Ok(true); + } + let at = now(); + Ok(Self::query(DB::TASK_CANDIDATES_SQL) + .bind(self.namespace.as_ref()) + .bind(queue) + .bind(kind) + .bind(DB::bind_time(at)) + .bind(DB::bind_time(at)) + .fetch_optional(&self.pool) + .await? + .is_some()) + } + async fn claim( &self, body: &Value,