@temporalio/core-bridge 0.23.0 → 1.0.0

package/sdk-core/core/src/abstractions.rs

@@ -1,13 +1,19 @@
 //! This module contains very generic helpers that can be used codebase-wide
 
 use crate::MetricsContext;
-use tokio::sync::{AcquireError, Semaphore, SemaphorePermit};
+use futures::{stream, Stream, StreamExt};
+use std::{
+    fmt::{Debug, Formatter},
+    future::Future,
+    sync::Arc,
+};
+use tokio::sync::{AcquireError, Notify, OwnedSemaphorePermit, Semaphore, TryAcquireError};
 
 /// Wraps a [Semaphore] with a function call that is fed the available permits any time a permit is
 /// acquired or restored through the provided methods
+#[derive(Clone)]
 pub(crate) struct MeteredSemaphore {
-    pub sem: Semaphore,
-    max_permits: usize,
+    sem: Arc<Semaphore>,
     metrics_ctx: MetricsContext,
     record_fn: fn(&MetricsContext, usize),
 }
@@ -19,27 +25,142 @@ impl MeteredSemaphore {
         record_fn: fn(&MetricsContext, usize),
     ) -> Self {
         Self {
-            sem: Semaphore::new(inital_permits),
-            max_permits: inital_permits,
+            sem: Arc::new(Semaphore::new(inital_permits)),
             metrics_ctx,
             record_fn,
         }
     }
 
-    pub async fn acquire(&self) -> Result<SemaphorePermit<'_>, AcquireError> {
-        let res = self.sem.acquire().await;
+    pub fn available_permits(&self) -> usize {
+        self.sem.available_permits()
+    }
+
+    pub async fn acquire_owned(&self) -> Result<OwnedMeteredSemPermit, AcquireError> {
+        let res = self.sem.clone().acquire_owned().await?;
+        self.record();
+        Ok(OwnedMeteredSemPermit {
+            inner: res,
+            record_fn: self.record_drop_owned(),
+        })
+    }
+
+    pub fn try_acquire_owned(&self) -> Result<OwnedMeteredSemPermit, TryAcquireError> {
+        let res = self.sem.clone().try_acquire_owned()?;
+        self.record();
+        Ok(OwnedMeteredSemPermit {
+            inner: res,
+            record_fn: self.record_drop_owned(),
+        })
+    }
+
+    fn record(&self) {
         (self.record_fn)(&self.metrics_ctx, self.sem.available_permits());
-        res
     }
 
