@temporalio/core-bridge 1.9.2 → 1.10.0

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

@@ -1,6 +1,6 @@
 //! This module contains very generic helpers that can be used codebase-wide
 
-pub mod take_cell;
+pub(crate) mod take_cell;
 
 use crate::MetricsContext;
 use derive_more::DebugCustom;
@@ -11,102 +11,149 @@ use std::{
         Arc,
     },
 };
-use tokio::sync::{AcquireError, OwnedSemaphorePermit, Semaphore, TryAcquireError};
+use temporal_sdk_core_api::worker::{
+    SlotKind, SlotReservationContext, SlotSupplier, SlotSupplierPermit,
+};
+use tokio::sync::watch;
 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
+/// Wraps a [SlotSupplier] and turns successful slot reservations into permit structs, as well
+/// as handling associated metrics tracking.
 #[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.
+pub(crate) struct MeteredPermitDealer<SK: SlotKind> {
+    supplier: Arc<dyn SlotSupplier<SlotKind = SK> + Send + Sync>,
+    /// The number of permit owners who have acquired a permit, 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>,
+    /// The number of permits that have been handed out
+    extant_permits: (watch::Sender<usize>, watch::Receiver<usize>),
+    /// The maximum number of extant permits which are allowed. Once the number of extant permits
+    /// is at this number, no more permits will be requested from the supplier until one is freed.
+    /// This avoids requesting slots when we are at the workflow cache size limit. If and when
+    /// we add user-defined cache sizing, that logic will need to live with the supplier and
+    /// there will need to be some associated refactoring.
+    max_permits: Option<usize>,
     metrics_ctx: MetricsContext,
-    record_fn: fn(&MetricsContext, usize),
 }
 
