Fix possibly sending LA marker after WF Complete command (#678)

Author: Sushisource

core/src/core_tests/local_activities.rs

@@ -2,8 +2,8 @@ use crate::{
     prost_dur,
     replay::{default_wes_attribs, TestHistoryBuilder, DEFAULT_WORKFLOW_TYPE},
     test_help::{
-        hist_to_poll_resp, mock_sdk, mock_sdk_cfg, mock_worker, single_hist_mock_sg, MockPollCfg,
-        ResponseType, WorkerExt,
+        build_mock_pollers, hist_to_poll_resp, mock_sdk, mock_sdk_cfg, mock_worker,
+        single_hist_mock_sg, MockPollCfg, ResponseType, WorkerExt,
     },
     worker::{client::mocks::mock_workflow_client, LEGACY_QUERY_ID},
 };
@@ -38,7 +38,7 @@ use temporal_sdk_core_protos::{
     },
     temporal::api::{
         common::v1::RetryPolicy,
-        enums::v1::{EventType, TimeoutType, WorkflowTaskFailedCause},
+        enums::v1::{CommandType, EventType, TimeoutType, WorkflowTaskFailedCause},
         failure::v1::{failure::FailureInfo, Failure},
         query::v1::WorkflowQuery,
     },
@@ -1295,3 +1295,98 @@ async fn queries_can_be_received_while_heartbeating() {
 
     core.drain_pollers_and_shutdown().await;
 }
+
+#[tokio::test]
+async fn local_activity_after_wf_complete_is_discarded() {
+    let wfid = "fake_wf_id";
+    let mut t = TestHistoryBuilder::default();
+    t.add_wfe_started_with_wft_timeout(Duration::from_millis(200));
+    t.add_full_wf_task();
+    t.add_workflow_task_scheduled_and_started();
+
+    let mock = mock_workflow_client();
+    let mut mock_cfg = MockPollCfg::from_resp_batches(
+        wfid,
+        t,
+        [ResponseType::ToTaskNum(1), ResponseType::ToTaskNum(2)],
+        mock,
+    );
+    mock_cfg.make_poll_stream_interminable = true;
+    mock_cfg.completion_asserts_from_expectations(|mut asserts| {
+        asserts
+            .then(move |wft| {
+                assert_eq!(wft.commands.len(), 0);
+            })
+            .then(move |wft| {
+                assert_eq!(wft.commands.len(), 2);
+                assert_eq!(wft.commands[0].command_type(), CommandType::RecordMarker);
+                assert_eq!(
+                    wft.commands[1].command_type(),
+                    CommandType::CompleteWorkflowExecution
+                );
+            });
+    });
+    let mut mock = build_mock_pollers(mock_cfg);
+    mock.worker_cfg(|wc| wc.max_cached_workflows = 1);
+    let core = mock_worker(mock);
+
+    let barr = Barrier::new(2);
+
+    let task = core.poll_workflow_activation().await.unwrap();
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
+        task.run_id,
+        vec![
+            ScheduleLocalActivity {
+                seq: 1,
+                activity_id: "1".to_string(),
+                activity_type: "test_act".to_string(),
+                start_to_close_timeout: Some(prost_dur!(from_secs(30))),
+                ..Default::default()
+            }
+            .into(),
+            ScheduleLocalActivity {
+                seq: 2,
+                activity_id: "2".to_string(),
+                activity_type: "test_act".to_string(),
+                start_to_close_timeout: Some(prost_dur!(from_secs(30))),
+                ..Default::default()
+            }
+            .into(),
+        ],
+    ))
+    .await
+    .unwrap();
+
+    let wf_poller = async {
+        let task = core.poll_workflow_activation().await.unwrap();
+        assert_matches!(
+            task.jobs.as_slice(),
+            [WorkflowActivationJob {
+                variant: Some(workflow_activation_job::Variant::ResolveActivity(_)),
+            }]
+        );
+        barr.wait().await;
+        core.complete_execution(&task.run_id).await;
+    };
+
+    let at_poller = async {
+        let act_task = core.poll_activity_task().await.unwrap();
+        core.complete_activity_task(ActivityTaskCompletion {
+            task_token: act_task.task_token,
+            result: Some(ActivityExecutionResult::ok(vec![1].into())),
+        })
+        .await
+        .unwrap();
+        let act_task = core.poll_activity_task().await.unwrap();
+        barr.wait().await;
+        core.complete_activity_task(ActivityTaskCompletion {
+            task_token: act_task.task_token,
+            result: Some(ActivityExecutionResult::ok(vec![2].into())),
+        })
+        .await
+        .unwrap();
+    };
+
+    join!(wf_poller, at_poller);
+    core.drain_pollers_and_shutdown().await;
+}

core/src/worker/mod.rs

@@ -179,8 +179,14 @@ impl WorkerTrait for Worker {
         // Let the manager know that shutdown has been initiated to try to unblock the local
         // activity poll in case this worker is an activity-only worker.
         self.local_act_mgr.shutdown_initiated();
+
         if !self.workflows.ever_polled() {
             self.local_act_mgr.workflows_have_shutdown();
+        } else {
+            // 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();
         }
     }
 
@@ -189,7 +195,6 @@ impl WorkerTrait for Worker {
     }
 
     async fn finalize_shutdown(self) {
-        self.shutdown().await;
         self.finalize_shutdown().await
     }
 }

core/src/worker/workflow/machines/workflow_machines.rs

@@ -1467,6 +1467,9 @@ impl WorkflowMachines {
         let cwfm = self.add_new_command_machine(machine);
         self.workflow_end_time = Some(SystemTime::now());
         self.current_wf_task_commands.push_back(cwfm);
+        // Wipe out any pending / executing local activity data since we're about to terminate
+        // and there's nothing to be done with them.
+        self.local_activity_data.indicate_terminating();
     }
 
     /// Add a new command/machines for that command to the current workflow task

