add explicit WorkerShutdown input (#1055)

Author: yuandrew

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

@@ -266,6 +266,10 @@ impl WorkerTrait for Worker {
                 .unregister_worker(self.worker_instance_key);
         }
 
+        // Push a Worker Shutdown message to the workflow activation queue. This ensures that
+        // any pending workflow activation polls will resolve, even if there are no other inputs.
+        self.workflows.send_worker_shutdown();
+
         // Second, we want to stop polling of both activity and workflow tasks
         if let Some(atm) = self.at_task_mgr.as_ref() {
             atm.initiate_shutdown();
@@ -274,13 +278,12 @@ impl WorkerTrait for Worker {
         // activity poll in case this worker is an activity-only worker.
         self.local_act_mgr.shutdown_initiated();
 
+        // If workflows have never been polled, immediately tell the local activity manager
+        // that workflows have shut down, so it can proceed with shutdown without waiting.
+        // This is particularly important for activity-only workers.
         if !self.workflows.ever_polled() {
             self.local_act_mgr.workflows_have_shutdown();
         }
-        // Bump the workflow stream with a pointless input, since if a client initiates shutdown
-        // and then immediately blocks waiting on a workflow activation poll, it's possible that
-        // there may not be any more inputs ever, and that poll will never resolve.
-        self.workflows.send_get_state_info_msg();
     }
 
     async fn shutdown(&self) {

crates/sdk-core/src/worker/workflow/mod.rs

@@ -626,6 +626,12 @@ impl Workflows {
         rx
     }
 
+    /// Send a `WorkerShutdown` message to the workflow stream. This ensures that any pending
+    /// workflow activation polls will resolve during shutdown, even if there are no other inputs.
+    pub(super) fn send_worker_shutdown(&self) {
+        self.send_local(LocalInputs::WorkerShutdown);
+    }
+
     /// Query the state of workflow management. Can return `None` if workflow state is shut down.
     pub(super) fn get_state_info(&self) -> impl Future<Output = Option<WorkflowStateInfo>> {
         let rx = self.send_get_state_info_msg();

crates/sdk-core/src/worker/workflow/workflow_stream.rs

@@ -147,6 +147,11 @@ impl WFStream {
                                 });
                                 None
                             }
+                            LocalInputs::WorkerShutdown => {
+                                // Worker shutdown message is used to bump the stream to ensure
+                                // pending workflow activation polls resolve during shutdown
+                                None
+                            }
                         }
                     }
                     WFStreamInput::FailedFetch {
@@ -633,6 +638,7 @@ pub(super) enum LocalInputs {
     RequestEviction(RequestEvictMsg),
     HeartbeatTimeout(String),
     GetStateInfo(GetStateInfoMsg),
+    WorkerShutdown,
 }
 impl LocalInputs {
     fn run_id(&self) -> Option<&str> {
@@ -643,7 +649,7 @@ impl LocalInputs {
             LocalInputs::PostActivation(pa) => &pa.run_id,
             LocalInputs::RequestEviction(re) => &re.run_id,
             LocalInputs::HeartbeatTimeout(hb) => hb,
-            LocalInputs::GetStateInfo(_) => return None,
+            LocalInputs::GetStateInfo(_) | LocalInputs::WorkerShutdown => return None,
         })
     }
 }