-impl MeteredSemaphore {
-    pub fn new(
-        inital_permits: usize,
+impl<SK> MeteredPermitDealer<SK>
+where
+    SK: SlotKind + 'static,
+{
+    pub(crate) fn new(
+        supplier: Arc<dyn SlotSupplier<SlotKind = SK> + Send + Sync>,
         metrics_ctx: MetricsContext,
-        record_fn: fn(&MetricsContext, usize),
+        max_permits: Option<usize>,
     ) -> Self {
         Self {
-            sem: Arc::new(Semaphore::new(inital_permits)),
+            supplier,
             unused_claimants: Arc::new(AtomicUsize::new(0)),
+            extant_permits: watch::channel(0),
             metrics_ctx,
-            record_fn,
+            max_permits,
         }
     }
 
-    pub fn available_permits(&self) -> usize {
-        self.sem.available_permits()
+    pub(crate) fn available_permits(&self) -> Option<usize> {
+        self.supplier.available_slots()
     }
 
     #[cfg(test)]
-    pub fn unused_permits(&self) -> usize {
-        self.sem.available_permits() + self.unused_claimants.load(Ordering::Acquire)
+    pub(crate) fn unused_permits(&self) -> Option<usize> {
+        self.available_permits()
+            .map(|ap| ap + self.unused_claimants.load(Ordering::Acquire))
     }
 
-    pub async fn acquire_owned(&self) -> Result<OwnedMeteredSemPermit, AcquireError> {
-        let res = self.sem.clone().acquire_owned().await?;
+    pub(crate) async fn acquire_owned(&self) -> Result<OwnedMeteredSemPermit<SK>, ()> {
+        if let Some(max) = self.max_permits {
+            self.extant_permits
+                .1
+                .clone()
+                .wait_for(|&ep| ep < max)
+                .await
+                .expect("Extant permit channel is never closed");
+        }
+        let res = self.supplier.reserve_slot(self).await;
         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))
+    pub(crate) fn try_acquire_owned(&self) -> Result<OwnedMeteredSemPermit<SK>, ()> {
+        if let Some(max) = self.max_permits {
+            if *self.extant_permits.1.borrow() >= max {
+                return Err(());
+            }
+        }
+        if let Some(res) = self.supplier.try_reserve_slot(self) {
+            Ok(self.build_owned(res))
+        } else {
+            Err(())
+        }
     }
 
-    fn build_owned(&self, res: OwnedSemaphorePermit) -> OwnedMeteredSemPermit {
+    fn build_owned(&self, res: SlotSupplierPermit) -> OwnedMeteredSemPermit<SK> {
         self.unused_claimants.fetch_add(1, Ordering::Release);
-        self.record();
+        self.extant_permits.0.send_modify(|ep| *ep += 1);
+        // Eww
+        let uc_c = self.unused_claimants.clone();
+        let ep_rx_c = self.extant_permits.1.clone();
+        let ep_tx_c = self.extant_permits.0.clone();
+        let supp = self.supplier.clone();
+        let supp_c = self.supplier.clone();
+        let supp_c_c = self.supplier.clone();
+        let mets = self.metrics_ctx.clone();
+        let metric_rec =
+            // When being called from the drop impl, the permit isn't actually dropped yet, so
+            // account for that with the `add_one` parameter.
+            move |add_one: bool| {
+                let extra = usize::from(add_one);
+                let unused = uc_c.load(Ordering::Acquire);
+                if let Some(avail) = supp.available_slots() {
+                    mets.available_task_slots(avail + unused + extra);
+                }
+                mets.task_slots_used((ep_rx_c.borrow().saturating_sub(unused) + extra) as u64);
+            };
+        let mrc = metric_rec.clone();
+        mrc(false);
+
         OwnedMeteredSemPermit {
-            inner: res,
+            _inner: res,
             unused_claimants: Some(self.unused_claimants.clone()),
-            record_fn: self.record_owned(),
+            use_fn: Box::new(move |info| {
+                supp_c.mark_slot_used(info);
+                metric_rec(false)
+            }),
+            release_fn: Box::new(move || {
+                supp_c_c.release_slot();
+                ep_tx_c.send_modify(|ep| *ep -= 1);
+                mrc(true)
+            }),
         }
     }
+}
 
-    fn record(&self) {
-        (self.record_fn)(
-            &self.metrics_ctx,
-            self.sem.available_permits() + self.unused_claimants.load(Ordering::Acquire),
-        );
-    }
-
-    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();
-        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,
-            )
-        })
+impl<SK: SlotKind> SlotReservationContext for MeteredPermitDealer<SK> {
+    fn num_issued_slots(&self) -> usize {
+        *self.extant_permits.1.borrow()
     }
 }
 
-/// A version of [MeteredSemaphore] that can be closed and supports waiting for close to complete.
+/// A version of [MeteredPermitDealer] 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>,
+pub(crate) struct ClosableMeteredPermitDealer<SK: SlotKind> {
+    inner: Arc<MeteredPermitDealer<SK>>,
     outstanding_permits: AtomicUsize,
     close_requested: AtomicBool,
     close_complete_token: CancellationToken,
 }
 
-impl ClosableMeteredSemaphore {
-    pub fn new_arc(sem: Arc<MeteredSemaphore>) -> Arc<Self> {
+impl<SK> ClosableMeteredPermitDealer<SK>
+where
+    SK: SlotKind,
+{
+    pub(crate) fn new_arc(sem: Arc<MeteredPermitDealer<SK>>) -> Arc<Self> {
         Arc::new(Self {
             inner: sem,
             outstanding_permits: Default::default(),
@@ -116,14 +163,17 @@ impl ClosableMeteredSemaphore {
     }
 }
 
-impl ClosableMeteredSemaphore {
+impl<SK> ClosableMeteredPermitDealer<SK>
+where
+    SK: SlotKind + 'static,
+{
     #[cfg(test)]
-    pub fn unused_permits(&self) -> usize {
+    pub(crate) fn unused_permits(&self) -> Option<usize> {
         self.inner.unused_permits()
     }
 
     /// Request to close the semaphore and prevent new permits from being acquired.
-    pub fn close(&self) {
+    pub(crate) fn close(&self) {
         self.close_requested.store(true, Ordering::Release);
         if self.outstanding_permits.load(Ordering::Acquire) == 0 {
             self.close_complete_token.cancel();
@@ -131,16 +181,16 @@ impl ClosableMeteredSemaphore {
     }
 
     /// Returns after close has been requested and all outstanding permits have been returned.
-    pub async fn close_complete(&self) {
+    pub(crate) 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(
+    pub(crate) fn try_acquire_owned(
         self: &Arc<Self>,
-    ) -> Result<TrackedOwnedMeteredSemPermit, TryAcquireError> {
+    ) -> Result<TrackedOwnedMeteredSemPermit<SK>, ()> {
         if self.close_requested.load(Ordering::Acquire) {
-            return Err(TryAcquireError::Closed);
+            return Err(());
         }
         self.outstanding_permits.fetch_add(1, Ordering::Release);
         let res = self.inner.try_acquire_owned();
@@ -168,72 +218,64 @@ impl ClosableMeteredSemaphore {
 
 /// Tracks an OwnedMeteredSemPermit and calls on_drop when dropped.
 #[derive(DebugCustom)]
+#[clippy::has_significant_drop]
 #[debug(fmt = "Tracked({inner:?})")]
-pub(crate) struct TrackedOwnedMeteredSemPermit {
-    inner: Option<OwnedMeteredSemPermit>,
+pub(crate) struct TrackedOwnedMeteredSemPermit<SK: SlotKind> {
+    inner: Option<OwnedMeteredSemPermit<SK>>,
     on_drop: Box<dyn Fn() + Send + Sync>,
 }
-impl From<TrackedOwnedMeteredSemPermit> for OwnedMeteredSemPermit {
-    fn from(mut value: TrackedOwnedMeteredSemPermit) -> Self {
+impl<SK: SlotKind> From<TrackedOwnedMeteredSemPermit<SK>> for OwnedMeteredSemPermit<SK> {
+    fn from(mut value: TrackedOwnedMeteredSemPermit<SK>) -> Self {
         value
             .inner
             .take()
             .expect("Inner permit should be available")
     }
 }
-impl Drop for TrackedOwnedMeteredSemPermit {
+impl<SK: SlotKind> Drop for TrackedOwnedMeteredSemPermit<SK> {
     fn drop(&mut self) {
         (self.on_drop)();
     }
 }
 
-/// Wraps an [OwnedSemaphorePermit] to update metrics when it's dropped
-pub(crate) struct OwnedMeteredSemPermit {
-    inner: OwnedSemaphorePermit,
-    /// See [MeteredSemaphore::unused_claimants]. If present when dropping, used to decrement the
+/// Wraps an [SlotSupplierPermit] to update metrics & when it's dropped
+#[clippy::has_significant_drop]
+pub(crate) struct OwnedMeteredSemPermit<SK: SlotKind> {
+    _inner: SlotSupplierPermit,
+    /// See [MeteredPermitDealer::unused_claimants]. If present when dropping, used to decrement the
     /// count.
     unused_claimants: Option<Arc<AtomicUsize>>,
-    record_fn: Box<dyn Fn(bool) + Send + Sync>,
+    #[allow(clippy::type_complexity)] // not really tho, bud
+    use_fn: Box<dyn Fn(SK::Info<'_>) + Send + Sync>,
+    release_fn: Box<dyn Fn() + Send + Sync>,
 }
-impl Drop for OwnedMeteredSemPermit {
+impl<SK: SlotKind> Drop for OwnedMeteredSemPermit<SK> {
     fn drop(&mut self) {
         if let Some(uc) = self.unused_claimants.take() {
             uc.fetch_sub(1, Ordering::Release);
         }
-        (self.record_fn)(true)
+        (self.release_fn)()
     }
 }
-impl Debug for OwnedMeteredSemPermit {
+impl<SK: SlotKind> Debug for OwnedMeteredSemPermit<SK> {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        self.inner.fmt(f)
+        f.write_str("OwnedMeteredSemPermit()")
     }
 }
-impl OwnedMeteredSemPermit {
+impl<SK: SlotKind> OwnedMeteredSemPermit<SK> {
     /// 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 {
+    pub(crate) fn into_used(mut self, info: SK::Info<'_>) -> UsedMeteredSemPermit<SK> {
         if let Some(uc) = self.unused_claimants.take() {
             uc.fetch_sub(1, Ordering::Release);
-            (self.record_fn)(false)
         }
+        (self.use_fn)(info);
         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(OwnedMeteredSemPermit {
-            inner,
-            unused_claimants: None,
-            record_fn: Box::new(|_| {}),
-        })
-    }
-}
+pub(crate) struct UsedMeteredSemPermit<SK: SlotKind>(#[allow(dead_code)] OwnedMeteredSemPermit<SK>);
 
 macro_rules! dbg_panic {
   ($($arg:tt)*) => {
@@ -244,13 +286,26 @@ macro_rules! dbg_panic {
 pub(crate) use dbg_panic;
 
 #[cfg(test)]
-mod tests {
+pub(crate) mod tests {
     use super::*;
+    use crate::{advance_fut, worker::tuner::FixedSizeSlotSupplier};
+    use futures_util::FutureExt;
+    use temporal_sdk_core_api::worker::WorkflowSlotKind;
 
-    #[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));
+    pub(crate) fn fixed_size_permit_dealer<SK: SlotKind + Send + Sync + 'static>(
+        size: usize,
+    ) -> MeteredPermitDealer<SK> {
+        MeteredPermitDealer::new(
+            Arc::new(FixedSizeSlotSupplier::new(size)),
+            MetricsContext::no_op(),
+            None,
+        )
+    }
+
+    #[test]
+    fn closable_semaphore_permit_drop_returns_permit() {
+        let inner = fixed_size_permit_dealer::<WorkflowSlotKind>(2);
+        let sem = ClosableMeteredPermitDealer::new_arc(Arc::new(inner));
         let perm = sem.try_acquire_owned().unwrap();
         let permits = sem.outstanding_permits.load(Ordering::Acquire);
         assert_eq!(permits, 1);
@@ -261,8 +316,8 @@ mod tests {
 
     #[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 inner = fixed_size_permit_dealer::<WorkflowSlotKind>(2);
+        let sem = ClosableMeteredPermitDealer::new_arc(Arc::new(inner));
         let perm = sem.try_acquire_owned().unwrap();
         sem.close();
         drop(perm);
@@ -271,18 +326,31 @@ mod tests {
 
     #[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));
+        let inner = fixed_size_permit_dealer::<WorkflowSlotKind>(2);
+        let sem = ClosableMeteredPermitDealer::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));
+    #[test]
+    fn closable_semaphore_does_not_hand_out_permits_after_closed() {
+        let inner = fixed_size_permit_dealer::<WorkflowSlotKind>(2);
+        let sem = ClosableMeteredPermitDealer::new_arc(Arc::new(inner));
         sem.close();
-        let perm = sem.try_acquire_owned().unwrap_err();
-        assert_matches!(perm, TryAcquireError::Closed);
+        sem.try_acquire_owned().unwrap_err();
+    }
+
+    #[tokio::test]
+    async fn respects_max_extant_permits() {
+        let mut sem = fixed_size_permit_dealer::<WorkflowSlotKind>(2);
+        sem.max_permits = Some(1);
+        let perm = sem.try_acquire_owned().unwrap();
+        sem.try_acquire_owned().unwrap_err();
+        let acquire_fut = sem.acquire_owned();
+        // Will be pending
+        advance_fut!(acquire_fut);
+        drop(perm);
+        // Now it'll proceed
+        acquire_fut.await.unwrap();
     }
 }

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

@@ -7,7 +7,7 @@ use crate::{
         QueueResponse, ResponseType, WorkerExt, WorkflowCachingPolicy, TEST_Q,
     },
     worker::client::mocks::{mock_manual_workflow_client, mock_workflow_client},
-    ActivityHeartbeat, Worker, WorkerConfigBuilder,
+    ActivityHeartbeat, Worker,
 };
 use futures::FutureExt;
 use itertools::Itertools;
@@ -618,11 +618,8 @@ async fn max_tq_acts_set_passed_to_poll_properly() {
             })
         });
 
-    let cfg = WorkerConfigBuilder::default()
-        .namespace("enchi")
-        .task_queue("cat")
+    let cfg = test_worker_cfg()
         .max_concurrent_at_polls(1_usize)
-        .worker_build_id("test_bin_id")
         .max_task_queue_activities_per_second(rate)
         .build()
         .unwrap();
@@ -921,7 +918,7 @@ async fn activity_tasks_from_completion_reserve_slots() {
     let mut mock = build_mock_pollers(mh);
     mock.worker_cfg(|cfg| {
         cfg.max_cached_workflows = 2;
-        cfg.max_outstanding_activities = 2;
+        cfg.max_outstanding_activities = Some(2);
     });
     mock.set_act_poller(mock_poller_from_resps(act_tasks));
     let core = Arc::new(mock_worker(mock));
@@ -1030,13 +1027,7 @@ async fn cant_complete_activity_with_unset_result_payload() {
             })
         });
 
-    let cfg = WorkerConfigBuilder::default()
-        .namespace("enchi")
-        .task_queue("cat")
-        .worker_build_id("enchi_loves_salmon")
-        .build()
-        .unwrap();
-    let worker = Worker::new_test(cfg, mock_client);
+    let worker = Worker::new_test(test_worker_cfg().build().unwrap(), mock_client);
     let t = worker.poll_activity_task().await.unwrap();
     let res = worker
         .complete_activity_task(ActivityTaskCompletion {

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

@@ -21,7 +21,8 @@ use temporal_sdk_core_protos::{
 };
 
 static DID_FAIL: AtomicBool = AtomicBool::new(false);
-pub async fn timer_wf_fails_once(ctx: WfContext) -> WorkflowResult<()> {
+
+pub(crate) async fn timer_wf_fails_once(ctx: WfContext) -> WorkflowResult<()> {
     ctx.timer(Duration::from_secs(1)).await;
     if DID_FAIL
         .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)

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

@@ -112,7 +112,7 @@ async fn local_act_two_wfts_before_marker(#[case] replay: bool, #[case] cached:
     worker.run_until_done().await.unwrap();
 }
 
-pub async fn local_act_fanout_wf(ctx: WfContext) -> WorkflowResult<()> {
+pub(crate) async fn local_act_fanout_wf(ctx: WfContext) -> WorkflowResult<()> {
     let las: Vec<_> = (1..=50)
         .map(|i| {
             ctx.local_activity(LocalActivityOptions {
@@ -187,7 +187,7 @@ async fn local_act_heartbeat(#[case] shutdown_middle: bool) {
     mh.enforce_correct_number_of_polls = false;
     let mut worker = mock_sdk_cfg(mh, |wc| {
         wc.max_cached_workflows = 1;
-        wc.max_outstanding_workflow_tasks = 1;
+        wc.max_outstanding_workflow_tasks = Some(1);
     });
     let core = worker.core_worker.clone();
 
@@ -1084,7 +1084,7 @@ async fn local_act_records_nonfirst_attempts_ok() {
     }));
     let mut worker = mock_sdk_cfg(mh, |wc| {
         wc.max_cached_workflows = 1;
-        wc.max_outstanding_workflow_tasks = 1;
+        wc.max_outstanding_workflow_tasks = Some(1);
     });
 
     worker.register_wf(

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

@@ -16,6 +16,7 @@ use crate::{
     Worker,
 };
 use futures::FutureExt;
+use once_cell::sync::Lazy;
 use std::time::Duration;
 use temporal_sdk_core_api::Worker as WorkerTrait;
 use temporal_sdk_core_protos::coresdk::workflow_completion::WorkflowActivationCompletion;
@@ -45,9 +46,8 @@ async fn after_shutdown_server_is_not_polled() {
 }
 
 // Better than cloning a billion arcs...
-lazy_static::lazy_static! {
-    static ref BARR: Barrier = Barrier::new(3);
-}
+static BARR: Lazy<Barrier> = Lazy::new(|| Barrier::new(3));
+
 #[tokio::test]
 async fn shutdown_interrupts_both_polls() {
     let mut mock_client = mock_manual_workflow_client();

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

@@ -893,8 +893,9 @@ async fn build_id_set_properly_on_query_on_first_task() {
     core.drain_pollers_and_shutdown().await;
 }
 
+#[rstest::rstest]
 #[tokio::test]
-async fn queries_arent_lost_in_buffer_void() {
+async fn queries_arent_lost_in_buffer_void(#[values(false, true)] buffered_because_cache: bool) {
     let wfid = "fake_wf_id";
     let mut t = TestHistoryBuilder::default();
     t.add_by_type(EventType::WorkflowExecutionStarted);
@@ -902,7 +903,24 @@ async fn queries_arent_lost_in_buffer_void() {
     t.add_we_signaled("sig", vec![]);
     t.add_full_wf_task();
     t.add_workflow_execution_completed();
-    let tasks = [
+
+    let mut tasks = if buffered_because_cache {
+        // A history for another wf which will occupy the only cache slot initially
+        let mut t1 = TestHistoryBuilder::default();
+        t1.add_by_type(EventType::WorkflowExecutionStarted);
+        t1.add_full_wf_task();
+        t1.add_workflow_execution_completed();
+
+        vec![hist_to_poll_resp(
+            &t1,
+            "cache_occupier".to_owned(),
+            1.into(),
+        )]
+    } else {
+        vec![]
+    };
+
+    tasks.extend([
         hist_to_poll_resp(&t, wfid.to_owned(), 1.into()),
         {
             let mut pr = hist_to_poll_resp(&t, wfid.to_owned(), 1.into());
@@ -923,8 +941,7 @@ async fn queries_arent_lost_in_buffer_void() {
             pr
         },
         hist_to_poll_resp(&t, wfid.to_owned(), 2.into()),
-    ]
-    .map(|r| r.resp);
+    ]);
 
     let mut mock = mock_workflow_client();
     mock.expect_complete_workflow_task()
@@ -932,10 +949,30 @@ async fn queries_arent_lost_in_buffer_void() {
     mock.expect_respond_legacy_query()
         .times(2)
         .returning(|_, _| Ok(Default::default()));
-    let mut mock = MocksHolder::from_wft_stream(mock, stream::iter(tasks));
+    let mut mock =
+        MocksHolder::from_wft_stream(mock, stream::iter(tasks.into_iter().map(|r| r.resp)));
     mock.worker_cfg(|wc| wc.max_cached_workflows = 1);
     let core = mock_worker(mock);
 
+    if buffered_because_cache {
+        // Poll for and complete the occupying workflow
+        let task = core.poll_workflow_activation().await.unwrap();
+        core.complete_execution(&task.run_id).await;
+        // Get the cache removal task
+        let task = core.poll_workflow_activation().await.unwrap();
+        assert_matches!(
+            task.jobs.as_slice(),
+            [WorkflowActivationJob {
+                variant: Some(workflow_activation_job::Variant::RemoveFromCache(_)),
+            }]
+        );
+        // Wait a beat to ensure the other task(s) have a chance to be buffered
+        tokio::time::sleep(Duration::from_millis(100)).await;
+        core.complete_workflow_activation(WorkflowActivationCompletion::empty(task.run_id))
+            .await
+            .unwrap();
+    }
+
     let task = core.poll_workflow_activation().await.unwrap();
     core.complete_workflow_activation(WorkflowActivationCompletion::empty(task.run_id))
         .await

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

@@ -175,7 +175,7 @@ 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
@@ -183,7 +183,6 @@ async fn complete_with_task_not_found_during_shutdown() {
             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]