-    /// Adds just one permit. Will not add if already at the initial/max capacity.
-    pub fn add_permit(&self) {
-        if self.sem.available_permits() < self.max_permits {
-            self.sem.add_permits(1);
-            (self.record_fn)(&self.metrics_ctx, self.sem.available_permits());
-        } else if cfg!(debug_assertions) {
-            // Panic only during debug mode if this happens
-            panic!("Tried to add permit to a semaphore that already was at capacity!");
+    fn record_drop_owned(&self) -> Box<dyn Fn() + Send + Sync> {
+        let rcf = self.record_fn;
+        let mets = self.metrics_ctx.clone();
+        let sem = self.sem.clone();
+        Box::new(move || rcf(&mets, sem.available_permits() + 1))
+    }
+}
+
+/// Wraps an [OwnedSemaphorePermit] to update metrics when it's dropped
+pub(crate) struct OwnedMeteredSemPermit {
+    inner: OwnedSemaphorePermit,
+    record_fn: Box<dyn Fn() + Send + Sync>,
+}
+impl Drop for OwnedMeteredSemPermit {
+    fn drop(&mut self) {
+        (self.record_fn)()
+    }
+}
+impl Debug for OwnedMeteredSemPermit {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        self.inner.fmt(f)
+    }
+}
+
+/// From the input stream, create a new stream which only pulls from the input stream when allowed.
+/// When allowed is determined by the passed in `proceeder` which must return a future every time
+/// it's called. The input stream is only pulled from when that future resolves.
+pub(crate) fn stream_when_allowed<S, F, FF>(
+    input: S,
+    proceeder: FF,
+) -> impl Stream<Item = (S::Item, F::Output)>
+where
+    S: Stream + Send + 'static,
+    F: Future,
+    FF: FnMut() -> F,
+{
+    let acceptable_notify = Arc::new(Notify::new());
+    acceptable_notify.notify_one();
+    let stream = stream::unfold(
+        (proceeder, input.boxed()),
+        |(mut proceeder, mut input)| async {
+            let v = proceeder().await;
+            input.next().await.map(|i| ((i, v), (proceeder, input)))
+        },
+    );
+    stream
+}
+
+macro_rules! dbg_panic {
+  ($($arg:tt)*) => {
+      error!($($arg)*);
+      debug_assert!(true, $($arg)*);
+  };
+}
+pub(crate) use dbg_panic;
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use futures::pin_mut;
+    use std::{cell::RefCell, task::Poll};
+    use tokio::sync::mpsc::unbounded_channel;
+
+    // This is fine. Test only / guaranteed to happen serially.
+    #[allow(clippy::await_holding_refcell_ref)]
+    #[test]
+    fn stream_when_allowed_works() {
+        let inputs = stream::iter([1, 2, 3]);
+        let (allow_tx, allow_rx) = unbounded_channel();
+        let allow_rx = RefCell::new(allow_rx);
+        let when_allowed = stream_when_allowed(inputs, || async {
+            allow_rx.borrow_mut().recv().await.unwrap()
+        });
+
+        let waker = futures::task::noop_waker_ref();
+        let mut cx = std::task::Context::from_waker(waker);
+        pin_mut!(when_allowed);
+
+        allow_tx.send(()).unwrap();
+        assert_eq!(
+            when_allowed.poll_next_unpin(&mut cx),
+            Poll::Ready(Some((1, ())))
+        );
+        // Now, it won't be ready
+        for _ in 1..10 {
+            assert_eq!(when_allowed.poll_next_unpin(&mut cx), Poll::Pending);
+        }
+        allow_tx.send(()).unwrap();
+        assert_eq!(
+            when_allowed.poll_next_unpin(&mut cx),
+            Poll::Ready(Some((2, ())))
+        );
+        for _ in 1..10 {
+            assert_eq!(when_allowed.poll_next_unpin(&mut cx), Poll::Pending);
+        }
+        allow_tx.send(()).unwrap();
+        assert_eq!(
+            when_allowed.poll_next_unpin(&mut cx),
+            Poll::Ready(Some((3, ())))
+        );
+        for _ in 1..10 {
+            assert_eq!(when_allowed.poll_next_unpin(&mut cx), Poll::Pending);
         }
+        allow_tx.send(()).unwrap();
+        assert_eq!(when_allowed.poll_next_unpin(&mut cx), Poll::Ready(None));
     }
 }

package/sdk-core/core/src/core_tests/activity_tasks.rs

@@ -1,11 +1,12 @@
 use crate::{
-    job_assert,
+    advance_fut, job_assert,
     test_help::{
-        build_fake_worker, canned_histories, gen_assert_and_reply, mock_manual_poller, mock_poller,
-        mock_worker, poll_and_reply, test_worker_cfg, MockWorker, MocksHolder,
+        build_fake_worker, build_mock_pollers, canned_histories, gen_assert_and_reply,
+        mock_manual_poller, mock_poller, mock_poller_from_resps, mock_worker, poll_and_reply,
+        single_hist_mock_sg, test_worker_cfg, MockPollCfg, MockWorkerInputs, MocksHolder,
+        ResponseType, WorkflowCachingPolicy, TEST_Q,
     },
     worker::client::mocks::{mock_manual_workflow_client, mock_workflow_client},
-    workflow::WorkflowCachingPolicy::NonSticky,
     ActivityHeartbeat, Worker, WorkerConfigBuilder,
 };
 use futures::FutureExt;
