@temporalio/core-bridge 1.13.1 → 1.14.0

package/sdk-core/crates/sdk-core/tests/integ_tests/worker_tests.rs

@@ -0,0 +1,920 @@
+use crate::{
+    common::{
+        CoreWfStarter, fake_grpc_server::fake_server, get_integ_runtime_options,
+        get_integ_server_options, get_integ_telem_options, mock_sdk_cfg,
+    },
+    shared_tests,
+};
+use assert_matches::assert_matches;
+use futures_util::FutureExt;
+use std::{
+    cell::Cell,
+    sync::{
+        Arc, Mutex,
+        atomic::{
+            AtomicBool, AtomicU8,
+            Ordering::{self, Relaxed},
+        },
+    },
+    time::Duration,
+};
+use temporalio_client::WorkflowOptions;
+use temporalio_common::{
+    Worker,
+    errors::WorkerValidationError,
+    protos::{
+        DEFAULT_WORKFLOW_TYPE, TestHistoryBuilder, canned_histories,
+        coresdk::{
+            ActivityTaskCompletion,
+            activity_result::ActivityExecutionResult,
+            workflow_completion::{
+                Failure, WorkflowActivationCompletion, workflow_activation_completion::Status,
+            },
+        },
+        temporal::api::{
+            command::v1::command::Attributes,
+            common::v1::WorkerVersionStamp,
+            enums::v1::{
+                EventType, WorkflowTaskFailedCause, WorkflowTaskFailedCause::GrpcMessageTooLarge,
+            },
+            failure::v1::Failure as InnerFailure,
+            history::v1::{
+                ActivityTaskScheduledEventAttributes, history_event,
+                history_event::Attributes::{
+                    self as EventAttributes, WorkflowTaskFailedEventAttributes,
+                },
+            },
+            workflowservice::v1::{
+                GetWorkflowExecutionHistoryResponse, PollActivityTaskQueueResponse,
+                RespondActivityTaskCompletedResponse,
+            },
+        },
+    },
+    worker::{
+        ActivitySlotKind, LocalActivitySlotKind, PollerBehavior, SlotInfo, SlotInfoTrait,
+        SlotMarkUsedContext, SlotReleaseContext, SlotReservationContext, SlotSupplier,
+        SlotSupplierPermit, WorkerConfigBuilder, WorkerTaskTypes, WorkerVersioningStrategy,
+        WorkflowSlotKind,
+    },
+};
+use temporalio_sdk::{
+    ActivityOptions, LocalActivityOptions, WfContext, interceptors::WorkerInterceptor,
+};
+use temporalio_sdk_core::{
+    CoreRuntime, ResourceBasedTuner, ResourceSlotOptions, TunerBuilder, init_worker,
+    test_help::{
+        FakeWfResponses, MockPollCfg, ResponseType, build_mock_pollers, drain_pollers_and_shutdown,
+        hist_to_poll_resp, mock_worker, mock_worker_client,
+    },
+};
+use tokio::sync::{Barrier, Notify, Semaphore};
+use tokio_util::sync::CancellationToken;
+use uuid::Uuid;
+
+#[tokio::test]
+async fn worker_validation_fails_on_nonexistent_namespace() {
+    let opts = get_integ_server_options();
+    let runtime =
+        CoreRuntime::new_assume_tokio(get_integ_runtime_options(get_integ_telem_options()))
+            .unwrap();
+    let retrying_client = opts
+        .connect_no_namespace(runtime.telemetry().get_temporal_metric_meter())
+        .await
+        .unwrap();
+
+    let worker = init_worker(
+        &runtime,
+        WorkerConfigBuilder::default()
+            .namespace("i_dont_exist")
+            .task_queue("Wheee!")
+            .versioning_strategy(WorkerVersioningStrategy::None {
+                build_id: "blah".to_owned(),
+            })
+            .task_types(WorkerTaskTypes::all())
+            .build()
+            .unwrap(),
+        retrying_client,
+    )
+    .unwrap();
+
+    let res = worker.validate().await;
+    assert_matches!(
+        res,
+        Err(WorkerValidationError::NamespaceDescribeError { .. })
+    );
+}
+
+#[tokio::test]
+async fn worker_handles_unknown_workflow_types_gracefully() {
+    let wf_type = "worker_handles_unknown_workflow_types_gracefully";
+    let mut starter = CoreWfStarter::new(wf_type);
+    let mut worker = starter.worker().await;
+
+    let run_id = worker
+        .submit_wf(
+            format!("wce-{}", Uuid::new_v4()),
+            "unregistered".to_string(),
+            vec![],
+            WorkflowOptions::default(),
+        )
+        .await
+        .unwrap();
+
+    struct GracefulAsserter {
+        notify: Arc<Notify>,
+        run_id: String,
+        unregistered_failure_seen: Cell<bool>,
+    }
+    #[async_trait::async_trait(?Send)]
+    impl WorkerInterceptor for GracefulAsserter {
+        async fn on_workflow_activation_completion(
+            &self,
+            completion: &WorkflowActivationCompletion,
+        ) {
+            if matches!(
+                completion,
+                WorkflowActivationCompletion {
+                    status: Some(Status::Failed(Failure {
+                        failure: Some(InnerFailure { message, .. }),
+                        ..
+                    })),
+                    run_id,
+                } if message == "Workflow type unregistered not found" && *run_id == self.run_id
+            ) {
+                self.unregistered_failure_seen.set(true);
+            }
+            // If we've seen the failure, and the completion is a success for the same run, we're done
+            if matches!(
+                completion,
+                WorkflowActivationCompletion {
+                    status: Some(Status::Successful(..)),
+                    run_id,
+                } if self.unregistered_failure_seen.get() && *run_id == self.run_id
+            ) {
+                // Shutdown the worker
+                self.notify.notify_one();
+            }
+        }
+        fn on_shutdown(&self, _: &temporalio_sdk::Worker) {}
+    }
+
+    let inner = worker.inner_mut();
+    let notify = Arc::new(Notify::new());
+    inner.set_worker_interceptor(GracefulAsserter {
+        notify: notify.clone(),
+        run_id,
+        unregistered_failure_seen: Cell::new(false),
+    });
+    tokio::join!(async { inner.run().await.unwrap() }, async move {
+        notify.notified().await;
+        let worker = starter.get_worker().await.clone();
+        drain_pollers_and_shutdown(&worker).await;
+    });
+}
+
+#[tokio::test]
+async fn resource_based_few_pollers_guarantees_non_sticky_poll() {
+    let wf_name = "resource_based_few_pollers_guarantees_non_sticky_poll";
+    let mut starter = CoreWfStarter::new(wf_name);
+    starter
+        .worker_config
+        .clear_max_outstanding_opts()
+        .task_types(WorkerTaskTypes::workflow_only())
+        // 3 pollers so the minimum slots of 2 can both be handed out to a sticky poller
+        .workflow_task_poller_behavior(PollerBehavior::SimpleMaximum(3_usize));
+    // Set the limits to zero so it's essentially unwilling to hand out slots
+    let mut tuner = ResourceBasedTuner::new(0.0, 0.0);
+    tuner.with_workflow_slots_options(ResourceSlotOptions::new(2, 10, Duration::from_millis(0)));
+    starter.worker_config.tuner(Arc::new(tuner));
+    let mut worker = starter.worker().await;
+
+    // Workflow doesn't actually need to do anything. We just need to see that we don't get stuck
+    // by assigning all slots to sticky pollers.
+    worker.register_wf(
+        wf_name.to_owned(),
+        |_: WfContext| async move { Ok(().into()) },
+    );
+    for i in 0..20 {
+        worker
+            .submit_wf(
+                format!("{wf_name}_{i}"),
+                wf_name.to_owned(),
+                vec![],
+                WorkflowOptions::default(),
+            )
+            .await
+            .unwrap();
+    }
+    worker.run_until_done().await.unwrap();
+}
+
+#[tokio::test]
+async fn oversize_grpc_message() {
+    use crate::common::{NAMESPACE, prom_metrics};
+    let wf_name = "oversize_grpc_message";
+    // Enable Prometheus metrics for this test and capture the address
+    let (telemopts, addr, _aborter) = prom_metrics(None);
+    let runtime = CoreRuntime::new_assume_tokio(get_integ_runtime_options(telemopts)).unwrap();
+    let mut starter = CoreWfStarter::new_with_runtime(wf_name, runtime);
+    starter
+        .worker_config
+        .task_types(WorkerTaskTypes::workflow_only());
+    let mut core = starter.worker().await;
+
+    static OVERSIZE_GRPC_MESSAGE_RUN: AtomicBool = AtomicBool::new(false);
+    core.register_wf(wf_name.to_owned(), |_ctx: WfContext| async move {
+        if OVERSIZE_GRPC_MESSAGE_RUN.load(Relaxed) {
+            Ok(vec![].into())
+        } else {
+            OVERSIZE_GRPC_MESSAGE_RUN.store(true, Relaxed);
+            let result: Vec<u8> = vec![0; 5000000];
+            Ok(result.into())
+        }
+    });
+    starter.start_with_worker(wf_name, &mut core).await;
+    core.run_until_done().await.unwrap();
+
+    assert!(starter.get_history().await.events.iter().any(|e| {
+        e.event_type == EventType::WorkflowTaskFailed as i32
+            && if let WorkflowTaskFailedEventAttributes(attr) = e.attributes.as_ref().unwrap() {
+                attr.cause == GrpcMessageTooLarge as i32
+                    && attr.failure.as_ref().unwrap().message == "GRPC Message too large"
+            } else {
+                false
+            }
+    }));
+
+    // Verify the workflow task failure metric includes the GrpcMessageTooLarge reason
+    let tq = starter.get_task_queue();
+    crate::common::eventually(
+        || async {
+            let body = crate::integ_tests::metrics_tests::get_text(format!("http://{addr}/metrics")).await;
+            if body.contains(&format!(
+                "temporal_workflow_task_execution_failed{{failure_reason=\"GrpcMessageTooLarge\",namespace=\"{NAMESPACE}\",service_name=\"temporal-core-sdk\",task_queue=\"{tq}\"}} 1"
+            )) {
+                Ok(())
+            } else {
+                Err(())
+            }
+        },
+        Duration::from_secs(2),
+    )
+    .await
+    .unwrap();
+}
+
+#[tokio::test]
+async fn grpc_message_too_large_test() {
+    shared_tests::grpc_message_too_large().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(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(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();
+    t.add_workflow_execution_completed();
+
+    let mut mock = mock_worker_client();
+    // Set up two tasks to be returned via normal activity polling
+    let act_tasks = 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_mock_fn = 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
+            .first()
+            .and_then(|cmd| cmd.attributes.as_ref())
+        {
+            if attrs.activity_id == "1" {
+                assert!(!attrs.request_eager_execution);
+            } else {
+                assert!(attrs.request_eager_execution);
+            }
+        }
+        Ok(Default::default())
+    }));
+    mh.activity_responses = Some(act_tasks);
+    let mut mock = build_mock_pollers(mh);
+    mock.worker_cfg(|cfg| {
+        cfg.max_cached_workflows = 2;
+        cfg.max_outstanding_activities = Some(2);
+    });
+    let core = Arc::new(mock_worker(mock));
+    let mut worker = crate::common::TestWorker::new(core.clone());
+
+    // First poll for activities twice, occupying both slots
+    let at1 = core.poll_activity_task().await.unwrap();
+    let at2 = core.poll_activity_task().await.unwrap();
+    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
+        .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;
+        // 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, temporalio_common::errors::PollError::ShutDown);
+    };
+    // 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);
+}
+
+#[tokio::test]
+async fn max_wft_respected() {
+    let total_wfs = 100;
+    let wf_ids: Vec<_> = (0..total_wfs).map(|i| format!("fake-wf-{i}")).collect();
+    let hists = wf_ids.iter().map(|wf_id| {
+        let hist = canned_histories::single_timer("1");
+        FakeWfResponses {
+            wf_id: wf_id.to_string(),
+            hist,
+            response_batches: vec![1.into(), 2.into()],
+        }
+    });
+    let mh = MockPollCfg::new(hists.into_iter().collect(), true, 0);
+    let mut worker = mock_sdk_cfg(mh, |cfg| {
+        cfg.max_cached_workflows = total_wfs as usize;
+        cfg.max_outstanding_workflow_tasks = Some(1);
+    });
+    let active_count: &'static _ = Box::leak(Box::new(Semaphore::new(1)));
+    worker.register_wf(DEFAULT_WORKFLOW_TYPE, move |ctx: WfContext| async move {
+        drop(
+            active_count
+                .try_acquire()
+                .expect("No multiple concurrent workflow tasks!"),
+        );
+        ctx.timer(Duration::from_secs(1)).await;
+        Ok(().into())
+    });
+
+    for wf_id in wf_ids {
+        worker
+            .submit_wf(wf_id, DEFAULT_WORKFLOW_TYPE, vec![], Default::default())
+            .await
+            .unwrap();
+    }
+    worker.run_until_done().await.unwrap();
+}
+
+#[rstest]
+#[tokio::test]
+async fn history_length_with_fail_and_timeout(
+    #[values(true, false)] use_cache: bool,
+    #[values(1, 2, 3)] history_responses_case: u8,
+) {
+    if !use_cache && history_responses_case == 3 {
+        eprintln!(
+            "Skipping history_length_with_fail_and_timeout::use_cache_2_false::history_responses_case_3_3 due to flaky hang"
+        );
+        return;
+    }
+    let wfid = "fake_wf_id";
+    let mut t = TestHistoryBuilder::default();
+    t.add_by_type(EventType::WorkflowExecutionStarted);
+    t.add_full_wf_task();
+    let timer_started_event_id = t.add_by_type(EventType::TimerStarted);
+    t.add_timer_fired(timer_started_event_id, "1".to_string());
+    t.add_workflow_task_scheduled_and_started();
+    t.add_workflow_task_failed_with_failure(WorkflowTaskFailedCause::Unspecified, "ahh".into());
+    t.add_workflow_task_scheduled_and_started();
+    t.add_workflow_task_timed_out();
+    t.add_full_wf_task();
+    let timer_started_event_id = t.add_by_type(EventType::TimerStarted);
+    t.add_timer_fired(timer_started_event_id, "2".to_string());
+    t.add_full_wf_task();
+    t.add_workflow_execution_completed();
+
+    let mut mock_client = mock_worker_client();
+    let history_responses = match history_responses_case {
+        1 => vec![ResponseType::AllHistory],
+        2 => vec![
+            ResponseType::ToTaskNum(1),
+            ResponseType::ToTaskNum(2),
+            ResponseType::AllHistory,
+        ],
+        3 => {
+            let mut needs_fetch = hist_to_poll_resp(&t, wfid, ResponseType::ToTaskNum(2)).resp;
+            needs_fetch.next_page_token = vec![1];
+            // Truncate the history a bit in order to force incomplete WFT
+            needs_fetch.history.as_mut().unwrap().events.truncate(6);
+            let needs_fetch_resp = ResponseType::Raw(needs_fetch);
+            let mut empty_fetch_resp: GetWorkflowExecutionHistoryResponse =
+                t.get_history_info(1).unwrap().into();
+            empty_fetch_resp.history.as_mut().unwrap().events = vec![];
+            mock_client
+                .expect_get_workflow_execution_history()
+                .returning(move |_, _, _| Ok(empty_fetch_resp.clone()))
+                .times(1);
+            vec![
+                ResponseType::ToTaskNum(1),
+                needs_fetch_resp,
+                ResponseType::ToTaskNum(2),
+                ResponseType::AllHistory,
+            ]
+        }
+        _ => unreachable!(),
+    };
+
+    let mut mh = MockPollCfg::from_resp_batches(wfid, t, history_responses, mock_client);
+    if history_responses_case == 3 {
+        // Expect the failed pagination fetch
+        mh.num_expected_fails = 1;
+    }
+    let mut worker = mock_sdk_cfg(mh, |wc| {
+        if use_cache {
+            wc.max_cached_workflows = 1;
+        }
+    });
+    worker.register_wf(DEFAULT_WORKFLOW_TYPE, |ctx: WfContext| async move {
+        assert_eq!(ctx.history_length(), 3);
+        ctx.timer(Duration::from_secs(1)).await;
+        assert_eq!(ctx.history_length(), 14);
+        ctx.timer(Duration::from_secs(1)).await;
+        assert_eq!(ctx.history_length(), 19);
+        Ok(().into())
+    });
+    worker
+        .submit_wf(
+            wfid.to_owned(),
+            DEFAULT_WORKFLOW_TYPE.to_owned(),
+            vec![],
+            WorkflowOptions::default(),
+        )
+        .await
+        .unwrap();
+    worker.run_until_done().await.unwrap();
+}
+
+#[allow(deprecated)]
+#[tokio::test]
+async fn sets_build_id_from_wft_complete() {
+    let wfid = "fake_wf_id";
+
+    let mut t = TestHistoryBuilder::default();
+    t.add_by_type(EventType::WorkflowExecutionStarted);
+    t.add_full_wf_task();
+    let timer_started_event_id = t.add_by_type(EventType::TimerStarted);
+    t.add_timer_fired(timer_started_event_id, "1".to_string());
+    t.add_full_wf_task();
+    t.modify_event(t.current_event_id(), |he| {
+        if let history_event::Attributes::WorkflowTaskCompletedEventAttributes(a) =
+            he.attributes.as_mut().unwrap()
+        {
+            a.worker_version = Some(WorkerVersionStamp {
+                build_id: "enchi-cat".to_string(),
+                ..Default::default()
+            });
+        }
+    });
+    let timer_started_event_id = t.add_by_type(EventType::TimerStarted);
+    t.add_timer_fired(timer_started_event_id, "2".to_string());
+    t.add_workflow_task_scheduled_and_started();
+
+    let mock = mock_worker_client();
+    let mut worker = mock_sdk_cfg(
+        MockPollCfg::from_resp_batches(wfid, t, [ResponseType::AllHistory], mock),
+        |cfg| {
+            cfg.versioning_strategy = WorkerVersioningStrategy::None {
+                build_id: "fierce-predator".to_string(),
+            };
+            cfg.max_cached_workflows = 1;
+        },
+    );
+
+    worker.register_wf(DEFAULT_WORKFLOW_TYPE, |ctx: WfContext| async move {
+        // First task, it should be empty, since replaying and nothing in first WFT completed
+        assert_eq!(ctx.current_deployment_version(), None);
+        ctx.timer(Duration::from_secs(1)).await;
+        assert_eq!(
+            ctx.current_deployment_version().unwrap().build_id,
+            "enchi-cat"
+        );
+        ctx.timer(Duration::from_secs(1)).await;
+        // Not replaying at this point, so we should see the worker's build id
+        assert_eq!(
+            ctx.current_deployment_version().unwrap().build_id,
+            "fierce-predator"
+        );
+        ctx.timer(Duration::from_secs(1)).await;
+        assert_eq!(
+            ctx.current_deployment_version().unwrap().build_id,
+            "fierce-predator"
+        );
+        Ok(().into())
+    });
+    worker
+        .submit_wf(wfid, DEFAULT_WORKFLOW_TYPE, vec![], Default::default())
+        .await
+        .unwrap();
+    worker.run_until_done().await.unwrap();
+}
+
+#[derive(Debug, Clone)]
+enum SlotEvent {
+    ReserveSlot {
+        slot_type: &'static str,
+    },
+    TryReserveSlot {
+        slot_type: &'static str,
+    },
+    MarkSlotUsed {
+        slot_type: &'static str,
+        is_sticky: bool,
+        workflow_type: Option<String>,
+        activity_type: Option<String>,
+    },
+    ReleaseSlot {
+        slot_type: &'static str,
+    },
+}
+
+struct TrackingSlotSupplier<SK> {
+    events: Arc<Mutex<Vec<SlotEvent>>>,
+    slot_type: &'static str,
+    _phantom: std::marker::PhantomData<SK>,
+}
+
+impl<SK> TrackingSlotSupplier<SK> {
+    fn new(slot_type: &'static str) -> Self {
+        Self {
+            events: Arc::new(Mutex::new(Vec::new())),
+            slot_type,
+            _phantom: std::marker::PhantomData,
+        }
+    }
+
+    fn get_events(&self) -> Vec<SlotEvent> {
+        self.events.lock().unwrap().clone()
+    }
+
+    fn add_event(&self, event: SlotEvent) {
+        self.events.lock().unwrap().push(event);
+    }
+
+    fn extract_slot_info(info: &dyn SlotInfoTrait) -> (bool, Option<String>, Option<String>) {
+        match info.downcast() {
+            SlotInfo::Workflow(w) => (w.is_sticky, Some(w.workflow_type.clone()), None),
+            SlotInfo::Activity(a) => (false, None, Some(a.activity_type.clone())),
+            SlotInfo::LocalActivity(a) => (false, None, Some(a.activity_type.clone())),
+            SlotInfo::Nexus(_) => (false, None, None),
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl<SK> SlotSupplier for TrackingSlotSupplier<SK>
+where
+    SK: temporalio_common::worker::SlotKind + Send + Sync,
+    SK::Info: SlotInfoTrait,
+{
+    type SlotKind = SK;
+
+    async fn reserve_slot(&self, _ctx: &dyn SlotReservationContext) -> SlotSupplierPermit {
+        self.add_event(SlotEvent::ReserveSlot {
+            slot_type: self.slot_type,
+        });
+        SlotSupplierPermit::with_user_data(())
+    }
+
+    fn try_reserve_slot(&self, _ctx: &dyn SlotReservationContext) -> Option<SlotSupplierPermit> {
+        self.add_event(SlotEvent::TryReserveSlot {
+            slot_type: self.slot_type,
+        });
+        Some(SlotSupplierPermit::with_user_data(()))
+    }
+
+    fn mark_slot_used(&self, ctx: &dyn SlotMarkUsedContext<SlotKind = Self::SlotKind>) {
+        let (is_sticky, workflow_type, activity_type) = Self::extract_slot_info(ctx.info());
+        self.add_event(SlotEvent::MarkSlotUsed {
+            slot_type: self.slot_type,
+            is_sticky,
+            workflow_type,
+            activity_type,
+        });
+    }
+
+    fn release_slot(&self, _ctx: &dyn SlotReleaseContext<SlotKind = Self::SlotKind>) {
+        self.add_event(SlotEvent::ReleaseSlot {
+            slot_type: self.slot_type,
+        });
+    }
+}
+
+#[tokio::test]
+async fn test_custom_slot_supplier_simple() {
+    let wf_supplier = Arc::new(TrackingSlotSupplier::<WorkflowSlotKind>::new("workflow"));
+    let activity_supplier = Arc::new(TrackingSlotSupplier::<ActivitySlotKind>::new("activity"));
+    let local_activity_supplier = Arc::new(TrackingSlotSupplier::<LocalActivitySlotKind>::new(
+        "local_activity",
+    ));
+
+    let mut starter = CoreWfStarter::new("test_custom_slot_supplier_simple");
+    starter.worker_config.clear_max_outstanding_opts();
+
+    let mut tb = TunerBuilder::default();
+    tb.workflow_slot_supplier(wf_supplier.clone());
+    tb.activity_slot_supplier(activity_supplier.clone());
+    tb.local_activity_slot_supplier(local_activity_supplier.clone());
+    starter.worker_config.tuner(Arc::new(tb.build()));
+
+    let mut worker = starter.worker().await;
+
+    worker.register_activity(
+        "SlotSupplierActivity",
+        |_: temporalio_sdk::ActContext, _: ()| async move { Ok(()) },
+    );
+    worker.register_wf(
+        "SlotSupplierWorkflow".to_owned(),
+        |ctx: WfContext| async move {
+            let _result = ctx
+                .activity(ActivityOptions {
+                    activity_type: "SlotSupplierActivity".to_string(),
+                    start_to_close_timeout: Some(Duration::from_secs(10)),
+                    ..Default::default()
+                })
+                .await;
+            let _result = ctx
+                .local_activity(LocalActivityOptions {
+                    activity_type: "SlotSupplierActivity".to_string(),
+                    start_to_close_timeout: Some(Duration::from_secs(10)),
+                    ..Default::default()
+                })
+                .await;
+            Ok(().into())
+        },
+    );
+
+    worker
+        .submit_wf(
+            "test-wf".to_owned(),
+            "SlotSupplierWorkflow".to_owned(),
+            vec![],
+            Default::default(),
+        )
+        .await
+        .unwrap();
+
+    worker.run_until_done().await.unwrap();
+
+    // Collect all events
+    let wf_events = wf_supplier.get_events();
+    let activity_events = activity_supplier.get_events();
+    let local_activity_events = local_activity_supplier.get_events();
+
+    // Verify workflow slot events - should have reserve, mark used, and release events
+    assert!(wf_events.iter().any(
+        |e| matches!(e, SlotEvent::ReserveSlot { slot_type, .. } if *slot_type == "workflow")
+    ));
+    assert!(wf_events.iter().any(
+        |e| matches!(e, SlotEvent::MarkSlotUsed { slot_type, .. } if *slot_type == "workflow")
+    ));
+    assert!(
+        wf_events
+            .iter()
+            .any(|e| matches!(e, SlotEvent::ReleaseSlot { slot_type } if *slot_type == "workflow"))
+    );
+
+    // Verify activity slot events - should have reserve, try_reserve (for eager execution), mark
+    // used, and release
+    assert!(activity_events.iter().any(
+        |e| matches!(e, SlotEvent::ReserveSlot { slot_type, .. } if *slot_type == "activity")
+    ));
+    assert!(
+        activity_events.iter().any(
+            |e| matches!(e, SlotEvent::TryReserveSlot { slot_type } if *slot_type == "activity")
+        )
+    );
+    assert!(activity_events.iter().any(
+        |e| matches!(e, SlotEvent::MarkSlotUsed { slot_type, .. } if *slot_type == "activity")
+    ));
+    assert!(
+        activity_events
+            .iter()
+            .any(|e| matches!(e, SlotEvent::ReleaseSlot { slot_type } if *slot_type == "activity"))
+    );
+
+    // Verify local activity slot events
+    assert!(local_activity_events.iter().any(
+        |e| matches!(e, SlotEvent::ReserveSlot { slot_type, .. } if *slot_type == "local_activity")
+    ));
+    assert!(local_activity_events.iter().any(
+        |e| matches!(e, SlotEvent::MarkSlotUsed { slot_type, .. } if *slot_type == "local_activity")
+    ));
+    assert!(local_activity_events.iter().any(
+        |e| matches!(e, SlotEvent::ReleaseSlot { slot_type } if *slot_type == "local_activity")
+    ));
+
+    assert!(
+        wf_events
+            .iter()
+            .any(|e| matches!(e, SlotEvent::MarkSlotUsed {
+                                    slot_type: "workflow",
+                                    workflow_type: Some(wf_type),
+                                    ..
+                                } if wf_type == "SlotSupplierWorkflow"))
+    );
+    assert!(
+        activity_events
+            .iter()
+            .any(|e| matches!(e, SlotEvent::MarkSlotUsed {
+                                    slot_type: "activity",
+                                    activity_type: Some(act_type),
+                                    ..
+                                } if act_type == "SlotSupplierActivity"))
+    );
+    assert!(
+        local_activity_events
+            .iter()
+            .any(|e| matches!(e, SlotEvent::MarkSlotUsed {
+                                    slot_type: "local_activity",
+                                    activity_type: Some(act_type),
+                                    ..
+                                } if act_type == "SlotSupplierActivity"))
+    );
+    assert!(wf_events.iter().any(|e| matches!(
+        e,
+        SlotEvent::MarkSlotUsed {
+            slot_type: "workflow",
+            is_sticky: false,
+            ..
+        }
+    )));
+
+    // Verify that the number of reserve/try_reserve events matches the number of release events
+    let total_reserves = wf_events
+        .iter()
+        .filter(|e| {
+            matches!(
+                e,
+                SlotEvent::ReserveSlot { .. } | SlotEvent::TryReserveSlot { .. }
+            )
+        })
+        .count()
+        + activity_events
+            .iter()
+            .filter(|e| {
+                matches!(
+                    e,
+                    SlotEvent::ReserveSlot { .. } | SlotEvent::TryReserveSlot { .. }
+                )
+            })
+            .count()
+        + local_activity_events
+            .iter()
+            .filter(|e| {
+                matches!(
+                    e,
+                    SlotEvent::ReserveSlot { .. } | SlotEvent::TryReserveSlot { .. }
+                )
+            })
+            .count();
+
+    let total_releases = wf_events
+        .iter()
+        .filter(|e| matches!(e, SlotEvent::ReleaseSlot { .. }))
+        .count()
+        + activity_events
+            .iter()
+            .filter(|e| matches!(e, SlotEvent::ReleaseSlot { .. }))
+            .count()
+        + local_activity_events
+            .iter()
+            .filter(|e| matches!(e, SlotEvent::ReleaseSlot { .. }))
+            .count();
+
+    assert_eq!(
+        total_reserves, total_releases,
+        "Number of reserves should equal number of releases"
+    );
+}
+
+#[tokio::test]
+async fn shutdown_worker_not_retried() {
+    let shutdown_call_count = Arc::new(AtomicU8::new(0));
+    let scc = shutdown_call_count.clone();
+    let fs = fake_server(move |req| {
+        if req.uri().to_string().contains("ShutdownWorker") {
+            scc.fetch_add(1, Ordering::Relaxed);
+        }
+        let s = tonic::Status::new(tonic::Code::Unknown, "bla").into_http();
+        async { s }.boxed()
+    })
+    .await;
+
+    let mut opts = get_integ_server_options();
+    let uri = format!("http://localhost:{}", fs.addr.port())
+        .parse()
+        .unwrap();
+    opts.target_url = uri;
+    opts.skip_get_system_info = true;
+    let client = opts.connect("ns", None).await.unwrap();
+
+    let wf_type = "shutdown_worker_not_retried";
+    let mut starter = CoreWfStarter::new_with_overrides(wf_type, None, Some(client));
+    let worker = starter.get_worker().await;
+    drain_pollers_and_shutdown(&worker).await;
+    assert_eq!(shutdown_call_count.load(Ordering::Relaxed), 1);
+}