@temporalio/core-bridge 1.6.0 → 1.7.1

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

@@ -1,18 +1,30 @@
 //! This module contains very generic helpers that can be used codebase-wide
 
+pub mod take_cell;
+
 use crate::MetricsContext;
+use derive_more::DebugCustom;
 use futures::{stream, Stream, StreamExt};
 use std::{
     fmt::{Debug, Formatter},
-    sync::Arc,
+    sync::{
+        atomic::{AtomicBool, AtomicUsize, Ordering},
+        Arc,
+    },
 };
 use tokio::sync::{AcquireError, OwnedSemaphorePermit, Semaphore, TryAcquireError};
+use tokio_util::sync::CancellationToken;
 
 /// 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 {
     sem: Arc<Semaphore>,
+    /// The number of permit owners who have acquired a permit from the semaphore, but are not yet
+    /// meaningfully using that permit. This is useful for giving a more semantically accurate count
+    /// of used task slots, since we typically wait for a permit first before polling, but that slot
+    /// isn't used in the sense the user expects until we actually also get the corresponding task.
+    unused_claimants: Arc<AtomicUsize>,
     metrics_ctx: MetricsContext,
     record_fn: fn(&MetricsContext, usize),
 }
@@ -25,6 +37,7 @@ impl MeteredSemaphore {
     ) -> Self {
         Self {
             sem: Arc::new(Semaphore::new(inital_permits)),
+            unused_claimants: Arc::new(AtomicUsize::new(0)),
             metrics_ctx,
             record_fn,
         }
@@ -36,42 +49,154 @@ impl MeteredSemaphore {
 
     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(),
-        })
+        Ok(self.build_owned(res))
     }
 
     pub fn try_acquire_owned(&self) -> Result<OwnedMeteredSemPermit, TryAcquireError> {
         let res = self.sem.clone().try_acquire_owned()?;
+        Ok(self.build_owned(res))
+    }
+
+    fn build_owned(&self, res: OwnedSemaphorePermit) -> OwnedMeteredSemPermit {
+        self.unused_claimants.fetch_add(1, Ordering::Release);
         self.record();
-        Ok(OwnedMeteredSemPermit {
+        OwnedMeteredSemPermit {
             inner: res,
-            record_fn: self.record_drop_owned(),
-        })
+            unused_claimants: Some(self.unused_claimants.clone()),
+            record_fn: self.record_owned(),
+        }
     }
 
     fn record(&self) {
-        (self.record_fn)(&self.metrics_ctx, self.sem.available_permits());
+        (self.record_fn)(
+            &self.metrics_ctx,
+            self.sem.available_permits() + self.unused_claimants.load(Ordering::Acquire),
+        );
     }
 