core/src/worker/workflow/machines/workflow_machines/local_acts.rs

@@ -20,6 +20,8 @@ pub(super) struct LocalActivityData {
     /// Maps local activity sequence numbers to their resolutions as found when looking ahead at
     /// next WFT
     preresolutions: HashMap<u32, ResolveDat>,
+    /// Set true if the workflow is terminating
+    am_terminating: bool,
 }
 
 impl LocalActivityData {
@@ -45,6 +47,10 @@ impl LocalActivityData {
         wf_id: &str,
         run_id: &str,
     ) -> Vec<LocalActRequest> {
+        if self.am_terminating {
+            return vec![LocalActRequest::CancelAllInRun(run_id.to_string())];
+        }
+
         self.cancel_requests
             .drain(..)
             .map(LocalActRequest::Cancel)
@@ -65,6 +71,9 @@ impl LocalActivityData {
 
     /// Returns all outstanding local activities, whether executing or requested and in the queue
     pub(super) fn outstanding_la_count(&self) -> usize {
+        if self.am_terminating {
+            return 0;
+        }
         self.executing.len() + self.new_requests.len()
     }
 
@@ -82,4 +91,10 @@ impl LocalActivityData {
             .position(|req| req.seq == seq)
             .map(|i| self.new_requests.remove(i))
     }
+
+    /// Store that this workflow is terminating, and thus no new LA requests need be processed,
+    /// and any executing LAs should not prevent us from shutting down.
+    pub(super) fn indicate_terminating(&mut self) {
+        self.am_terminating = true;
+    }
 }

core/src/worker/workflow/mod.rs

@@ -180,58 +180,60 @@ impl Workflows {
         // We must spawn a task to constantly poll the activation stream, because otherwise
         // activation completions would not cause anything to happen until the next poll.
         let tracing_sub = telem_instance.and_then(|ti| ti.trace_subscriber());
-        let processing_task = thread::spawn(move || {
-            if let Some(ts) = tracing_sub {
-                set_trace_subscriber_for_current_thread(ts);
-            }
-            let rt = tokio::runtime::Builder::new_current_thread()
-                .enable_all()
-                .thread_name("workflow-processing")
-                .build()
-                .unwrap();
-            let local = LocalSet::new();
-            local.block_on(&rt, async move {
-                let mut stream = WFStream::build(
-                    basics,
-                    extracted_wft_stream,
-                    locals_stream,
-                    local_activity_request_sink,
-                );
+        let processing_task = thread::Builder::new()
+            .name("workflow-processing".to_string())
+            .spawn(move || {
+                if let Some(ts) = tracing_sub {
+                    set_trace_subscriber_for_current_thread(ts);
+                }
+                let rt = tokio::runtime::Builder::new_current_thread()
+                    .enable_all()
+                    .build()
+                    .unwrap();
+                let local = LocalSet::new();
+                local.block_on(&rt, async move {
+                    let mut stream = WFStream::build(
+                        basics,
+                        extracted_wft_stream,
+                        locals_stream,
+                        local_activity_request_sink,
+                    );
 
-                // However, we want to avoid plowing ahead until we've been asked to poll at least
-                // once. This supports activity-only workers.
-                let do_poll = tokio::select! {
-                    sp = start_polling_rx => {
-                        sp.is_ok()
-                    }
-                    _ = shutdown_tok.cancelled() => {
-                        false
+                    // However, we want to avoid plowing ahead until we've been asked to poll at least
+                    // once. This supports activity-only workers.
+                    let do_poll = tokio::select! {
+                        sp = start_polling_rx => {
+                            sp.is_ok()
+                        }
+                        _ = shutdown_tok.cancelled() => {
+                            false
+                        }
+                    };
+                    if !do_poll {
+                        return;
                     }
-                };
-                if !do_poll {
-                    return;
-                }
-                while let Some(output) = stream.next().await {
-                    match output {
-                        Ok(o) => {
-                            for fetchreq in o.fetch_histories {
-                                fetch_tx
-                                    .send(fetchreq)
-                                    .expect("Fetch channel must not be dropped");
-                            }
-                            for act in o.activations {
-                                activation_tx
-                                    .send(Ok(act))
-                                    .expect("Activation processor channel not dropped");
+                    while let Some(output) = stream.next().await {
+                        match output {
+                            Ok(o) => {
+                                for fetchreq in o.fetch_histories {
+                                    fetch_tx
+                                        .send(fetchreq)
+                                        .expect("Fetch channel must not be dropped");
+                                }
+                                for act in o.activations {
+                                    activation_tx
+                                        .send(Ok(act))
+                                        .expect("Activation processor channel not dropped");
+                                }
                             }
+                            Err(e) => activation_tx
+                                .send(Err(e))
+                                .expect("Activation processor channel not dropped"),
                         }
-                        Err(e) => activation_tx
-                            .send(Err(e))
-                            .expect("Activation processor channel not dropped"),
                     }
-                }
-            });
-        });
+                });
+            })
+            .expect("Must be able to spawn workflow processing thread");
         Self {
             task_queue,
             local_tx,
@@ -498,10 +500,17 @@ impl Workflows {
         });
     }
 
-    /// 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>> {
+    /// Send a `GetStateInfoMsg` to the workflow stream. Can be used to bump the stream if there
+    /// would otherwise be no new inputs.
+    pub(super) fn send_get_state_info_msg(&self) -> oneshot::Receiver<WorkflowStateInfo> {
         let (tx, rx) = oneshot::channel();
         self.send_local(GetStateInfoMsg { response_tx: tx });
+        rx
+    }
+
+    /// 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();
         async move { rx.await.ok() }
     }