diff --git a/docs/rust-server-transition-locking.md b/docs/rust-server-transition-locking.md index e593eab1a..392f7743c 100644 --- a/docs/rust-server-transition-locking.md +++ b/docs/rust-server-transition-locking.md @@ -26,6 +26,16 @@ 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 (`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 +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..b31a8349d 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 + // (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/mod.rs b/rust/src/runtime/mod.rs index 2fa6aa8fe..9f5766b20 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..6c0265829 --- /dev/null +++ b/rust/src/runtime/poll_readiness.rs @@ -0,0 +1,37 @@ +//! 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 (`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}; + +/// 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), + } + } +} diff --git a/rust/src/runtime/queries.rs b/rust/src/runtime/queries.rs index 861532c37..9a6b17ec6 100644 --- a/rust/src/runtime/queries.rs +++ b/rust/src/runtime/queries.rs @@ -354,14 +354,20 @@ 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; + let mut pace = super::poll_readiness::ReadPace::new(); loop { let wake = self.wake.notified(); - match self.claim_query(&body, config, &proof).await? { - QueryClaim::Task(task) => return Ok(poll_result(Some(task), "leased")), - 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), "leased")), + 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, "empty")); @@ -375,7 +381,7 @@ where json!({"task":null,"poll_status":"task_queue_changed","server_capabilities":super::store::capabilities()}), ); } - tokio::select! { _ = wake => {}, _ = tokio::time::sleep(Duration::from_millis(50)) => {} } + pace.wait(wake, deadline).await; } } @@ -431,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 88da7f9e7..32e675e98 100644 --- a/rust/src/runtime/store.rs +++ b/rust/src/runtime/store.rs @@ -935,11 +935,18 @@ 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; + let mut pace = super::poll_readiness::ReadPace::new(); 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": "leased", "server_capabilities": capabilities()}), ); @@ -958,10 +965,47 @@ where json!({"task":null,"poll_status":"task_queue_changed","server_capabilities":capabilities()}), ); } - tokio::select! { _ = wake => {}, _ = tokio::time::sleep(Duration::from_millis(50)) => {} } + pace.wait(wake, deadline).await; } } + /// 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, diff --git a/rust/tests/execution.rs b/rust/tests/execution.rs index 0c1f6a19e..a9097d460 100644 --- a/rust/tests/execution.rs +++ b/rust/tests/execution.rs @@ -64,6 +64,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")