Replace client test passes, need to write unit tests in worker_registry

Author: yuandrew

client/src/lib.rs

@@ -33,7 +33,7 @@ pub use temporal_sdk_core_protos::temporal::api::{
 };
 pub use tonic;
 pub use worker_registry::{
-    ClientWorkerSet, SharedNamespaceWorkerTrait, Slot, SlotProvider, WorkerKey,
+    ClientWorker, ClientWorkerSet, HeartbeatCallback, SharedNamespaceWorkerTrait, Slot, WorkerKey,
 };
 pub use workflow_handle::{
     GetWorkflowResultOpts, WorkflowExecutionInfo, WorkflowExecutionResult, WorkflowHandle,

client/src/worker_registry/mod.rs

@@ -2996,6 +2996,7 @@ async fn both_normal_and_sticky_pollers_poll_concurrently() {
         Arc::new(mock_client),
         None,
         None,
+        false,
     );
 
     for _ in 1..50 {

core/src/core_tests/workflow_tasks.rs

@@ -70,7 +70,6 @@ use temporal_sdk_core_api::{
     telemetry::TelemetryOptions,
 };
 use temporal_sdk_core_protos::coresdk::ActivityHeartbeat;
-use uuid::Uuid;
 
 /// Initialize a worker bound to a task queue.
 ///
@@ -124,6 +123,7 @@ where
         client_bag.clone(),
         Some(&runtime.telemetry),
         runtime.heartbeat_interval,
+        false,
     );
 
     Ok(worker)

core/src/lib.rs

@@ -114,7 +114,7 @@ where
                 hist_allow_tx.send("Failed".to_string()).unwrap();
                 async move { Ok(RespondWorkflowTaskFailedResponse::default()) }.boxed()
             });
-        let mut worker = Worker::new(self.config, None, Arc::new(client), None, None);
+        let mut worker = Worker::new(self.config, None, Arc::new(client), None, None, false);
         worker.set_post_activate_hook(post_activate);
         shutdown_tok(worker.shutdown_token());
         Ok(worker)

core/src/replay/mod.rs

@@ -205,6 +205,7 @@ pub fn mock_worker(mocks: MocksHolder) -> Worker {
         },
         None,
         None,
+        false,
     )
 }
 

core/src/test_help/integ_helpers.rs

@@ -12,6 +12,7 @@ use temporal_client::SharedNamespaceWorkerTrait;
 use temporal_sdk_core_api::worker::{WorkerConfigBuilder, WorkerVersioningStrategy};
 use temporal_sdk_core_protos::temporal::api::worker::v1::WorkerHeartbeat;
 use tokio::sync::Notify;
+use tokio_util::sync::CancellationToken;
 
 /// Callback used to collect heartbeat data from each worker at the time of heartbeat
 pub(crate) type HeartbeatFn = Box<dyn Fn() -> WorkerHeartbeat + Send + Sync>;
@@ -22,6 +23,7 @@ pub(crate) type HeartbeatFn = Box<dyn Fn() -> WorkerHeartbeat + Send + Sync>;
 pub(crate) struct SharedNamespaceWorker {
     heartbeat_map: Arc<Mutex<HashMap<String, HeartbeatFn>>>,
     namespace: String,
+    cancel: CancellationToken,
 }
 
 impl SharedNamespaceWorker {
@@ -51,12 +53,16 @@ impl SharedNamespaceWorker {
             TaskPollers::Real,
             telemetry,
             None,
+            true,
         );
 
         let last_heartbeat_time_map = Mutex::new(HashMap::new());
 
         // heartbeat task
         let reset_notify = Arc::new(Notify::new());
+        let cancel = CancellationToken::new();
+        let child = cancel.child_token();
+
         let client_clone = client;
         let namespace_clone = namespace.clone();
 
@@ -67,15 +73,6 @@ impl SharedNamespaceWorker {
             let mut ticker = tokio::time::interval(heartbeat_interval);
             ticker.tick().await;
             loop {
-                // TODO: Race condition here, can technically shut down before anything is ever initialized
-                if heartbeat_map_clone.lock().is_empty() {
-                    println!(
-                        "// TODO: Race condition here, can technically shut down before anything is ever initialized"
-                    );
-                    worker.shutdown().await;
-                    return;
-                }
-
                 tokio::select! {
                     _ = ticker.tick() => {
                         let mut hb_to_send = Vec::new();
@@ -94,10 +91,18 @@ impl SharedNamespaceWorker {
                             } else {
                                 None
                             };
+                            let sdk_name_and_ver = client_clone.sdk_name_and_version();
 
                             heartbeat.elapsed_since_last_heartbeat = elapsed_since_last_heartbeat;
                             heartbeat.heartbeat_time = Some(now.into());
 
+                            // All of these heartbeat details rely on a client. Due to type limitations
+                            // from the dependency graph of worker_registry, this must be populated
+                            // from within SharedNamespaceWorker to get the current and proper client
+                            heartbeat.worker_identity = client_clone.identity();
+                            heartbeat.sdk_name = sdk_name_and_ver.0;
+                            heartbeat.sdk_version = sdk_name_and_ver.1;
+
                             hb_to_send.push(heartbeat);
 
                             last_heartbeat_time_map.insert(instance_key.clone(), now);
@@ -115,13 +120,18 @@ impl SharedNamespaceWorker {
                     _ = reset_notify.notified() => {
                         ticker.reset();
                     }
+                    _ = child.cancelled() => {
+                        worker.shutdown().await;
+                        return;
+                    }
                 }
             }
         });
 
         Self {
             heartbeat_map,
             namespace,
+            cancel,
         }
     }
 }
@@ -136,10 +146,11 @@ impl SharedNamespaceWorkerTrait for SharedNamespaceWorker {
         worker_instance_key: String,
     ) -> (Option<Box<dyn Fn() -> WorkerHeartbeat + Send + Sync>>, bool) {
         let mut heartbeat_map = self.heartbeat_map.lock();
-        (
-            heartbeat_map.remove(&worker_instance_key),
-            heartbeat_map.is_empty(),
-        )
+        let heartbeat_callback = heartbeat_map.remove(&worker_instance_key);
+        if heartbeat_map.is_empty() {
+            self.cancel.cancel();
+        }
+        (heartbeat_callback, heartbeat_map.is_empty())
     }
     fn register_callback(
         &self,
@@ -230,6 +241,7 @@ mod tests {
             client.clone(),
             None,
             Some(Duration::from_millis(100)),
+            false,
         );
 
         tokio::time::sleep(Duration::from_millis(250)).await;