@temporalio/core-bridge 1.9.2 → 1.10.0

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

@@ -1,11 +1,14 @@
 mod activities;
 pub(crate) mod client;
 mod slot_provider;
+pub(crate) mod tuner;
 mod workflow;
 
 pub use temporal_sdk_core_api::worker::{WorkerConfig, WorkerConfigBuilder};
-#[cfg(feature = "save_wf_inputs")]
-pub use workflow::replay_wf_state_inputs;
+pub use tuner::{
+    FixedSizeSlotSupplier, RealSysInfo, ResourceBasedSlots, ResourceBasedTuner,
+    ResourceSlotOptions, TunerBuilder, TunerHolder,
+};
 
 pub(crate) use activities::{
     ExecutingLAId, LocalActRequest, LocalActivityExecutionResult, LocalActivityResolution,
@@ -13,10 +16,10 @@ pub(crate) use activities::{
 };
 pub(crate) use workflow::{wft_poller::new_wft_poller, LEGACY_QUERY_ID};
 
-use temporal_client::WorkerKey;
+use temporal_client::{ConfiguredClient, TemporalServiceClientWithMetrics, WorkerKey};
 
 use crate::{
-    abstractions::{dbg_panic, MeteredSemaphore},
+    abstractions::{dbg_panic, MeteredPermitDealer},
     errors::CompleteWfError,
     pollers::{
         new_activity_task_buffer, new_workflow_task_buffer, BoxedActPoller, WorkflowTaskPoller,
@@ -38,6 +41,7 @@ use crate::{
 };
 use activities::WorkerActivityTasks;
 use futures_util::{stream, StreamExt};
+use parking_lot::Mutex;
 use slot_provider::SlotProvider;
 use std::{
     convert::TryInto,
@@ -81,7 +85,7 @@ pub struct Worker {
     config: WorkerConfig,
     wf_client: Arc<dyn WorkerClient>,
     /// Registration key to enable eager workflow start for this worker
-    worker_key: Option<WorkerKey>,
+    worker_key: Mutex<Option<WorkerKey>>,
     /// Manages all workflows and WFT processing
     workflows: Workflows,
     /// Manages activity tasks for this worker/task queue
@@ -169,7 +173,7 @@ impl WorkerTrait for Worker {
         }
         self.shutdown_token.cancel();
         // First, disable Eager Workflow Start
-        if let Some(key) = self.worker_key {
+        if let Some(key) = *self.worker_key.lock() {
             self.wf_client.workers().unregister(key);
         }
         // Second, we want to stop polling of both activity and workflow tasks
@@ -217,12 +221,23 @@ impl Worker {
         )
     }
 
+    /// Replace client and return a new client. For eager workflow purposes, this new client will
+    /// now apply to future eager start requests and the older client will not.
+    pub fn replace_client(&self, new_client: ConfiguredClient<TemporalServiceClientWithMetrics>) {
+        // Unregister worker from current client, register in new client at the end
+        let mut worker_key = self.worker_key.lock();
+        let slot_provider = (*worker_key).and_then(|k| self.wf_client.workers().unregister(k));
+        self.wf_client
+            .replace_client(super::init_worker_client(&self.config, new_client));
+        *worker_key = slot_provider
+            .and_then(|slot_provider| self.wf_client.workers().register(slot_provider));
+    }
+
     #[cfg(test)]
     pub(crate) fn new_test(config: WorkerConfig, client: impl WorkerClient + 'static) -> Self {
         Self::new(config, None, Arc::new(client), None)
     }
 
-    #[allow(clippy::too_many_arguments)] // Not much worth combining here
     pub(crate) fn new_with_pollers(
         config: WorkerConfig,
         sticky_queue_name: Option<String>,
@@ -230,22 +245,40 @@ impl Worker {
         task_pollers: TaskPollers,
         telem_instance: Option<&TelemetryInstance>,
     ) -> Self {
-        let metrics = if let Some(ti) = telem_instance {
-            MetricsContext::top_level(config.namespace.clone(), config.task_queue.clone(), ti)
+        let (metrics, meter) = if let Some(ti) = telem_instance {
+            (
+                MetricsContext::top_level(config.namespace.clone(), config.task_queue.clone(), ti),
+                ti.get_metric_meter(),
+            )
         } else {
-            MetricsContext::no_op()
+            (MetricsContext::no_op(), None)
         };
+        let tuner = config
+            .tuner
+            .as_ref()
+            .cloned()
+            .unwrap_or_else(|| TunerBuilder::from_config(&config).build());
+
         metrics.worker_registered();
+        if let Some(meter) = meter {
+            tuner.attach_metrics(meter.clone());
+        }
         let shutdown_token = CancellationToken::new();
-        let wft_semaphore = Arc::new(MeteredSemaphore::new(
-            config.max_outstanding_workflow_tasks,
+        let wft_slots = Arc::new(MeteredPermitDealer::new(
+            tuner.workflow_task_slot_supplier(),
             metrics.with_new_attrs([workflow_worker_type()]),
-            MetricsContext::available_task_slots,
+            if config.max_cached_workflows > 0 {
+                // Since we always need to be able to poll the normal task queue as well as the
+                // sticky queue, we need a value of at least 2 here.
+                Some(std::cmp::max(2, config.max_cached_workflows))
+            } else {
+                None
+            },
         ));
-        let act_semaphore = Arc::new(MeteredSemaphore::new(
-            config.max_outstanding_activities,
+        let act_slots = Arc::new(MeteredPermitDealer::new(
+            tuner.activity_task_slot_supplier(),
             metrics.with_new_attrs([activity_worker_type()]),
-            MetricsContext::available_task_slots,
+            None,
         ));
         let (external_wft_tx, external_wft_rx) = unbounded_channel();
         let (wft_stream, act_poller) = match task_pollers {
@@ -265,7 +298,7 @@ impl Worker {
                         normal_name: "".to_string(),
                     },
                     max_nonsticky_polls,
-                    wft_semaphore.clone(),
+                    wft_slots.clone(),
                     shutdown_token.child_token(),
                     Some(move |np| {
                         wft_metrics.record_num_pollers(np);
@@ -281,7 +314,7 @@ impl Worker {
                             normal_name: config.task_queue.clone(),
                         },
                         max_sticky_polls,
-                        wft_semaphore.clone(),
+                        wft_slots.clone(),
                         shutdown_token.child_token(),
                         Some(move |np| {
                             sticky_metrics.record_num_pollers(np);
@@ -296,7 +329,7 @@ impl Worker {
                         client.clone(),
                         config.task_queue.clone(),
                         config.max_concurrent_at_polls,
-                        act_semaphore.clone(),
+                        act_slots.clone(),
                         config.max_task_queue_activities_per_second,
                         shutdown_token.child_token(),
                         Some(move |np| act_metrics.record_num_pollers(np)),
@@ -327,9 +360,8 @@ impl Worker {
                 wft_stream,
                 act_poller,
             } => {
-                let ap =
-                    act_poller.map(|ap| MockPermittedPollBuffer::new(act_semaphore.clone(), ap));
-                let wft_semaphore = wft_semaphore.clone();
+                let ap = act_poller.map(|ap| MockPermittedPollBuffer::new(act_slots.clone(), ap));
+                let wft_semaphore = wft_slots.clone();
                 let wfs = wft_stream.then(move |s| {
                     let wft_semaphore = wft_semaphore.clone();
                     async move {
@@ -347,14 +379,14 @@ impl Worker {
 
         let (hb_tx, hb_rx) = unbounded_channel();
         let local_act_mgr = Arc::new(LocalActivityManager::new(
-            config.max_outstanding_local_activities,
+            tuner.local_activity_slot_supplier(),
             config.namespace.clone(),
             hb_tx,
             metrics.with_new_attrs([local_activity_worker_type()]),
         ));
         let at_task_mgr = act_poller.map(|ap| {
             WorkerActivityTasks::new(
-                act_semaphore,
+                act_slots,
                 ap,
                 client.clone(),
                 metrics.clone(),
@@ -368,14 +400,14 @@ impl Worker {
         if !poll_on_non_local_activities {
             info!("Activity polling is disabled for this worker");
         };
-        let la_sink = LAReqSink::new(local_act_mgr.clone(), config.wf_state_inputs.clone());
+        let la_sink = LAReqSink::new(local_act_mgr.clone());
         let provider = SlotProvider::new(
             config.namespace.clone(),
             config.task_queue.clone(),
-            wft_semaphore.clone(),
+            wft_slots.clone(),
             external_wft_tx,
         );
-        let worker_key = client.workers().register(Box::new(provider));
+        let worker_key = Mutex::new(client.workers().register(Box::new(provider)));
         Self {
             worker_key,
             wf_client: client.clone(),
@@ -384,7 +416,7 @@ impl Worker {
                     config.clone(),
                     metrics,
                     shutdown_token.child_token(),
-                    client.capabilities().cloned().unwrap_or_default(),
+                    client.capabilities().clone().unwrap_or_default(),
                 ),
                 sticky_queue_name.map(|sq| StickyExecutionAttributes {
                     worker_task_queue: Some(TaskQueue {
@@ -400,7 +432,7 @@ impl Worker {
                     ),
                 }),
                 client,
-                wft_semaphore,
+                wft_slots,
                 wft_stream,
                 la_sink,
                 local_act_mgr.clone(),
@@ -473,11 +505,11 @@ impl Worker {
     }
 
     #[allow(unused)]
-    pub(crate) fn available_wft_permits(&self) -> usize {
+    pub(crate) fn available_wft_permits(&self) -> Option<usize> {
         self.workflows.available_wft_permits()
     }
     #[cfg(test)]
-    pub(crate) fn unused_wft_permits(&self) -> usize {
+    pub(crate) fn unused_wft_permits(&self) -> Option<usize> {
         self.workflows.unused_wft_permits()
     }
 
@@ -659,10 +691,9 @@ impl Worker {
     }
 }
 
-pub struct PostActivateHookData<'a> {
-    pub run_id: &'a str,
-    pub most_recent_event: usize,
-    pub replaying: bool,
+pub(crate) struct PostActivateHookData<'a> {
+    pub(crate) run_id: &'a str,
+    pub(crate) replaying: bool,
 }
 
 fn build_wf_basics(
@@ -718,7 +749,7 @@ mod tests {
                 .as_ref()
                 .unwrap()
                 .remaining_activity_capacity(),
-            5
+            Some(5)
         );
     }
 
@@ -735,16 +766,15 @@ mod tests {
             .unwrap();
         let worker = Worker::new_test(cfg, mock_client);
         assert!(worker.activity_poll().await.is_err());
-        assert_eq!(worker.at_task_mgr.unwrap().remaining_activity_capacity(), 5);
+        assert_eq!(
+            worker.at_task_mgr.unwrap().remaining_activity_capacity(),
+            Some(5)
+        );
     }
 
     #[test]
     fn max_polls_calculated_properly() {
-        let mut wcb = WorkerConfigBuilder::default();
-        let cfg = wcb
-            .namespace("default")
-            .task_queue("whatever")
-            .worker_build_id("test_bin_id")
+        let cfg = test_worker_cfg()
             .max_concurrent_wft_polls(5_usize)
             .build()
             .unwrap();

package/sdk-core/core/src/worker/slot_provider.rs

@@ -3,27 +3,38 @@
 //! This enables latency optimizations such as Eager Workflow Start.
 
 use crate::{
-    abstractions::{MeteredSemaphore, OwnedMeteredSemPermit},
+    abstractions::{MeteredPermitDealer, OwnedMeteredSemPermit},
     protosext::ValidPollWFTQResponse,
     worker::workflow::wft_poller::validate_wft,
 };
 
 use std::sync::Arc;
 use temporal_client::{Slot as SlotTrait, SlotProvider as SlotProviderTrait};
+use temporal_sdk_core_api::worker::WorkflowSlotKind;
 use temporal_sdk_core_protos::temporal::api::workflowservice::v1::PollWorkflowTaskQueueResponse;
 use tokio::sync::mpsc::UnboundedSender;
 use tonic::Status;
 
-type WFTStreamSender =
-    UnboundedSender<Result<(ValidPollWFTQResponse, OwnedMeteredSemPermit), Status>>;
-
-pub struct Slot {
-    permit: OwnedMeteredSemPermit,
+type WFTStreamSender = UnboundedSender<
+    Result<
+        (
+            ValidPollWFTQResponse,
+            OwnedMeteredSemPermit<WorkflowSlotKind>,
+        ),
+        Status,
+    >,
+>;
+
+struct Slot {
+    permit: OwnedMeteredSemPermit<WorkflowSlotKind>,
     external_wft_tx: WFTStreamSender,
 }
 
 impl Slot {
-    fn new(permit: OwnedMeteredSemPermit, external_wft_tx: WFTStreamSender) -> Self {
+    fn new(
+        permit: OwnedMeteredSemPermit<WorkflowSlotKind>,
+        external_wft_tx: WFTStreamSender,
+    ) -> Self {
         Self {
             permit,
             external_wft_tx,
@@ -44,18 +55,18 @@ impl SlotTrait for Slot {
 
 #[derive(derive_more::DebugCustom)]
 #[debug(fmt = "SlotProvider {{ namespace:{namespace}, task_queue: {task_queue} }}")]
-pub struct SlotProvider {
+pub(super) struct SlotProvider {
     namespace: String,
     task_queue: String,
-    wft_semaphore: Arc<MeteredSemaphore>,
+    wft_semaphore: Arc<MeteredPermitDealer<WorkflowSlotKind>>,
     external_wft_tx: WFTStreamSender,
 }
 
 impl SlotProvider {
-    pub(crate) fn new(
+    pub(super) fn new(
         namespace: String,
         task_queue: String,
-        wft_semaphore: Arc<MeteredSemaphore>,
+        wft_semaphore: Arc<MeteredPermitDealer<WorkflowSlotKind>>,
         external_wft_tx: WFTStreamSender,
     ) -> Self {
         Self {
@@ -86,6 +97,7 @@ impl SlotProviderTrait for SlotProvider {
 mod tests {
     use super::*;
 
+    use crate::abstractions::tests::fixed_size_permit_dealer;
     use temporal_sdk_core_protos::temporal::api::{
         common::v1::{WorkflowExecution, WorkflowType},
         history::v1::History,
@@ -106,11 +118,7 @@ mod tests {
 
     #[tokio::test]
     async fn slot_propagates_through_channel() {
-        let wft_semaphore = Arc::new(MeteredSemaphore::new(
-            2,
-            crate::MetricsContext::no_op(),
-            |_, _| {},
-        ));
+        let wft_semaphore = Arc::new(fixed_size_permit_dealer(2));
         let (external_wft_tx, mut external_wft_rx) = unbounded_channel();
 
         let provider = SlotProvider::new(
@@ -133,11 +141,7 @@ mod tests {
         let (external_wft_tx, mut external_wft_rx) = unbounded_channel();
         {
             let external_wft_tx = external_wft_tx;
-            let wft_semaphore = Arc::new(MeteredSemaphore::new(
-                2,
-                crate::MetricsContext::no_op(),
-                |_, _| {},
-            ));
+            let wft_semaphore = Arc::new(fixed_size_permit_dealer(2));
             let provider = SlotProvider::new(
                 "my_namespace".to_string(),
                 "my_queue".to_string(),
@@ -151,11 +155,7 @@ mod tests {
 
     #[tokio::test]
     async fn unused_slots_reclaimed() {
-        let wft_semaphore = Arc::new(MeteredSemaphore::new(
-            2,
-            crate::MetricsContext::no_op(),
-            |_, _| {},
-        ));
+        let wft_semaphore = Arc::new(fixed_size_permit_dealer(2));
         {
             let wft_semaphore = wft_semaphore.clone();
             let (external_wft_tx, _) = unbounded_channel();
@@ -167,9 +167,9 @@ mod tests {
             );
             let slot = provider.try_reserve_wft_slot();
             assert!(slot.is_some());
-            assert_eq!(wft_semaphore.available_permits(), 1);
+            assert_eq!(wft_semaphore.available_permits(), Some(1));
             // drop slot without using it
         }
-        assert_eq!(wft_semaphore.available_permits(), 2);
+        assert_eq!(wft_semaphore.available_permits(), Some(2));
     }
 }

package/sdk-core/core/src/worker/slot_supplier.rs

@@ -0,0 +1 @@
+

package/sdk-core/core/src/worker/tuner/fixed_size.rs

@@ -0,0 +1,52 @@
+use std::{marker::PhantomData, sync::Arc};
+use temporal_sdk_core_api::worker::{
+    SlotKind, SlotReservationContext, SlotSupplier, SlotSupplierPermit,
+};
+use tokio::sync::Semaphore;
+
+/// Implements [SlotSupplier] with a fixed number of slots
+pub struct FixedSizeSlotSupplier<SK> {
+    sem: Arc<Semaphore>,
+    _pd: PhantomData<SK>,
+}
+
+impl<SK> FixedSizeSlotSupplier<SK> {
+    /// Create a slot supplier which will only hand out at most the provided number of slots
+    pub fn new(size: usize) -> Self {
+        Self {
+            sem: Arc::new(Semaphore::new(size)),
+            _pd: Default::default(),
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl<SK> SlotSupplier for FixedSizeSlotSupplier<SK>
+where
+    SK: SlotKind + Send + Sync,
+{
+    type SlotKind = SK;
+
+    async fn reserve_slot(&self, _: &dyn SlotReservationContext) -> SlotSupplierPermit {
+        let perm = self
+            .sem
+            .clone()
+            .acquire_owned()
+            .await
+            .expect("inner semaphore is never closed");
+        SlotSupplierPermit::with_user_data(perm)
+    }
+
+    fn try_reserve_slot(&self, _: &dyn SlotReservationContext) -> Option<SlotSupplierPermit> {
+        let perm = self.sem.clone().try_acquire_owned();
+        perm.ok().map(SlotSupplierPermit::with_user_data)
+    }
+
+    fn mark_slot_used(&self, _info: SK::Info<'_>) {}
+
+    fn release_slot(&self) {}
+
+    fn available_slots(&self) -> Option<usize> {
+        Some(self.sem.available_permits())
+    }
+}