-    fn record_drop_owned(&self) -> Box<dyn Fn() + Send + Sync> {
+    fn record_owned(&self) -> Box<dyn Fn(bool) + 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))
+        let uc = self.unused_claimants.clone();
+        // When being called from the drop impl, the semaphore permit isn't actually dropped yet,
+        // so account for that.
+        Box::new(move |add_one: bool| {
+            let extra = usize::from(add_one);
+            rcf(
+                &mets,
+                sem.available_permits() + uc.load(Ordering::Acquire) + extra,
+            )
+        })
+    }
+}
+
+/// A version of [MeteredSemaphore] that can be closed and supports waiting for close to complete.
+/// Once closed, no permits will be handed out.
+/// Close completes when all permits have been returned.
+pub(crate) struct ClosableMeteredSemaphore {
+    inner: Arc<MeteredSemaphore>,
+    outstanding_permits: AtomicUsize,
+    close_requested: AtomicBool,
+    close_complete_token: CancellationToken,
+}
+
+impl ClosableMeteredSemaphore {
+    pub fn new_arc(sem: Arc<MeteredSemaphore>) -> Arc<Self> {
+        Arc::new(Self {
+            inner: sem,
+            outstanding_permits: Default::default(),
+            close_requested: AtomicBool::new(false),
+            close_complete_token: CancellationToken::new(),
+        })
+    }
+}
+
+impl ClosableMeteredSemaphore {
+    #[cfg(test)]
+    pub fn available_permits(&self) -> usize {
+        self.inner.available_permits()
+    }
+
+    /// Request to close the semaphore and prevent new permits from being acquired.
+    pub fn close(&self) {
+        self.close_requested.store(true, Ordering::Release);
+        if self.outstanding_permits.load(Ordering::Acquire) == 0 {
+            self.close_complete_token.cancel();
+        }
+    }
+
+    /// Returns after close has been requested and all outstanding permits have been returned.
+    pub async fn close_complete(&self) {
+        self.close_complete_token.cancelled().await;
+    }
+
+    /// Acquire a permit if one is available and close was not requested.
+    pub fn try_acquire_owned(
+        self: &Arc<Self>,
+    ) -> Result<TrackedOwnedMeteredSemPermit, TryAcquireError> {
+        if self.close_requested.load(Ordering::Acquire) {
+            return Err(TryAcquireError::Closed);
+        }
+        self.outstanding_permits.fetch_add(1, Ordering::Release);
+        let res = self.inner.try_acquire_owned();
+        if res.is_err() {
+            self.outstanding_permits.fetch_sub(1, Ordering::Release);
+        }
+        res.map(|permit| TrackedOwnedMeteredSemPermit {
+            inner: Some(permit),
+            on_drop: self.on_permit_dropped(),
+        })
+    }
+
+    fn on_permit_dropped(self: &Arc<Self>) -> Box<dyn Fn() + Send + Sync> {
+        let sem = self.clone();
+        Box::new(move || {
+            sem.outstanding_permits.fetch_sub(1, Ordering::Release);
+            if sem.close_requested.load(Ordering::Acquire)
+                && sem.outstanding_permits.load(Ordering::Acquire) == 0
+            {
+                sem.close_complete_token.cancel();
+            }
+        })
+    }
+}
+
+/// Tracks an OwnedMeteredSemPermit and calls on_drop when dropped.
+#[derive(DebugCustom)]
+#[debug(fmt = "Tracked({inner:?})")]
+pub(crate) struct TrackedOwnedMeteredSemPermit {
+    inner: Option<OwnedMeteredSemPermit>,
+    on_drop: Box<dyn Fn() + Send + Sync>,
+}
+impl From<TrackedOwnedMeteredSemPermit> for OwnedMeteredSemPermit {
+    fn from(mut value: TrackedOwnedMeteredSemPermit) -> Self {
+        value
+            .inner
+            .take()
+            .expect("Inner permit should be available")
+    }
+}
+impl Drop for TrackedOwnedMeteredSemPermit {
+    fn drop(&mut self) {
+        (self.on_drop)();
     }
 }
 
 /// Wraps an [OwnedSemaphorePermit] to update metrics when it's dropped
 pub(crate) struct OwnedMeteredSemPermit {
     inner: OwnedSemaphorePermit,
-    record_fn: Box<dyn Fn() + Send + Sync>,
+    /// See [MeteredSemaphore::unused_claimants]. If present when dropping, used to decrement the
+    /// count.
+    unused_claimants: Option<Arc<AtomicUsize>>,
+    record_fn: Box<dyn Fn(bool) + Send + Sync>,
 }
 impl Drop for OwnedMeteredSemPermit {
     fn drop(&mut self) {
-        (self.record_fn)()
+        if let Some(uc) = self.unused_claimants.take() {
+            uc.fetch_sub(1, Ordering::Release);
+        }
+        (self.record_fn)(true)
     }
 }
 impl Debug for OwnedMeteredSemPermit {
@@ -79,15 +204,30 @@ impl Debug for OwnedMeteredSemPermit {
         self.inner.fmt(f)
     }
 }
