temporalio 0.0.1 → 0.1.0

data/bridge/sdk-core/core/src/worker/client.rs

@@ -2,16 +2,18 @@
 
 pub(crate) mod mocks;
 
-use temporal_client::{Client, RetryClient, WorkflowService, WorkflowTaskCompletion};
+use temporal_client::{Client, RetryClient, WorkflowService};
 use temporal_sdk_core_protos::{
     coresdk::workflow_commands::QueryResult,
     temporal::api::{
-        common::v1::{Payloads, WorkflowExecution},
+        command::v1::Command,
+        common::v1::{MeteringMetadata, Payloads, WorkflowExecution},
         enums::v1::{TaskQueueKind, WorkflowTaskFailedCause},
         failure::v1::Failure,
         query::v1::WorkflowQueryResult,
-        taskqueue::v1::{TaskQueue, TaskQueueMetadata},
-        workflowservice::v1::*,
+        sdk::v1::WorkflowTaskCompletedMetadata,
+        taskqueue::v1::{StickyExecutionAttributes, TaskQueue, TaskQueueMetadata, VersionId},
+        workflowservice::v1::{get_system_info_response::Capabilities, *},
     },
     TaskToken,
 };
@@ -108,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]
@@ -133,7 +138,9 @@ impl WorkerClient for WorkerClientBag {
             } else {
                 self.worker_build_id.clone()
             },
-            worker_versioning_build_id: self.versioning_build_id(),
+            worker_versioning_id: Some(VersionId {
+                worker_build_id: self.versioning_build_id(),
+            }),
         };
 
         Ok(self
@@ -159,7 +166,9 @@ impl WorkerClient for WorkerClientBag {
             task_queue_metadata: max_tasks_per_sec.map(|tps| TaskQueueMetadata {
                 max_tasks_per_second: Some(tps),
             }),
-            worker_versioning_build_id: self.versioning_build_id(),
+            worker_versioning_id: Some(VersionId {
+                worker_build_id: self.versioning_build_id(),
+            }),
         };
 
         Ok(self
@@ -181,6 +190,10 @@ impl WorkerClient for WorkerClientBag {
             sticky_attributes: request.sticky_attributes,
             return_new_workflow_task: request.return_new_workflow_task,
             force_create_new_workflow_task: request.force_create_new_workflow_task,
+            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
@@ -198,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
@@ -294,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
@@ -344,4 +360,30 @@ 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
+/// server client
+#[derive(Debug, Clone, PartialEq)]
+pub(crate) struct WorkflowTaskCompletion {
+    /// The task token that would've been received from polling for a workflow activation
+    pub task_token: TaskToken,
+    /// A list of new commands to send to the server, such as starting a timer.
+    pub commands: Vec<Command>,
+    /// If set, indicate that next task should be queued on sticky queue with given attributes.
+    pub sticky_attributes: Option<StickyExecutionAttributes>,
+    /// Responses to queries in the `queries` field of the workflow task.
+    pub query_responses: Vec<QueryResult>,
+    /// Indicate that the task completion should return a new WFT if one is available
+    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,
 }

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]
@@ -74,7 +91,7 @@ impl WorkerTrait for Worker {
         self.next_workflow_activation().await
     }
 
-    #[instrument(level = "debug", skip(self))]
+    #[instrument(skip(self))]
     async fn poll_activity_task(&self) -> Result<ActivityTask, PollActivityError> {
         loop {
             match self.activity_poll().await.transpose() {
@@ -94,7 +111,6 @@ impl WorkerTrait for Worker {
         self.complete_workflow_activation(completion).await
     }
 
-    #[instrument(level = "debug", skip(self, completion), fields(completion=%&completion))]
     async fn complete_activity_task(
         &self,
         completion: ActivityTaskCompletion,
@@ -135,7 +151,14 @@ impl WorkerTrait for Worker {
         if let Some(atm) = self.at_task_mgr.as_ref() {
             atm.notify_shutdown();
         }
-        info!("Initiated 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,
+            "Initiated shutdown",
+        );
     }
 
     async fn shutdown(&self) {
@@ -153,9 +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, "Initializing worker");
+        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();
@@ -218,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), Default::default())
+        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,
@@ -254,17 +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(),
-                },
+                    shutdown_token.child_token(),
+                    client.capabilities().cloned().unwrap_or_default(),
+                ),
                 sticky_queue_name.map(|sq| StickyExecutionAttributes {
                     worker_task_queue: Some(TaskQueue {
                         name: sq,
@@ -279,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(),
         }
     }
 
@@ -298,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;
         }
     }
 
@@ -317,6 +361,10 @@ impl Worker {
         }
     }
 
+    pub(crate) fn shutdown_token(&self) -> CancellationToken {
+        self.shutdown_token.clone()
+    }
+
     /// Returns number of currently cached workflows
     pub async fn cached_workflows(&self) -> usize {
         self.workflows
@@ -336,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
@@ -351,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,
         }
     }
@@ -394,6 +466,9 @@ impl Worker {
         }
     }
 
+    #[instrument(skip(self, task_token, status),
+                 fields(task_token=%&task_token, status=%&status,
+                        task_queue=%self.config.task_queue, workflow_id, run_id))]
     pub(crate) async fn complete_activity(
         &self,
         task_token: TaskToken,
@@ -433,22 +508,35 @@ impl Worker {
         Ok(())
     }
 
-    #[instrument(level = "debug", skip(self), fields(run_id))]
+    #[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(level = "debug", skip(self, completion),
-      fields(completion=%&completion, run_id=%completion.run_id))]
+    #[instrument(skip(self, completion),
+                 fields(completion=%&completion, run_id=%completion.run_id, workflow_id,
+                        task_queue=%self.config.task_queue))]
     pub(crate) async fn complete_workflow_activation(
         &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(())
     }
 
@@ -470,16 +558,6 @@ impl Worker {
         self.post_activate_hook = Some(Box::new(callback))
     }
 
-    /// Used for replay workers - causes the worker to shutdown when the given run reaches the
-    /// given event number
-    pub(crate) fn set_shutdown_on_run_reaches_event(&mut self, run_id: String, last_event: i64) {
-        self.set_post_activate_hook(move |worker, activated_run_id, last_processed_event| {
-            if activated_run_id == run_id && last_processed_event >= last_event as usize {
-                worker.initiate_shutdown();
-            }
-        });
-    }
-
     fn complete_local_act(
         &self,
         la_res: LocalActivityExecutionResult,
@@ -494,7 +572,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,
             }),
         )
     }
@@ -504,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()
@@ -522,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