@temporalio/core-bridge 1.5.2 → 1.7.0

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

@@ -3,6 +3,8 @@ pub(crate) mod client;
 mod workflow;
 
 pub use temporal_sdk_core_api::worker::{WorkerConfig, WorkerConfigBuilder};
+#[cfg(feature = "save_wf_inputs")]
+pub use workflow::replay_wf_state_inputs;
 
 pub(crate) use activities::{
     ExecutingLAId, LocalActRequest, LocalActivityExecutionResult, LocalActivityResolution,
@@ -19,20 +21,30 @@ 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},
         client::WorkerClient,
-        workflow::{LocalResolution, WorkflowBasics, Workflows},
+        workflow::{LAReqSink, LocalResolution, WorkflowBasics, Workflows},
     },
     ActivityHeartbeat, CompleteActivityError, PollActivityError, PollWfError, WorkerTrait,
 };
 use activities::{LocalInFlightActInfo, WorkerActivityTasks};
 use futures::Stream;
-use std::{convert::TryInto, sync::Arc};
+use std::{
+    convert::TryInto,
+    future,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+};
 use temporal_sdk_core_protos::{
     coresdk::{
         activity_result::activity_execution_result,
@@ -44,10 +56,11 @@ 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,
 };
+use tokio::sync::mpsc::unbounded_channel;
 use tokio_util::sync::CancellationToken;
 
 /// A worker polls on a certain task queue
@@ -65,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]
@@ -132,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,
@@ -156,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();
@@ -223,31 +252,34 @@ 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(
-        config: WorkerConfig,
+        mut config: WorkerConfig,
         sticky_queue_name: Option<String>,
         client: Arc<dyn WorkerClient>,
         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();
         let local_act_mgr = Arc::new(LocalActivityManager::new(
             config.max_outstanding_local_activities,
             config.namespace.clone(),
+            hb_tx,
             metrics.with_new_attrs([local_activity_worker_type()]),
         ));
-        let lam_clone = local_act_mgr.clone();
-        let local_act_req_sink = move |requests| lam_clone.enqueue(requests);
         let at_task_mgr = act_poller.map(|ap| {
             WorkerActivityTasks::new(
                 config.max_outstanding_activities,
@@ -257,20 +289,23 @@ impl Worker {
                 metrics.clone(),
                 config.max_heartbeat_throttle_interval,
                 config.default_heartbeat_throttle_interval,
+                config.graceful_shutdown_period,
             )
         });
+        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(local_act_mgr.clone(), config.wf_state_inputs.clone());
         Self {
             wf_client: client.clone(),
             workflows: Workflows::new(
-                WorkflowBasics {
-                    max_cached_workflows: config.max_cached_workflows,
-                    max_outstanding_wfts: config.max_outstanding_workflow_tasks,
-                    shutdown_token: shutdown_token.child_token(),
+                build_wf_basics(
+                    &mut config,
                     metrics,
-                    namespace: config.namespace.clone(),
-                    task_queue: config.task_queue.clone(),
-                    ignore_evicts_on_shutdown: config.ignore_evicts_on_shutdown,
-                },
+                    shutdown_token.child_token(),
+                    client.capabilities().cloned().unwrap_or_default(),
+                ),
                 sticky_queue_name.map(|sq| StickyExecutionAttributes {
                     worker_task_queue: Some(TaskQueue {
                         name: sq,
@@ -285,16 +320,22 @@ impl Worker {
                 }),
                 client,
                 wft_stream,
-                local_act_req_sink,
+                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(),
         }
     }
 
@@ -302,9 +343,11 @@ 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
-        self.local_act_mgr.shutdown_and_wait_all_finished().await;
+        // 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;
         // Wait for workflows to finish
         self.workflows
             .shutdown()
@@ -312,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;
         }
@@ -346,13 +390,9 @@ impl Worker {
             .unwrap_or_default()
     }
 
-    #[cfg(test)]
+    #[allow(unused)]
     pub(crate) async fn available_wft_permits(&self) -> usize {
-        self.workflows
-            .get_state_info()
-            .await
-            .expect("You can only check for available permits before shutdown")
-            .available_wft_permits
+        self.workflows.available_wft_permits()
     }
 
     /// Get new activity tasks (may be local or nonlocal). Local activities are returned first
@@ -361,35 +401,63 @@ 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 {
-                info!("Activity polling is disabled for this worker");
-                self.shutdown_token.cancelled().await;
-                Err(PollActivityError::ShutDown)
+                // We expect the local activity branch below to produce shutdown when appropriate if
+                // there are no activity pollers.
+                future::pending::<()>().await;
+                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,
         }
     }
@@ -448,7 +516,16 @@ impl Worker {
 
     #[instrument(skip(self), fields(run_id, workflow_id, task_queue=%self.config.task_queue))]
     pub(crate) async fn next_workflow_activation(&self) -> Result<WorkflowActivation, PollWfError> {
-        self.workflows.next_workflow_activation().await
+        let r = self.workflows.next_workflow_activation().await;
+        // In the event workflows are shutdown, begin shutdown of everything else, since that's
+        // 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();
+        }
+        r
     }
 
     #[instrument(skip(self, completion),
@@ -458,11 +535,14 @@ 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,
+                self.post_activate_hook
+                    .as_ref()
+                    .map(|h| |data: PostActivateHookData| h(self, data)),
+            )
+            .await?;
         Ok(())
     }
 
@@ -479,7 +559,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))
     }