-#[cfg(feature = "save_wf_inputs")]
 impl OwnedMeteredSemPermit {
+    /// Should be called once this permit is actually being "used" for the work it was meant to
+    /// permit.
+    pub(crate) fn into_used(mut self) -> UsedMeteredSemPermit {
+        if let Some(uc) = self.unused_claimants.take() {
+            uc.fetch_sub(1, Ordering::Release);
+            (self.record_fn)(false)
+        }
+        UsedMeteredSemPermit(self)
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct UsedMeteredSemPermit(OwnedMeteredSemPermit);
+impl UsedMeteredSemPermit {
+    #[cfg(feature = "save_wf_inputs")]
     pub(crate) fn fake_deserialized() -> Self {
         let sem = Arc::new(Semaphore::new(1));
         let inner = sem.try_acquire_owned().unwrap();
-        Self {
+        Self(OwnedMeteredSemPermit {
             inner,
-            record_fn: Box::new(|| {}),
-        }
+            unused_claimants: None,
+            record_fn: Box::new(|_| {}),
+        })
     }
 }
 
@@ -174,4 +314,43 @@ mod tests {
         allow_tx.send(()).unwrap();
         assert_eq!(when_allowed.poll_next_unpin(&mut cx), Poll::Ready(None));
     }
+
+    #[tokio::test]
+    async fn closable_semaphore_permit_drop_returns_permit() {
+        let inner = MeteredSemaphore::new(2, MetricsContext::no_op(), |_, _| {});
+        let sem = ClosableMeteredSemaphore::new_arc(Arc::new(inner));
+        let perm = sem.try_acquire_owned().unwrap();
+        let permits = sem.outstanding_permits.load(Ordering::Acquire);
+        assert_eq!(permits, 1);
+        drop(perm);
+        let permits = sem.outstanding_permits.load(Ordering::Acquire);
+        assert_eq!(permits, 0);
+    }
+
+    #[tokio::test]
+    async fn closable_semaphore_permit_drop_after_close_resolves_close_complete() {
+        let inner = MeteredSemaphore::new(2, MetricsContext::no_op(), |_, _| {});
+        let sem = ClosableMeteredSemaphore::new_arc(Arc::new(inner));
+        let perm = sem.try_acquire_owned().unwrap();
+        sem.close();
+        drop(perm);
+        sem.close_complete().await;
+    }
+
+    #[tokio::test]
+    async fn closable_semaphore_close_complete_ready_if_unused() {
+        let inner = MeteredSemaphore::new(2, MetricsContext::no_op(), |_, _| {});
+        let sem = ClosableMeteredSemaphore::new_arc(Arc::new(inner));
+        sem.close();
+        sem.close_complete().await;
+    }
+
+    #[tokio::test]
+    async fn closable_semaphore_does_not_hand_out_permits_after_closed() {
+        let inner = MeteredSemaphore::new(2, MetricsContext::no_op(), |_, _| {});
+        let sem = ClosableMeteredSemaphore::new_arc(Arc::new(inner));
+        sem.close();
+        let perm = sem.try_acquire_owned().unwrap_err();
+        assert_matches!(perm, TryAcquireError::Closed);
+    }
 }

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

@@ -4,7 +4,7 @@ use crate::{
         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,
+        QueueResponse, ResponseType, WorkerExt, WorkflowCachingPolicy, TEST_Q,
     },
     worker::client::mocks::{mock_manual_workflow_client, mock_workflow_client},
     ActivityHeartbeat, Worker, WorkerConfigBuilder,
@@ -13,7 +13,8 @@ use futures::FutureExt;
 use itertools::Itertools;
 use std::{
     cell::RefCell,
-    collections::{hash_map::Entry, HashMap, VecDeque},
+    collections::{hash_map::Entry, HashMap, HashSet, VecDeque},
+    future,
     rc::Rc,
     sync::{
         atomic::{AtomicUsize, Ordering},
@@ -23,14 +24,17 @@ use std::{
 };
 use temporal_client::WorkflowOptions;
 use temporal_sdk::{ActivityOptions, WfContext};
-use temporal_sdk_core_api::{errors::CompleteActivityError, Worker as WorkerTrait};
+use temporal_sdk_core_api::{
+    errors::{CompleteActivityError, PollActivityError},
+    Worker as WorkerTrait,
+};
 use temporal_sdk_core_protos::{
     coresdk::{
         activity_result::{
             activity_execution_result, activity_resolution, ActivityExecutionResult,
             ActivityResolution, Success,
         },
-        activity_task::{activity_task, ActivityTask},
+        activity_task::{activity_task, ActivityCancelReason, ActivityTask, Cancel},
         workflow_activation::{workflow_activation_job, ResolveActivity, WorkflowActivationJob},
         workflow_commands::{
             ActivityCancellationType, CompleteWorkflowExecution, RequestCancelActivity,
@@ -42,6 +46,9 @@ use temporal_sdk_core_protos::{
     temporal::api::{
         command::v1::{command::Attributes, ScheduleActivityTaskCommandAttributes},
         enums::v1::EventType,
+        history::v1::{
+            history_event::Attributes as EventAttributes, ActivityTaskScheduledEventAttributes,
+        },
         workflowservice::v1::{
             PollActivityTaskQueueResponse, RecordActivityTaskHeartbeatResponse,
             RespondActivityTaskCanceledResponse, RespondActivityTaskCompletedResponse,
@@ -52,11 +59,10 @@ use temporal_sdk_core_protos::{
 };
 use temporal_sdk_core_test_utils::{fanout_tasks, start_timer_cmd, TestWorker};
 use tokio::{sync::Barrier, time::sleep};
+use tokio_util::sync::CancellationToken;
 
-#[tokio::test]
-async fn max_activities_respected() {
-    let _task_q = "q";
-    let mut tasks = VecDeque::from(vec![
+fn three_tasks() -> VecDeque<PollActivityTaskQueueResponse> {
+    VecDeque::from(vec![
         PollActivityTaskQueueResponse {
             task_token: vec![1],
             activity_id: "act1".to_string(),
@@ -72,7 +78,13 @@ async fn max_activities_respected() {
             activity_id: "act3".to_string(),
             ..Default::default()
         },
-    ]);
+    ])
+}
+
+#[tokio::test]
+async fn max_activities_respected() {
+    let _task_q = "q";
+    let mut tasks = three_tasks();
     let mut mock_client = mock_workflow_client();
     mock_client
         .expect_poll_activity_task()
@@ -121,7 +133,7 @@ async fn activity_not_found_returns_ok() {
     })
     .await
     .unwrap();
-    core.shutdown().await;
+    core.drain_activity_poller_and_shutdown().await;
 }
 
 #[tokio::test]
@@ -217,12 +229,14 @@ async fn heartbeats_report_cancels_only_once() {
     })
     .await
     .unwrap();
-    core.shutdown().await;
+    core.drain_activity_poller_and_shutdown().await;
 }
 
 #[tokio::test]
 async fn activity_cancel_interrupts_poll() {
     let mut mock_poller = mock_manual_poller();
+    let shutdown_token = CancellationToken::new();
+    let shutdown_token_clone = shutdown_token.clone();
     let mut poll_resps = VecDeque::from(vec![
         async {
             Some(Ok(PollActivityTaskQueueResponse {
@@ -237,10 +251,15 @@ async fn activity_cancel_interrupts_poll() {
             Some(Ok(Default::default()))
         }
         .boxed(),
+        async move {
+            shutdown_token.cancelled().await;
+            None
+        }
+        .boxed(),
     ]);
     mock_poller
         .expect_poll()
-        .times(2)
+        .times(3)
         .returning(move || poll_resps.pop_front().unwrap());
 
     let mut mock_client = mock_manual_workflow_client();
@@ -289,11 +308,12 @@ async fn activity_cancel_interrupts_poll() {
                 }
             ).await.unwrap();
             last_finisher.store(2, Ordering::SeqCst);
+            shutdown_token_clone.cancel();
         }
     };
     // So that we know we blocked
     assert_eq!(last_finisher.load(Ordering::Acquire), 2);
-    core.shutdown().await;
+    core.drain_activity_poller_and_shutdown().await;
 }
 
 #[tokio::test]
@@ -342,13 +362,10 @@ async fn many_concurrent_heartbeat_cancels() {
             })
             .collect::<Vec<_>>(),
     );
-    // Because the mock is so fast, it's possible it can return before the cancel channel in
-    // the activity task poll selector. So, the final poll when there are no more tasks must
-    // take a while.
     poll_resps.push_back(
         async {
-            sleep(Duration::from_secs(10)).await;
-            unreachable!("Long poll")
+            future::pending::<()>().await;
+            unreachable!()
         }
         .boxed(),
     );
@@ -431,7 +448,7 @@ async fn many_concurrent_heartbeat_cancels() {
     })
     .await;
 
-    worker.shutdown().await;
+    worker.drain_activity_poller_and_shutdown().await;
 }
 
 #[tokio::test]
@@ -483,7 +500,7 @@ async fn activity_timeout_no_double_resolve() {
     )
     .await;
 
-    core.shutdown().await;
+    core.drain_pollers_and_shutdown().await;
 }
 
 #[tokio::test]
@@ -529,7 +546,7 @@ async fn can_heartbeat_acts_during_shutdown() {
     })
     .await
     .unwrap();
-    shutdown_fut.await;
+    core.drain_activity_poller_and_shutdown().await;
 }
 
 /// Verifies that if a user has tried to record a heartbeat and then immediately after failed the
@@ -580,7 +597,7 @@ async fn complete_act_with_fail_flushes_heartbeat() {
     })
     .await
     .unwrap();
-    core.shutdown().await;
+    core.drain_activity_poller_and_shutdown().await;
 
     // Verify the last seen call to record a heartbeat had the last detail payload
     let last_seen_payload = &last_seen_payload.take().unwrap().payloads[0];
@@ -686,7 +703,7 @@ async fn no_eager_activities_requested_when_worker_options_disable_remote_activi
     .await
     .unwrap();
 
-    core.shutdown().await;
+    core.drain_pollers_and_shutdown().await;
 
     assert_eq!(num_eager_requested.load(Ordering::Relaxed), 0);
 }
