@temporalio/core-bridge 1.6.0 → 1.7.1

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

@@ -1,15 +1,33 @@
 use super::*;
 use futures::Future;
 
+pub(crate) static DEFAULT_TEST_CAPABILITIES: &Capabilities = &Capabilities {
+    signal_and_query_header: true,
+    internal_error_differentiation: true,
+    activity_failure_include_heartbeat: true,
+    supports_schedules: true,
+    encoded_failure_attributes: true,
+    build_id_based_versioning: true,
+    upsert_memo: true,
+    eager_workflow_start: true,
+    sdk_metadata: true,
+};
+
 #[cfg(test)]
 /// Create a mock client primed with basic necessary expectations
 pub(crate) fn mock_workflow_client() -> MockWorkerClient {
-    MockWorkerClient::new()
+    let mut r = MockWorkerClient::new();
+    r.expect_capabilities()
+        .returning(|| Some(DEFAULT_TEST_CAPABILITIES));
+    r
 }
 
 /// Create a mock manual client primed with basic necessary expectations
 pub(crate) fn mock_manual_workflow_client() -> MockManualWorkerClient {
-    MockManualWorkerClient::new()
+    let mut r = MockManualWorkerClient::new();
+    r.expect_capabilities()
+        .returning(|| Some(DEFAULT_TEST_CAPABILITIES));
+    r
 }
 
 // Need a version of the mock that can return futures so we can return potentially pending
@@ -83,5 +101,7 @@ mockall::mock! {
             query_result: QueryResult,
         ) -> impl Future<Output = Result<RespondQueryTaskCompletedResponse>> + Send + 'b
             where 'a: 'b, Self: 'b;
+
+        fn capabilities(&self) -> Option<&'static get_system_info_response::Capabilities>;
     }
 }

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

@@ -7,12 +7,13 @@ use temporal_sdk_core_protos::{
     coresdk::workflow_commands::QueryResult,
     temporal::api::{
         command::v1::Command,
-        common::v1::{Payloads, WorkflowExecution},
+        common::v1::{MeteringMetadata, Payloads, WorkflowExecution},
         enums::v1::{TaskQueueKind, WorkflowTaskFailedCause},
         failure::v1::Failure,
         query::v1::WorkflowQueryResult,
+        sdk::v1::WorkflowTaskCompletedMetadata,
         taskqueue::v1::{StickyExecutionAttributes, TaskQueue, TaskQueueMetadata, VersionId},
-        workflowservice::v1::*,
+        workflowservice::v1::{get_system_info_response::Capabilities, *},
     },
     TaskToken,
 };
@@ -109,6 +110,9 @@ pub(crate) trait WorkerClient: Sync + Send {
         task_token: TaskToken,
         query_result: QueryResult,
     ) -> Result<RespondQueryTaskCompletedResponse>;
+
+    #[allow(clippy::needless_lifetimes)] // Clippy is wrong here
+    fn capabilities<'a>(&'a self) -> Option<&'a get_system_info_response::Capabilities>;
 }
 
 #[async_trait::async_trait]
@@ -189,6 +193,7 @@ impl WorkerClient for WorkerClientBag {
             worker_versioning_id: Some(VersionId {
                 worker_build_id: self.versioning_build_id(),
             }),
+            messages: vec![],
             binary_checksum: self.worker_build_id.clone(),
             query_results: request
                 .query_responses
@@ -206,6 +211,8 @@ impl WorkerClient for WorkerClientBag {
                 })
                 .collect(),
             namespace: self.namespace.clone(),
+            sdk_metadata: Some(request.sdk_metadata),
+            metering_metadata: Some(request.metering_metadata),
         };
         Ok(self
             .client
@@ -302,6 +309,7 @@ impl WorkerClient for WorkerClientBag {
             identity: self.identity.clone(),
             binary_checksum: self.worker_build_id.clone(),
             namespace: self.namespace.clone(),
+            messages: vec![],
         };
         Ok(self
             .client
@@ -352,6 +360,10 @@ impl WorkerClient for WorkerClientBag {
             .await?
             .into_inner())
     }
+
+    fn capabilities(&self) -> Option<&Capabilities> {
+        self.client.get_client().inner().capabilities()
+    }
 }
 
 /// A version of [RespondWorkflowTaskCompletedRequest] that will finish being filled out by the
@@ -370,4 +382,8 @@ pub(crate) struct WorkflowTaskCompletion {
     pub return_new_workflow_task: bool,
     /// Force a new WFT to be created after this completion
     pub force_create_new_workflow_task: bool,
+    /// SDK-specific metadata to send
+    pub sdk_metadata: WorkflowTaskCompletedMetadata,
+    /// Metering info
+    pub metering_metadata: MeteringMetadata,
 }

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

