Ensure waiters are always notified in Nexus worker

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

@@ -10,11 +10,14 @@ use crate::{
         self, PollerBehavior,
         client::{
             MockWorkerClient,
-            mocks::{DEFAULT_TEST_CAPABILITIES, DEFAULT_WORKERS_REGISTRY, mock_worker_client},
+            mocks::{
+                DEFAULT_TEST_CAPABILITIES, DEFAULT_WORKERS_REGISTRY, mock_manual_worker_client,
+                mock_worker_client,
+            },
         },
     },
 };
-use futures_util::{stream, stream::StreamExt};
+use futures_util::{FutureExt, future, stream, stream::StreamExt};
 use std::{
     cell::RefCell,
     collections::HashMap,
@@ -1326,3 +1329,83 @@ async fn graceful_shutdown_sends_shutdown_worker_rpc_during_initiate() {
 
     worker.finalize_shutdown().await;
 }
+
+/// Verifies that even if the a nexus task completion is dropped, the nexus worker
+/// is able to trigger PollShutdown.
+#[tokio::test]
+async fn nexus_shutdown_does_not_hang_when_pending_completion_is_cancelled() {
+    let mut client = mock_manual_worker_client();
+    let completion_rpc_started = Arc::new(Barrier::new(2));
+    let completion_rpc_started_clone = completion_rpc_started.clone();
+
+    // Create a client that will hang on complete_nexus_task
+    client
+        .expect_complete_nexus_task()
+        .times(1)
+        .returning(move |_, _| {
+            let completion_rpc_started = completion_rpc_started_clone.clone();
+            async move {
+                completion_rpc_started.wait().await;
+                future::pending::<Result<RespondNexusTaskCompletedResponse, tonic::Status>>().await
+            }
+            .boxed()
+        });
+
+    let mut mocks = MocksHolder::from_client_with_custom(
+        client,
+        None::<stream::Empty<PollWorkflowTaskQueueResponse>>,
+        None::<Vec<QueueResponse<PollActivityTaskQueueResponse>>>,
+        Some(vec![QueueResponse::from(create_test_nexus_task(
+            None,
+            Some(nexus::v1::request::Capabilities {
+                temporal_failure_responses: true,
+            }),
+        ))]),
+    );
+    mocks.worker_cfg(|w| {
+        w.task_queue = "nexus-shutdown-cancelled-completion".to_string();
+        w.task_types = WorkerTaskTypes {
+            enable_workflows: false,
+            enable_local_activities: false,
+            enable_remote_activities: false,
+            enable_nexus: true,
+        };
+        w.skip_client_worker_set_check = true;
+    });
+    let worker = mock_worker(mocks);
+
+    // Poll the first task from the stream and initiate shutdown
+    let nexus_task = worker.poll_nexus_task().await.unwrap();
+    worker.initiate_shutdown();
+
+    // Poll the stream again and expect it to wait for the task to complete
+    let mut poll_future = Box::pin(worker.poll_nexus_task());
+    tokio::time::timeout(Duration::from_millis(50), poll_future.as_mut())
+        .await
+        .expect_err("poll should wait for the outstanding Nexus task");
+
+    // Send completion that we know will hang indefinitely before notifying waiters
+    let mut completion =
+        Box::pin(worker.complete_nexus_task(create_test_nexus_completion(nexus_task.task_token())));
+
+    // Wait for the completion to start then drop it before notify_waiters is triggered
+    // Use select so the completion future is polled and completion actually starts
+    tokio::select! {
+        _ = completion_rpc_started.wait() => {}
+        result = completion.as_mut() => {
+            panic!("completion should stay pending, got {result:?}");
+        }
+    }
+    drop(completion);
+
+    // Polling again should now return PollError::ShutDown because the outstanding task map is empty
+    // and waiters should have been notified.
+    let poll_result = tokio::time::timeout(Duration::from_secs(1), poll_future.as_mut())
+        .await
+        .expect("shutdown poll should finish when a pending completion future is cancelled");
+    assert_matches!(poll_result, Err(PollError::ShutDown));
+    drop(poll_future);
+
+    worker.shutdown().await;
+    worker.finalize_shutdown().await;
+}

crates/sdk-core/src/worker/nexus.rs

@@ -120,6 +120,8 @@ impl NexusManager {
     ) -> Result<(), CompleteNexusError> {
         let removed = self.outstanding_task_map.lock().remove(&tt);
         if let Some(task_info) = removed {
+            let task_completed_notify =
+                TaskCompletedNotify::new(self.task_completed_notify.clone());
             self.metrics
                 .nexus_task_execution_latency(task_info.start_time.elapsed());
             task_info.timeout_task.inspect(|jh| jh.abort());
@@ -294,7 +296,7 @@ impl NexusManager {
                 }
             };
 
-            self.task_completed_notify.notify_waiters();
+            task_completed_notify.notify_waiters();
 
             if let Some(e) = maybe_net_err {
                 if e.code() == tonic::Code::NotFound {
@@ -326,6 +328,37 @@ impl NexusManager {
     }
 }
 
+/// TaskCompleteNotify is used to ensure that waiters are notified when a task
+/// is removed from the outstanding task map even in the event that the running
+/// future is dropped.
+struct TaskCompletedNotify {
+    inner: Option<Arc<Notify>>,
+}
+
+impl TaskCompletedNotify {
+    fn new(notify: Arc<Notify>) -> Self {
+        Self {
+            inner: Some(notify),
+        }
+    }
+
+    fn notify_inner(&mut self) {
+        if let Some(notify) = self.inner.take() {
+            notify.notify_waiters();
+        }
+    }
+
+    fn notify_waiters(mut self) {
+        self.notify_inner();
+    }
+}
+
+impl Drop for TaskCompletedNotify {
+    fn drop(&mut self) {
+        self.notify_inner();
+    }
+}
+
 struct NexusTaskStream<S> {
     source_stream: S,
     outstanding_task_map: OutstandingTaskMap,