@@ -753,11 +770,8 @@ async fn activity_tasks_from_completion_are_delivered() {
         .times(3)
         .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));
+    let act_tasks: Vec<QueueResponse<PollActivityTaskQueueResponse>> = vec![];
+    mock.set_act_poller(mock_poller_from_resps(act_tasks));
     mock.worker_cfg(|wc| wc.max_cached_workflows = 2);
     let core = mock_worker(mock);
 
@@ -816,7 +830,7 @@ async fn activity_tasks_from_completion_are_delivered() {
         .unwrap();
     }
 
-    core.shutdown().await;
+    core.drain_pollers_and_shutdown().await;
 
     // Verify only a single eager activity was scheduled (the one on our worker's task queue)
     assert_eq!(num_eager_requested.load(Ordering::Relaxed), 3);
@@ -828,11 +842,23 @@ async fn activity_tasks_from_completion_reserve_slots() {
     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 schedid = t.add(EventAttributes::ActivityTaskScheduledEventAttributes(
+        ActivityTaskScheduledEventAttributes {
+            activity_id: "1".to_string(),
+            activity_type: Some("act1".into()),
+            ..Default::default()
+        },
+    ));
     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 schedid = t.add(EventAttributes::ActivityTaskScheduledEventAttributes(
+        ActivityTaskScheduledEventAttributes {
+            activity_id: "2".to_string(),
+            activity_type: Some("act2".into()),
+            ..Default::default()
+        },
+    ));
     let startid = t.add_activity_task_started(schedid);
     t.add_activity_task_completed(schedid, startid, b"hi".into());
     t.add_full_wf_task();
@@ -902,19 +928,25 @@ async fn activity_tasks_from_completion_reserve_slots() {
     // 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())
+    let workflow_complete_token = CancellationToken::new();
+    let workflow_complete_token_clone = workflow_complete_token.clone();
+
+    worker.register_wf(DEFAULT_WORKFLOW_TYPE, move |ctx: WfContext| {
+        let complete_token = workflow_complete_token.clone();
+        async move {
+            ctx.activity(ActivityOptions {
+                activity_type: "act1".to_string(),
+                ..Default::default()
+            })
+            .await;
+            ctx.activity(ActivityOptions {
+                activity_type: "act2".to_string(),
+                ..Default::default()
+            })
+            .await;
+            complete_token.cancel();
+            Ok(().into())
+        }
     });
 
     worker
@@ -941,6 +973,13 @@ async fn activity_tasks_from_completion_reserve_slots() {
         .await
         .unwrap();
         barr.wait().await;
+        // Wait for workflow to complete in order for all eager activities to be requested before shutting down.
+        // After shutdown, no eager activities slots can be allocated.
+        workflow_complete_token_clone.cancelled().await;
+        core.initiate_shutdown();
+        // Even though this test requests eager activity tasks, none are returned in poll responses.
+        let err = core.poll_activity_task().await.unwrap_err();
+        assert_matches!(err, PollActivityError::ShutDown);
     };
     // This wf poll should *not* set the flag that it wants tasks back since both slots are
     // occupied
@@ -974,7 +1013,7 @@ async fn retryable_net_error_exhaustion_is_nonfatal() {
     })
     .await
     .unwrap();
-    core.shutdown().await;
+    core.drain_activity_poller_and_shutdown().await;
 }
 
 #[tokio::test]
