Implements fix & changes WFT boundaries

* One UT fixed to show, others fail

Author: Sushisource

core/src/worker/workflow/history_update.rs

@@ -695,7 +695,8 @@ fn find_end_index_of_next_wft_seq(
         return NextWFTSeqEndIndex::Incomplete(0);
     }
     let mut last_index = 0;
-    let mut saw_any_command_event = false;
+    let mut saw_command_or_started = false;
+    let mut saw_command = false;
     let mut wft_started_event_id_to_index = vec![];
     for (ix, e) in events.iter().enumerate() {
         last_index = ix;
@@ -706,8 +707,11 @@ fn find_end_index_of_next_wft_seq(
             continue;
         }
 
-        if e.is_command_event() || e.event_type() == EventType::WorkflowExecutionStarted {
-            saw_any_command_event = true;
+        if e.is_command_event() {
+            saw_command = true;
+            if e.event_type() == EventType::WorkflowExecutionStarted {
+                saw_command_or_started = true;
+            }
         }
         if e.is_final_wf_execution_event() {
             return NextWFTSeqEndIndex::Complete(last_index);
@@ -728,13 +732,14 @@ fn find_end_index_of_next_wft_seq(
                         | EventType::WorkflowExecutionCanceled
                 ) {
                     continue;
-                }
-                // If we've never seen an interesting event and the next two events are a completion
-                // followed immediately again by scheduled, then this is a WFT heartbeat and also
-                // doesn't conclude the sequence.
-                else if next_event_type == EventType::WorkflowTaskCompleted {
+                } else if next_event_type == EventType::WorkflowTaskCompleted {
                     if let Some(next_next_event) = events.get(ix + 2) {
-                        if next_next_event.event_type() == EventType::WorkflowTaskScheduled {
+                        if !saw_command
+                            && next_next_event.event_type() == EventType::WorkflowTaskScheduled
+                        {
+                            // If we've never seen an interesting event and the next two events are
+                            // a completion followed immediately again by scheduled, then this is a
+                            // WFT heartbeat and also doesn't conclude the sequence.
                             continue;
                         } else {
                             // If we see an update accepted command after WFT completed, we want to
@@ -766,18 +771,18 @@ fn find_end_index_of_next_wft_seq(
                             }
                             return NextWFTSeqEndIndex::Complete(ix);
                         }
-                    } else if !has_last_wft && !saw_any_command_event {
+                    } else if !has_last_wft && !saw_command_or_started {
                         // Don't have enough events to look ahead of the WorkflowTaskCompleted. Need
                         // to fetch more.
                         continue;
                     }
                 }
-            } else if !has_last_wft && !saw_any_command_event {
+            } else if !has_last_wft && !saw_command_or_started {
                 // Don't have enough events to look ahead of the WorkflowTaskStarted. Need to fetch
                 // more.
                 continue;
             }
-            if saw_any_command_event {
+            if saw_command_or_started {
                 return NextWFTSeqEndIndex::Complete(ix);
             }
         }
@@ -921,7 +926,9 @@ mod tests {
         let seq = next_check_peek(&mut update, 0);
         assert_eq!(seq.len(), 6);
         let seq = next_check_peek(&mut update, 6);
-        assert_eq!(seq.len(), 13);
+        assert_eq!(seq.len(), 4);
+        let seq = next_check_peek(&mut update, 10);
+        assert_eq!(seq.len(), 9);
         let seq = next_check_peek(&mut update, 19);
         assert_eq!(seq.len(), 4);
         let seq = next_check_peek(&mut update, 23);

sdk-core-protos/src/lib.rs

@@ -628,8 +628,8 @@ pub mod coresdk {
                     workflow_activation_job::Variant::ResolveRequestCancelExternalWorkflow(_) => {
                         write!(f, "ResolveRequestCancelExternalWorkflow")
                     }
-                    workflow_activation_job::Variant::DoUpdate(_) => {
-                        write!(f, "DoUpdate")
+                    workflow_activation_job::Variant::DoUpdate(u) => {
+                        write!(f, "DoUpdate({})", u.id)
                     }
                     workflow_activation_job::Variant::ResolveNexusOperationStart(_) => {
                         write!(f, "ResolveNexusOperationStart")

sdk/src/workflow_context.rs

@@ -151,6 +151,11 @@ impl WfContext {
         self.shared.read().random_seed
     }
 
+    /// Returns true if the current workflow task is happening under replay
+    pub fn is_replaying(&self) -> bool {
+        self.shared.read().is_replaying
+    }
+
     /// Return various information that the workflow was initialized with. Will eventually become
     /// a proper non-proto workflow info struct.
     pub fn workflow_initial_info(&self) -> &InitializeWorkflow {

tests/integ_tests/update_tests.rs

@@ -1154,23 +1154,33 @@ async fn update_lost_on_activity_mismatch() {
     let mut worker = starter.worker().await;
     let client = starter.get_client().await;
 
-    static CAN_RUN: AtomicBool = AtomicBool::new(true);
     static DID_FAIL: AtomicBool = AtomicBool::new(false);
     worker.register_wf(wf_name.to_owned(), move |ctx: WfContext| async move {
+        let can_run = Arc::new(AtomicBool::new(true));
+        let cr = can_run.clone();
         ctx.update_handler(
             "update",
             |_: &_, _: ()| Ok(()),
-            move |_: UpdateContext, _: ()| async move {
-                dbg!("Running update!!");
-                CAN_RUN.store(true, Ordering::Relaxed);
-                Ok(())
+            move |c: UpdateContext, _: ()| {
+                let cr = cr.clone();
+                async move {
+                    dbg!(
+                        "Running update!!",
+                        c.wf_ctx.is_replaying(),
+                        &c.info.update_id
+                    );
+                    cr.store(true, Ordering::Relaxed);
+                    Ok(())
+                }
             },
         );
         for i in 1..=3 {
-            ctx.wait_condition(|| CAN_RUN.load(Ordering::Relaxed)).await;
-            if i == 2 {
-                ctx.force_task_fail(anyhow!("Fail on purpose"));
+            let cr = can_run.clone();
+            ctx.wait_condition(|| cr.load(Ordering::Relaxed)).await;
+            if i == 2 && !DID_FAIL.load(Ordering::Relaxed) {
                 DID_FAIL.store(true, Ordering::Relaxed);
+                // Pre-fix, test completes without this panic, but hangs with it
+                // panic!("Die on purpose");
             }
             ctx.activity(ActivityOptions {
                 activity_type: "echo".to_string(),
@@ -1179,18 +1189,22 @@ async fn update_lost_on_activity_mismatch() {
                 ..Default::default()
             })
             .await;
-            CAN_RUN.store(false, Ordering::Release);
+            can_run.store(false, Ordering::Release);
         }
         Ok(().into())
     });
     worker.register_activity("echo", |_ctx: ActContext, echo_me: String| async move {
         Ok(echo_me)
     });
 
+    let core_worker = worker.core_worker.clone();
     let handle = starter.start_with_worker(wf_name, &mut worker).await;
 
     let wf_id = starter.get_task_queue().to_string();
     let update = async {
+        tokio::time::sleep(Duration::from_millis(500)).await;
+        // Evict wf
+        core_worker.request_workflow_eviction(&handle.info().run_id.as_ref().unwrap());
         for _ in 1..=2 {
             // TODO: Without this, hangs, verify it doesn't after fix
             tokio::time::sleep(Duration::from_millis(500)).await;
@@ -1213,9 +1227,6 @@ async fn update_lost_on_activity_mismatch() {
         worker.run_until_done().await.unwrap();
     };
     join!(update, runner);
-    // This triggers the NDE before bugfix
-    CAN_RUN.store(true, Ordering::Relaxed);
-    dbg!("Replaying");
     handle
         .fetch_history_and_replay(worker.inner_mut())
         .await

tests/integ_tests/workflow_tests/determinism.rs

@@ -1,9 +1,10 @@
 use std::{
-    sync::atomic::{AtomicUsize, Ordering},
+    sync::atomic::{AtomicBool, AtomicUsize, Ordering},
     time::Duration,
 };
-use temporal_sdk::{ActivityOptions, WfContext, WorkflowResult};
-use temporal_sdk_core_test_utils::CoreWfStarter;
+use temporal_sdk::{ActContext, ActivityOptions, WfContext, WorkflowResult};
+use temporal_sdk_core_protos::coresdk::AsJsonPayloadExt;
+use temporal_sdk_core_test_utils::{CoreWfStarter, WorkflowHandleExt};
 
 static RUN_CT: AtomicUsize = AtomicUsize::new(1);
 
@@ -45,3 +46,40 @@ async fn test_determinism_error_then_recovers() {
     // 4 because we still add on the 3rd and final attempt
     assert_eq!(RUN_CT.load(Ordering::Relaxed), 4);
 }
+
+#[tokio::test]
+async fn task_fail_causes_replay_unset_too_soon() {
+    let wf_name = "task_fail_causes_replay_unset_too_soon";
+    let mut starter = CoreWfStarter::new(wf_name);
+    let mut worker = starter.worker().await;
+
+    static DID_FAIL: AtomicBool = AtomicBool::new(false);
+    worker.register_wf(wf_name.to_owned(), move |ctx: WfContext| async move {
+        if DID_FAIL.load(Ordering::Relaxed) {
+            assert!(ctx.is_replaying());
+        }
+        ctx.activity(ActivityOptions {
+            activity_type: "echo".to_string(),
+            input: "hi!".as_json_payload().expect("serializes fine"),
+            start_to_close_timeout: Some(Duration::from_secs(2)),
+            ..Default::default()
+        })
+        .await;
+        if !DID_FAIL.load(Ordering::Relaxed) {
+            DID_FAIL.store(true, Ordering::Relaxed);
+            panic!("Die on purpose");
+        }
+        Ok(().into())
+    });
+    worker.register_activity("echo", |_ctx: ActContext, echo_me: String| async move {
+        Ok(echo_me)
+    });
+
+    let handle = starter.start_with_worker(wf_name, &mut worker).await;
+
+    worker.run_until_done().await.unwrap();
+    handle
+        .fetch_history_and_replay(worker.inner_mut())
+        .await
+        .unwrap();
+}