@temporalio/core-bridge 1.13.2 → 1.14.0

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

@@ -1,42 +1,49 @@
 //! Worker-specific client needs
 
 pub(crate) mod mocks;
-use crate::{
-    abstractions::dbg_panic, protosext::legacy_query_failure, worker::heartbeat::HeartbeatFn,
-};
+use crate::protosext::legacy_query_failure;
+use parking_lot::Mutex;
+use prost_types::Duration as PbDuration;
 use std::{
-    sync::{Arc, OnceLock},
-    time::Duration,
+    collections::HashMap,
+    sync::Arc,
+    time::{Duration, SystemTime},
 };
-use temporal_client::{
-    Client, IsWorkerTaskLongPoll, Namespace, NamespacedClient, NoRetryOnMatching, RetryClient,
-    SharedReplaceableClient, SlotManager, WorkflowService,
+use temporalio_client::{
+    Client, Namespace, NamespacedClient, RetryClient, RetryOptions, SharedReplaceableClient,
+    WorkflowService,
+    request_extensions::{IsWorkerTaskLongPoll, NoRetryOnMatching, RetryConfigForCall},
+    worker::ClientWorkerSet,
 };
-use temporal_sdk_core_api::worker::WorkerVersioningStrategy;
-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_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},
-        worker::v1::WorkerHeartbeat,
-        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>;
 
