From 85d0500621e28d51fabd2b69ce56cabc0b764137 Mon Sep 17 00:00:00 2001 From: Andrew Yuan Date: Tue, 12 May 2026 19:02:59 -0400 Subject: [PATCH 1/3] Avoid losing shutdown RPC handle --- crates/sdk-core/src/worker/mod.rs | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/crates/sdk-core/src/worker/mod.rs b/crates/sdk-core/src/worker/mod.rs index 0a5ecd5e1..a40c1e263 100644 --- a/crates/sdk-core/src/worker/mod.rs +++ b/crates/sdk-core/src/worker/mod.rs @@ -1376,6 +1376,7 @@ impl Worker { /// /// You can then wait on `shutdown` or [Worker::finalize_shutdown]. pub fn initiate_shutdown(&self) { + let mut shutdown_rpc_handle = self.shutdown_rpc_handle.lock(); if !self.shutdown_token.is_cancelled() { info!( task_queue=%self.config.task_queue, @@ -1419,10 +1420,7 @@ impl Worker { } } - // Spawn the ShutdownWorker RPC so the server can complete in-flight polls. - // The handle is stored and awaited in shutdown() to ensure completion. - let mut guard = self.shutdown_rpc_handle.lock(); - if guard.is_some() || already_initiated_shutdown { + if shutdown_rpc_handle.is_some() || already_initiated_shutdown { return; } @@ -1458,7 +1456,7 @@ impl Worker { _ => {} } }); - *guard = Some(handle); + *shutdown_rpc_handle = Some(handle); } /// Unique identifier for this worker instance. From 38bbd964fc3f29d9b3db485c9b645430fb206dff Mon Sep 17 00:00:00 2001 From: Andrew Yuan Date: Fri, 17 Jul 2026 14:44:12 -0400 Subject: [PATCH 2/3] handle multiple shutdown() callers and cancellation of original handle owner --- crates/sdk-core/src/core_tests/workers.rs | 168 +++++++++++++++++++++- crates/sdk-core/src/worker/mod.rs | 46 ++++-- 2 files changed, 200 insertions(+), 14 deletions(-) diff --git a/crates/sdk-core/src/core_tests/workers.rs b/crates/sdk-core/src/core_tests/workers.rs index db6ec59a5..772ba9541 100644 --- a/crates/sdk-core/src/core_tests/workers.rs +++ b/crates/sdk-core/src/core_tests/workers.rs @@ -10,14 +10,18 @@ use crate::{ self, PollerBehavior, client::{ MockWorkerClient, - mocks::{DEFAULT_TEST_CAPABILITIES, DEFAULT_WORKERS_REGISTRY, mock_worker_client}, + mocks::{ + DEFAULT_TEST_CAPABILITIES, DEFAULT_WORKERS_REGISTRY, MockManualWorkerClient, + mock_worker_client, + }, }, }, }; -use futures_util::{stream, stream::StreamExt}; +use futures_util::{FutureExt, stream, stream::StreamExt}; use std::{ cell::RefCell, collections::HashMap, + future::{Future, poll_fn}, sync::{ Arc, atomic::{AtomicBool, Ordering}, @@ -61,7 +65,7 @@ use temporalio_common::{ }, worker::WorkerTaskTypes, }; -use tokio::sync::{Barrier, Notify, watch}; +use tokio::sync::{Barrier, Notify, oneshot, watch}; use uuid::Uuid; #[tokio::test] @@ -1326,3 +1330,161 @@ async fn graceful_shutdown_sends_shutdown_worker_rpc_during_initiate() { worker.finalize_shutdown().await; } + +fn worker_with_blocked_shutdown_rpc() -> (Arc, Arc, watch::Sender) { + let rpc_started = Arc::new(Barrier::new(2)); + let rpc_started_for_client = rpc_started.clone(); + let (release_rpc, release_rpc_rx) = watch::channel(false); + + let mut mock_client = MockManualWorkerClient::new(); + mock_client + .expect_capabilities() + .returning(|| Some(*DEFAULT_TEST_CAPABILITIES)); + mock_client + .expect_workers() + .returning(|| DEFAULT_WORKERS_REGISTRY.clone()); + mock_client.expect_is_mock().returning(|| true); + mock_client + .expect_sdk_name_and_version() + .returning(|| ("test-core".to_string(), "0.0.0".to_string())); + mock_client + .expect_identity() + .returning(|| "test-identity".to_string()); + mock_client + .expect_worker_grouping_key() + .returning(Uuid::new_v4); + mock_client + .expect_worker_instance_key() + .returning(Uuid::new_v4); + mock_client + .expect_shutdown_worker() + .times(1) + .returning(move |_, _, _, _| { + let rpc_started = rpc_started_for_client.clone(); + let mut release_rpc = release_rpc_rx.clone(); + async move { + rpc_started.wait().await; + release_rpc.wait_for(|released| *released).await.unwrap(); + Ok(ShutdownWorkerResponse {}) + } + .boxed() + }); + + let mw = MockWorkerInputs::new(stream::pending().boxed()); + let mut mocks = MocksHolder::from_mock_worker(mock_client, mw); + // No task managers are needed here, allowing an incorrectly unblocked shutdown caller to + // return immediately instead of being hidden by unrelated drain waits. + mocks.worker_cfg(|config| { + config.task_types = WorkerTaskTypes::nexus_only(); + }); + (Arc::new(mock_worker(mocks)), rpc_started, release_rpc) +} + +#[tokio::test] +async fn concurrent_shutdown_waits_for_shared_rpc_and_only_starts_once() { + let (worker, rpc_started, release_rpc) = worker_with_blocked_shutdown_rpc(); + let start = Arc::new(Barrier::new(4)); + + let spawn_shutdown = |worker: Arc, start: Arc| { + let (first_poll_tx, first_poll_rx) = oneshot::channel(); + let shutdown = tokio::spawn(async move { + start.wait().await; + let shutdown = worker.shutdown(); + tokio::pin!(shutdown); + let mut first_poll_tx = Some(first_poll_tx); + poll_fn(|cx| { + let result = shutdown.as_mut().poll(cx); + if let Some(first_poll_tx) = first_poll_tx.take() { + let _ = first_poll_tx.send(result.is_pending()); + } + result + }) + .await; + }); + (shutdown, first_poll_rx) + }; + + let initiator = { + let worker = worker.clone(); + let start = start.clone(); + tokio::spawn(async move { + start.wait().await; + worker.initiate_shutdown(); + }) + }; + let (shutdown_one, shutdown_one_first_poll) = spawn_shutdown(worker.clone(), start.clone()); + let (shutdown_two, shutdown_two_first_poll) = spawn_shutdown(worker.clone(), start.clone()); + + start.wait().await; + tokio::time::timeout(Duration::from_secs(5), rpc_started.wait()) + .await + .expect("shutdown RPC should start"); + assert!( + shutdown_one_first_poll.await.unwrap(), + "first shutdown caller completed while the RPC was blocked" + ); + assert!( + shutdown_two_first_poll.await.unwrap(), + "second shutdown caller completed while the RPC was blocked" + ); + + release_rpc.send(true).unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + shutdown_one.await.unwrap(); + shutdown_two.await.unwrap(); + }) + .await + .expect("concurrent shutdown callers should complete after the RPC"); + + let worker = Arc::try_unwrap(worker).unwrap_or_else(|_| panic!("worker still shared")); + worker.finalize_shutdown().await; +} + +#[tokio::test] +async fn shutdown_waits_for_rpc_after_handle_owner_is_cancelled() { + let (worker, rpc_started, release_rpc) = worker_with_blocked_shutdown_rpc(); + + let first_shutdown = { + let worker = worker.clone(); + tokio::spawn(async move { + worker.shutdown().await; + }) + }; + tokio::time::timeout(Duration::from_secs(5), rpc_started.wait()) + .await + .expect("shutdown RPC should start"); + first_shutdown.abort(); + assert!(first_shutdown.await.unwrap_err().is_cancelled()); + + worker.initiate_shutdown(); + let (replacement_first_poll_tx, replacement_first_poll_rx) = oneshot::channel(); + let replacement_shutdown = { + let worker = worker.clone(); + tokio::spawn(async move { + let shutdown = worker.shutdown(); + tokio::pin!(shutdown); + let mut replacement_first_poll_tx = Some(replacement_first_poll_tx); + poll_fn(|cx| { + let result = shutdown.as_mut().poll(cx); + if let Some(replacement_first_poll_tx) = replacement_first_poll_tx.take() { + let _ = replacement_first_poll_tx.send(result.is_pending()); + } + result + }) + .await; + }) + }; + assert!( + replacement_first_poll_rx.await.unwrap(), + "replacement shutdown completed while the original RPC was blocked" + ); + + release_rpc.send(true).unwrap(); + tokio::time::timeout(Duration::from_secs(5), replacement_shutdown) + .await + .expect("replacement shutdown should complete after the RPC") + .unwrap(); + + let worker = Arc::try_unwrap(worker).unwrap_or_else(|_| panic!("worker still shared")); + worker.finalize_shutdown().await; +} diff --git a/crates/sdk-core/src/worker/mod.rs b/crates/sdk-core/src/worker/mod.rs index 53fdbb460..9572f5ae4 100644 --- a/crates/sdk-core/src/worker/mod.rs +++ b/crates/sdk-core/src/worker/mod.rs @@ -440,8 +440,18 @@ pub struct Worker { /// Capabilities as returned by a describe namespace rpc. Not set until after validate() is /// called. capabilities: Arc, - /// Handle for the spawned ShutdownWorker RPC task, awaited during shutdown. - shutdown_rpc_handle: Mutex>>, + /// Keeps initiation atomic even after a shutdown caller has taken the RPC handle. + shutdown_rpc_state: Mutex, + /// Allows a subsequent shutdown caller to wait for the in-flight RPC when the caller that took + /// its join handle is cancelled. + shutdown_rpc_complete: CancellationToken, + /// Serializes shutdown because several downstream shutdown routines consume join handles. + shutdown_complete: tokio::sync::Mutex, +} + +enum ShutdownRpcState { + NotStarted, + Started(Option>), } /// Namespace capabilities discovered via `describe_namespace` during worker validation. @@ -989,7 +999,9 @@ impl Worker { client_worker_registrator, status: worker_status, capabilities, - shutdown_rpc_handle: Mutex::new(None), + shutdown_rpc_state: Mutex::new(ShutdownRpcState::NotStarted), + shutdown_rpc_complete: CancellationToken::new(), + shutdown_complete: tokio::sync::Mutex::new(false), }) } @@ -1006,13 +1018,23 @@ impl Worker { /// Lang implementations should use [Worker::initiate_shutdown] followed by /// [Worker::finalize_shutdown]. pub async fn shutdown(&self) { + let mut shutdown_complete = self.shutdown_complete.lock().await; + if *shutdown_complete { + return; + } + self.initiate_shutdown(); // Ensure the ShutdownWorker RPC completes before waiting for polls to drain, // otherwise graceful poll shutdown deadlocks. - let handle = self.shutdown_rpc_handle.lock().take(); + let handle = match &mut *self.shutdown_rpc_state.lock() { + ShutdownRpcState::NotStarted => unreachable!("shutdown RPC must have been started"), + ShutdownRpcState::Started(handle) => handle.take(), + }; if let Some(handle) = handle { let _ = handle.await; + } else { + self.shutdown_rpc_complete.cancelled().await; } // We need to wait for all local activities to finish so no more workflow task heartbeats @@ -1042,6 +1064,7 @@ impl Worker { dbg_panic!("Waiting for all slot permits to release took too long!"); } } + *shutdown_complete = true; } /// Completes shutdown and frees all resources. You should avoid simply dropping workers, as @@ -1421,7 +1444,11 @@ impl Worker { /// /// You can then wait on `shutdown` or [Worker::finalize_shutdown]. pub fn initiate_shutdown(&self) { - let mut shutdown_rpc_handle = self.shutdown_rpc_handle.lock(); + let mut shutdown_rpc_state = self.shutdown_rpc_state.lock(); + if matches!(*shutdown_rpc_state, ShutdownRpcState::Started(_)) { + return; + } + if !self.shutdown_token.is_cancelled() { info!( task_queue=%self.config.task_queue, @@ -1429,7 +1456,6 @@ impl Worker { "Initiated shutdown", ); } - let already_initiated_shutdown = self.shutdown_token.is_cancelled(); self.shutdown_token.cancel(); { *self.status.write() = WorkerStatus::ShuttingDown; @@ -1465,10 +1491,6 @@ impl Worker { } } - if shutdown_rpc_handle.is_some() || already_initiated_shutdown { - return; - } - let client = self.client.clone(); let sticky_name = self .workflows @@ -1491,7 +1513,9 @@ impl Worker { .heartbeat_manager .as_ref() .map(|hm| hm.heartbeat_callback.clone()()); + let shutdown_rpc_complete = self.shutdown_rpc_complete.clone(); let handle = tokio::spawn(async move { + let _complete_on_drop = shutdown_rpc_complete.drop_guard(); match client .shutdown_worker(sticky_name, task_queue, task_queue_types, heartbeat) .await @@ -1510,7 +1534,7 @@ impl Worker { _ => {} } }); - *shutdown_rpc_handle = Some(handle); + *shutdown_rpc_state = ShutdownRpcState::Started(Some(handle)); } /// Unique identifier for this worker instance. From 684ced82d259e252cb33659a3cca400e90377f38 Mon Sep 17 00:00:00 2001 From: Andrew Yuan Date: Wed, 22 Jul 2026 15:51:51 -0400 Subject: [PATCH 3/3] Fix shutdown test deadlock --- crates/sdk-core/src/core_tests/activity_tasks.rs | 6 +++++- crates/sdk-core/src/core_tests/workers.rs | 1 + 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/crates/sdk-core/src/core_tests/activity_tasks.rs b/crates/sdk-core/src/core_tests/activity_tasks.rs index 1d61a3cb6..af6e3ac17 100644 --- a/crates/sdk-core/src/core_tests/activity_tasks.rs +++ b/crates/sdk-core/src/core_tests/activity_tasks.rs @@ -649,7 +649,11 @@ async fn can_heartbeat_acts_during_shutdown() { }) .await .unwrap(); - core.drain_activity_poller_and_shutdown().await; + assert_matches!( + core.poll_activity_task().await.unwrap_err(), + crate::PollError::ShutDown + ); + shutdown_fut.await; } /// Verifies that if a user has tried to record a heartbeat and then immediately after failed the diff --git a/crates/sdk-core/src/core_tests/workers.rs b/crates/sdk-core/src/core_tests/workers.rs index 772ba9541..9ab43e68a 100644 --- a/crates/sdk-core/src/core_tests/workers.rs +++ b/crates/sdk-core/src/core_tests/workers.rs @@ -1427,6 +1427,7 @@ async fn concurrent_shutdown_waits_for_shared_rpc_and_only_starts_once() { shutdown_two_first_poll.await.unwrap(), "second shutdown caller completed while the RPC was blocked" ); + initiator.await.unwrap(); release_rpc.send(true).unwrap(); tokio::time::timeout(Duration::from_secs(5), async {