@@ -1012,3 +1051,56 @@ async fn cant_complete_activity_with_unset_result_payload() {
         Err(CompleteActivityError::MalformedActivityCompletion { .. })
     )
 }
+
+#[tokio::test]
+async fn graceful_shutdown() {
+    let _task_q = "q";
+    let mut tasks = three_tasks();
+    let mut mock_client = mock_workflow_client();
+    mock_client
+        .expect_poll_activity_task()
+        .times(3)
+        .returning(move |_, _| Ok(tasks.pop_front().unwrap()));
+    // They shall all be reported as failed
+    mock_client
+        .expect_fail_activity_task()
+        .times(3)
+        .returning(|_, _| Ok(Default::default()));
+
+    let worker = Worker::new_test(
+        test_worker_cfg()
+            .graceful_shutdown_period(Duration::from_millis(500))
+            .build()
+            .unwrap(),
+        mock_client,
+    );
+
+    let _1 = worker.poll_activity_task().await.unwrap();
+    let _2 = worker.poll_activity_task().await.unwrap();
+    let _3 = worker.poll_activity_task().await.unwrap();
+
+    worker.initiate_shutdown();
+    let expected_tts = HashSet::from([vec![1], vec![2], vec![3]]);
+    let mut seen_tts = HashSet::new();
+    for _ in 1..=3 {
+        let cancel = worker.poll_activity_task().await.unwrap();
+        assert_matches!(
+            cancel.variant,
+            Some(activity_task::Variant::Cancel(Cancel {
+                reason: r
+            })) if r == ActivityCancelReason::WorkerShutdown as i32
+        );
+        seen_tts.insert(cancel.task_token);
+    }
+    assert_eq!(expected_tts, seen_tts);
+    for tt in seen_tts {
+        worker
+            .complete_activity_task(ActivityTaskCompletion {
+                task_token: tt,
+                result: Some(ActivityExecutionResult::cancel_from_details(None)),
+            })
+            .await
+            .unwrap();
+    }
+    worker.drain_pollers_and_shutdown().await;
+}

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

@@ -88,11 +88,12 @@ async fn parent_cancels_child_wf(ctx: WfContext) -> WorkflowResult<()> {
         .await
         .into_started()
         .expect("Child should get started");
-    let cancel_fut = start_res.cancel(&ctx);
-    let resfut = start_res.result();
-    let (cancel_res, res) = join!(cancel_fut, resfut);
-    cancel_res.expect("cancel result is ok");
-    let stat = res.status.expect("child wf result is ok");
+    start_res.cancel(&ctx);
+    let stat = start_res
+        .result()
+        .await
+        .status
+        .expect("child wf result is ok");
     assert_matches!(stat, child_workflow_result::Status::Cancelled(_));
     Ok(().into())
 }