@@ -21,9 +21,12 @@ use crate::{
         WorkflowTaskPoller,
     },
     protosext::{validate_activity_completion, ValidPollWFTQResponse},
-    telemetry::metrics::{
-        activity_poller, local_activity_worker_type, workflow_poller, workflow_sticky_poller,
-        MetricsContext,
+    telemetry::{
+        metrics::{
+            activity_poller, local_activity_worker_type, workflow_poller, workflow_sticky_poller,
+            MetricsContext,
+        },
+        TelemetryInstance,
     },
     worker::{
         activities::{DispatchOrTimeoutLA, LACompleteAction, LocalActivityManager},
@@ -34,7 +37,14 @@ use crate::{
 };
 use activities::{LocalInFlightActInfo, WorkerActivityTasks};
 use futures::Stream;
-use std::{convert::TryInto, future, sync::Arc};
+use std::{
+    convert::TryInto,
+    future,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+};
 use temporal_sdk_core_protos::{
     coresdk::{
         activity_result::activity_execution_result,
@@ -46,7 +56,7 @@ use temporal_sdk_core_protos::{
     temporal::api::{
         enums::v1::TaskQueueKind,
         taskqueue::v1::{StickyExecutionAttributes, TaskQueue},
-        workflowservice::v1::PollActivityTaskQueueResponse,
+        workflowservice::v1::{get_system_info_response, PollActivityTaskQueueResponse},
     },
     TaskToken,
 };
@@ -68,7 +78,11 @@ pub struct Worker {
     shutdown_token: CancellationToken,
     /// Will be called at the end of each activation completion
     #[allow(clippy::type_complexity)] // Sorry clippy, there's no simple way to re-use here.
-    post_activate_hook: Option<Box<dyn Fn(&Self, &str, usize) + Send + Sync>>,
+    post_activate_hook: Option<Box<dyn Fn(&Self, PostActivateHookData) + Send + Sync>>,
+    /// Set when non-local activities are complete and should stop being polled
+    non_local_activities_complete: Arc<AtomicBool>,
+    /// Set when local activities are complete and should stop being polled
+    local_activities_complete: Arc<AtomicBool>,
 }
 
 #[async_trait::async_trait]
@@ -135,7 +149,13 @@ impl WorkerTrait for Worker {
         self.shutdown_token.cancel();
         // First, we want to stop polling of both activity and workflow tasks
         if let Some(atm) = self.at_task_mgr.as_ref() {
-            atm.notify_shutdown();
+            atm.initiate_shutdown();
+        }
+        // Let the manager know that shutdown has been initiated to try to unblock the local
+        // activity poll in case this worker is an activity-only worker.
+        self.local_act_mgr.shutdown_initiated();
+        if !self.workflows.ever_polled() {
+            self.local_act_mgr.workflows_have_shutdown();
         }
         info!(
             task_queue=%self.config.task_queue,
@@ -159,11 +179,17 @@ impl Worker {
         config: WorkerConfig,
         sticky_queue_name: Option<String>,
         client: Arc<dyn WorkerClient>,
-        metrics: MetricsContext,
+        telem_instance: Option<&TelemetryInstance>,
     ) -> Self {
         info!(task_queue=%config.task_queue,
               namespace=%config.namespace,
               "Initializing worker");
+        let metrics = if let Some(ti) = telem_instance {
+            MetricsContext::top_level(config.namespace.clone(), ti)
+                .with_task_q(config.task_queue.clone())
+        } else {
+            MetricsContext::no_op()
+        };
         metrics.worker_registered();
 
         let shutdown_token = CancellationToken::new();
@@ -226,15 +252,17 @@ impl Worker {
             wft_stream,
             act_poll_buffer,
             metrics,
+            telem_instance,
             shutdown_token,
         )
     }
 
     #[cfg(test)]
     pub(crate) fn new_test(config: WorkerConfig, client: impl WorkerClient + 'static) -> Self {
-        Self::new(config, None, Arc::new(client), MetricsContext::no_op())
+        Self::new(config, None, Arc::new(client), None)
     }
 
+    #[allow(clippy::too_many_arguments)] // Not much worth combining here
     pub(crate) fn new_with_pollers(
         mut config: WorkerConfig,
         sticky_queue_name: Option<String>,
@@ -242,6 +270,7 @@ impl Worker {
         wft_stream: impl Stream<Item = Result<ValidPollWFTQResponse, tonic::Status>> + Send + 'static,
         act_poller: Option<BoxedActPoller>,
         metrics: MetricsContext,
+        telem_instance: Option<&TelemetryInstance>,
         shutdown_token: CancellationToken,
     ) -> Self {
         let (hb_tx, hb_rx) = unbounded_channel();
@@ -251,7 +280,6 @@ impl Worker {
             hb_tx,
             metrics.with_new_attrs([local_activity_worker_type()]),
         ));
-        let lam_clone = local_act_mgr.clone();
         let at_task_mgr = act_poller.map(|ap| {
             WorkerActivityTasks::new(
                 config.max_outstanding_activities,
@@ -261,16 +289,23 @@ impl Worker {
                 metrics.clone(),
                 config.max_heartbeat_throttle_interval,
                 config.default_heartbeat_throttle_interval,
+                config.graceful_shutdown_period,
             )
         });
-        if at_task_mgr.is_none() {
+        let poll_on_non_local_activities = at_task_mgr.is_some();
+        if !poll_on_non_local_activities {
             info!("Activity polling is disabled for this worker");
-        }
-        let la_sink = LAReqSink::new(lam_clone, config.wf_state_inputs.clone());
+        };
+        let la_sink = LAReqSink::new(local_act_mgr.clone(), config.wf_state_inputs.clone());
         Self {
             wf_client: client.clone(),
             workflows: Workflows::new(
-                build_wf_basics(&mut config, metrics, shutdown_token.child_token()),
+                build_wf_basics(
+                    &mut config,
+                    metrics,
+                    shutdown_token.child_token(),
+                    client.capabilities().cloned().unwrap_or_default(),
+                ),
                 sticky_queue_name.map(|sq| StickyExecutionAttributes {
                     worker_task_queue: Some(TaskQueue {
                         name: sq,
@@ -286,16 +321,21 @@ impl Worker {
                 client,
                 wft_stream,
                 la_sink,
+                local_act_mgr.clone(),
                 hb_rx,
                 at_task_mgr
                     .as_ref()
                     .map(|mgr| mgr.get_handle_for_workflows()),
+                telem_instance,
             ),
             at_task_mgr,
             local_act_mgr,
             config,
             shutdown_token,
             post_activate_hook: None,
+            // Complete if there configured not to poll on non-local activities.
+            non_local_activities_complete: Arc::new(AtomicBool::new(!poll_on_non_local_activities)),
+            local_activities_complete: Default::default(),
         }
     }
 
@@ -303,8 +343,8 @@ impl Worker {
     /// completed
     async fn shutdown(&self) {
         self.initiate_shutdown();
-        // Next we need to wait for all local activities to finish so no more workflow task
-        // heartbeats will be generated
+        // We need to wait for all local activities to finish so no more workflow task heartbeats
+        // will be generated
         self.local_act_mgr
             .wait_all_outstanding_tasks_finished()
             .await;
@@ -315,12 +355,13 @@ impl Worker {
             .expect("Workflow processing terminates cleanly");
         // Wait for activities to finish
         if let Some(acts) = self.at_task_mgr.as_ref() {
-            acts.wait_all_finished().await;
+            acts.shutdown().await;
         }
     }
 
     /// Finish shutting down by consuming the background pollers and freeing all resources
     async fn finalize_shutdown(self) {
+        self.shutdown().await;
         if let Some(b) = self.at_task_mgr {
             b.shutdown().await;
         }
@@ -360,9 +401,27 @@ impl Worker {
     /// Returns `Ok(None)` in the event of a poll timeout or if the polling loop should otherwise
     /// be restarted
     async fn activity_poll(&self) -> Result<Option<ActivityTask>, PollActivityError> {
+        let local_activities_complete = self.local_activities_complete.load(Ordering::Relaxed);
+        let non_local_activities_complete =
+            self.non_local_activities_complete.load(Ordering::Relaxed);
+        if local_activities_complete && non_local_activities_complete {
+            return Err(PollActivityError::ShutDown);
+        }
         let act_mgr_poll = async {
+            if non_local_activities_complete {
+                future::pending::<()>().await;
+                unreachable!()
+            }
             if let Some(ref act_mgr) = self.at_task_mgr {
-                act_mgr.poll().await
+                let res = act_mgr.poll().await;
+                if let Err(err) = res.as_ref() {
+                    if matches!(err, PollActivityError::ShutDown) {
+                        self.non_local_activities_complete
+                            .store(true, Ordering::Relaxed);
+                        return Ok(None);
+                    }
+                };
+                res.map(Some)
             } else {
                 // We expect the local activity branch below to produce shutdown when appropriate if
                 // there are no activity pollers.
@@ -370,26 +429,35 @@ impl Worker {
                 unreachable!()
             }
         };
+        let local_activities_poll = async {
+            if local_activities_complete {
+                future::pending::<()>().await;
+                unreachable!()
+            }
+            match self.local_act_mgr.next_pending().await {
+                Some(DispatchOrTimeoutLA::Dispatch(r)) => Ok(Some(r)),
+                Some(DispatchOrTimeoutLA::Timeout {
+                    run_id,
+                    resolution,
+                    task,
+                }) => {
+                    self.notify_local_result(&run_id, LocalResolution::LocalActivity(resolution));
+                    Ok(task)
+                }
+                None => {
+                    if self.shutdown_token.is_cancelled() {
+                        self.local_activities_complete
+                            .store(true, Ordering::Relaxed);
+                    }
+                    Ok(None)
+                }
+            }
+        };
 
         tokio::select! {
             biased;
 
-            r = self.local_act_mgr.next_pending() => {
-                match r {
-                    Some(DispatchOrTimeoutLA::Dispatch(r)) => Ok(Some(r)),
-                    Some(DispatchOrTimeoutLA::Timeout { run_id, resolution, task }) => {
-                        self.notify_local_result(
-                            &run_id, LocalResolution::LocalActivity(resolution));
-                        Ok(task)
-                    },
-                    None => {
-                        if self.shutdown_token.is_cancelled() {
-                            return Err(PollActivityError::ShutDown);
-                        }
-                        Ok(None)
-                    }
-                }
-            },
+            r = local_activities_poll => r,
             r = act_mgr_poll => r,
         }
     }
@@ -453,6 +521,7 @@ impl Worker {
         // about to happen anyway. Tell the local activity manager that, so that it can know to
         // cancel any remaining outstanding LAs and shutdown.
         if matches!(r, Err(PollWfError::ShutDown)) {
+            // This is covering the situation where WFT pollers dying is the reason for shutdown
             self.initiate_shutdown();
             self.local_act_mgr.workflows_have_shutdown();
         }
@@ -466,11 +535,15 @@ impl Worker {
         &self,
         completion: WorkflowActivationCompletion,
     ) -> Result<(), CompleteWfError> {
-        let run_id = completion.run_id.clone();
-        let most_recent_event = self.workflows.activation_completed(completion).await?;
-        if let Some(h) = &self.post_activate_hook {
-            h(self, &run_id, most_recent_event);
-        }
+        self.workflows
+            .activation_completed(
+                completion,
+                false,
+                self.post_activate_hook
+                    .as_ref()
+                    .map(|h| |data: PostActivateHookData| h(self, data)),
+            )
+            .await?;
         Ok(())
     }
 
@@ -487,7 +560,7 @@ impl Worker {
     /// Sets a function to be called at the end of each activation completion
     pub(crate) fn set_post_activate_hook(
         &mut self,
-        callback: impl Fn(&Self, &str, usize) + Send + Sync + 'static,
+        callback: impl Fn(&Self, PostActivateHookData) + Send + Sync + 'static,
     ) {
         self.post_activate_hook = Some(Box::new(callback))
     }
@@ -516,10 +589,17 @@ impl Worker {
     }
 }
 
+pub struct PostActivateHookData<'a> {
+    pub run_id: &'a str,
+    pub most_recent_event: usize,
+    pub replaying: bool,
+}
+
 fn build_wf_basics(
     config: &mut WorkerConfig,
     metrics: MetricsContext,
     shutdown_token: CancellationToken,
+    server_capabilities: get_system_info_response::Capabilities,
 ) -> WorkflowBasics {
     WorkflowBasics {
         max_cached_workflows: config.max_cached_workflows,
@@ -530,6 +610,7 @@ fn build_wf_basics(
         task_queue: config.task_queue.clone(),
         ignore_evicts_on_shutdown: config.ignore_evicts_on_shutdown,
         fetching_concurrency: config.fetching_concurrency,
+        server_capabilities,
         #[cfg(feature = "save_wf_inputs")]
         wf_state_inputs: config.wf_state_inputs.take(),
     }
@@ -538,11 +619,15 @@ fn build_wf_basics(
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::{test_help::test_worker_cfg, worker::client::mocks::mock_workflow_client};
+    use crate::{
+        advance_fut, test_help::test_worker_cfg, worker::client::mocks::mock_workflow_client,
+    };
+    use futures::FutureExt;
+
     use temporal_sdk_core_protos::temporal::api::workflowservice::v1::PollActivityTaskQueueResponse;
 
     #[tokio::test]
-    async fn activity_timeouts_dont_eat_permits() {
+    async fn activity_timeouts_maintain_permit() {
         let mut mock_client = mock_workflow_client();
         mock_client
             .expect_poll_activity_task()
@@ -553,8 +638,16 @@ mod tests {
             .build()
             .unwrap();
         let worker = Worker::new_test(cfg, mock_client);
-        assert_eq!(worker.activity_poll().await.unwrap(), None);
-        assert_eq!(worker.at_task_mgr.unwrap().remaining_activity_capacity(), 5);
+        let fut = worker.poll_activity_task();
+        advance_fut!(fut);
+        assert_eq!(
+            worker
+                .at_task_mgr
+                .as_ref()
+                .unwrap()
+                .remaining_activity_capacity(),
+            4
+        );
     }
 
     #[tokio::test]