@temporalio/core-bridge 1.13.1 → 1.14.0

package/sdk-core/{core → crates/sdk-core}/src/worker/client.rs

@@ -1,41 +1,49 @@
 //! Worker-specific client needs
 
 pub(crate) mod mocks;
-use crate::abstractions::dbg_panic;
 use crate::protosext::legacy_query_failure;
-use crate::worker::heartbeat::HeartbeatFn;
-use parking_lot::RwLock;
-use std::sync::OnceLock;
-use std::{sync::Arc, time::Duration};
-use temporal_client::{
-    Client, IsWorkerTaskLongPoll, Namespace, NamespacedClient, NoRetryOnMatching, RetryClient,
-    SlotManager, WorkflowService,
+use parking_lot::Mutex;
+use prost_types::Duration as PbDuration;
+use std::{
+    collections::HashMap,
+    sync::Arc,
+    time::{Duration, SystemTime},
 };
-use temporal_sdk_core_api::worker::WorkerVersioningStrategy;
-use temporal_sdk_core_protos::temporal::api::worker::v1::WorkerHeartbeat;
-use temporal_sdk_core_protos::{
-    TaskToken,
-    coresdk::{workflow_commands::QueryResult, workflow_completion},
-    temporal::api::{
-        command::v1::Command,
-        common::v1::{
-            MeteringMetadata, Payloads, WorkerVersionCapabilities, WorkerVersionStamp,
-            WorkflowExecution,
-        },
-        deployment,
-        enums::v1::{
-            TaskQueueKind, VersioningBehavior, WorkerVersioningMode, WorkflowTaskFailedCause,
+use temporalio_client::{
+    Client, Namespace, NamespacedClient, RetryClient, RetryOptions, SharedReplaceableClient,
+    WorkflowService,
+    request_extensions::{IsWorkerTaskLongPoll, NoRetryOnMatching, RetryConfigForCall},
+    worker::ClientWorkerSet,
+};
+use temporalio_common::{
+    protos::{
+        TaskToken,
+        coresdk::{workflow_commands::QueryResult, workflow_completion},
+        temporal::api::{
+            command::v1::Command,
+            common::v1::{
+                MeteringMetadata, Payloads, WorkerVersionCapabilities, WorkerVersionStamp,
+                WorkflowExecution,
+            },
+            deployment,
+            enums::v1::{
+                TaskQueueKind, VersioningBehavior, WorkerStatus, WorkerVersioningMode,
+                WorkflowTaskFailedCause,
+            },
+            failure::v1::Failure,
+            nexus,
+            protocol::v1::Message as ProtocolMessage,
+            query::v1::WorkflowQueryResult,
+            sdk::v1::WorkflowTaskCompletedMetadata,
+            taskqueue::v1::{StickyExecutionAttributes, TaskQueue, TaskQueueMetadata},
+            worker::v1::{WorkerHeartbeat, WorkerSlotsInfo},
+            workflowservice::v1::{get_system_info_response::Capabilities, *},
         },
-        failure::v1::Failure,
-        nexus,
-        protocol::v1::Message as ProtocolMessage,
-        query::v1::WorkflowQueryResult,
-        sdk::v1::WorkflowTaskCompletedMetadata,
-        taskqueue::v1::{StickyExecutionAttributes, TaskQueue, TaskQueueMetadata},
-        workflowservice::v1::{get_system_info_response::Capabilities, *},
     },
+    worker::WorkerVersioningStrategy,
 };
 use tonic::IntoRequest;
+use uuid::Uuid;
 
 type Result<T, E = tonic::Status> = std::result::Result<T, E>;
 
@@ -46,34 +54,29 @@ pub enum LegacyQueryResult {
 
 /// Contains everything a worker needs to interact with the server
 pub(crate) struct WorkerClientBag {
-    replaceable_client: RwLock<RetryClient<Client>>,
+    client: RetryClient<SharedReplaceableClient<Client>>,
     namespace: String,
     identity: String,
     worker_versioning_strategy: WorkerVersioningStrategy,
-    heartbeat_data: Option<Arc<OnceLock<HeartbeatFn>>>,
+    worker_heartbeat_map: Arc<Mutex<HashMap<String, ClientHeartbeatData>>>,
 }
 
 impl WorkerClientBag {
     pub(crate) fn new(
-        client: RetryClient<Client>,
+        client: RetryClient<SharedReplaceableClient<Client>>,
         namespace: String,
         identity: String,
         worker_versioning_strategy: WorkerVersioningStrategy,
-        heartbeat_data: Option<Arc<OnceLock<HeartbeatFn>>>,
     ) -> Self {
         Self {
-            replaceable_client: RwLock::new(client),
+            client,
             namespace,
             identity,
             worker_versioning_strategy,
-            heartbeat_data,
+            worker_heartbeat_map: Arc::new(Mutex::new(HashMap::new())),
         }
     }
 
-    fn cloned_client(&self) -> RetryClient<Client> {
-        self.replaceable_client.read().clone()
-    }
-
     fn default_capabilities(&self) -> Capabilities {
         self.capabilities().unwrap_or_default()
     }
@@ -127,24 +130,11 @@ impl WorkerClientBag {
             None
         }
     }
-
-    fn capture_heartbeat(&self) -> Option<WorkerHeartbeat> {
-        if let Some(heartbeat_data) = self.heartbeat_data.as_ref() {
-            if let Some(hb) = heartbeat_data.get() {
-                hb()
-            } else {
-                dbg_panic!("Heartbeat function never set");
-                None
-            }
-        } else {
-            None
-        }
-    }
 }
 
 /// This trait contains everything workers need to interact with Temporal, and hence provides a
-/// minimal mocking surface. Delegates to [WorkflowClientTrait] so see that for details.
-#[cfg_attr(test, mockall::automock)]
+/// minimal mocking surface.
+#[cfg_attr(any(feature = "test-utilities", test), mockall::automock)]
 #[async_trait::async_trait]
 pub trait WorkerClient: Sync + Send {
     /// Poll workflow tasks
@@ -163,6 +153,7 @@ pub trait WorkerClient: Sync + Send {
     async fn poll_nexus_task(
         &self,
         poll_options: PollOptions,
+        send_heartbeat: bool,
     ) -> Result<PollNexusTaskQueueResponse>;
     /// Complete a workflow task
     async fn complete_workflow_task(
@@ -228,25 +219,35 @@ pub trait WorkerClient: Sync + Send {
     /// Describe the namespace
     async fn describe_namespace(&self) -> Result<DescribeNamespaceResponse>;
     /// Shutdown the worker
-    async fn shutdown_worker(&self, sticky_task_queue: String) -> Result<ShutdownWorkerResponse>;
+    async fn shutdown_worker(
+        &self,
+        sticky_task_queue: String,
+        final_heartbeat: Option<WorkerHeartbeat>,
+    ) -> Result<ShutdownWorkerResponse>;
     /// Record a worker heartbeat
     async fn record_worker_heartbeat(
         &self,
-        heartbeat: WorkerHeartbeat,
+        namespace: String,
+        worker_heartbeat: Vec<WorkerHeartbeat>,
     ) -> Result<RecordWorkerHeartbeatResponse>;
 
     /// Replace the underlying client
-    fn replace_client(&self, new_client: RetryClient<Client>);
+    fn replace_client(&self, new_client: Client);
     /// Return server capabilities
     fn capabilities(&self) -> Option<Capabilities>;
     /// Return workers using this client
-    fn workers(&self) -> Arc<SlotManager>;
+    fn workers(&self) -> Arc<ClientWorkerSet>;
     /// Indicates if this is a mock client
     fn is_mock(&self) -> bool;
     /// Return name and version of the SDK
     fn sdk_name_and_version(&self) -> (String, String);
     /// Get worker identity
-    fn get_identity(&self) -> String;
+    fn identity(&self) -> String;
+    /// Get worker grouping key
+    fn worker_grouping_key(&self) -> Uuid;
+    /// Sets the client-reliant fields for WorkerHeartbeat. This also updates client-level tracking
+    /// of heartbeat fields, like last heartbeat timestamp.
+    fn set_heartbeat_client_fields(&self, heartbeat: &mut WorkerHeartbeat);
 }
 
 /// Configuration options shared by workflow, activity, and Nexus polling calls
@@ -312,7 +313,8 @@ impl WorkerClient for WorkerClientBag {
         }
 
         Ok(self
-            .cloned_client()
+            .client
+            .clone()
             .poll_workflow_task_queue(request)
             .await?
             .into_inner())
@@ -349,7 +351,8 @@ impl WorkerClient for WorkerClientBag {
         }
 
         Ok(self
-            .cloned_client()
+            .client
+            .clone()
             .poll_activity_task_queue(request)
             .await?
             .into_inner())
@@ -358,6 +361,7 @@ impl WorkerClient for WorkerClientBag {
     async fn poll_nexus_task(
         &self,
         poll_options: PollOptions,
+        _send_heartbeat: bool,
     ) -> Result<PollNexusTaskQueueResponse> {
         #[allow(deprecated)] // want to list all fields explicitly
         let mut request = PollNexusTaskQueueRequest {
@@ -370,7 +374,7 @@ impl WorkerClient for WorkerClientBag {
             identity: self.identity.clone(),
             worker_version_capabilities: self.worker_version_capabilities(),
             deployment_options: self.deployment_options(),
-            worker_heartbeat: self.capture_heartbeat().into_iter().collect(),
+            worker_heartbeat: Vec::new(),
         }
         .into_request();
         request.extensions_mut().insert(IsWorkerTaskLongPoll);
@@ -382,7 +386,8 @@ impl WorkerClient for WorkerClientBag {
         }
 
         Ok(self
-            .cloned_client()
+            .client
+            .clone()
             .poll_nexus_task_queue(request)
             .await?
             .into_inner())
@@ -432,8 +437,9 @@ impl WorkerClient for WorkerClientBag {
             deployment_options: self.deployment_options(),
         };
         Ok(self
-            .cloned_client()
-            .respond_workflow_task_completed(request)
+            .client
+            .clone()
+            .respond_workflow_task_completed(request.into_request())
             .await?
             .into_inner())
     }
@@ -444,7 +450,8 @@ impl WorkerClient for WorkerClientBag {
         result: Option<Payloads>,
     ) -> Result<RespondActivityTaskCompletedResponse> {
         Ok(self
-            .cloned_client()
+            .client
+            .clone()
             .respond_activity_task_completed(
                 #[allow(deprecated)] // want to list all fields explicitly
                 RespondActivityTaskCompletedRequest {
@@ -456,7 +463,8 @@ impl WorkerClient for WorkerClientBag {
                     // Will never be set, deprecated.
                     deployment: None,
                     deployment_options: self.deployment_options(),
-                },
+                }
+                .into_request(),
             )
             .await?
             .into_inner())
@@ -468,13 +476,17 @@ impl WorkerClient for WorkerClientBag {
         response: nexus::v1::Response,
     ) -> Result<RespondNexusTaskCompletedResponse> {
         Ok(self
-            .cloned_client()
-            .respond_nexus_task_completed(RespondNexusTaskCompletedRequest {
-                namespace: self.namespace.clone(),
-                identity: self.identity.clone(),
-                task_token: task_token.0,
-                response: Some(response),
-            })
+            .client
+            .clone()
+            .respond_nexus_task_completed(
+                RespondNexusTaskCompletedRequest {
+                    namespace: self.namespace.clone(),
+                    identity: self.identity.clone(),
+                    task_token: task_token.0,
+                    response: Some(response),
+                }
+                .into_request(),
+            )
             .await?
             .into_inner())
     }
@@ -485,13 +497,17 @@ impl WorkerClient for WorkerClientBag {
         details: Option<Payloads>,
     ) -> Result<RecordActivityTaskHeartbeatResponse> {
         Ok(self
-            .cloned_client()
-            .record_activity_task_heartbeat(RecordActivityTaskHeartbeatRequest {
-                task_token: task_token.0,
-                details,
-                identity: self.identity.clone(),
-                namespace: self.namespace.clone(),
-            })
+            .client
+            .clone()
+            .record_activity_task_heartbeat(
+                RecordActivityTaskHeartbeatRequest {
+                    task_token: task_token.0,
+                    details,
+                    identity: self.identity.clone(),
+                    namespace: self.namespace.clone(),
+                }
+                .into_request(),
+            )
             .await?
             .into_inner())
     }
@@ -502,7 +518,8 @@ impl WorkerClient for WorkerClientBag {
         details: Option<Payloads>,
     ) -> Result<RespondActivityTaskCanceledResponse> {
         Ok(self
-            .cloned_client()
+            .client
+            .clone()
             .respond_activity_task_canceled(
                 #[allow(deprecated)] // want to list all fields explicitly
                 RespondActivityTaskCanceledRequest {
@@ -514,7 +531,8 @@ impl WorkerClient for WorkerClientBag {
                     // Will never be set, deprecated.
                     deployment: None,
                     deployment_options: self.deployment_options(),
-                },
+                }
+                .into_request(),
             )
             .await?
             .into_inner())
@@ -526,7 +544,8 @@ impl WorkerClient for WorkerClientBag {
         failure: Option<Failure>,
     ) -> Result<RespondActivityTaskFailedResponse> {
         Ok(self
-            .cloned_client()
+            .client
+            .clone()
             .respond_activity_task_failed(
                 #[allow(deprecated)] // want to list all fields explicitly
                 RespondActivityTaskFailedRequest {
@@ -540,7 +559,8 @@ impl WorkerClient for WorkerClientBag {
                     // Will never be set, deprecated.
                     deployment: None,
                     deployment_options: self.deployment_options(),
-                },
+                }
+                .into_request(),
             )
             .await?
             .into_inner())
@@ -567,8 +587,9 @@ impl WorkerClient for WorkerClientBag {
             deployment_options: self.deployment_options(),
         };
         Ok(self
-            .cloned_client()
-            .respond_workflow_task_failed(request)
+            .client
+            .clone()
+            .respond_workflow_task_failed(request.into_request())
             .await?
             .into_inner())
     }
@@ -579,13 +600,17 @@ impl WorkerClient for WorkerClientBag {
         error: nexus::v1::HandlerError,
     ) -> Result<RespondNexusTaskFailedResponse> {
         Ok(self
-            .cloned_client()
-            .respond_nexus_task_failed(RespondNexusTaskFailedRequest {
-                namespace: self.namespace.clone(),
-                identity: self.identity.clone(),
-                task_token: task_token.0,
-                error: Some(error),
-            })
+            .client
+            .clone()
+            .respond_nexus_task_failed(
+                RespondNexusTaskFailedRequest {
+                    namespace: self.namespace.clone(),
+                    identity: self.identity.clone(),
+                    task_token: task_token.0,
+                    error: Some(error),
+                }
+                .into_request(),
+            )
             .await?
             .into_inner())
     }
@@ -597,16 +622,20 @@ impl WorkerClient for WorkerClientBag {
         page_token: Vec<u8>,
     ) -> Result<GetWorkflowExecutionHistoryResponse> {
         Ok(self
-            .cloned_client()
-            .get_workflow_execution_history(GetWorkflowExecutionHistoryRequest {
-                namespace: self.namespace.clone(),
-                execution: Some(WorkflowExecution {
-                    workflow_id,
-                    run_id: run_id.unwrap_or_default(),
-                }),
-                next_page_token: page_token,
-                ..Default::default()
-            })
+            .client
+            .clone()
+            .get_workflow_execution_history(
+                GetWorkflowExecutionHistoryRequest {
+                    namespace: self.namespace.clone(),
+                    execution: Some(WorkflowExecution {
+                        workflow_id,
+                        run_id: run_id.unwrap_or_default(),
+                    }),
+                    next_page_token: page_token,
+                    ..Default::default()
+                }
+                .into_request(),
+            )
             .await?
             .into_inner())
     }
@@ -629,74 +658,99 @@ impl WorkerClient for WorkerClientBag {
         let (_, completed_type, query_result, error_message) = query_result.into_components();
 
         Ok(self
-            .cloned_client()
-            .respond_query_task_completed(RespondQueryTaskCompletedRequest {
-                task_token: task_token.into(),
-                completed_type: completed_type as i32,
-                query_result,
-                error_message,
-                namespace: self.namespace.clone(),
-                failure,
-                cause: cause.into(),
-            })
+            .client
+            .clone()
+            .respond_query_task_completed(
+                RespondQueryTaskCompletedRequest {
+                    task_token: task_token.into(),
+                    completed_type: completed_type as i32,
+                    query_result,
+                    error_message,
+                    namespace: self.namespace.clone(),
+                    failure,
+                    cause: cause.into(),
+                }
+                .into_request(),
+            )
             .await?
             .into_inner())
     }
 
     async fn describe_namespace(&self) -> Result<DescribeNamespaceResponse> {
         Ok(self
-            .cloned_client()
+            .client
+            .clone()
             .describe_namespace(
-                Namespace::Name(self.namespace.clone()).into_describe_namespace_request(),
+                Namespace::Name(self.namespace.clone())
+                    .into_describe_namespace_request()
+                    .into_request(),
             )
             .await?
             .into_inner())
     }
 
-    async fn shutdown_worker(&self, sticky_task_queue: String) -> Result<ShutdownWorkerResponse> {
-        let request = ShutdownWorkerRequest {
+    async fn shutdown_worker(
+        &self,
+        sticky_task_queue: String,
+        final_heartbeat: Option<WorkerHeartbeat>,
+    ) -> Result<ShutdownWorkerResponse> {
+        let mut final_heartbeat = final_heartbeat;
+        if let Some(w) = final_heartbeat.as_mut() {
+            w.status = WorkerStatus::Shutdown.into();
+            self.set_heartbeat_client_fields(w);
+        }
+        let mut request = ShutdownWorkerRequest {
             namespace: self.namespace.clone(),
             identity: self.identity.clone(),
             sticky_task_queue,
             reason: "graceful shutdown".to_string(),
-            worker_heartbeat: self.capture_heartbeat(),
-        };
+            worker_heartbeat: final_heartbeat,
+        }
+        .into_request();
+        request
+            .extensions_mut()
+            .insert(RetryConfigForCall(RetryOptions::no_retries()));
 
         Ok(
-            WorkflowService::shutdown_worker(&mut self.cloned_client(), request)
+            WorkflowService::shutdown_worker(&mut self.client.clone(), request)
                 .await?
                 .into_inner(),
         )
     }
 
-    fn replace_client(&self, new_client: RetryClient<Client>) {
-        let mut replaceable_client = self.replaceable_client.write();
-        *replaceable_client = new_client;
-    }
-
     async fn record_worker_heartbeat(
         &self,
-        heartbeat: WorkerHeartbeat,
+        namespace: String,
+        worker_heartbeat: Vec<WorkerHeartbeat>,
     ) -> Result<RecordWorkerHeartbeatResponse> {
+        let request = RecordWorkerHeartbeatRequest {
+            namespace,
+            identity: self.identity.clone(),
+            worker_heartbeat,
+        };
         Ok(self
-            .cloned_client()
-            .record_worker_heartbeat(RecordWorkerHeartbeatRequest {
-                namespace: self.namespace.clone(),
-                identity: self.identity.clone(),
-                worker_heartbeat: vec![heartbeat],
-            })
+            .client
+            .clone()
+            .record_worker_heartbeat(request.into_request())
             .await?
             .into_inner())
     }
 
+    fn replace_client(&self, new_client: Client) {
+        self.client.get_client().replace_client(new_client);
+    }
+
     fn capabilities(&self) -> Option<Capabilities> {
-        let client = self.replaceable_client.read();
-        client.get_client().inner().capabilities().cloned()
+        self.client
+            .get_client()
+            .inner_cow()
+            .inner()
+            .capabilities()
+            .cloned()
     }
 
-    fn workers(&self) -> Arc<SlotManager> {
-        let client = self.replaceable_client.read();
-        client.get_client().inner().workers()
+    fn workers(&self) -> Arc<ClientWorkerSet> {
+        self.client.get_client().inner_cow().inner().workers()
     }
 
     fn is_mock(&self) -> bool {
@@ -704,23 +758,71 @@ impl WorkerClient for WorkerClientBag {
     }
 
     fn sdk_name_and_version(&self) -> (String, String) {
-        let lock = self.replaceable_client.read();
-        let opts = lock.get_client().inner().options();
+        let inner = self.client.get_client().inner_cow();
+        let opts = inner.options();
         (opts.client_name.clone(), opts.client_version.clone())
     }
 
-    fn get_identity(&self) -> String {
+    fn identity(&self) -> String {
         self.identity.clone()
     }
+
+    fn worker_grouping_key(&self) -> Uuid {
+        self.client.get_client().inner_cow().worker_grouping_key()
+    }
+
+    fn set_heartbeat_client_fields(&self, heartbeat: &mut WorkerHeartbeat) {
+        if let Some(host_info) = heartbeat.host_info.as_mut() {
+            host_info.process_key = self.worker_grouping_key().to_string();
+        }
+        heartbeat.worker_identity = WorkerClient::identity(self);
+        let sdk_name_and_ver = self.sdk_name_and_version();
+        heartbeat.sdk_name = sdk_name_and_ver.0;
+        heartbeat.sdk_version = sdk_name_and_ver.1;
+
+        let now = SystemTime::now();
+        heartbeat.heartbeat_time = Some(now.into());
+        let mut heartbeat_map = self.worker_heartbeat_map.lock();
+        let client_heartbeat_data = heartbeat_map
+            .entry(heartbeat.worker_instance_key.clone())
+            .or_default();
+        let elapsed_since_last_heartbeat =
+            client_heartbeat_data.last_heartbeat_time.map(|hb_time| {
+                let dur = now.duration_since(hb_time).unwrap_or(Duration::ZERO);
+                PbDuration {
+                    seconds: dur.as_secs() as i64,
+                    nanos: dur.subsec_nanos() as i32,
+                }
+            });
+        heartbeat.elapsed_since_last_heartbeat = elapsed_since_last_heartbeat;
+        client_heartbeat_data.last_heartbeat_time = Some(now);
+
+        update_slots(
+            &mut heartbeat.workflow_task_slots_info,
+            &mut client_heartbeat_data.workflow_task_slots_info,
+        );
+        update_slots(
+            &mut heartbeat.activity_task_slots_info,
+            &mut client_heartbeat_data.activity_task_slots_info,
+        );
+        update_slots(
+            &mut heartbeat.nexus_task_slots_info,
+            &mut client_heartbeat_data.nexus_task_slots_info,
+        );
+        update_slots(
+            &mut heartbeat.local_activity_slots_info,
+            &mut client_heartbeat_data.local_activity_slots_info,
+        );
+    }
 }
 
 impl NamespacedClient for WorkerClientBag {
-    fn namespace(&self) -> &str {
-        &self.namespace
+    fn namespace(&self) -> String {
+        self.namespace.clone()
     }
 
-    fn get_identity(&self) -> &str {
-        &self.identity
+    fn identity(&self) -> String {
+        self.identity.clone()
     }
 }
 
@@ -749,3 +851,31 @@ pub struct WorkflowTaskCompletion {
     /// Versioning behavior of the workflow, if any.
     pub versioning_behavior: VersioningBehavior,
 }
+
+#[derive(Clone, Default)]
+struct SlotsInfo {
+    total_processed_tasks: i32,
+    total_failed_tasks: i32,
+}
+
+#[derive(Clone, Default)]
+struct ClientHeartbeatData {
+    last_heartbeat_time: Option<SystemTime>,
+
+    workflow_task_slots_info: SlotsInfo,
+    activity_task_slots_info: SlotsInfo,
+    nexus_task_slots_info: SlotsInfo,
+    local_activity_slots_info: SlotsInfo,
+}
+
+fn update_slots(slots_info: &mut Option<WorkerSlotsInfo>, client_heartbeat_data: &mut SlotsInfo) {
+    if let Some(wft_slot_info) = slots_info.as_mut() {
+        wft_slot_info.last_interval_processed_tasks =
+            wft_slot_info.total_processed_tasks - client_heartbeat_data.total_processed_tasks;
+        wft_slot_info.last_interval_failure_tasks =
+            wft_slot_info.total_failed_tasks - client_heartbeat_data.total_failed_tasks;
+
+        client_heartbeat_data.total_processed_tasks = wft_slot_info.total_processed_tasks;
+        client_heartbeat_data.total_failed_tasks = wft_slot_info.total_failed_tasks;
+    }
+}