@@ -498,7 +578,7 @@ impl Worker {
                 runtime: info.dispatch_time.elapsed(),
                 attempt: info.attempt,
                 backoff,
-                original_schedule_time: Some(info.la_info.schedule_time),
+                original_schedule_time: info.la_info.schedule_cmd.original_schedule_time,
             }),
         )
     }
@@ -508,14 +588,45 @@ 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,
+        max_outstanding_wfts: config.max_outstanding_workflow_tasks,
+        shutdown_token,
+        metrics,
+        namespace: config.namespace.clone(),
+        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(),
+    }
+}
+
 #[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()
@@ -526,8 +637,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]

package/sdk-core/core/src/worker/workflow/bridge.rs

@@ -25,11 +25,9 @@ impl WorkflowBridge {
     }
 }
 
-#[async_trait::async_trait]
 impl WorkflowFetcher for WorkflowBridge {
-    async fn fetch_workflow_iteration_output(&mut self) -> Vec<WFCommand> {
+    fn fetch_workflow_iteration_output(&mut self) -> Vec<WFCommand> {
         let in_cmds = self.incoming_commands.try_recv();
-
         let in_cmds = in_cmds.unwrap_or_else(|_| vec![WFCommand::NoCommandsFromLang]);
         debug!(in_cmds = %in_cmds.display(), "wf bridge iteration fetch");
         in_cmds

package/sdk-core/core/src/worker/workflow/driven_workflow.rs

@@ -78,16 +78,14 @@ impl DrivenWorkflow {
     }
 }
 
-#[async_trait::async_trait]
 impl WorkflowFetcher for DrivenWorkflow {
-    async fn fetch_workflow_iteration_output(&mut self) -> Vec<WFCommand> {
-        self.fetcher.fetch_workflow_iteration_output().await
+    fn fetch_workflow_iteration_output(&mut self) -> Vec<WFCommand> {
+        self.fetcher.fetch_workflow_iteration_output()
     }
 }
 
 /// Implementors of this trait represent a way to fetch output from executing/iterating some
 /// workflow code (or a mocked workflow).
-#[async_trait::async_trait]
 pub trait WorkflowFetcher: Send {
     /// Obtain any output from the workflow's recent execution(s). Because the lang sdk is
     /// responsible for calling workflow code as a result of receiving tasks from
@@ -97,5 +95,5 @@ pub trait WorkflowFetcher: Send {
     ///
     /// In the case of the real [WorkflowBridge] implementation, commands are simply pulled from
     /// a buffer that the language side sinks into when it calls [crate::Core::complete_task]
-    async fn fetch_workflow_iteration_output(&mut self) -> Vec<WFCommand>;
+    fn fetch_workflow_iteration_output(&mut self) -> Vec<WFCommand>;
 }