Let held long polls wait without the transition lock

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

rust/src/runtime/backend.rs

@@ -337,8 +337,8 @@ impl Backend for Postgres {
     async fn begin(pool: &Pool<Self>) -> Result<Transaction<'static, Self>> {
         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?;

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;

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),
+        }
+    }
+}

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<bool> {
+        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,

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<bool> {
+        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,

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;
 

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;
+}

rust/tests/execution/read_snapshots.rs

@@ -3,15 +3,15 @@
 use super::*;
 use sqlx::Connection;
 
-enum Writer {
+pub(super) enum Writer {
     Sqlite(sqlx::SqliteConnection),
     Postgres(sqlx::PgConnection),
     MySql(sqlx::MySqlConnection),
 }
 
 // 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")