@temporalio/core-bridge 1.5.2 → 1.6.0

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,
@@ -26,13 +28,13 @@ use crate::{
     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::Arc};
 use temporal_sdk_core_protos::{
     coresdk::{
         activity_result::activity_execution_result,
@@ -48,6 +50,7 @@ use temporal_sdk_core_protos::{
     },
     TaskToken,
 };
+use tokio::sync::mpsc::unbounded_channel;
 use tokio_util::sync::CancellationToken;
 
 /// A worker polls on a certain task queue
@@ -233,7 +236,7 @@ impl Worker {
     }
 
     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,
@@ -241,13 +244,14 @@ impl Worker {
         metrics: MetricsContext,
         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,
@@ -259,18 +263,14 @@ impl Worker {
                 config.default_heartbeat_throttle_interval,
             )
         });
+        if at_task_mgr.is_none() {
+            info!("Activity polling is disabled for this worker");
+        }
+        let la_sink = LAReqSink::new(lam_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(),
-                    metrics,
-                    namespace: config.namespace.clone(),
-                    task_queue: config.task_queue.clone(),
-                    ignore_evicts_on_shutdown: config.ignore_evicts_on_shutdown,
-                },
+                build_wf_basics(&mut config, metrics, shutdown_token.child_token()),
                 sticky_queue_name.map(|sq| StickyExecutionAttributes {
                     worker_task_queue: Some(TaskQueue {
                         name: sq,
@@ -285,7 +285,8 @@ impl Worker {
                 }),
                 client,
                 wft_stream,
-                local_act_req_sink,
+                la_sink,
+                hb_rx,
                 at_task_mgr
                     .as_ref()
                     .map(|mgr| mgr.get_handle_for_workflows()),
@@ -304,7 +305,9 @@ impl Worker {
         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;
+        self.local_act_mgr
+            .wait_all_outstanding_tasks_finished()
+            .await;
         // Wait for workflows to finish
         self.workflows
             .shutdown()
@@ -346,13 +349,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
@@ -365,9 +364,10 @@ impl Worker {
             if let Some(ref act_mgr) = self.at_task_mgr {
                 act_mgr.poll().await
             } 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!()
             }
         };
 
@@ -448,7 +448,15 @@ 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)) {
+            self.initiate_shutdown();
+            self.local_act_mgr.workflows_have_shutdown();
+        }
+        r
     }
 
     #[instrument(skip(self, completion),
@@ -498,7 +506,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,6 +516,25 @@ impl Worker {
     }
 }
 
+fn build_wf_basics(
+    config: &mut WorkerConfig,
+    metrics: MetricsContext,
+    shutdown_token: CancellationToken,
+) -> 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,
+        #[cfg(feature = "save_wf_inputs")]
+        wf_state_inputs: config.wf_state_inputs.take(),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

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>;
 }