Switch to worker_set_key

Author: yuandrew

client/src/lib.rs

@@ -443,6 +443,11 @@ impl<C> ConfiguredClient<C> {
     pub fn workers(&self) -> Arc<ClientWorkerSet> {
         self.workers.clone()
     }
+
+    /// Returns the worker set key, this should be unique across each client
+    pub fn worker_set_key(&self) -> Uuid {
+        self.workers.worker_set_key()
+    }
 }
 
 #[derive(Debug)]
@@ -499,10 +504,9 @@ impl ClientOptions {
         &self,
         namespace: impl Into<String>,
         metrics_meter: Option<TemporalMeter>,
-        process_key: Uuid,
     ) -> Result<RetryClient<Client>, ClientInitError> {
         let client = self.connect_no_namespace(metrics_meter).await?.into_inner();
-        let client = Client::new(client, namespace.into(), process_key);
+        let client = Client::new(client, namespace.into());
         let retry_client = RetryClient::new(client, self.retry_config.clone());
         Ok(retry_client)
     }
@@ -851,21 +855,17 @@ pub struct Client {
     inner: ConfiguredClient<TemporalServiceClientWithMetrics>,
     /// The namespace this client interacts with
     namespace: String,
-    /// Process-wide key, used for worker heartbeating
-    process_key: Uuid,
 }
 
 impl Client {
     /// Create a new client from an existing configured lower level client and a namespace
     pub fn new(
         client: ConfiguredClient<TemporalServiceClientWithMetrics>,
         namespace: String,
-        process_key: Uuid,
     ) -> Self {
         Client {
             inner: client,
             namespace,
-            process_key,
         }
     }
 
@@ -910,8 +910,8 @@ impl Client {
     }
 
     /// Returns the process-wide key
-    pub fn process_key(&self) -> Uuid {
-        self.process_key
+    pub fn worker_set_key(&self) -> Uuid {
+        self.inner.worker_set_key()
     }
 }
 

client/src/worker_registry/mod.rs

@@ -7,6 +7,7 @@ use slotmap::SlotMap;
 use std::collections::{HashMap, hash_map::Entry::Vacant};
 use temporal_sdk_core_protos::temporal::api::worker::v1::WorkerHeartbeat;
 use temporal_sdk_core_protos::temporal::api::workflowservice::v1::PollWorkflowTaskQueueResponse;
+use uuid::Uuid;
 
 slotmap::new_key_type! {
     /// Registration key for a worker
@@ -131,12 +132,15 @@ pub trait SharedNamespaceWorkerTrait: std::fmt::Debug {
 }
 
 /// Enables local workers to make themselves visible to a shared client instance.
-/// There can only be one worker registered per namespace+queue_name+client, others will get ignored.
+///
+/// For slot managing, there can only be one worker registered per
+/// namespace+queue_name+client, others will get ignored.
 /// It also provides a convenient method to find compatible slots within the collection.
 #[derive(Default, Debug)]
 pub struct ClientWorkerSet {
     slot_manager: RwLock<SlotManagerImpl>,
     heartbeat_manager: Mutex<HashMap<String, Box<dyn SharedNamespaceWorkerTrait + Send + Sync>>>,
+    worker_set_key: Uuid,
 }
 
 impl ClientWorkerSet {
@@ -145,6 +149,7 @@ impl ClientWorkerSet {
         Self {
             slot_manager: RwLock::new(SlotManagerImpl::new()),
             heartbeat_manager: Mutex::new(HashMap::new()),
+            worker_set_key: Uuid::new_v4(),
         }
     }
 
@@ -208,6 +213,11 @@ impl ClientWorkerSet {
         }
     }
 
+    /// Returns the worker set key, which is unique for each client. Used for worker heartbeating
+    pub fn worker_set_key(&self) -> Uuid {
+        self.worker_set_key
+    }
+
     #[cfg(test)]
     /// Returns (num_providers, num_buckets), where a bucket key is namespace+task_queue.
     /// There is only one provider per bucket so `num_providers` should be equal to `num_buckets`.

core/src/core_tests/workers.rs

@@ -315,7 +315,7 @@ async fn worker_shutdown_api(#[case] use_cache: bool, #[case] api_success: bool)
     mock.expect_is_mock().returning(|| true);
     mock.expect_sdk_name_and_version()
         .returning(|| ("test-core".to_string(), "0.0.0".to_string()));
-    mock.expect_get_identity()
+    mock.expect_identity()
         .returning(|| "test-identity".to_string());
     if use_cache {
         if api_success {

core/src/lib.rs

@@ -96,7 +96,6 @@ where
         worker_config.namespace.clone(),
         worker_config.client_identity_override.clone(),
         client_inner,
-        runtime.process_key(),
     );
     let namespace = worker_config.namespace.clone();
     if client.namespace() != namespace {
@@ -151,9 +150,8 @@ pub(crate) fn init_worker_client(
     namespace: String,
     client_identity_override: Option<String>,
     client: ConfiguredClient<TemporalServiceClientWithMetrics>,
-    process_key: Uuid,
 ) -> RetryClient<Client> {
-    let mut client = Client::new(client, namespace, process_key);
+    let mut client = Client::new(client, namespace);
     if let Some(ref id_override) = client_identity_override {
         client.options_mut().identity.clone_from(id_override);
     }
@@ -228,7 +226,6 @@ pub struct CoreRuntime {
     telemetry: TelemetryInstance,
     runtime: Option<tokio::runtime::Runtime>,
     runtime_handle: tokio::runtime::Handle,
-    process_key: Uuid,
     heartbeat_interval: Option<Duration>,
 }
 
@@ -335,7 +332,6 @@ impl CoreRuntime {
             telemetry,
             runtime: None,
             runtime_handle,
-            process_key: Uuid::new_v4(),
             heartbeat_interval,
         }
     }
@@ -350,11 +346,6 @@ impl CoreRuntime {
         &self.telemetry
     }
 
-    /// Return a process-wide unique key
-    pub fn process_key(&self) -> Uuid {
-        self.process_key
-    }
-
     /// Return a mutable reference to the owned [TelemetryInstance]
     pub fn telemetry_mut(&mut self) -> &mut TelemetryInstance {
         &mut self.telemetry

core/src/worker/client.rs

@@ -230,9 +230,9 @@ pub trait WorkerClient: Sync + Send {
     /// Return name and version of the SDK
     fn sdk_name_and_version(&self) -> (String, String);
     /// Get worker identity
-    fn get_identity(&self) -> String;
-    /// Get process key
-    fn get_process_key(&self) -> Uuid;
+    fn identity(&self) -> String;
+    /// Get worker set key
+    fn worker_set_key(&self) -> Uuid;
 }
 
 /// Configuration options shared by workflow, activity, and Nexus polling calls
@@ -698,12 +698,12 @@ 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 get_process_key(&self) -> Uuid {
-        self.replaceable_client.read().get_client().process_key()
+    fn worker_set_key(&self) -> Uuid {
+        self.replaceable_client.read().get_client().worker_set_key()
     }
 }
 

core/src/worker/client/mocks.rs

@@ -33,9 +33,9 @@ pub fn mock_worker_client() -> MockWorkerClient {
         .returning(|_| Ok(ShutdownWorkerResponse {}));
     r.expect_sdk_name_and_version()
         .returning(|| ("test-core".to_string(), "0.0.0".to_string()));
-    r.expect_get_identity()
+    r.expect_identity()
         .returning(|| "test-identity".to_string());
-    r.expect_get_process_key().returning(Uuid::new_v4);
+    r.expect_worker_set_key().returning(Uuid::new_v4);
     r
 }
 
@@ -49,7 +49,7 @@ pub(crate) fn mock_manual_worker_client() -> MockManualWorkerClient {
     r.expect_is_mock().returning(|| true);
     r.expect_sdk_name_and_version()
         .returning(|| ("test-core".to_string(), "0.0.0".to_string()));
-    r.expect_get_identity()
+    r.expect_identity()
         .returning(|| "test-identity".to_string());
     r
 }
@@ -162,7 +162,7 @@ mockall::mock! {
         fn workers(&self) -> Arc<ClientWorkerSet>;
         fn is_mock(&self) -> bool;
         fn sdk_name_and_version(&self) -> (String, String);
-        fn get_identity(&self) -> String;
-        fn get_process_key(&self) -> Uuid;
+        fn identity(&self) -> String;
+        fn worker_set_key(&self) -> Uuid;
     }
 }

core/src/worker/heartbeat.rs

@@ -35,7 +35,7 @@ impl SharedNamespaceWorker {
             .namespace(namespace.clone())
             .task_queue(format!(
                 "temporal-sys/worker-commands/{namespace}/{}",
-                client.get_process_key()
+                client.worker_set_key(),
             ))
             .no_remote_activities(true)
             .max_outstanding_nexus_tasks(5_usize)

core/src/worker/mod.rs

@@ -342,7 +342,6 @@ impl Worker {
             self.config.namespace.clone(),
             self.config.client_identity_override.clone(),
             new_client,
-            self.client.get_process_key(),
         );
 
         self.client.replace_client(new_worker_client);
@@ -428,7 +427,7 @@ impl Worker {
         let shutdown_token = CancellationToken::new();
         let slot_context_data = Arc::new(PermitDealerContextData {
             task_queue: config.task_queue.clone(),
-            worker_identity: client.get_identity(),
+            worker_identity: client.identity(),
             worker_deployment_version: config.computed_deployment_version(),
         });
         let wft_slots = MeteredPermitDealer::new(
@@ -592,7 +591,7 @@ impl Worker {
         let worker_heartbeat = if let Some(heartbeat_interval) = worker_heartbeat_interval {
             let worker_instance_key = Uuid::new_v4();
             let worker_instance_key_clone = worker_instance_key.to_string();
-            let worker_identity = client.get_identity();
+            let worker_identity = client.identity();
             let task_queue = config.task_queue.clone();
             let sdk_name_and_ver = sdk_name_and_ver.clone();
 

sdk/src/lib.rs

@@ -25,7 +25,7 @@
 //!     let runtime_options = RuntimeOptionsBuilder::default().telemetry_options(telemetry_options).build().unwrap();
 //!     let runtime = CoreRuntime::new_assume_tokio(runtime_options)?;
 //!
-//!     let client = server_options.connect("default", None, runtime.process_key()).await?;
+//!     let client = server_options.connect("default", None, runtime.runtime_key()).await?;
 //!
 //!     let worker_config = WorkerConfigBuilder::default()
 //!         .namespace("default")

tests/common/mod.rs

@@ -204,7 +204,6 @@ pub(crate) async fn get_cloud_client() -> RetryClient<Client> {
     sgo.connect(
         env::var("TEMPORAL_NAMESPACE").expect("TEMPORAL_NAMESPACE must be set"),
         None,
-        Uuid::new_v4(),
     )
     .await
     .unwrap()
@@ -456,7 +455,6 @@ impl CoreWfStarter {
                             .connect(
                                 cfg.namespace.clone(),
                                 rt.telemetry().get_temporal_metric_meter(),
-                                rt.process_key(),
                             )
                             .await
                             .expect("Must connect"),