Shutdown rpc initiate shutdown (#1224)

Author: Sushisource

.cargo/config.toml

@@ -1,6 +1,6 @@
 [env]
 # This temporarily overrides the version of the CLI used for integration tests, locally and in CI
-#CLI_VERSION_OVERRIDE = "v1.4.1-cloud-v1-29-0-139-2.0"
+CLI_VERSION_OVERRIDE = "v1.6.3-serverless"
 
 [alias]
 # Not sure why --all-features doesn't work

crates/sdk-core-c-bridge/src/worker.rs

@@ -930,6 +930,7 @@ pub extern "C" fn temporal_core_worker_request_workflow_eviction(
 #[unsafe(no_mangle)]
 pub extern "C" fn temporal_core_worker_initiate_shutdown(worker: *mut Worker) {
     let worker = unsafe { &*worker };
+    enter_sync!(worker.runtime);
     worker.worker.as_ref().unwrap().initiate_shutdown();
 }
 

crates/sdk-core/src/core_tests/workers.rs

@@ -13,7 +13,15 @@ use crate::{
     },
 };
 use futures_util::{stream, stream::StreamExt};
+use std::sync::{
+    Arc,
+    atomic::{AtomicBool, Ordering},
+};
 use std::{cell::RefCell, collections::HashMap, time::Duration};
+use temporalio_common::protos::temporal::api::{
+    namespace::v1::{NamespaceInfo, namespace_info::Capabilities},
+    workflowservice::v1::DescribeNamespaceResponse,
+};
 use temporalio_common::{
     protos::{
         canned_histories,
@@ -51,6 +59,7 @@ use temporalio_common::{
     },
     worker::WorkerTaskTypes,
 };
+use tokio::sync::Notify;
 use tokio::sync::{Barrier, watch};
 use uuid::Uuid;
 
@@ -1209,3 +1218,106 @@ async fn nexus_start_operation_failure_converts_to_legacy_for_old_server(
     worker.shutdown().await;
     worker.finalize_shutdown().await;
 }
+
+/// Verifies that `initiate_shutdown` sends the `ShutdownWorker` RPC so that the server can
+/// complete in-flight polls. Without this, graceful poll shutdown deadlocks: the SDK waits for
+/// polls to drain, but the server was never told to flush them.
+#[tokio::test]
+async fn graceful_shutdown_sends_shutdown_worker_rpc_during_initiate() {
+    let shutdown_rpc_called = Arc::new(AtomicBool::new(false));
+    let shutdown_rpc_called_clone = shutdown_rpc_called.clone();
+    // When the shutdown_worker RPC fires, it signals polls to complete (simulating server
+    // behavior where ShutdownWorker causes the server to return empty poll responses).
+    let poll_releaser = Arc::new(Notify::new());
+    let poll_releaser_for_rpc = poll_releaser.clone();
+
+    let mut mock_client = MockWorkerClient::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_set_heartbeat_client_fields()
+        .returning(|hb| {
+            hb.sdk_name = "test-core".to_string();
+            hb.sdk_version = "0.0.0".to_string();
+            hb.worker_identity = "test-identity".to_string();
+            hb.heartbeat_time = Some(std::time::SystemTime::now().into());
+        });
+    // Return the worker_poll_complete_on_shutdown capability so validate() enables graceful mode
+    mock_client.expect_describe_namespace().returning(move || {
+        Ok(DescribeNamespaceResponse {
+            namespace_info: Some(NamespaceInfo {
+                capabilities: Some(Capabilities {
+                    worker_poll_complete_on_shutdown: true,
+                    ..Capabilities::default()
+                }),
+                ..NamespaceInfo::default()
+            }),
+            ..DescribeNamespaceResponse::default()
+        })
+    });
+    // When shutdown_worker RPC is called, mark it and release polls
+    mock_client
+        .expect_shutdown_worker()
+        .returning(move |_, _, _, _| {
+            shutdown_rpc_called_clone.store(true, Ordering::SeqCst);
+            poll_releaser_for_rpc.notify_waiters();
+            Ok(ShutdownWorkerResponse {})
+        });
+    mock_client
+        .expect_complete_workflow_task()
+        .returning(|_| Ok(RespondWorkflowTaskCompletedResponse::default()));
+
+    // Polls block until shutdown_worker RPC releases them (simulating server holding polls
+    // open until it receives the ShutdownWorker signal)
+    let poll_releaser_for_stream = poll_releaser.clone();
+    let stream = stream::unfold(poll_releaser_for_stream, |releaser| async move {
+        releaser.notified().await;
+        Some((
+            Ok(PollWorkflowTaskQueueResponse::default().try_into().unwrap()),
+            releaser,
+        ))
+    });
+
+    let mw = MockWorkerInputs::new(stream.boxed());
+    let worker = mock_worker(MocksHolder::from_mock_worker(mock_client, mw));
+
+    // validate() reads describe_namespace and sets capabilities.graceful_poll_shutdown = true
+    worker.validate().await.unwrap();
+
+    let poll_fut = worker.poll_workflow_activation();
+    let shutdown_fut = async {
+        // initiate_shutdown must send the ShutdownWorker RPC, which releases the polls
+        worker.initiate_shutdown();
+    };
+
+    let (poll_result, _) = tokio::time::timeout(Duration::from_secs(5), async {
+        tokio::join!(poll_fut, shutdown_fut)
+    })
+    .await
+    .expect("Shutdown should complete within 5s -- if it hangs, the ShutdownWorker RPC was not sent during initiate_shutdown");
+
+    assert_matches!(poll_result.unwrap_err(), PollError::ShutDown);
+    assert!(
+        shutdown_rpc_called.load(Ordering::SeqCst),
+        "ShutdownWorker RPC must be called during initiate_shutdown"
+    );
+
+    worker.finalize_shutdown().await;
+}

crates/sdk-core/src/lib.rs

@@ -42,13 +42,13 @@ pub use temporalio_common::protos::TaskToken;
 pub use url::Url;
 pub use worker::{
     ActivitySlotKind, CompleteActivityError, CompleteNexusError, CompleteWfError,
-    FixedSizeSlotSupplier, LocalActivitySlotKind, NexusSlotKind, PollError, PollerBehavior,
-    ResourceBasedSlotsOptions, ResourceBasedSlotsOptionsBuilder, ResourceBasedTuner,
-    ResourceSlotOptions, SlotInfo, SlotInfoTrait, SlotKind, SlotKindType, SlotMarkUsedContext,
-    SlotReleaseContext, SlotReservationContext, SlotSupplier, SlotSupplierOptions,
-    SlotSupplierPermit, TunerBuilder, TunerHolder, TunerHolderOptions, TunerHolderOptionsBuilder,
-    Worker, WorkerConfig, WorkerConfigBuilder, WorkerTuner, WorkerValidationError,
-    WorkerVersioningStrategy, WorkflowErrorType, WorkflowSlotKind,
+    FixedSizeSlotSupplier, LocalActivitySlotKind, NamespaceCapabilities, NexusSlotKind, PollError,
+    PollerBehavior, ResourceBasedSlotsOptions, ResourceBasedSlotsOptionsBuilder,
+    ResourceBasedTuner, ResourceSlotOptions, SlotInfo, SlotInfoTrait, SlotKind, SlotKindType,
+    SlotMarkUsedContext, SlotReleaseContext, SlotReservationContext, SlotSupplier,
+    SlotSupplierOptions, SlotSupplierPermit, TunerBuilder, TunerHolder, TunerHolderOptions,
+    TunerHolderOptionsBuilder, Worker, WorkerConfig, WorkerConfigBuilder, WorkerTuner,
+    WorkerValidationError, WorkerVersioningStrategy, WorkflowErrorType, WorkflowSlotKind,
 };
 
 use crate::{

crates/sdk-core/src/pollers/mod.rs

@@ -157,6 +157,11 @@ where
                         return match state.poller.poll().await {
                             Some(Ok((task, permit))) => {
                                 if task == Default::default() {
+                                    if state.poller_was_shutdown {
+                                        // Server sent an empty response after we initiated
+                                        // shutdown — this is the graceful shutdown signal.
+                                        return None;
+                                    }
                                     // We get the default proto in the event that the long poll
                                     // times out.
                                     debug!("Poll {} task timeout", T::task_name());
@@ -276,3 +281,93 @@ pub(crate) fn new_nexus_task_poller(
     )
     .into_stream()
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{
+        abstractions::tests::fixed_size_permit_dealer, pollers::MockPermittedPollBuffer,
+        test_help::mock_poller, worker::ActivitySlotKind,
+    };
+    use futures_util::{StreamExt, pin_mut};
+    use std::sync::{
+        Arc,
+        atomic::{AtomicUsize, Ordering},
+    };
+
+    /// Verify that empty responses after shutdown are not treated as poll timeout and retried
+    /// indefinitely
+    #[tokio::test]
+    async fn empty_response_after_shutdown_terminates_stream() {
+        let poll_count = Arc::new(AtomicUsize::new(0));
+        let poll_count_clone = poll_count.clone();
+
+        let mut mock_poller = mock_poller();
+        mock_poller.expect_poll().returning(move || {
+            poll_count_clone.fetch_add(1, Ordering::SeqCst);
+            Some(Ok(PollActivityTaskQueueResponse::default()))
+        });
+
+        let sem = Arc::new(fixed_size_permit_dealer::<ActivitySlotKind>(10));
+        let shutdown_token = CancellationToken::new();
+
+        let stream = new_activity_task_poller(
+            Box::new(MockPermittedPollBuffer::new(sem, mock_poller)),
+            MetricsContext::no_op(),
+            shutdown_token.clone(),
+        );
+        pin_mut!(stream);
+
+        shutdown_token.cancel();
+
+        let result = tokio::time::timeout(std::time::Duration::from_secs(2), stream.next()).await;
+        assert!(
+            result.is_ok(),
+            "Stream should terminate promptly after shutdown, not hang"
+        );
+        assert!(
+            result.unwrap().is_none(),
+            "Stream should return None (terminated) on empty response after shutdown"
+        );
+
+        let total = poll_count.load(Ordering::SeqCst);
+        assert!(
+            total < 5,
+            "Expected stream to terminate quickly, but poller was called {total} times"
+        );
+    }
+
+    #[tokio::test]
+    async fn empty_response_before_shutdown_retries() {
+        let mut mock_poller = mock_poller();
+        let call_count = Arc::new(AtomicUsize::new(0));
+        let call_count_clone = call_count.clone();
+        mock_poller.expect_poll().returning(move || {
+            let n = call_count_clone.fetch_add(1, Ordering::SeqCst);
+            if n < 2 {
+                Some(Ok(PollActivityTaskQueueResponse::default()))
+            } else {
+                None
+            }
+        });
+
+        let sem = Arc::new(fixed_size_permit_dealer::<ActivitySlotKind>(10));
+        let shutdown_token = CancellationToken::new();
+
+        let stream = new_activity_task_poller(
+            Box::new(MockPermittedPollBuffer::new(sem, mock_poller)),
+            MetricsContext::no_op(),
+            shutdown_token,
+        );
+        pin_mut!(stream);
+
+        // Without shutdown, empty responses should be skipped and the stream terminates
+        // only when the poller returns None.
+        let result = stream.next().await;
+        assert!(
+            result.is_none(),
+            "Stream should end when poller returns None"
+        );
+        assert_eq!(call_count.load(Ordering::SeqCst), 3);
+    }
+}