Fix replay hanging on empty WFT complete @ end of history (#522)

Author: Sushisource

core/src/replay/mod.rs

@@ -3,14 +3,16 @@
 //! users during testing.
 
 use crate::{
-    worker::client::{mocks::mock_manual_workflow_client, WorkerClient},
+    worker::{
+        client::{mocks::mock_manual_workflow_client, WorkerClient},
+        PostActivateHookData,
+    },
     Worker,
 };
 use futures::{FutureExt, Stream, StreamExt};
 use once_cell::sync::OnceCell;
 use parking_lot::Mutex;
 use std::{
-    collections::HashMap,
     pin::Pin,
     sync::Arc,
     task::{Context, Poll},
@@ -147,25 +149,18 @@ impl Historator {
 
     /// Returns a callback that can be used as the post-activation hook for a worker to indicate
     /// we're ready to replay the next history, or whatever else.
-    pub(crate) fn get_post_activate_hook(&self) -> impl Fn(&Worker, &str, usize) + Send + Sync {
-        let dat = self.dat.clone();
+    pub(crate) fn get_post_activate_hook(
+        &self,
+    ) -> impl Fn(&Worker, PostActivateHookData) + Send + Sync {
         let done_tx = self.replay_done_tx.clone();
-        move |worker, activated_run_id, last_processed_event| {
-            // We can't hold the lock while evaluating the hook, or we'd deadlock.
-            let last_event_in_hist = dat
-                .lock()
-                .run_id_to_last_event_num
-                .get(activated_run_id)
-                .cloned();
-            if let Some(le) = last_event_in_hist {
-                if last_processed_event >= le {
-                    worker.request_wf_eviction(
-                        activated_run_id,
-                        "Always evict workflows after replay",
-                        EvictionReason::LangRequested,
-                    );
-                    done_tx.send(activated_run_id.to_string()).unwrap();
-                }
+        move |worker, data| {
+            if !data.replaying {
+                worker.request_wf_eviction(
+                    data.run_id,
+                    "Always evict workflows after replay",
+                    EvictionReason::LangRequested,
+                );
+                done_tx.send(data.run_id.to_string()).unwrap();
             }
         }
     }
@@ -184,19 +179,14 @@ impl Stream for Historator {
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         match self.iter.poll_next_unpin(cx) {
             Poll::Ready(Some(history)) => {
-                let run_id = history
+                history
                     .hist
                     .extract_run_id_from_start()
                     .expect(
                         "Histories provided for replay must contain run ids in their workflow \
                                  execution started events",
                     )
                     .to_string();
-                let last_event = history.hist.last_event_id();
-                self.dat
-                    .lock()
-                    .run_id_to_last_event_num
-                    .insert(run_id, last_event as usize);
                 Poll::Ready(Some(history))
             }
             Poll::Ready(None) => {
@@ -210,6 +200,5 @@ impl Stream for Historator {
 
 #[derive(Default)]
 struct HistoratorDat {
-    run_id_to_last_event_num: HashMap<String, usize>,
     all_dispatched: bool,
 }

core/src/worker/mod.rs

@@ -78,7 +78,7 @@ pub struct Worker {
     shutdown_token: CancellationToken,
     /// Will be called at the end of each activation completion
     #[allow(clippy::type_complexity)] // Sorry clippy, there's no simple way to re-use here.
-    post_activate_hook: Option<Box<dyn Fn(&Self, &str, usize) + Send + Sync>>,
+    post_activate_hook: Option<Box<dyn Fn(&Self, PostActivateHookData) + Send + Sync>>,
     /// Set when non-local activities are complete and should stop being polled
     non_local_activities_complete: Arc<AtomicBool>,
     /// Set when local activities are complete and should stop being polled
@@ -538,9 +538,9 @@ impl Worker {
         self.workflows
             .activation_completed(
                 completion,
-                self.post_activate_hook.as_ref().map(|h| {
-                    |run_id: &str, most_recent_event: usize| h(self, run_id, most_recent_event)
-                }),
+                self.post_activate_hook
+                    .as_ref()
+                    .map(|h| |data: PostActivateHookData| h(self, data)),
             )
             .await?;
         Ok(())
@@ -559,7 +559,7 @@ impl Worker {
     /// Sets a function to be called at the end of each activation completion
     pub(crate) fn set_post_activate_hook(
         &mut self,
-        callback: impl Fn(&Self, &str, usize) + Send + Sync + 'static,
+        callback: impl Fn(&Self, PostActivateHookData) + Send + Sync + 'static,
     ) {
         self.post_activate_hook = Some(Box::new(callback))
     }
@@ -588,6 +588,12 @@ impl Worker {
     }
 }
 
+pub struct PostActivateHookData<'a> {
+    pub run_id: &'a str,
+    pub most_recent_event: usize,
+    pub replaying: bool,
+}
+
 fn build_wf_basics(
     config: &mut WorkerConfig,
     metrics: MetricsContext,

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

@@ -420,6 +420,8 @@ impl WorkflowMachines {
         // then we don't need to do anything here, and in fact we need to avoid re-applying the
         // final WFT.
         if self.have_seen_terminal_event {
+            // Replay clearly counts as done now, since we return here and never do anything else.
+            self.replaying = false;
             return Ok(0);
         }
 

core/src/worker/workflow/managed_run.rs

@@ -988,7 +988,8 @@ impl ManagedRun {
         let should_respond = !(self.wfm.machines.has_pending_jobs()
             || outgoing_cmds.replaying
             || is_query_playback
-            || data.activation_was_only_eviction);
+            || data.activation_was_only_eviction
+            || self.wfm.machines.have_seen_terminal_event);
         // If there are pending LA resolutions, and we're responding to a query here,
         // we want to make sure to force a new task, as otherwise once we tell lang about
         // the LA resolution there wouldn't be any task to reply to with the result of iterating
@@ -1011,10 +1012,7 @@ impl ManagedRun {
             ActivationCompleteOutcome::DoNothing
         };
         FulfillableActivationComplete {
-            result: ActivationCompleteResult {
-                most_recently_processed_event: self.wfm.machines.last_processed_event as usize,
-                outcome,
-            },
+            result: self.build_activation_complete_result(outcome),
             resp_chan,
         }
     }
@@ -1042,11 +1040,19 @@ impl ManagedRun {
         chan: Option<oneshot::Sender<ActivationCompleteResult>>,
     ) {
         if let Some(chan) = chan {
-            chan.send(ActivationCompleteResult {
-                most_recently_processed_event: self.most_recently_processed_event_number() as usize,
-                outcome,
-            })
-            .expect("Rcv half of activation reply not dropped");
+            chan.send(self.build_activation_complete_result(outcome))
+                .expect("Rcv half of activation reply not dropped");
+        }
+    }
+
+    fn build_activation_complete_result(
+        &self,
+        outcome: ActivationCompleteOutcome,
+    ) -> ActivationCompleteResult {
+        ActivationCompleteResult {
+            outcome,
+            most_recently_processed_event: self.most_recently_processed_event_number() as usize,
+            replaying: self.wfm.machines.replaying,
         }
     }
 

core/src/worker/workflow/mod.rs

@@ -43,6 +43,7 @@ use crate::{
             workflow_stream::{LocalInput, LocalInputs, WFStream},
         },
         LocalActRequest, LocalActivityExecutionResult, LocalActivityResolution,
+        PostActivateHookData,
     },
     MetricsContext,
 };
@@ -295,7 +296,7 @@ impl Workflows {
                         },
                         // We need to say a type, but the type is irrelevant, so imagine some
                         // boxed function we'll never call.
-                        Option::<Box<dyn Fn(&str, usize) + Send>>::None,
+                        Option::<Box<dyn Fn(PostActivateHookData) + Send>>::None,
                     )
                     .await?;
                 }
@@ -308,7 +309,7 @@ impl Workflows {
                             run_id,
                             status: Some(auto_fail_to_complete_status(machines_err)),
                         },
-                        Option::<Box<dyn Fn(&str, usize) + Send>>::None,
+                        Option::<Box<dyn Fn(PostActivateHookData) + Send>>::None,
                     )
                     .await?;
                 }
@@ -323,7 +324,7 @@ impl Workflows {
     pub(super) async fn activation_completed(
         &self,
         completion: WorkflowActivationCompletion,
-        post_activate_hook: Option<impl Fn(&str, usize)>,
+        post_activate_hook: Option<impl Fn(PostActivateHookData)>,
     ) -> Result<usize, CompleteWfError> {
         let is_empty_completion = completion.is_empty();
         let completion = validate_completion(completion)?;
@@ -448,7 +449,11 @@ impl Workflows {
         };
 
         if let Some(h) = post_activate_hook {
-            h(&run_id, completion_outcome.most_recently_processed_event);
+            h(PostActivateHookData {
+                run_id: &run_id,
+                most_recent_event: completion_outcome.most_recently_processed_event,
+                replaying: completion_outcome.replaying,
+            });
         }
 
         self.post_activation(PostActivationMsg {
@@ -956,6 +961,7 @@ struct GetStateInfoMsg {
 #[derive(Debug)]
 struct ActivationCompleteResult {
     most_recently_processed_event: usize,
+    replaying: bool,
     outcome: ActivationCompleteOutcome,
 }
 /// What needs to be done after calling [Workflows::activation_completed]

histories/ends_empty_wft_complete.bin

@@ -11,6 +11,7 @@ use temporal_sdk_core_protos::{
         workflow_commands::{ScheduleActivity, StartTimer},
         workflow_completion::WorkflowActivationCompletion,
     },
+    temporal::api::enums::v1::EventType,
     TestHistoryBuilder, DEFAULT_WORKFLOW_TYPE,
 };
 use temporal_sdk_core_test_utils::{
@@ -130,6 +131,23 @@ async fn replay_using_wf_function() {
     worker.run().await.unwrap();
 }
 
+#[tokio::test]
+async fn replay_ending_wft_complete_with_commands_but_no_scheduled_started() {
+    let mut t = TestHistoryBuilder::default();
+    t.add_by_type(EventType::WorkflowExecutionStarted);
+    t.add_full_wf_task();
+
+    for i in 1..=2 {
+        let timer_started_event_id = t.add_by_type(EventType::TimerStarted);
+        t.add_timer_fired(timer_started_event_id, i.to_string());
+        t.add_full_wf_task();
+    }
+    let func = timers_wf(3);
+    let mut worker = replay_sdk_worker([test_hist_to_replay(t)]);
+    worker.register_wf(DEFAULT_WORKFLOW_TYPE, func);
+    worker.run().await.unwrap();
+}
+
 async fn replay_abrupt_ending(t: TestHistoryBuilder) {
     let func = timers_wf(1);
     let mut worker = replay_sdk_worker([test_hist_to_replay(t)]);
@@ -219,6 +237,38 @@ async fn replay_old_patch_format() {
     worker.run().await.unwrap();
 }
 
+#[tokio::test]
+async fn replay_ends_with_empty_wft() {
+    let core = init_core_replay_preloaded(
+        "SayHelloWorkflow",
+        [HistoryForReplay::new(
+            history_from_proto_binary("histories/ends_empty_wft_complete.bin")
+                .await
+                .unwrap(),
+            "fake".to_owned(),
+        )],
+    );
+    let task = core.poll_workflow_activation().await.unwrap();
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
+        task.run_id,
+        vec![ScheduleActivity {
+            seq: 1,
+            activity_id: "1".to_string(),
+            activity_type: "say_hello".to_string(),
+            ..Default::default()
+        }
+        .into()],
+    ))
+    .await
+    .unwrap();
+    let task = core.poll_workflow_activation().await.unwrap();
+    core.complete_workflow_activation(WorkflowActivationCompletion::empty(task.run_id))
+        .await
+        .unwrap();
+    let task = core.poll_workflow_activation().await.unwrap();
+    assert!(task.eviction_reason().is_some());
+}
+
 fn timers_wf(num_timers: u32) -> WorkflowFunction {
     WorkflowFunction::new(move |ctx: WfContext| async move {
         for _ in 1..=num_timers {