@temporalio/core-bridge 1.11.2 → 1.11.4

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

@@ -25,7 +25,7 @@ use crate::{
 };
 use activity_heartbeat_manager::ActivityHeartbeatManager;
 use dashmap::DashMap;
-use futures::{
+use futures_util::{
     stream,
     stream::{BoxStream, PollNext},
     Stream, StreamExt,
@@ -37,14 +37,14 @@ use std::{
         atomic::{AtomicBool, Ordering},
         Arc,
     },
-    time::{Duration, Instant},
+    time::{Duration, Instant, SystemTime},
 };
-use temporal_sdk_core_api::worker::{ActivitySlotInfo, ActivitySlotKind};
+use temporal_sdk_core_api::worker::ActivitySlotKind;
 use temporal_sdk_core_protos::{
     coresdk::{
         activity_result::{self as ar, activity_execution_result as aer},
         activity_task::{ActivityCancelReason, ActivityTask},
-        ActivityHeartbeat,
+        ActivityHeartbeat, ActivitySlotInfo,
     },
     temporal::api::{
         failure::v1::{failure::FailureInfo, ApplicationFailureInfo, CanceledFailureInfo, Failure},
@@ -93,6 +93,7 @@ struct InFlightActInfo {
     /// Only kept for logging reasons
     workflow_run_id: String,
     start_time: Instant,
+    scheduled_time: Option<SystemTime>,
 }
 
 /// Augments [InFlightActInfo] with details specific to remote activities
@@ -127,6 +128,7 @@ impl RemoteInFlightActInfo {
                 workflow_id: wec.workflow_id,
                 workflow_run_id: wec.run_id,
                 start_time: Instant::now(),
+                scheduled_time: poll_resp.scheduled_time.and_then(|i| i.try_into().ok()),
             },
             heartbeat_timeout: poll_resp.heartbeat_timeout,
             issued_cancel_to_lang: None,
@@ -180,7 +182,7 @@ enum ActivityTaskSource {
 impl WorkerActivityTasks {
     #[allow(clippy::too_many_arguments)]
     pub(crate) fn new(
-        semaphore: Arc<MeteredPermitDealer<ActivitySlotKind>>,
+        semaphore: MeteredPermitDealer<ActivitySlotKind>,
         poller: BoxedActPoller,
         client: Arc<dyn WorkerClient>,
         metrics: MetricsContext,
@@ -194,7 +196,7 @@ impl WorkerActivityTasks {
         let server_poller_stream =
             new_activity_task_poller(poller, metrics.clone(), shutdown_initiated_token.clone());
         let (eager_activities_tx, eager_activities_rx) = unbounded_channel();
-        let eager_activities_semaphore = ClosableMeteredPermitDealer::new_arc(semaphore);
+        let eager_activities_semaphore = ClosableMeteredPermitDealer::new_arc(Arc::new(semaphore));
 
         let start_tasks_stream_complete = CancellationToken::new();
         let starts_stream = Self::merge_start_task_sources(
@@ -277,7 +279,7 @@ impl WorkerActivityTasks {
         // Prefer eager activities over polling the server
         stream::select_with_strategy(non_poll_stream, poller_stream, |_: &mut ()| PollNext::Left)
             .map(|res| Some(res.map_err(Into::into)))
-            .chain(futures::stream::once(async move {
+            .chain(futures_util::stream::once(async move {
                 on_complete_token.cancel();
                 None
             }))
@@ -336,10 +338,19 @@ impl WorkerActivityTasks {
                 let _flushing_guard = self.completers_lock.read().await;
                 let maybe_net_err = match status {
                     aer::Status::WillCompleteAsync(_) => None,
-                    aer::Status::Completed(ar::Success { result }) => client
-                        .complete_activity_task(task_token.clone(), result.map(Into::into))
-                        .await
-                        .err(),
+                    aer::Status::Completed(ar::Success { result }) => {
+                        if let Some(sched_time) = act_info
+                            .base
+                            .scheduled_time
+                            .and_then(|st| st.elapsed().ok())
+                        {
+                            act_metrics.act_execution_succeeded(sched_time);
+                        }
+                        client
+                            .complete_activity_task(task_token.clone(), result.map(Into::into))
+                            .await
+                            .err()
+                    }
                     aer::Status::Failed(ar::Failure { failure }) => {
                         act_metrics.act_execution_failed();
                         client
@@ -542,7 +553,7 @@ where
                                 outstanding_entry.insert(RemoteInFlightActInfo::new(
                                     &task.resp,
                                     task.permit.into_used(ActivitySlotInfo {
-                                        activity_type: activity_type_name,
+                                        activity_type: activity_type_name.to_string(),
                                     }),
                                 ));
                             // If we have already waited the grace period and issued cancels,
@@ -739,7 +750,7 @@ mod tests {
             .times(2)
             .returning(|_, _| Ok(Default::default()));
         let mock_client = Arc::new(mock_client);
-        let sem = Arc::new(fixed_size_permit_dealer(10));
+        let sem = fixed_size_permit_dealer(10);
         let shutdown_token = CancellationToken::new();
         let ap = new_activity_task_buffer(
             mock_client.clone(),
@@ -828,7 +839,7 @@ mod tests {
                 })
             });
         let mock_client = Arc::new(mock_client);
-        let sem = Arc::new(fixed_size_permit_dealer(1));
+        let sem = fixed_size_permit_dealer(1);
         let shutdown_token = CancellationToken::new();
         let ap = new_activity_task_buffer(
             mock_client.clone(),
@@ -899,7 +910,7 @@ mod tests {
             .times(2)
             .returning(|_, _| Ok(Default::default()));
         let mock_client = Arc::new(mock_client);
-        let sem = Arc::new(fixed_size_permit_dealer(1));
+        let sem = fixed_size_permit_dealer(1);
         let shutdown_token = CancellationToken::new();
         let ap = new_activity_task_buffer(
             mock_client.clone(),

package/sdk-core/core/src/worker/client/mocks.rs

@@ -1,11 +1,11 @@
 use super::*;
-use futures::Future;
-use once_cell::sync::Lazy;
+use futures_util::Future;
 use std::sync::Arc;
+use std::sync::LazyLock;
 use temporal_client::SlotManager;
 
-static DEFAULT_WORKERS_REGISTRY: Lazy<Arc<SlotManager>> =
-    Lazy::new(|| Arc::new(SlotManager::new()));
+pub(crate) static DEFAULT_WORKERS_REGISTRY: LazyLock<Arc<SlotManager>> =
+    LazyLock::new(|| Arc::new(SlotManager::new()));
 
 pub(crate) static DEFAULT_TEST_CAPABILITIES: &Capabilities = &Capabilities {
     signal_and_query_header: true,
@@ -18,6 +18,7 @@ pub(crate) static DEFAULT_TEST_CAPABILITIES: &Capabilities = &Capabilities {
     eager_workflow_start: true,
     sdk_metadata: true,
     count_group_by_execution_status: false,
+    nexus: false,
 };
 
 #[cfg(test)]
@@ -29,6 +30,8 @@ pub(crate) fn mock_workflow_client() -> MockWorkerClient {
     r.expect_workers()
         .returning(|| DEFAULT_WORKERS_REGISTRY.clone());
     r.expect_is_mock().returning(|| true);
+    r.expect_shutdown_worker()
+        .returning(|_| Ok(ShutdownWorkerResponse {}));
     r
 }
 
@@ -119,6 +122,9 @@ mockall::mock! {
           impl Future<Output = Result<DescribeNamespaceResponse>> + Send + 'b
           where 'a: 'b, Self: 'b;
 
+        fn shutdown_worker<'a, 'b>(&self, sticky_task_queue: String) -> impl Future<Output = Result<ShutdownWorkerResponse>> + Send + 'b
+            where 'a: 'b, Self: 'b;
+
         fn replace_client(&self, new_client: RetryClient<Client>);
         fn capabilities(&self) -> Option<Capabilities>;
         fn workers(&self) -> Arc<SlotManager>;

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

@@ -146,6 +146,7 @@ pub(crate) trait WorkerClient: Sync + Send {
         query_result: QueryResult,
     ) -> Result<RespondQueryTaskCompletedResponse>;
     async fn describe_namespace(&self) -> Result<DescribeNamespaceResponse>;
+    async fn shutdown_worker(&self, sticky_task_queue: String) -> Result<ShutdownWorkerResponse>;
 
     fn replace_client(&self, new_client: RetryClient<Client>);
     fn capabilities(&self) -> Option<Capabilities>;
@@ -232,6 +233,7 @@ impl WorkerClient for WorkerClientBag {
             namespace: self.namespace.clone(),
             sdk_metadata: Some(request.sdk_metadata),
             metering_metadata: Some(request.metering_metadata),
+            capabilities: None,
         };
         Ok(self
             .cloned_client()
@@ -384,6 +386,21 @@ impl WorkerClient for WorkerClientBag {
         .await
     }
 
+    async fn shutdown_worker(&self, sticky_task_queue: String) -> Result<ShutdownWorkerResponse> {
+        let request = ShutdownWorkerRequest {
+            namespace: self.namespace.clone(),
+            identity: self.identity.clone(),
+            sticky_task_queue,
+            reason: "graceful shutdown".to_string(),
+        };
+
+        Ok(
+            WorkflowService::shutdown_worker(&mut self.cloned_client(), request)
+                .await?
+                .into_inner(),
+        )
+    }
+
     fn replace_client(&self, new_client: RetryClient<Client>) {
         let mut replaceable_client = self.replaceable_client.write();
         *replaceable_client = new_client;

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

@@ -26,8 +26,8 @@ use crate::{
     protosext::validate_activity_completion,
     telemetry::{
         metrics::{
-            activity_poller, activity_worker_type, local_activity_worker_type, workflow_poller,
-            workflow_sticky_poller, workflow_worker_type, MetricsContext,
+            activity_poller, activity_worker_type, workflow_poller, workflow_sticky_poller,
+            workflow_worker_type, MetricsContext,
         },
         TelemetryInstance,
     },
@@ -49,6 +49,7 @@ use std::{
         atomic::{AtomicBool, Ordering},
         Arc,
     },
+    time::Duration,
 };
 use temporal_client::{ConfiguredClient, TemporalServiceClientWithMetrics, WorkerKey};
 use temporal_sdk_core_protos::{
@@ -66,10 +67,13 @@ use temporal_sdk_core_protos::{
     },
     TaskToken,
 };
-use tokio::sync::mpsc::unbounded_channel;
+use tokio::sync::{mpsc::unbounded_channel, watch};
 use tokio_stream::wrappers::UnboundedReceiverStream;
 use tokio_util::sync::CancellationToken;
 
+use crate::{
+    abstractions::PermitDealerContextData, telemetry::metrics::local_activity_worker_type,
+};
 use temporal_sdk_core_api::errors::WorkerValidationError;
 #[cfg(test)]
 use {
@@ -77,7 +81,7 @@ use {
         pollers::{BoxedPoller, MockPermittedPollBuffer},
         protosext::ValidPollWFTQResponse,
     },
-    futures::stream::BoxStream,
+    futures_util::stream::BoxStream,
     temporal_sdk_core_protos::temporal::api::workflowservice::v1::PollActivityTaskQueueResponse,
 };
 
@@ -102,6 +106,22 @@ pub struct Worker {
     non_local_activities_complete: Arc<AtomicBool>,
     /// Set when local activities are complete and should stop being polled
     local_activities_complete: Arc<AtomicBool>,
+    /// Used to track all permits have been released
+    all_permits_tracker: tokio::sync::Mutex<AllPermitsTracker>,
+}
+
+struct AllPermitsTracker {
+    wft_permits: watch::Receiver<usize>,
+    act_permits: watch::Receiver<usize>,
+    la_permits: watch::Receiver<usize>,
+}
+
+impl AllPermitsTracker {
+    async fn all_done(&mut self) {
+        let _ = self.wft_permits.wait_for(|x| *x == 0).await;
+        let _ = self.act_permits.wait_for(|x| *x == 0).await;
+        let _ = self.la_permits.wait_for(|x| *x == 0).await;
+    }
 }
 
 #[async_trait::async_trait]
@@ -270,7 +290,12 @@ impl Worker {
             tuner.attach_metrics(meter.clone());
         }
         let shutdown_token = CancellationToken::new();
-        let wft_slots = Arc::new(MeteredPermitDealer::new(
+        let slot_context_data = Arc::new(PermitDealerContextData {
+            task_queue: config.task_queue.clone(),
+            worker_identity: config.client_identity_override.clone().unwrap_or_default(),
+            worker_build_id: config.worker_build_id.clone(),
+        });
+        let wft_slots = MeteredPermitDealer::new(
             tuner.workflow_task_slot_supplier(),
             metrics.with_new_attrs([workflow_worker_type()]),
             if config.max_cached_workflows > 0 {
@@ -280,12 +305,16 @@ impl Worker {
             } else {
                 None
             },
-        ));
-        let act_slots = Arc::new(MeteredPermitDealer::new(
+            slot_context_data.clone(),
+        );
+        let wft_permits = wft_slots.get_extant_count_rcv();
+        let act_slots = MeteredPermitDealer::new(
             tuner.activity_task_slot_supplier(),
             metrics.with_new_attrs([activity_worker_type()]),
             None,
-        ));
+            slot_context_data.clone(),
+        );
+        let act_permits = act_slots.get_extant_count_rcv();
         let (external_wft_tx, external_wft_rx) = unbounded_channel();
         let (wft_stream, act_poller) = match task_pollers {
             TaskPollers::Real => {
@@ -320,7 +349,7 @@ impl Worker {
                             normal_name: config.task_queue.clone(),
                         },
                         max_sticky_polls,
-                        wft_slots.clone(),
+                        wft_slots.clone().into_sticky(),
                         shutdown_token.child_token(),
                         Some(move |np| {
                             sticky_metrics.record_num_pollers(np);
@@ -366,7 +395,8 @@ impl Worker {
                 wft_stream,
                 act_poller,
             } => {
-                let ap = act_poller.map(|ap| MockPermittedPollBuffer::new(act_slots.clone(), ap));
+                let ap = act_poller
+                    .map(|ap| MockPermittedPollBuffer::new(Arc::new(act_slots.clone()), ap));
                 let wft_semaphore = wft_slots.clone();
                 let wfs = wft_stream.then(move |s| {
                     let wft_semaphore = wft_semaphore.clone();
@@ -381,11 +411,18 @@ impl Worker {
         };
 
         let (hb_tx, hb_rx) = unbounded_channel();
-        let local_act_mgr = Arc::new(LocalActivityManager::new(
+        let la_pemit_dealer = MeteredPermitDealer::new(
             tuner.local_activity_slot_supplier(),
+            metrics.with_new_attrs([local_activity_worker_type()]),
+            None,
+            slot_context_data,
+        );
+        let la_permits = la_pemit_dealer.get_extant_count_rcv();
+        let local_act_mgr = Arc::new(LocalActivityManager::new(
             config.namespace.clone(),
+            la_pemit_dealer,
             hb_tx,
-            metrics.with_new_attrs([local_activity_worker_type()]),
+            metrics.clone(),
         ));
         let at_task_mgr = act_poller.map(|ap| {
             WorkerActivityTasks::new(
@@ -453,6 +490,11 @@ impl Worker {
             // Non-local activities are already complete if configured not to poll for them.
             non_local_activities_complete: Arc::new(AtomicBool::new(!poll_on_non_local_activities)),
             local_activities_complete: Default::default(),
+            all_permits_tracker: tokio::sync::Mutex::new(AllPermitsTracker {
+                wft_permits,
+                act_permits,
+                la_permits,
+            }),
         }
     }
 
@@ -460,6 +502,15 @@ impl Worker {
     /// completed
     async fn shutdown(&self) {
         self.initiate_shutdown();
+        if let Some(name) = self.workflows.get_sticky_queue_name() {
+            // This is a best effort call and we can still shutdown the worker if it fails
+            match self.client.shutdown_worker(name).await {
+                Err(err) if err.code() != tonic::Code::Unavailable => {
+                    warn!("Failed to shutdown sticky queue  {:?}", err);
+                }
+                _ => {}
+            }
+        }
         // We need to wait for all local activities to finish so no more workflow task heartbeats
         // will be generated
         self.local_act_mgr
@@ -474,6 +525,13 @@ impl Worker {
         if let Some(acts) = self.at_task_mgr.as_ref() {
             acts.shutdown().await;
         }
+        // Wait for all permits to be released, but don't totally hang real-world shutdown.
+        tokio::select! {
+            _ = async { self.all_permits_tracker.lock().await.all_done().await } => {},
+            _ = tokio::time::sleep(Duration::from_secs(1)) => {
+                dbg_panic!("Waiting for all slot permits to release took too long!");
+            }
+        };
     }
 
     /// Finish shutting down by consuming the background pollers and freeing all resources
@@ -750,7 +808,7 @@ mod tests {
     use crate::{
         advance_fut, test_help::test_worker_cfg, worker::client::mocks::mock_workflow_client,
     };
-    use futures::FutureExt;
+    use futures_util::FutureExt;
 
     use temporal_sdk_core_protos::temporal::api::workflowservice::v1::PollActivityTaskQueueResponse;
 

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

@@ -7,8 +7,6 @@ use crate::{
     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;
@@ -58,7 +56,7 @@ impl SlotTrait for Slot {
 pub(super) struct SlotProvider {
     namespace: String,
     task_queue: String,
-    wft_semaphore: Arc<MeteredPermitDealer<WorkflowSlotKind>>,
+    wft_semaphore: MeteredPermitDealer<WorkflowSlotKind>,
     external_wft_tx: WFTStreamSender,
 }
 
@@ -66,7 +64,7 @@ impl SlotProvider {
     pub(super) fn new(
         namespace: String,
         task_queue: String,
-        wft_semaphore: Arc<MeteredPermitDealer<WorkflowSlotKind>>,
+        wft_semaphore: MeteredPermitDealer<WorkflowSlotKind>,
         external_wft_tx: WFTStreamSender,
     ) -> Self {
         Self {
@@ -119,7 +117,7 @@ mod tests {
 
     #[tokio::test]
     async fn slot_propagates_through_channel() {
-        let wft_semaphore = Arc::new(fixed_size_permit_dealer(2));
+        let wft_semaphore = fixed_size_permit_dealer(2);
         let (external_wft_tx, mut external_wft_rx) = unbounded_channel();
 
         let provider = SlotProvider::new(
@@ -142,7 +140,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(fixed_size_permit_dealer(2));
+            let wft_semaphore = fixed_size_permit_dealer(2);
             let provider = SlotProvider::new(
                 "my_namespace".to_string(),
                 "my_queue".to_string(),
@@ -156,7 +154,7 @@ mod tests {
 
     #[tokio::test]
     async fn unused_slots_reclaimed() {
-        let wft_semaphore = Arc::new(fixed_size_permit_dealer(2));
+        let wft_semaphore = fixed_size_permit_dealer(2);
         {
             let wft_semaphore = wft_semaphore.clone();
             let (external_wft_tx, _) = unbounded_channel();

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

@@ -1,6 +1,7 @@
 use std::{marker::PhantomData, sync::Arc};
 use temporal_sdk_core_api::worker::{
-    SlotKind, SlotReservationContext, SlotSupplier, SlotSupplierPermit,
+    SlotKind, SlotMarkUsedContext, SlotReleaseContext, SlotReservationContext, SlotSupplier,
+    SlotSupplierPermit,
 };
 use tokio::sync::Semaphore;
 
@@ -42,9 +43,9 @@ where
         perm.ok().map(SlotSupplierPermit::with_user_data)
     }
 
-    fn mark_slot_used(&self, _info: SK::Info<'_>) {}
+    fn mark_slot_used(&self, _ctx: &dyn SlotMarkUsedContext<SlotKind = Self::SlotKind>) {}
 
-    fn release_slot(&self) {}
+    fn release_slot(&self, _ctx: &dyn SlotReleaseContext<SlotKind = Self::SlotKind>) {}
 
     fn available_slots(&self) -> Option<usize> {
         Some(self.sem.available_permits())