Worker heartbeat: New in-memory metrics mechism, plumb rest of heartbeat data

client/src/lib.rs

@@ -1227,6 +1227,14 @@ pub trait WorkflowClientTrait: NamespacedClient {
         query: String,
     ) -> Result<ListArchivedWorkflowExecutionsResponse>;
 
+    /// List workers registered with server through worker heartbeats
+    async fn list_workers(
+        &self,
+        page_size: i32,
+        next_page_token: Vec<u8>,
+        query: String,
+    ) -> Result<ListWorkersResponse>;
+
     /// Get Cluster Search Attributes
     async fn get_search_attributes(&self) -> Result<GetSearchAttributesResponse>;
 
@@ -1793,6 +1801,21 @@ where
         .into_inner())
     }
 
+    async fn list_workers(
+        &self,
+        page_size: i32,
+        next_page_token: Vec<u8>,
+        query: String,
+    ) -> Result<ListWorkersResponse> {
+        Ok(WorkflowService::list_workers(&mut self.clone(),
+         ListWorkersRequest {
+             namespace: self.namespace().to_owned(),
+             page_size,
+             next_page_token,
+             query,
+         }).await?.into_inner())
+    }
+
     async fn get_search_attributes(&self) -> Result<GetSearchAttributesResponse> {
         Ok(WorkflowService::get_search_attributes(&mut self.clone(),
             GetSearchAttributesRequest {},

client/src/worker_registry/mod.rs

@@ -133,14 +133,8 @@ impl ClientWorkerSetImpl {
 
         if let Some(w) = self.shared_worker.get_mut(worker.namespace()) {
             let (callback, is_empty) = w.unregister_callback(worker.worker_instance_key());
-            if let Some(cb) = callback {
-                if is_empty {
-                    self.shared_worker.remove(worker.namespace());
-                }
-
-                // To maintain single ownership of the callback, we must re-register the callback
-                // back to the ClientWorker
-                worker.register_callback(cb);
+            if callback.is_some() && is_empty {
+                self.shared_worker.remove(worker.namespace());
             }
         }
 
@@ -212,14 +206,6 @@ impl ClientWorkerSet {
             .try_reserve_wft_slot(namespace, task_queue)
     }
 
-    /// Unregisters a local worker, typically when that worker starts shutdown.
-    pub fn unregister_worker(
-        &self,
-        worker_instance_key: Uuid,
-    ) -> Result<Arc<dyn ClientWorker + Send + Sync>, anyhow::Error> {
-        self.worker_manager.write().unregister(worker_instance_key)
-    }
-
     /// Register a local worker that can provide WFT processing slots and potentially worker heartbeating.
     pub fn register_worker(
         &self,
@@ -228,6 +214,14 @@ impl ClientWorkerSet {
         self.worker_manager.write().register(worker)
     }
 
+    /// Unregisters a local worker, typically when that worker starts shutdown.
+    pub fn unregister_worker(
+        &self,
+        worker_instance_key: Uuid,
+    ) -> Result<Arc<dyn ClientWorker + Send + Sync>, anyhow::Error> {
+        self.worker_manager.write().unregister(worker_instance_key)
+    }
+
     /// Returns the worker grouping key, which is unique for each worker.
     pub fn worker_grouping_key(&self) -> Uuid {
         self.worker_grouping_key
@@ -256,7 +250,7 @@ impl std::fmt::Debug for ClientWorkerSet {
 }
 
 /// Contains a worker heartbeat callback, wrapped for mocking
-pub type HeartbeatCallback = Box<dyn Fn() -> WorkerHeartbeat + Send + Sync>;
+pub type HeartbeatCallback = Arc<dyn Fn() -> WorkerHeartbeat + Send + Sync>;
 
 /// Represents a complete worker that can handle both slot management
 /// and worker heartbeat functionality.
@@ -289,9 +283,6 @@ pub trait ClientWorker: Send + Sync {
     fn new_shared_namespace_worker(
         &self,
     ) -> Result<Box<dyn SharedNamespaceWorkerTrait + Send + Sync>, anyhow::Error>;
-
-    /// Registers a worker heartbeat callback, typically when a worker is unregistered from a client
-    fn register_callback(&self, callback: HeartbeatCallback);
 }
 
 #[cfg(test)]
@@ -453,7 +444,7 @@ mod tests {
         if heartbeat_enabled {
             mock_provider
                 .expect_heartbeat_callback()
-                .returning(|| Some(Box::new(WorkerHeartbeat::default)));
+                .returning(|| Some(Arc::new(WorkerHeartbeat::default)));
 
             let namespace_clone = namespace.clone();
             mock_provider
@@ -463,8 +454,6 @@ mod tests {
                         namespace_clone.clone(),
                     )))
                 });
-
-            mock_provider.expect_register_callback().returning(|_| {});
         }
 
         mock_provider

core-api/Cargo.toml

@@ -31,6 +31,7 @@ tonic = { workspace = true }
 tracing = "0.1"
 tracing-core = "0.1"
 url = "2.5"
+uuid = { version = "1.18.1", features = ["v4"] }
 
 [dependencies.temporal-sdk-core-protos]
 path = "../sdk-core-protos"

core-api/src/lib.rs

@@ -19,6 +19,7 @@ use temporal_sdk_core_protos::coresdk::{
     workflow_activation::WorkflowActivation,
     workflow_completion::WorkflowActivationCompletion,
 };
+use uuid::Uuid;
 
 /// This trait is the primary way by which language specific SDKs interact with the core SDK.
 /// It represents one worker, which has a (potentially shared) client for connecting to the service
@@ -138,6 +139,10 @@ pub trait Worker: Send + Sync {
     /// This should be called only after [Worker::shutdown] has resolved and/or both polling
     /// functions have returned `ShutDown` errors.
     async fn finalize_shutdown(self);
+
+    /// Unique identifier for this worker instance.
+    /// This must be stable across the worker's lifetime but unique per instance.
+    fn worker_instance_key(&self) -> Uuid;
 }
 
 #[async_trait::async_trait]
@@ -205,6 +210,10 @@ where
     async fn finalize_shutdown(self) {
         panic!("Can't finalize shutdown on Arc'd worker")
     }
+
+    fn worker_instance_key(&self) -> Uuid {
+        (**self).worker_instance_key()
+    }
 }
 
 macro_rules! dbg_panic {