temporalio 0.0.2 → 0.1.0

data/bridge/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
@@ -66,6 +79,10 @@ pub struct Worker {
     /// 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>>,
+    /// 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]
@@ -134,6 +151,9 @@ impl WorkerTrait for Worker {
         if let Some(atm) = self.at_task_mgr.as_ref() {
             atm.notify_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();
         info!(
             task_queue=%self.config.task_queue,
             namespace=%self.config.namespace,
@@ -156,11 +176,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 +249,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,
@@ -259,18 +288,20 @@ impl Worker {
                 config.default_heartbeat_throttle_interval,
             )
         });
+        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 +316,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(),
         }
     }
 
@@ -304,15 +341,16 @@ 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;
         // Wait for workflows to finish
         self.workflows
             .shutdown()
             .await
             .expect("Workflow processing terminates cleanly");
+        let lam = self.local_act_mgr.clone();
+        lam.wait_all_outstanding_tasks_finished().await;
         // Wait for activities to finish
         if let Some(acts) = self.at_task_mgr.as_ref() {
-            acts.wait_all_finished().await;
+            acts.shutdown().await;
         }
     }
 
@@ -346,13 +384,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 +395,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 +510,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 +529,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| {
+                    |run_id: &str, most_recent_event: usize| h(self, run_id, most_recent_event)
+                }),
+            )
+            .await?;
         Ok(())
     }
 
@@ -508,14 +582,39 @@ impl Worker {
     }
 }
 
+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 +625,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]

data/bridge/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

data/bridge/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>;
 }