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

.cargo/config.toml

@@ -1,6 +1,6 @@
 [env]
 # This temporarily overrides the version of the CLI used for integration tests, locally and in CI
-CLI_VERSION_OVERRIDE = "v1.4.1-cloud-v1-29-0-139-2.0"
+#CLI_VERSION_OVERRIDE = "v1.4.1-cloud-v1-29-0-139-2.0"
 
 [alias]
 # Not sure why --all-features doesn't work

client/src/raw.rs

@@ -1345,6 +1345,15 @@ proxier! {
             r.extensions_mut().insert(labels);
         }
     );
+    (
+        describe_worker,
+        DescribeWorkerRequest,
+        DescribeWorkerResponse,
+        |r| {
+            let labels = namespaced_request!(r);
+            r.extensions_mut().insert(labels);
+        }
+    );
     (
         record_worker_heartbeat,
         RecordWorkerHeartbeatRequest,
@@ -1382,6 +1391,15 @@ proxier! {
             r.extensions_mut().insert(labels);
         }
     );
+    (
+        set_worker_deployment_manager,
+        SetWorkerDeploymentManagerRequest,
+        SetWorkerDeploymentManagerResponse,
+        |r| {
+            let labels = namespaced_request!(r);
+            r.extensions_mut().insert(labels);
+        }
+    );
 }
 
 proxier! {

client/src/worker_registry/mod.rs

@@ -76,12 +76,13 @@ impl ClientWorkerSetImpl {
     fn register(
         &mut self,
         worker: Arc<dyn ClientWorker + Send + Sync>,
+        skip_client_worker_set_check: bool,
     ) -> Result<(), anyhow::Error> {
         let slot_key = SlotKey::new(
             worker.namespace().to_string(),
             worker.task_queue().to_string(),
         );
-        if self.slot_providers.contains_key(&slot_key) {
+        if self.slot_providers.contains_key(&slot_key) && !skip_client_worker_set_check {
             bail!(
                 "Registration of multiple workers on the same namespace and task queue for the same client not allowed: {slot_key:?}, worker_instance_key: {:?}.",
                 worker.worker_instance_key()
@@ -133,14 +134,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,6 +207,17 @@ impl ClientWorkerSet {
             .try_reserve_wft_slot(namespace, task_queue)
     }
 
+    /// Register a local worker that can provide WFT processing slots and potentially worker heartbeating.
+    pub fn register_worker(
+        &self,
+        worker: Arc<dyn ClientWorker + Send + Sync>,
+        skip_client_worker_set_check: bool,
+    ) -> Result<(), anyhow::Error> {
+        self.worker_manager
+            .write()
+            .register(worker, skip_client_worker_set_check)
+    }
+
     /// Unregisters a local worker, typically when that worker starts shutdown.
     pub fn unregister_worker(
         &self,
@@ -220,14 +226,6 @@ impl ClientWorkerSet {
         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,
-        worker: Arc<dyn ClientWorker + Send + Sync>,
-    ) -> Result<(), anyhow::Error> {
-        self.worker_manager.write().register(worker)
-    }
-
     /// Returns the worker grouping key, which is unique for each worker.
     pub fn worker_grouping_key(&self) -> Uuid {
         self.worker_grouping_key
@@ -256,7 +254,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.
@@ -276,7 +274,7 @@ pub trait ClientWorker: Send + Sync {
     fn try_reserve_wft_slot(&self) -> Option<Box<dyn Slot + Send>>;
 
     /// Unique identifier for this worker instance.
-    /// This must be stable across the worker's lifetime but unique per instance.
+    /// This must be stable across the worker's lifetime and unique per instance.
     fn worker_instance_key(&self) -> Uuid;
 
     /// Indicates if worker heartbeating is enabled for this client worker.
@@ -289,9 +287,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)]
@@ -350,7 +345,7 @@ mod tests {
                 new_mock_provider(namespace, "bar_q".to_string(), false, false, false);
             let worker_instance_key = mock_provider.worker_instance_key();
 
-            let result = manager.register_worker(Arc::new(mock_provider));
+            let result = manager.register_worker(Arc::new(mock_provider), false);
             if result.is_ok() {
                 successful_registrations += 1;
                 worker_keys.push(worker_instance_key);
@@ -453,7 +448,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 +458,6 @@ mod tests {
                         namespace_clone.clone(),
                     )))
                 });
-
-            mock_provider.expect_register_callback().returning(|_| {});
         }
 
         mock_provider
@@ -489,10 +482,10 @@ mod tests {
             Uuid::new_v4(),
         );
 
-        manager.register_worker(Arc::new(worker1)).unwrap();
+        manager.register_worker(Arc::new(worker1), false).unwrap();
 
         // second worker register should fail due to duplicate namespace+task_queue
-        let result = manager.register_worker(Arc::new(worker2));
+        let result = manager.register_worker(Arc::new(worker2), false);
         assert!(result.is_err());
         assert!(
             result
@@ -528,8 +521,8 @@ mod tests {
             Uuid::new_v4(),
         );
 
-        manager.register_worker(Arc::new(worker1)).unwrap();
-        manager.register_worker(Arc::new(worker2)).unwrap();
+        manager.register_worker(Arc::new(worker1), false).unwrap();
+        manager.register_worker(Arc::new(worker2), false).unwrap();
 
         assert_eq!(2, manager.num_providers());
         assert_eq!(manager.num_heartbeat_workers(), 2);
@@ -558,8 +551,8 @@ mod tests {
             Uuid::new_v4(),
         );
 
-        manager.register_worker(Arc::new(worker1)).unwrap();
-        manager.register_worker(Arc::new(worker2)).unwrap();
+        manager.register_worker(Arc::new(worker1), false).unwrap();
+        manager.register_worker(Arc::new(worker2), false).unwrap();
 
         assert_eq!(2, manager.num_providers());
         assert_eq!(manager.num_heartbeat_workers(), 2);
@@ -590,8 +583,10 @@ mod tests {
         let worker_instance_key1 = worker1.worker_instance_key();
         let worker_instance_key2 = worker2.worker_instance_key();
 
-        manager.register_worker(Arc::new(worker1)).unwrap();
-        manager.register_worker(Arc::new(worker2)).unwrap();
+        assert_ne!(worker_instance_key1, worker_instance_key2);
+
+        manager.register_worker(Arc::new(worker1), false).unwrap();
+        manager.register_worker(Arc::new(worker2), false).unwrap();
 
         // Verify initial state: 2 slot providers, 2 heartbeat workers, 1 shared worker
         assert_eq!(2, manager.num_providers());

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 and 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 {