@@ -13,9 +14,14 @@ use std::{
     cell::RefCell,
     collections::{hash_map::Entry, HashMap, VecDeque},
     rc::Rc,
-    sync::atomic::{AtomicUsize, Ordering},
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
     time::Duration,
 };
+use temporal_client::WorkflowOptions;
+use temporal_sdk::{ActivityOptions, WfContext};
 use temporal_sdk_core_api::Worker as WorkerTrait;
 use temporal_sdk_core_protos::{
     coresdk::{
@@ -26,16 +32,22 @@ use temporal_sdk_core_protos::{
             ActivityCancellationType, CompleteWorkflowExecution, RequestCancelActivity,
             ScheduleActivity,
         },
+        workflow_completion::WorkflowActivationCompletion,
         ActivityTaskCompletion,
     },
-    temporal::api::workflowservice::v1::{
-        PollActivityTaskQueueResponse, RecordActivityTaskHeartbeatResponse,
-        RespondActivityTaskCanceledResponse, RespondActivityTaskCompletedResponse,
-        RespondActivityTaskFailedResponse,
+    temporal::api::{
+        command::v1::command::Attributes,
+        enums::v1::EventType,
+        workflowservice::v1::{
+            PollActivityTaskQueueResponse, RecordActivityTaskHeartbeatResponse,
+            RespondActivityTaskCanceledResponse, RespondActivityTaskCompletedResponse,
+            RespondActivityTaskFailedResponse, RespondWorkflowTaskCompletedResponse,
+        },
     },
+    TestHistoryBuilder, DEFAULT_WORKFLOW_TYPE,
 };
-use temporal_sdk_core_test_utils::{fanout_tasks, start_timer_cmd};
-use tokio::{join, time::sleep};
+use temporal_sdk_core_test_utils::{fanout_tasks, start_timer_cmd, TestWorker};
+use tokio::{sync::Barrier, time::sleep};
 
 #[tokio::test]
 async fn max_activities_respected() {
@@ -77,23 +89,18 @@ async fn max_activities_respected() {
     // We allow two outstanding activities, therefore first two polls should return right away
     let r1 = worker.poll_activity_task().await.unwrap();
     let _r2 = worker.poll_activity_task().await.unwrap();
-    // Third should block until we complete one of the first two
-    let last_finisher = AtomicUsize::new(0);
-    tokio::join! {
-        async {
-            worker.complete_activity_task(ActivityTaskCompletion {
-                task_token: r1.task_token,
-                result: Some(ActivityExecutionResult::ok(vec![1].into()))
-            }).await.unwrap();
-            last_finisher.store(1, Ordering::SeqCst);
-        },
-        async {
-            worker.poll_activity_task().await.unwrap();
-            last_finisher.store(2, Ordering::SeqCst);
-        }
-    };
-    // So that we know we blocked
-    assert_eq!(last_finisher.load(Ordering::Acquire), 2);
+    // Third poll should block until we complete one of the first two. To ensure this, manually
+    // poll it a bunch to see it's not resolving.
+    let poll_fut = worker.poll_activity_task();
+    advance_fut!(poll_fut);
+    worker
+        .complete_activity_task(ActivityTaskCompletion {
+            task_token: r1.task_token,
+            result: Some(ActivityExecutionResult::ok(vec![1].into())),
+        })
+        .await
+        .unwrap();
+    poll_fut.await.unwrap();
 }
 
 #[tokio::test]
@@ -102,7 +109,7 @@ async fn activity_not_found_returns_ok() {
     // Mock won't even be called, since we weren't tracking activity
     mock_client.expect_complete_activity_task().times(0);
 
-    let core = mock_worker(MocksHolder::from_client_with_responses(mock_client, [], []));
+    let core = mock_worker(MocksHolder::from_client_with_activities(mock_client, []));
 
     core.complete_activity_task(ActivityTaskCompletion {
         task_token: vec![1],
@@ -133,22 +140,23 @@ async fn heartbeats_report_cancels_only_once() {
         .times(1)
         .returning(|_, _| Ok(RespondActivityTaskCanceledResponse::default()));
 
-    let core = mock_worker(MocksHolder::from_client_with_responses(
+    let core = mock_worker(MocksHolder::from_client_with_activities(
         mock_client,
-        [],
         [
             PollActivityTaskQueueResponse {
                 task_token: vec![1],
                 activity_id: "act1".to_string(),
                 heartbeat_timeout: Some(Duration::from_millis(1).into()),
                 ..Default::default()
-            },
+            }
+            .into(),
             PollActivityTaskQueueResponse {
                 task_token: vec![2],
                 activity_id: "act2".to_string(),
                 heartbeat_timeout: Some(Duration::from_millis(1).into()),
                 ..Default::default()
-            },
+            }
+            .into(),
         ],
     ));
 
@@ -248,7 +256,7 @@ async fn activity_cancel_interrupts_poll() {
         .times(1)
         .returning(|_, _| async { Ok(RespondActivityTaskCompletedResponse::default()) }.boxed());
 
-    let mw = MockWorker {
+    let mw = MockWorkerInputs {
         act_poller: Some(Box::from(mock_poller)),
         ..Default::default()
     };
@@ -300,7 +308,7 @@ async fn activity_poll_timeout_retries() {
             }))
         }
     });
-    let mw = MockWorker {
+    let mw = MockWorkerInputs {
         act_poller: Some(Box::from(mock_act_poller)),
         ..Default::default()
     };
@@ -430,7 +438,7 @@ async fn activity_timeout_no_double_resolve() {
 
     poll_and_reply(
         &core,
-        NonSticky,
+        WorkflowCachingPolicy::NonSticky,
         &[
             gen_assert_and_reply(
                 &job_assert!(workflow_activation_job::Variant::StartWorkflow(_)),
@@ -490,41 +498,34 @@ async fn can_heartbeat_acts_during_shutdown() {
         .times(1)
         .returning(|_, _| Ok(RespondActivityTaskCompletedResponse::default()));
 
-    let core = mock_worker(MocksHolder::from_client_with_responses(
+    let core = mock_worker(MocksHolder::from_client_with_activities(
         mock_client,
-        [],
         [PollActivityTaskQueueResponse {
             task_token: vec![1],
             activity_id: "act1".to_string(),
             heartbeat_timeout: Some(Duration::from_millis(1).into()),
             ..Default::default()
-        }],
+        }
+        .into()],
     ));
 
     let act = core.poll_activity_task().await.unwrap();
-    let complete_order = RefCell::new(vec![]);
-    // Start shutdown before completing the activity
-    let shutdown_fut = async {
-        core.shutdown().await;
-        complete_order.borrow_mut().push(1);
-    };
-    let complete_fut = async {
-        core.record_activity_heartbeat(ActivityHeartbeat {
-            task_token: act.task_token.clone(),
+    // Make sure shutdown has progressed before trying to record heartbeat / complete
+    let shutdown_fut = core.shutdown();
+    advance_fut!(shutdown_fut);
+    core.record_activity_heartbeat(ActivityHeartbeat {
+        task_token: act.task_token.clone(),
 
-            details: vec![vec![1_u8, 2, 3].into()],
-        });
-        core.complete_activity_task(ActivityTaskCompletion {
-            task_token: act.task_token,
+        details: vec![vec![1_u8, 2, 3].into()],
+    });
+    core.complete_activity_task(ActivityTaskCompletion {
+        task_token: act.task_token,
 
-            result: Some(ActivityExecutionResult::ok(vec![1].into())),
-        })
-        .await
-        .unwrap();
-        complete_order.borrow_mut().push(2);
-    };
-    join!(shutdown_fut, complete_fut);
-    assert_eq!(&complete_order.into_inner(), &[2, 1])
+        result: Some(ActivityExecutionResult::ok(vec![1].into())),
+    })
+    .await
+    .unwrap();
+    shutdown_fut.await;
 }
 
 /// Verifies that if a user has tried to record a heartbeat and then immediately after failed the
@@ -550,15 +551,15 @@ async fn complete_act_with_fail_flushes_heartbeat() {
         .times(1)
         .returning(|_, _| Ok(RespondActivityTaskFailedResponse::default()));
 
-    let core = mock_worker(MocksHolder::from_client_with_responses(
+    let core = mock_worker(MocksHolder::from_client_with_activities(
         mock_client,
-        [],
         [PollActivityTaskQueueResponse {
             task_token: vec![1],
             activity_id: "act1".to_string(),
             heartbeat_timeout: Some(Duration::from_secs(10).into()),
             ..Default::default()
-        }],
+        }
+        .into()],
     ));
 
     let act = core.poll_activity_task().await.unwrap();
@@ -600,9 +601,203 @@ async fn max_tq_acts_set_passed_to_poll_properly() {
         .namespace("enchi")
         .task_queue("cat")
         .max_concurrent_at_polls(1_usize)
+        .worker_build_id("test_bin_id")
         .max_task_queue_activities_per_second(rate)
         .build()
         .unwrap();
     let worker = Worker::new_test(cfg, mock_client);
     worker.poll_activity_task().await.unwrap();
 }
+
+/// This test verifies that activity tasks which come as replies to completing a WFT are properly
+/// delivered via polling.
+#[tokio::test]
+async fn activity_tasks_from_completion_are_delivered() {
+    let wfid = "fake_wf_id";
+    let mut t = TestHistoryBuilder::default();
+    t.add_by_type(EventType::WorkflowExecutionStarted);
+    t.add_full_wf_task();
+    let schedid = t.add_activity_task_scheduled("act_id");
+    let startid = t.add_activity_task_started(schedid);
+    t.add_activity_task_completed(schedid, startid, b"hi".into());
+    t.add_full_wf_task();
+    t.add_workflow_execution_completed();
+
+    let mut mock = mock_workflow_client();
+    mock.expect_complete_workflow_task()
+        .times(1)
+        .returning(move |_| {
+            Ok(RespondWorkflowTaskCompletedResponse {
+                workflow_task: None,
+                activity_tasks: vec![PollActivityTaskQueueResponse {
+                    task_token: vec![1],
+                    activity_id: "act1".to_string(),
+                    ..Default::default()
+                }],
+            })
+        });
+    mock.expect_complete_activity_task()
+        .times(1)
+        .returning(|_, _| Ok(RespondActivityTaskCompletedResponse::default()));
+    let mut mock = single_hist_mock_sg(wfid, t, [1], mock, true);
+    let mut mock_poller = mock_manual_poller();
+    mock_poller
+        .expect_poll()
+        .returning(|| futures::future::pending().boxed());
+    mock.set_act_poller(Box::new(mock_poller));
+    mock.worker_cfg(|wc| wc.max_cached_workflows = 2);
+    let core = mock_worker(mock);
+
+    let wf_task = core.poll_workflow_activation().await.unwrap();
+    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
+        wf_task.run_id,
+        ScheduleActivity {
+            seq: 1,
+            activity_id: "act_id".to_string(),
+            cancellation_type: ActivityCancellationType::TryCancel as i32,
+            ..Default::default()
+        }
+        .into(),
+    ))
+    .await
+    .unwrap();
+
+    // We should see the activity when we poll now
+    let act_task = core.poll_activity_task().await.unwrap();
+    assert_eq!(act_task.task_token, vec![1]);
+
+    core.complete_activity_task(ActivityTaskCompletion {
+        task_token: act_task.task_token.clone(),
+        result: Some(ActivityExecutionResult::ok("hi".into())),
+    })
+    .await
+    .unwrap();
+
+    core.shutdown().await;
+}
+
+#[tokio::test]
+async fn activity_tasks_from_completion_reserve_slots() {
+    let wf_id = "fake_wf_id";
+    let mut t = TestHistoryBuilder::default();
+    t.add_by_type(EventType::WorkflowExecutionStarted);
+    t.add_full_wf_task();
+    let schedid = t.add_activity_task_scheduled("1");
+    let startid = t.add_activity_task_started(schedid);
+    t.add_activity_task_completed(schedid, startid, b"hi".into());
+    t.add_full_wf_task();
+    let schedid = t.add_activity_task_scheduled("2");
+    let startid = t.add_activity_task_started(schedid);
+    t.add_activity_task_completed(schedid, startid, b"hi".into());
+    t.add_full_wf_task();
+    t.add_workflow_execution_completed();
+
+    let mut mock = mock_workflow_client();
+    // Set up two tasks to be returned via normal activity polling
+    let act_tasks = VecDeque::from(vec![
+        PollActivityTaskQueueResponse {
+            task_token: vec![1],
+            activity_id: "act1".to_string(),
+            ..Default::default()
+        }
+        .into(),
+        PollActivityTaskQueueResponse {
+            task_token: vec![2],
+            activity_id: "act2".to_string(),
+            ..Default::default()
+        }
+        .into(),
+    ]);
+    mock.expect_complete_activity_task()
+        .times(2)
+        .returning(|_, _| Ok(RespondActivityTaskCompletedResponse::default()));
+    let barr: &'static Barrier = Box::leak(Box::new(Barrier::new(2)));
+    let mut mh = MockPollCfg::from_resp_batches(
+        wf_id,
+        t,
+        [
+            ResponseType::ToTaskNum(1),
+            // We don't want the second task to be delivered until *after* the activity tasks
+            // have been completed, so that the second activity schedule will have slots available
+            ResponseType::UntilResolved(
+                async {
+                    barr.wait().await;
+                    barr.wait().await;
+                }
+                .boxed(),
+                2,
+            ),
+            ResponseType::AllHistory,
+        ],
+        mock,
+    );
+    mh.completion_asserts = Some(Box::new(|wftc| {
+        // Make sure when we see the completion with the schedule act command that it does
+        // not have the eager execution flag set the first time, and does the second.
+        if let Some(Attributes::ScheduleActivityTaskCommandAttributes(attrs)) =
+            wftc.commands.get(0).and_then(|cmd| cmd.attributes.as_ref())
+        {
+            if attrs.activity_id == "1" {
+                assert!(!attrs.request_eager_execution);
+            } else {
+                assert!(attrs.request_eager_execution);
+            }
+        }
+    }));
+    let mut mock = build_mock_pollers(mh);
+    mock.worker_cfg(|cfg| {
+        cfg.max_cached_workflows = 2;
+        cfg.max_outstanding_activities = 2;
+    });
+    mock.set_act_poller(mock_poller_from_resps(act_tasks));
+    let core = Arc::new(mock_worker(mock));
+    let mut worker = TestWorker::new(core.clone(), TEST_Q.to_string());
+
+    // First poll for activities twice, occupying both slots
+    let at1 = core.poll_activity_task().await.unwrap();
+    let at2 = core.poll_activity_task().await.unwrap();
+
+    worker.register_wf(DEFAULT_WORKFLOW_TYPE, move |ctx: WfContext| async move {
+        ctx.activity(ActivityOptions {
+            activity_type: "act1".to_string(),
+            ..Default::default()
+        })
+        .await;
+        ctx.activity(ActivityOptions {
+            activity_type: "act2".to_string(),
+            ..Default::default()
+        })
+        .await;
+        Ok(().into())
+    });
+
+    worker
+        .submit_wf(
+            wf_id.to_owned(),
+            DEFAULT_WORKFLOW_TYPE,
+            vec![],
+            WorkflowOptions::default(),
+        )
+        .await
+        .unwrap();
+    let act_completer = async {
+        barr.wait().await;
+        core.complete_activity_task(ActivityTaskCompletion {
+            task_token: at1.task_token,
+            result: Some(ActivityExecutionResult::ok("hi".into())),
+        })
+        .await
+        .unwrap();
+        core.complete_activity_task(ActivityTaskCompletion {
+            task_token: at2.task_token,
+            result: Some(ActivityExecutionResult::ok("hi".into())),
+        })
+        .await
+        .unwrap();
+        barr.wait().await;
+    };
+    // This wf poll should *not* set the flag that it wants tasks back since both slots are
+    // occupied
+    let run_fut = async { worker.run_until_done().await.unwrap() };
+    tokio::join!(run_fut, act_completer);
+}

package/sdk-core/core/src/core_tests/child_workflows.rs

@@ -1,8 +1,7 @@
 use crate::{
     replay::DEFAULT_WORKFLOW_TYPE,
     test_help::{canned_histories, mock_sdk, MockPollCfg, ResponseType},
-    worker::client::mocks::mock_workflow_client,
-    workflow::managed_wf::ManagedWFFunc,
+    worker::{client::mocks::mock_workflow_client, ManagedWFFunc},
 };
 use temporal_client::WorkflowOptions;
 use temporal_sdk::{ChildWorkflowOptions, Signal, WfContext, WorkflowFunction, WorkflowResult};

package/sdk-core/core/src/core_tests/determinism.rs

@@ -34,7 +34,7 @@ async fn test_panic_wf_task_rejected_properly() {
     let mut mh = MockPollCfg::from_resp_batches(wf_id, t, [1, 2, 2], mock);
     // We should see one wft failure which has unspecified cause, since panics don't have a defined
     // type.
-    mh.num_expected_fails = Some(1);
+    mh.num_expected_fails = 1;
     mh.expect_fail_wft_matcher =
         Box::new(|_, cause, _| matches!(cause, WorkflowTaskFailedCause::Unspecified));
     let mut worker = mock_sdk(mh);
@@ -71,7 +71,7 @@ async fn test_wf_task_rejected_properly_due_to_nondeterminism(#[case] use_cache:
         mock,
     );
     // We should see one wft failure which has nondeterminism cause
-    mh.num_expected_fails = Some(1);
+    mh.num_expected_fails = 1;
     mh.expect_fail_wft_matcher =
         Box::new(|_, cause, _| matches!(cause, WorkflowTaskFailedCause::NonDeterministicError));
     let mut worker = mock_sdk_cfg(mh, |cfg| {

package/sdk-core/core/src/core_tests/local_activities.rs

@@ -12,8 +12,8 @@ use std::{
 use temporal_client::WorkflowOptions;
 use temporal_sdk::{ActContext, LocalActivityOptions, WfContext, WorkflowResult};
 use temporal_sdk_core_protos::{
-    coresdk::{common::RetryPolicy, AsJsonPayloadExt},
-    temporal::api::{enums::v1::EventType, failure::v1::Failure},
+    coresdk::AsJsonPayloadExt,
+    temporal::api::{common::v1::RetryPolicy, enums::v1::EventType, failure::v1::Failure},
 };
 use tokio::sync::Barrier;
 
@@ -159,12 +159,13 @@ async fn local_act_heartbeat(#[case] shutdown_middle: bool) {
 
     let wf_id = "fakeid";
     let mock = mock_workflow_client();
-    // Allow returning incomplete history more than once, as the wft timeout can be timing sensitive
-    // and might poll an extra time
-    let mut mh = MockPollCfg::from_resp_batches(wf_id, t, [1, 2, 2, 2, 2], mock);
+    let mut mh = MockPollCfg::from_resp_batches(wf_id, t, [1, 2, 2, 2], mock);
     mh.enforce_correct_number_of_polls = false;
-    let mut worker = mock_sdk_cfg(mh, |wc| wc.max_cached_workflows = 1);
-    let core = worker.orig_core_worker.clone();
+    let mut worker = mock_sdk_cfg(mh, |wc| {
+        wc.max_cached_workflows = 1;
+        wc.max_outstanding_workflow_tasks = 1;
+    });
+    let core = worker.core_worker.clone();
 
     let shutdown_barr: &'static Barrier = Box::leak(Box::new(Barrier::new(2)));