@temporalio/core-bridge 1.9.2 → 1.10.0

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

@@ -10,7 +10,10 @@ use crate::{
         test_worker_cfg, FakeWfResponses, MockPollCfg, MocksHolder, ResponseType, WorkerExt,
         WorkflowCachingPolicy::{self, AfterEveryReply, NonSticky},
     },
-    worker::client::mocks::{mock_manual_workflow_client, mock_workflow_client},
+    worker::{
+        client::mocks::{mock_manual_workflow_client, mock_workflow_client},
+        TunerBuilder,
+    },
     Worker,
 };
 use futures::{stream, FutureExt};
@@ -26,7 +29,13 @@ use std::{
 };
 use temporal_client::WorkflowOptions;
 use temporal_sdk::{ActivityOptions, CancellableFuture, WfContext};
-use temporal_sdk_core_api::{errors::PollWfError, Worker as WorkerTrait};
+use temporal_sdk_core_api::{
+    errors::PollWfError,
+    worker::{
+        SlotKind, SlotReservationContext, SlotSupplier, SlotSupplierPermit, WorkflowSlotKind,
+    },
+    Worker as WorkerTrait,
+};
 use temporal_sdk_core_protos::{
     coresdk::{
         activity_result::{self as ar, activity_resolution, ActivityResolution},
@@ -911,7 +920,7 @@ async fn max_wft_respected() {
     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 = 1;
+        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 {
@@ -1506,7 +1515,7 @@ async fn failing_wft_doesnt_eat_permit_forever() {
     let mut mock = build_mock_pollers(mock);
     mock.worker_cfg(|cfg| {
         cfg.max_cached_workflows = 2;
-        cfg.max_outstanding_workflow_tasks = 2;
+        cfg.max_outstanding_workflow_tasks = Some(2);
     });
     let outstanding_mock_tasks = mock.outstanding_task_map.clone();
     let worker = mock_worker(mock);
@@ -1531,7 +1540,7 @@ async fn failing_wft_doesnt_eat_permit_forever() {
                 variant: Some(workflow_activation_job::Variant::RemoveFromCache(_)),
             },]
         );
-        run_id = activation.run_id.clone();
+        run_id.clone_from(&activation.run_id);
         worker
             .complete_workflow_activation(WorkflowActivationCompletion::empty(activation.run_id))
             .await
@@ -1546,7 +1555,7 @@ async fn failing_wft_doesnt_eat_permit_forever() {
     outstanding_mock_tasks.unwrap().release_run(&run_id);
     let activation = worker.poll_workflow_activation().await.unwrap();
     // There should be no change in permits, since this just unbuffered the buffered task
-    assert_eq!(worker.available_wft_permits(), 1);
+    assert_eq!(worker.available_wft_permits(), Some(1));
     worker
         .complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
             activation.run_id,
@@ -1555,7 +1564,7 @@ async fn failing_wft_doesnt_eat_permit_forever() {
         .await
         .unwrap();
     worker.shutdown().await;
-    assert_eq!(worker.available_wft_permits(), 2);
+    assert_eq!(worker.available_wft_permits(), Some(2));
 }
 
 #[tokio::test]
@@ -1582,7 +1591,7 @@ async fn cache_miss_will_fetch_history() {
     mock.worker_cfg(|cfg| {
         cfg.max_cached_workflows = 1;
         // Also verifies tying the WFT permit to the fetch request doesn't get us stuck
-        cfg.max_outstanding_workflow_tasks = 1;
+        cfg.max_outstanding_workflow_tasks = Some(1);
     });
     let worker = mock_worker(mock);
 
@@ -1808,7 +1817,7 @@ async fn poll_faster_than_complete_wont_overflow_cache() {
     let mut mock = build_mock_pollers(mock_cfg);
     mock.worker_cfg(|wc| {
         wc.max_cached_workflows = 3;
-        wc.max_outstanding_workflow_tasks = 3;
+        wc.max_outstanding_workflow_tasks = Some(3);
     });
     let core = mock_worker(mock);
     // Poll 4 times, completing once, such that max tasks are never exceeded
@@ -1915,7 +1924,7 @@ async fn poll_faster_than_complete_wont_overflow_cache() {
 
     join!(blocking_poll, complete_evict);
     // p5 outstanding and final poll outstanding -- hence one permit available
-    assert_eq!(core.available_wft_permits(), 1);
+    assert_eq!(core.available_wft_permits(), Some(1));
     assert_eq!(core.cached_workflows().await, 3);
 }
 
@@ -2663,8 +2672,8 @@ async fn poller_wont_run_ahead_of_task_slots() {
     }
 
     assert_eq!(worker.outstanding_workflow_tasks().await, 10);
-    assert_eq!(worker.available_wft_permits(), 0);
-    assert_eq!(worker.unused_wft_permits(), 0);
+    assert_eq!(worker.available_wft_permits(), Some(0));
+    assert_eq!(worker.unused_wft_permits(), Some(0));
 
     // This one should hang until we complete some tasks since we're at the limit
     let hung_poll = async {
@@ -2752,13 +2761,13 @@ async fn use_compatible_version_flag(
             let can_cmd = c.commands.pop().unwrap().attributes.unwrap();
             match can_cmd {
                 Attributes::ContinueAsNewWorkflowExecutionCommandAttributes(a) => {
-                    assert_eq!(a.use_compatible_version, compat_flag_expected);
+                    assert_eq!(a.inherit_build_id, compat_flag_expected);
                 }
                 Attributes::ScheduleActivityTaskCommandAttributes(a) => {
-                    assert_eq!(a.use_compatible_version, compat_flag_expected);
+                    assert_eq!(a.use_workflow_build_id, compat_flag_expected);
                 }
                 Attributes::StartChildWorkflowExecutionCommandAttributes(a) => {
-                    assert_eq!(a.use_compatible_version, compat_flag_expected);
+                    assert_eq!(a.inherit_build_id, compat_flag_expected);
                 }
                 _ => panic!("invalid attributes type"),
             }
@@ -2855,3 +2864,94 @@ async fn sets_build_id_from_wft_complete() {
         .unwrap();
     worker.run_until_done().await.unwrap();
 }
+
+#[tokio::test]
+async fn slot_provider_cant_hand_out_more_permits_than_cache_size() {
+    let popped_tasks = Arc::new(AtomicUsize::new(0));
+    let ptc = popped_tasks.clone();
+    let mut bunch_of_first_tasks = (1..50).map(move |i| {
+        ptc.fetch_add(1, Ordering::Relaxed);
+        hist_to_poll_resp(
+            &canned_histories::single_timer(&format!("{i}")),
+            format!("wf-{i}"),
+            1.into(),
+        )
+        .resp
+    });
+    let mut mock_client = mock_workflow_client();
+    mock_client
+        .expect_poll_workflow_task()
+        .returning(move |_| Ok(bunch_of_first_tasks.next().unwrap()));
+    mock_client
+        .expect_complete_workflow_task()
+        .returning(|_| Ok(Default::default()));
+
+    struct EndlessSupplier {}
+    #[async_trait::async_trait]
+    impl SlotSupplier for EndlessSupplier {
+        type SlotKind = WorkflowSlotKind;
+        async fn reserve_slot(&self, _: &dyn SlotReservationContext) -> SlotSupplierPermit {
+            SlotSupplierPermit::default()
+        }
+        fn try_reserve_slot(&self, _: &dyn SlotReservationContext) -> Option<SlotSupplierPermit> {
+            Some(SlotSupplierPermit::default())
+        }
+        fn mark_slot_used(&self, _: <Self::SlotKind as SlotKind>::Info<'_>) {}
+        fn release_slot(&self) {}
+        fn available_slots(&self) -> Option<usize> {
+            None
+        }
+    }
+
+    let worker = Worker::new_test(
+        test_worker_cfg()
+            .max_cached_workflows(10_usize)
+            .tuner(
+                TunerBuilder::default()
+                    .workflow_slot_supplier(Arc::new(EndlessSupplier {}))
+                    .build(),
+            )
+            .max_concurrent_wft_polls(10_usize)
+            .no_remote_activities(true)
+            .build()
+            .unwrap(),
+        mock_client,
+    );
+
+    // Should be able to get at 10 tasks
+    let mut tasks = vec![];
+    for _ in 0..10 {
+        tasks.push(worker.poll_workflow_activation().await.unwrap());
+    }
+    // 11th should hang
+
+    assert_eq!(worker.outstanding_workflow_tasks().await, 10);
+    // assert_eq!(worker.available_wft_permits(), Some(0));
+    // assert_eq!(worker.unused_wft_permits(), Some(0));
+
+    // This one should hang until we complete some tasks since we're at the limit
+    let hung_poll = async {
+        // This should end up getting shut down after the other routine finishes tasks
+        assert_matches!(
+            worker.poll_workflow_activation().await.unwrap_err(),
+            PollWfError::ShutDown
+        );
+    };
+    // Wait for a bit concurrently with above, verify no extra tasks got taken, shutdown
+    let ender = async {
+        time::sleep(Duration::from_millis(300)).await;
+        // initiate shutdown, then complete open tasks
+        worker.initiate_shutdown();
+        for t in tasks {
+            worker
+                .complete_workflow_activation(WorkflowActivationCompletion::empty(t.run_id))
+                .await
+                .unwrap();
+        }
+        worker.shutdown().await;
+    };
+    join!(hung_poll, ender);
+    // We shouldn't have got more than the 10 tasks from the poller -- verifying that the concurrent
+    // polling is not exceeding the task limit
+    assert_eq!(popped_tasks.load(Ordering::Relaxed), 10);
+}

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

@@ -7,6 +7,7 @@ use futures::StreamExt;
 use serde::Deserialize;
 use std::{
     fs::OpenOptions,
+    io,
     path::{Path, PathBuf},
 };
 use temporal_client::ClientOptionsBuilder;
@@ -22,99 +23,6 @@ use zip::read::read_zipfile_from_stream;
 use std::os::unix::fs::OpenOptionsExt;
 use std::process::Stdio;
 
-/// Configuration for Temporalite.
-/// Will be removed eventually as its successor, Temporal CLI matures.
-/// We don't care for the duplication between this struct and [TemporalDevServerConfig] and prefer that over another
-/// abstraction since the existence of this struct is temporary.
-#[derive(Debug, Clone, derive_builder::Builder)]
-pub struct TemporaliteConfig {
-    /// Required path to executable or download info.
-    pub exe: EphemeralExe,
-    /// Namespace to use.
-    #[builder(default = "\"default\".to_owned()")]
-    pub namespace: String,
-    /// IP to bind to.
-    #[builder(default = "\"127.0.0.1\".to_owned()")]
-    pub ip: String,
-    /// Port to use or obtains a free one if none given.
-    #[builder(default)]
-    pub port: Option<u16>,
-    /// Sqlite DB filename if persisting or non-persistent if none.
-    #[builder(default)]
-    pub db_filename: Option<String>,
-    /// Whether to enable the UI.
-    #[builder(default)]
-    pub ui: bool,
-    /// Log format and level
-    #[builder(default = "(\"pretty\".to_owned(), \"warn\".to_owned())")]
-    pub log: (String, String),
-    /// Additional arguments to Temporalite.
-    #[builder(default)]
-    pub extra_args: Vec<String>,
-}
-
-impl TemporaliteConfig {
-    /// Start a Temporalite server.
-    pub async fn start_server(&self) -> anyhow::Result<EphemeralServer> {
-        self.start_server_with_output(Stdio::inherit(), Stdio::inherit())
-            .await
-    }
-
-    /// Start a Temporalite server with configurable stdout destination.
-    pub async fn start_server_with_output(
-        &self,
-        output: Stdio,
-        err_output: Stdio,
-    ) -> anyhow::Result<EphemeralServer> {
-        // Get exe path
-        let exe_path = self
-            .exe
-            .get_or_download("temporalite", "temporalite", None)
-            .await?;
-
-        // Get free port if not already given
-        let port = self.port.unwrap_or_else(|| get_free_port(&self.ip));
-
-        // Build arg set
-        let mut args = vec![
-            "start".to_owned(),
-            "--port".to_owned(),
-            port.to_string(),
-            "--namespace".to_owned(),
-            self.namespace.clone(),
-            "--ip".to_owned(),
-            self.ip.clone(),
-            "--log-format".to_owned(),
-            self.log.0.clone(),
-            "--log-level".to_owned(),
-            self.log.1.clone(),
-            "--dynamic-config-value".to_owned(),
-            "frontend.enableServerVersionCheck=false".to_owned(),
-        ];
-        if let Some(db_filename) = &self.db_filename {
-            args.push("--filename".to_owned());
-            args.push(db_filename.clone());
-        } else {
-            args.push("--ephemeral".to_owned());
-        }
-        if !self.ui {
-            args.push("--headless".to_owned());
-        }
-        args.extend(self.extra_args.clone());
-
-        // Start
-        EphemeralServer::start(EphemeralServerConfig {
-            exe_path,
-            port,
-            args,
-            has_test_service: false,
-            output,
-            err_output,
-        })
-        .await
-    }
-}
-
 /// Configuration for Temporal CLI dev server.
 #[derive(Debug, Clone, derive_builder::Builder)]
 pub struct TemporalDevServerConfig {
@@ -163,7 +71,10 @@ impl TemporalDevServerConfig {
             .await?;
 
         // Get free port if not already given
-        let port = self.port.unwrap_or_else(|| get_free_port(&self.ip));
+        let port = match self.port {
+            Some(p) => p,
+            None => get_free_port(&self.ip)?,
+        };
 
         // Build arg set
         let mut args = vec![
@@ -241,7 +152,10 @@ impl TestServerConfig {
             .await?;
 
         // Get free port if not already given
-        let port = self.port.unwrap_or_else(|| get_free_port("0.0.0.0"));
+        let port = match self.port {
+            Some(p) => p,
+            None => get_free_port("0.0.0.0")?,
+        };
 
         // Build arg set
         let mut args = vec![port.to_string()];
@@ -307,11 +221,7 @@ impl EphemeralServer {
             .build()?;
         for _ in 0..50 {
             sleep(Duration::from_millis(100)).await;
-            if client_options
-                .connect_no_namespace(None, None)
-                .await
-                .is_ok()
-            {
+            if client_options.connect_no_namespace(None).await.is_ok() {
                 return success;
             }
         }
@@ -322,7 +232,6 @@ impl EphemeralServer {
     /// a kill if the child process appears completed, but such a check is not
     /// atomic so a kill could still fail as completed if completed just before
     /// kill.
-    #[cfg(not(target_family = "unix"))]
     pub async fn shutdown(&mut self) -> anyhow::Result<()> {
         // Only kill if there is a PID
         if self.child.id().is_some() {
@@ -332,32 +241,10 @@ impl EphemeralServer {
         }
     }
 
-    /// Shutdown the server (i.e. kill the child process). This does not attempt
-    /// a kill if the child process appears completed, but such a check is not
-    /// atomic so a kill could still fail as completed if completed just before
-    /// kill.
-    #[cfg(target_family = "unix")]
-    pub async fn shutdown(&mut self) -> anyhow::Result<()> {
-        // For whatever reason, Tokio is not properly waiting on result
-        // after sending kill in some cases which is causing defunct zombie
-        // processes to remain and kill() to hang. Therefore, we are sending
-        // SIGKILL and waiting on the process ourselves using a low-level call.
-        //
-        // WARNING: This is based on empirical evidence starting a Python test
-        // run on Linux with Python 3.7 (does not happen on Python 3.10 nor does
-        // it happen on Temporalite nor does it happen in Rust integration
-        // tests). Don't alter without running that scenario. EX: SIGINT works but not SIGKILL
-        if let Some(pid) = self.child.id() {
-            let nix_pid = nix::unistd::Pid::from_raw(pid as i32);
-            Ok(spawn_blocking(move || {
-                nix::sys::signal::kill(nix_pid, nix::sys::signal::Signal::SIGINT)?;
-                nix::sys::wait::waitpid(Some(nix_pid), None)
-            })
-            .await?
-            .map(|_| ())?)
-        } else {
-            Ok(())
-        }
+    /// Get the process ID of the child. This will be None if the process is
+    /// considered to be complete.
+    pub fn child_process_id(&self) -> Option<u32> {
+        self.child.id()
     }
 }
 
@@ -492,14 +379,53 @@ impl EphemeralExe {
     }
 }
 
-fn get_free_port(bind_ip: &str) -> u16 {
-    // Can just ask OS to give us a port then close socket. OS's don't give that
-    // port back to anyone else anytime soon.
-    std::net::TcpListener::bind(format!("{bind_ip}:0"))
-        .unwrap()
-        .local_addr()
-        .unwrap()
-        .port()
+/// Returns a TCP port that is available to listen on for the given local host.
+///
+/// This works by binding a new TCP socket on port 0, which requests the OS to
+/// allocate a free port. There is no strict guarantee that the port will remain
+/// available after this function returns, but it should be safe to assume that
+/// a given port will not be allocated again to any process on this machine
+/// within a few seconds.
+///
+/// On Unix-based systems, binding to the port returned by this function
+/// requires setting the `SO_REUSEADDR` socket option (Rust already does that by
+/// default, but other languages may not); otherwise, the OS may fail with a
+/// message such as "address already in use". Windows default behavior is
+/// already appropriate in this regard; on that platform, `SO_REUSEADDR` has a
+/// different meaning and should not be set (setting it may have unpredictable
+/// consequences).
+fn get_free_port(bind_ip: &str) -> io::Result<u16> {
+    let listen = std::net::TcpListener::bind((bind_ip, 0))?;
+    let addr = listen.local_addr()?;
+
+    // On Linux and some BSD variants, ephemeral ports are randomized, and may
+    // consequently repeat within a short time frame after the listenning end
+    // has been closed. To avoid this, we make a connection to the port, then
+    // close that connection from the server's side (this is very important),
+    // which puts the connection in TIME_WAIT state for some time (by default,
+    // 60s on Linux). While it remains in that state, the OS will not reallocate
+    // that port number for bind(:0) syscalls, yet we are not prevented from
+    // explicitly binding to it (thanks to SO_REUSEADDR).
+    //
+    // On macOS and Windows, the above technique is not necessary, as the OS
+    // allocates ephemeral ports sequentially, meaning a port number will only
+    // be reused after the entire range has been exhausted. Quite the opposite,
+    // given that these OSes use a significantly smaller range for ephemeral
+    // ports, making an extra connection just to reserve a port might actually
+    // be harmful (by hastening ephemeral port exhaustion).
+    #[cfg(not(any(target_os = "windows", target_os = "macos")))]
+    {
+        // Establish a connection to the bind_ip:port
+        let _stream = std::net::TcpStream::connect(addr)?;
+
+        // Accept the connection from the listening side
+        let (socket, _addr) = listen.accept()?;
+
+        // Explicitly drop the socket to close the connection from the listening side first
+        std::mem::drop(socket);
+    }
+
+    Ok(addr.port())
 }
 
 /// Returns false if we successfully waited for another download to complete, or
@@ -625,7 +551,7 @@ async fn download_and_extract(
                 // that requires Seek.
                 if let Some(mut file) = read_zipfile_from_stream(&mut reader)? {
                     // If this is the file we're expecting, extract it
-                    if file.enclosed_name() == Some(&file_to_extract) {
+                    if file.enclosed_name().as_ref() == Some(&file_to_extract) {
                         std::io::copy(&mut file, &mut dest)?;
                         return Ok(());
                     }
@@ -637,3 +563,50 @@ async fn download_and_extract(
     })
     .await?
 }
+
+#[cfg(test)]
+mod tests {
+    use super::get_free_port;
+    use std::collections::HashSet;
+    use std::net::{TcpListener, TcpStream};
+
+    #[test]
+    fn get_free_port_no_double() {
+        let host = "127.0.0.1";
+        let mut port_set = HashSet::new();
+
+        for _ in 0..2000 {
+            let port = get_free_port(host).unwrap();
+            assert!(
+                !port_set.contains(&port),
+                "Port {port} has been assigned more than once"
+            );
+
+            // Add port to the set
+            port_set.insert(port);
+        }
+    }
+
+    #[test]
+    fn get_free_port_can_bind_immediately() {
+        let host = "127.0.0.1";
+
+        for _ in 0..500 {
+            let port = get_free_port(host).unwrap();
+            try_listen_and_dial_on(host, port).expect("Failed to bind to port");
+        }
+    }
+
+    fn try_listen_and_dial_on(host: &str, port: u16) -> std::io::Result<()> {
+        let listener = TcpListener::bind((host, port))?;
+        let _stream = TcpStream::connect((host, port))?;
+
+        // Accept the connection from the listening side
+        let (socket, _addr) = listener.accept()?;
+
+        // Explicitly drop the socket to close the connection from the listening side first
+        std::mem::drop(socket);
+
+        Ok(())
+    }
+}

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

@@ -43,7 +43,7 @@ pub(crate) enum InternalFlags {
 }
 
 impl InternalFlags {
-    pub fn new(server_capabilities: &get_system_info_response::Capabilities) -> Self {
+    pub(crate) fn new(server_capabilities: &get_system_info_response::Capabilities) -> Self {
         match server_capabilities.sdk_metadata {
             true => Self::Enabled {
                 core: Default::default(),
@@ -55,7 +55,7 @@ impl InternalFlags {
         }
     }
 
-    pub fn add_from_complete(&mut self, e: &WorkflowTaskCompletedEventAttributes) {
+    pub(crate) fn add_from_complete(&mut self, e: &WorkflowTaskCompletedEventAttributes) {
         if let Self::Enabled { core, lang, .. } = self {
             if let Some(metadata) = e.sdk_metadata.as_ref() {
                 core.extend(
@@ -69,7 +69,7 @@ impl InternalFlags {
         }
     }
 
-    pub fn add_lang_used(&mut self, flags: impl IntoIterator<Item = u32>) {
+    pub(crate) fn add_lang_used(&mut self, flags: impl IntoIterator<Item = u32>) {
         if let Self::Enabled {
             lang_since_last_complete,
             ..
@@ -82,7 +82,7 @@ impl InternalFlags {
     /// Returns true if this flag may currently be used. If `should_record` is true, always returns
     /// true and records the flag as being used, for taking later via
     /// [Self::gather_for_wft_complete].
-    pub fn try_use(&mut self, core_patch: CoreInternalFlags, should_record: bool) -> bool {
+    pub(crate) fn try_use(&mut self, core_patch: CoreInternalFlags, should_record: bool) -> bool {
         match self {
             Self::Enabled {
                 core,
@@ -104,7 +104,7 @@ impl InternalFlags {
 
     /// Writes all known core flags to the set which should be recorded in the current WFT if not
     /// already known. Must only be called if not replaying.
-    pub fn write_all_known(&mut self) {
+    pub(crate) fn write_all_known(&mut self) {
         if let Self::Enabled {
             core_since_last_complete,
             ..
@@ -117,7 +117,7 @@ impl InternalFlags {
     /// Wipes the recorded flags used during the current WFT and returns a partially filled
     /// sdk metadata message that can be combined with any existing data before sending the WFT
     /// complete
-    pub fn gather_for_wft_complete(&mut self) -> WorkflowTaskCompletedMetadata {
+    pub(crate) fn gather_for_wft_complete(&mut self) -> WorkflowTaskCompletedMetadata {
         match self {
             Self::Enabled {
                 core_since_last_complete,
@@ -148,7 +148,7 @@ impl InternalFlags {
         }
     }
 
-    pub fn all_lang(&self) -> impl Iterator<Item = u32> + '_ {
+    pub(crate) fn all_lang(&self) -> impl Iterator<Item = u32> + '_ {
         match self {
             Self::Enabled { lang, .. } => Either::Left(lang.iter().copied()),
             Self::Disabled => Either::Right(iter::empty()),
@@ -165,7 +165,7 @@ impl CoreInternalFlags {
         }
     }
 
-    pub fn all_except_too_high() -> impl Iterator<Item = CoreInternalFlags> {
+    pub(crate) fn all_except_too_high() -> impl Iterator<Item = CoreInternalFlags> {
         enum_iterator::all::<CoreInternalFlags>()
             .filter(|f| !matches!(f, CoreInternalFlags::TooHigh))
     }

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

@@ -38,9 +38,10 @@ pub use temporal_sdk_core_api as api;
 pub use temporal_sdk_core_protos as protos;
 pub use temporal_sdk_core_protos::TaskToken;
 pub use url::Url;
-#[cfg(feature = "save_wf_inputs")]
-pub use worker::replay_wf_state_inputs;
-pub use worker::{Worker, WorkerConfig, WorkerConfigBuilder};
+pub use worker::{
+    FixedSizeSlotSupplier, RealSysInfo, ResourceBasedSlots, ResourceBasedTuner,
+    ResourceSlotOptions, TunerBuilder, TunerHolder, Worker, WorkerConfig, WorkerConfigBuilder,
+};
 
 use crate::{
     replay::{HistoryForReplay, ReplayWorkerInput},
@@ -79,14 +80,7 @@ pub fn init_worker<CT>(
 where
     CT: Into<sealed::AnyClient>,
 {
-    let client = {
-        let ll = client.into().into_inner();
-        let mut client = Client::new(*ll, worker_config.namespace.clone());
-        if let Some(ref id_override) = worker_config.client_identity_override {
-            client.options_mut().identity = id_override.clone();
-        }
-        RetryClient::new(client, RetryConfig::default())
-    };
+    let client = init_worker_client(&worker_config, *client.into().into_inner());
     if client.namespace() != worker_config.namespace {
         panic!("Passed in client is not bound to the same namespace as the worker");
     }
@@ -125,6 +119,17 @@ where
     rwi.into_core_worker()
 }
 
+pub(crate) fn init_worker_client(
+    config: &WorkerConfig,
+    client: ConfiguredClient<TemporalServiceClientWithMetrics>,
+) -> RetryClient<Client> {
+    let mut client = Client::new(client, config.namespace.clone());
+    if let Some(ref id_override) = config.client_identity_override {
+        client.options_mut().identity.clone_from(id_override);
+    }
+    RetryClient::new(client, RetryConfig::default())
+}
+
 /// Creates a unique sticky queue name for a worker, iff the config allows for 1 or more cached
 /// workflows.
 pub(crate) fn sticky_q_name_for_worker(

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

@@ -17,15 +17,16 @@ use temporal_sdk_core_protos::temporal::api::workflowservice::v1::{
 use futures::Future;
 #[cfg(test)]
 pub(crate) use poll_buffer::MockPermittedPollBuffer;
+use temporal_sdk_core_api::worker::{ActivitySlotKind, WorkflowSlotKind};
 
-pub type Result<T, E = tonic::Status> = std::result::Result<T, E>;
+pub(crate) type Result<T, E = tonic::Status> = std::result::Result<T, E>;
 
 /// A trait for things that poll the server. Hides complexity of concurrent polling or polling
 /// on sticky/nonsticky queues simultaneously.
 #[cfg_attr(test, mockall::automock)]
 #[cfg_attr(test, allow(unused))]
 #[async_trait::async_trait]
-pub trait Poller<PollResult>
+pub(crate) trait Poller<PollResult>
 where
     PollResult: Send + Sync + 'static,
 {
@@ -36,9 +37,14 @@ where
     async fn shutdown_box(self: Box<Self>);
 }
 pub(crate) type BoxedPoller<T> = Box<dyn Poller<T> + Send + Sync + 'static>;
-pub(crate) type BoxedWFPoller = BoxedPoller<(PollWorkflowTaskQueueResponse, OwnedMeteredSemPermit)>;
-pub(crate) type BoxedActPoller =
-    BoxedPoller<(PollActivityTaskQueueResponse, OwnedMeteredSemPermit)>;
+pub(crate) type BoxedWFPoller = BoxedPoller<(
+    PollWorkflowTaskQueueResponse,
+    OwnedMeteredSemPermit<WorkflowSlotKind>,
+)>;
+pub(crate) type BoxedActPoller = BoxedPoller<(
+    PollActivityTaskQueueResponse,
+    OwnedMeteredSemPermit<ActivitySlotKind>,
+)>;
 
 #[async_trait::async_trait]
 impl<T> Poller<T> for Box<dyn Poller<T> + Send + Sync>