@@ -51,7 +58,7 @@ pub(crate) struct WorkerClientBag {
     namespace: String,
     identity: String,
     worker_versioning_strategy: WorkerVersioningStrategy,
-    heartbeat_data: Option<Arc<OnceLock<HeartbeatFn>>>,
+    worker_heartbeat_map: Arc<Mutex<HashMap<String, ClientHeartbeatData>>>,
 }
 
 impl WorkerClientBag {
@@ -60,14 +67,13 @@ impl WorkerClientBag {
         namespace: String,
         identity: String,
         worker_versioning_strategy: WorkerVersioningStrategy,
-        heartbeat_data: Option<Arc<OnceLock<HeartbeatFn>>>,
     ) -> Self {
         Self {
             client,
             namespace,
             identity,
             worker_versioning_strategy,
-            heartbeat_data,
+            worker_heartbeat_map: Arc::new(Mutex::new(HashMap::new())),
         }
     }
 
@@ -124,23 +130,10 @@ 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.
+/// minimal mocking surface.
 #[cfg_attr(any(feature = "test-utilities", test), mockall::automock)]
 #[async_trait::async_trait]
 pub trait WorkerClient: Sync + Send {
@@ -160,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(
@@ -225,11 +219,16 @@ 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
@@ -237,13 +236,18 @@ pub trait WorkerClient: Sync + Send {
     /// 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
@@ -357,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 {
@@ -369,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);
@@ -684,45 +689,57 @@ impl WorkerClient for WorkerClientBag {
             .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.client.clone(), request.into_request())
+            WorkflowService::shutdown_worker(&mut self.client.clone(), request)
                 .await?
                 .into_inner(),
         )
     }
 
-    fn replace_client(&self, new_client: Client) {
-        self.client.get_client().replace_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
             .client
             .clone()
-            .record_worker_heartbeat(
-                RecordWorkerHeartbeatRequest {
-                    namespace: self.namespace.clone(),
-                    identity: self.identity.clone(),
-                    worker_heartbeat: vec![heartbeat],
-                }
-                .into_request(),
-            )
+            .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> {
         self.client
             .get_client()
@@ -732,7 +749,7 @@ impl WorkerClient for WorkerClientBag {
             .cloned()
     }
 
-    fn workers(&self) -> Arc<SlotManager> {
+    fn workers(&self) -> Arc<ClientWorkerSet> {
         self.client.get_client().inner_cow().inner().workers()
     }
 
@@ -746,9 +763,57 @@ impl WorkerClient for WorkerClientBag {
         (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 {
@@ -786,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;
+    }
+}

package/sdk-core/crates/sdk-core/src/worker/heartbeat.rs

@@ -0,0 +1,246 @@
+use crate::{
+    WorkerClient,
+    worker::{TaskPollers, WorkerTelemetry},
+};
+use parking_lot::RwLock;
+use std::{collections::HashMap, sync::Arc, time::Duration};
+use temporalio_client::worker::SharedNamespaceWorkerTrait;
+use temporalio_common::{
+    protos::temporal::api::worker::v1::WorkerHeartbeat,
+    worker::{PollerBehavior, WorkerConfigBuilder, WorkerTaskTypes, WorkerVersioningStrategy},
+};
+use tokio::sync::Notify;
+use tokio_util::sync::CancellationToken;
+use uuid::Uuid;
+
+/// Callback used to collect heartbeat data from each worker at the time of heartbeat
+pub(crate) type HeartbeatFn = Arc<dyn Fn() -> WorkerHeartbeat + Send + Sync>;
+
+/// SharedNamespaceWorker is responsible for polling nexus-delivered worker commands and sending
+/// worker heartbeats to the server. This invokes callbacks on all workers in the same process that
+/// share the same namespace.
+pub(crate) struct SharedNamespaceWorker {
+    heartbeat_map: Arc<RwLock<HashMap<Uuid, HeartbeatFn>>>,
+    namespace: String,
+    cancel: CancellationToken,
+}
+
+impl SharedNamespaceWorker {
+    pub(crate) fn new(
+        client: Arc<dyn WorkerClient>,
+        namespace: String,
+        heartbeat_interval: Duration,
+        telemetry: Option<WorkerTelemetry>,
+    ) -> Result<Self, anyhow::Error> {
+        let config = WorkerConfigBuilder::default()
+            .namespace(namespace.clone())
+            .task_queue(format!(
+                "temporal-sys/worker-commands/{namespace}/{}",
+                client.worker_grouping_key(),
+            ))
+            .task_types(WorkerTaskTypes::nexus_only())
+            .max_outstanding_nexus_tasks(5_usize)
+            .versioning_strategy(WorkerVersioningStrategy::None {
+                build_id: "1.0".to_owned(),
+            })
+            .nexus_task_poller_behavior(PollerBehavior::SimpleMaximum(1_usize))
+            .build()
+            .expect("all required fields should be implemented");
+        let worker = crate::worker::Worker::new_with_pollers(
+            config,
+            None,
+            client.clone(),
+            TaskPollers::Real,
+            telemetry,
+            None,
+            true,
+        )?;
+
+        let reset_notify = Arc::new(Notify::new());
+        let cancel = CancellationToken::new();
+        let cancel_clone = cancel.clone();
+
+        let client_clone = client;
+        let namespace_clone = namespace.clone();
+
+        let heartbeat_map = Arc::new(RwLock::new(HashMap::<Uuid, HeartbeatFn>::new()));
+        let heartbeat_map_clone = heartbeat_map.clone();
+
+        tokio::spawn(async move {
+            match client_clone.describe_namespace().await {
+                Ok(namespace_resp) => {
+                    if namespace_resp
+                        .namespace_info
+                        .and_then(|info| info.capabilities)
+                        .map(|caps| caps.worker_heartbeats)
+                        != Some(true)
+                    {
+                        warn!(
+                            "Worker heartbeating configured for runtime, but server version does not support it."
+                        );
+                        worker.shutdown().await;
+                        return;
+                    }
+                }
+                Err(e) => {
+                    warn!(error=?e, "Network error while describing namespace for heartbeat capabilities");
+                    worker.shutdown().await;
+                    return;
+                }
+            }
+            let mut ticker = tokio::time::interval(heartbeat_interval);
+            loop {
+                tokio::select! {
+                    _ = ticker.tick() => {
+                        let mut hb_to_send = Vec::new();
+                        let hb_callbacks = {
+                            heartbeat_map_clone.read().values().cloned().collect::<Vec<_>>()
+                        };
+                        for heartbeat_callback in hb_callbacks {
+                            let mut heartbeat = heartbeat_callback();
+                            // All of these heartbeat details rely on a client. To avoid circular
+                            // dependencies, this must be populated from within SharedNamespaceWorker
+                            // to get info from the current client
+                            client_clone.set_heartbeat_client_fields(&mut heartbeat);
+                            hb_to_send.push(heartbeat);
+                        }
+                        if let Err(e) = client_clone.record_worker_heartbeat(namespace_clone.clone(), hb_to_send).await {
+                            if matches!(e.code(), tonic::Code::Unimplemented) {
+                                worker.shutdown().await;
+                                return;
+                            }
+                            warn!(error=?e, "Network error while sending worker heartbeat");
+                        }
+                    }
+                    _ = reset_notify.notified() => {
+                        ticker.reset();
+                    }
+                    _ = cancel_clone.cancelled() => {
+                        worker.shutdown().await;
+                        return;
+                    }
+                }
+            }
+        });
+
+        Ok(Self {
+            heartbeat_map,
+            namespace,
+            cancel,
+        })
+    }
+}
+
+impl SharedNamespaceWorkerTrait for SharedNamespaceWorker {
+    fn namespace(&self) -> String {
+        self.namespace.clone()
+    }
+
+    fn register_callback(&self, worker_instance_key: Uuid, heartbeat_callback: HeartbeatFn) {
+        self.heartbeat_map
+            .write()
+            .insert(worker_instance_key, heartbeat_callback);
+    }
+
+    fn unregister_callback(&self, worker_instance_key: Uuid) -> (Option<HeartbeatFn>, bool) {
+        let mut heartbeat_map = self.heartbeat_map.write();
+        let heartbeat_callback = heartbeat_map.remove(&worker_instance_key);
+        if heartbeat_map.is_empty() {
+            self.cancel.cancel();
+        }
+        (heartbeat_callback, heartbeat_map.is_empty())
+    }
+
+    fn num_workers(&self) -> usize {
+        self.heartbeat_map.read().len()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        test_help::{WorkerExt, test_worker_cfg},
+        worker,
+        worker::client::mocks::mock_worker_client,
+    };
+    use std::{
+        sync::{
+            Arc,
+            atomic::{AtomicUsize, Ordering},
+        },
+        time::Duration,
+    };
+    use temporalio_common::{
+        protos::temporal::api::{
+            namespace::v1::{NamespaceInfo, namespace_info::Capabilities},
+            workflowservice::v1::{DescribeNamespaceResponse, RecordWorkerHeartbeatResponse},
+        },
+        worker::PollerBehavior,
+    };
+
+    #[tokio::test]
+    async fn worker_heartbeat_basic() {
+        let mut mock = mock_worker_client();
+        let heartbeat_count = Arc::new(AtomicUsize::new(0));
+        let heartbeat_count_clone = heartbeat_count.clone();
+        mock.expect_poll_workflow_task()
+            .returning(move |_namespace, _task_queue| Ok(Default::default()));
+        mock.expect_poll_nexus_task()
+            .returning(move |_poll_options, _send_heartbeat| Ok(Default::default()));
+        mock.expect_record_worker_heartbeat().times(3).returning(
+            move |_namespace, worker_heartbeat| {
+                assert_eq!(1, worker_heartbeat.len());
+                let heartbeat = worker_heartbeat[0].clone();
+                let host_info = heartbeat.host_info.clone().unwrap();
+                assert_eq!("test-identity", heartbeat.worker_identity);
+                assert!(!heartbeat.worker_instance_key.is_empty());
+                assert_eq!(
+                    host_info.host_name,
+                    gethostname::gethostname().to_string_lossy().to_string()
+                );
+                assert_eq!(host_info.process_id, std::process::id().to_string());
+                assert_eq!(heartbeat.sdk_name, "test-core");
+                assert_eq!(heartbeat.sdk_version, "0.0.0");
+                assert!(heartbeat.heartbeat_time.is_some());
+                assert!(heartbeat.start_time.is_some());
+
+                heartbeat_count_clone.fetch_add(1, Ordering::Relaxed);
+
+                Ok(RecordWorkerHeartbeatResponse {})
+            },
+        );
+        mock.expect_describe_namespace().returning(move || {
+            Ok(DescribeNamespaceResponse {
+                namespace_info: Some(NamespaceInfo {
+                    capabilities: Some(Capabilities {
+                        worker_heartbeats: true,
+                        ..Capabilities::default()
+                    }),
+                    ..NamespaceInfo::default()
+                }),
+                ..DescribeNamespaceResponse::default()
+            })
+        });
+
+        let config = test_worker_cfg()
+            .activity_task_poller_behavior(PollerBehavior::SimpleMaximum(1_usize))
+            .max_outstanding_activities(1_usize)
+            .build()
+            .unwrap();
+
+        let client = Arc::new(mock);
+        let worker = worker::Worker::new(
+            config,
+            None,
+            client.clone(),
+            None,
+            Some(Duration::from_millis(100)),
+        )
+        .unwrap();
+
+        tokio::time::sleep(Duration::from_millis(250)).await;
+        worker.drain_activity_poller_and_shutdown().await;
+
+        assert_eq!(3, heartbeat_count.load(Ordering::Relaxed));
+    }
+}