Fix some unit test flakes (#703)

Author: Sushisource

core/src/core_tests/workers.rs

@@ -175,15 +175,14 @@ async fn complete_with_task_not_found_during_shutdown() {
     // Initiate shutdown before completing the activation
     let shutdown_fut = async {
         core.shutdown().await;
-        complete_order.borrow_mut().push(3);
+        complete_order.borrow_mut().push(2);
     };
     let poll_fut = async {
         // This will return shutdown once the completion goes through
         assert_matches!(
             core.poll_workflow_activation().await.unwrap_err(),
             PollWfError::ShutDown
         );
-        complete_order.borrow_mut().push(2);
     };
     let complete_fut = async {
         core.complete_workflow_activation(WorkflowActivationCompletion::from_cmds(
@@ -195,7 +194,7 @@ async fn complete_with_task_not_found_during_shutdown() {
         complete_order.borrow_mut().push(1);
     };
     tokio::join!(shutdown_fut, poll_fut, complete_fut);
-    assert_eq!(&complete_order.into_inner(), &[1, 2, 3])
+    assert_eq!(&complete_order.into_inner(), &[1, 2])
 }
 
 #[tokio::test]

core/src/test_help/mod.rs

@@ -139,7 +139,10 @@ pub(crate) fn build_fake_worker(
 
 pub(crate) fn build_fake_sdk(mock_cfg: MockPollCfg) -> temporal_sdk::Worker {
     let mut mock = build_mock_pollers(mock_cfg);
-    mock.worker_cfg(|c| c.max_cached_workflows = 1);
+    mock.worker_cfg(|c| {
+        c.max_cached_workflows = 1;
+        c.ignore_evicts_on_shutdown = false;
+    });
     let core = mock_worker(mock);
     let mut worker = temporal_sdk::Worker::new_from_core(Arc::new(core), "replay_q".to_string());
     worker.set_worker_interceptor(FailOnNondeterminismInterceptor {});

core/src/worker/activities/local_activities.rs

@@ -378,7 +378,19 @@ impl LocalActivityManager {
         let (new_or_retry, permit) = match self.rcvs.lock().await.next().await? {
             NewOrCancel::Cancel(c) => {
                 return match c {
-                    CancelOrTimeout::Cancel(c) => Some(NextPendingLAAction::Dispatch(c)),
+                    CancelOrTimeout::Cancel(c) => {
+                        if self
+                            .dat
+                            .lock()
+                            .outstanding_activity_tasks
+                            .contains_key(c.task_token.as_slice())
+                        {
+                            Some(NextPendingLAAction::Dispatch(c))
+                        } else {
+                            // Don't dispatch cancels for things we've already stopped tracking
+                            None
+                        }
+                    }
                     CancelOrTimeout::Timeout { run_id, resolution } => {
                         let tt = self
                             .dat

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

@@ -310,8 +310,8 @@ mod test {
         mock_cfg.num_expected_fails = 1;
         mock_cfg.expect_fail_wft_matcher = Box::new(move |_, cause, f| {
             matches!(cause, WorkflowTaskFailedCause::NonDeterministicError)
-                && matches!(f, Some(Failure { source, .. }) 
-            if source.contains("Timer fired event did not have expected timer id 1"))
+                && matches!(f, Some(Failure {message, .. })
+            if message.contains("Timer fired event did not have expected timer id 1"))
         });
         let mut worker = build_fake_sdk(mock_cfg);
         worker.register_wf(DEFAULT_WORKFLOW_TYPE, |ctx: WfContext| async move {

core/src/worker/workflow/managed_run.rs

@@ -968,11 +968,10 @@ impl ManagedRun {
                     } else {
                         WorkflowTaskFailedCause::Unspecified
                     };
-                    let wft_fail_str = format!("{:?}", fail.source);
                     self.failed_completion(
                         fail_cause,
                         fail.source.evict_reason(),
-                        Failure::application_failure(wft_fail_str, false).into(),
+                        fail.source.as_failure(),
                         true,
                         Some(resp_chan),
                     )

core/src/worker/workflow/mod.rs

@@ -288,7 +288,7 @@ impl Workflows {
                     self.activation_completed(
                         WorkflowActivationCompletion {
                             run_id,
-                            status: Some(auto_fail_to_complete_status(machines_err)),
+                            status: Some(machines_err.as_failure().into()),
                         },
                         true,
                         Option::<Box<dyn Fn(PostActivateHookData) + Send>>::None,
@@ -1297,6 +1297,18 @@ impl WFMachinesError {
             WFMachinesError::Fatal(_) => EvictionReason::Fatal,
         }
     }
+
+    pub fn as_failure(&self) -> Failure {
+        Failure {
+            failure: Some(
+                temporal_sdk_core_protos::temporal::api::failure::v1::Failure::application_failure(
+                    self.to_string(),
+                    false,
+                ),
+            ),
+            force_cause: WorkflowTaskFailedCause::from(self.evict_reason()) as i32,
+        }
+    }
 }
 
 impl From<TimestampError> for WFMachinesError {
@@ -1311,22 +1323,6 @@ impl From<anyhow::Error> for WFMachinesError {
     }
 }
 
-fn auto_fail_to_complete_status(err: WFMachinesError) -> workflow_activation_completion::Status {
-    workflow_activation_completion::Status::Failed(Failure {
-        failure: Some(
-            temporal_sdk_core_protos::temporal::api::failure::v1::Failure {
-                message: "Error while processing workflow task".to_string(),
-                source: err.to_string(),
-                stack_trace: "".to_string(),
-                encoded_attributes: None,
-                cause: None,
-                failure_info: None,
-            },
-        ),
-        force_cause: WorkflowTaskFailedCause::from(err.evict_reason()) as i32,
-    })
-}
-
 pub(crate) trait LocalActivityRequestSink: Send + Sync + 'static {
     fn sink_reqs(&self, reqs: Vec<LocalActRequest>) -> Vec<LocalActivityResolution>;
 }

sdk-core-protos/src/task_token.rs

@@ -1,5 +1,8 @@
 use base64::{prelude::BASE64_STANDARD, Engine};
-use std::fmt::{Debug, Display, Formatter};
+use std::{
+    borrow::Borrow,
+    fmt::{Debug, Display, Formatter},
+};
 
 static LOCAL_ACT_TASK_TOKEN_PREFIX: &[u8] = b"local_act_";
 
@@ -43,6 +46,12 @@ impl Debug for TaskToken {
     }
 }
 
+impl Borrow<[u8]> for TaskToken {
+    fn borrow(&self) -> &[u8] {
+        self.0.as_slice()
+    }
+}
+
 pub fn fmt_tt(tt: &[u8]) -> String {
     BASE64_STANDARD.encode(tt)
 }

sdk/src/lib.rs

@@ -521,7 +521,10 @@ impl ActivityHalf {
                 });
             }
             Some(activity_task::Variant::Cancel(_)) => {
-                if let Some(ct) = self.task_tokens_to_cancels.get(&activity.task_token.into()) {
+                if let Some(ct) = self
+                    .task_tokens_to_cancels
+                    .get(activity.task_token.as_slice())
+                {
                     ct.cancel();
                 }
             }

test-utils/src/lib.rs

@@ -11,7 +11,7 @@ pub mod workflows;
 pub use temporal_sdk_core::replay::HistoryForReplay;
 
 use crate::stream::{Stream, TryStreamExt};
-use anyhow::Context;
+use anyhow::{Context, Error};
 use base64::{prelude::BASE64_STANDARD, Engine};
 use futures::{future, stream, stream::FuturesUnordered, StreamExt};
 use parking_lot::Mutex;
@@ -47,6 +47,7 @@ use temporal_sdk_core_api::{
 };
 use temporal_sdk_core_protos::{
     coresdk::{
+        workflow_activation::WorkflowActivation,
         workflow_commands::{
             workflow_command, ActivityCancellationType, CompleteWorkflowExecution, QueryResult,
             QuerySuccess, ScheduleActivity, ScheduleLocalActivity, StartTimer,
@@ -510,7 +511,7 @@ impl TestWorker {
         });
     }
 
-    /// Runs until all expected workflows have completed
+    /// Runs until all expected workflows have completed and then shuts down the worker
     pub async fn run_until_done(&mut self) -> Result<(), anyhow::Error> {
         self.run_until_done_intercepted(Option::<TestWorkerCompletionIceptor>::None)
             .await
@@ -604,6 +605,12 @@ impl WorkerInterceptor for TestWorkerCompletionIceptor {
             n.on_shutdown(sdk_worker);
         }
     }
+    async fn on_workflow_activation(&self, a: &WorkflowActivation) -> Result<(), Error> {
+        if let Some(n) = self.next.as_ref() {
+            n.on_workflow_activation(a).await?;
+        }
+        Ok(())
+    }
 }
 
 /// Returns the client options used to connect to the server used for integration tests.