@temporalio/core-bridge 1.14.2-canary-release-testing.0 → 1.16.0

package/sdk-core/crates/client/src/retry.rs

@@ -1,13 +1,21 @@
 use crate::{
-    ERROR_RETURNED_DUE_TO_SHORT_CIRCUIT, MESSAGE_TOO_LARGE_KEY, NamespacedClient, Result,
-    RetryOptions,
-    raw::IsUserLongPoll,
+    ERROR_RETURNED_DUE_TO_SHORT_CIRCUIT, MESSAGE_TOO_LARGE_KEY,
+    grpc::IsUserLongPoll,
     request_extensions::{IsWorkerTaskLongPoll, NoRetryOnMatching, RetryConfigForCall},
 };
-use backoff::{Clock, SystemClock, backoff::Backoff, exponential::ExponentialBackoff};
+use backoff::{
+    Clock, SystemClock,
+    backoff::Backoff,
+    exponential::{self, ExponentialBackoff},
+};
 use futures_retry::{ErrorHandler, FutureRetry, RetryPolicy};
-use std::{error::Error, fmt::Debug, future::Future, sync::Arc, time::Duration};
-use tonic::{Code, Request};
+use std::{
+    error::Error,
+    fmt::Debug,
+    future::Future,
+    time::{Duration, Instant},
+};
+use tonic::Code;
 
 /// List of gRPC error codes that client will retry.
 #[doc(hidden)]
@@ -22,44 +30,77 @@ pub const RETRYABLE_ERROR_CODES: [Code; 7] = [
 ];
 const LONG_POLL_FATAL_GRACE: Duration = Duration::from_secs(60);
 
-/// A wrapper for a [crate::WorkflowClientTrait] or [crate::WorkflowService] implementor which
-/// performs auto-retries
-#[derive(Debug, Clone)]
-pub struct RetryClient<SG> {
-    client: SG,
-    retry_config: Arc<RetryOptions>,
+/// Configuration for retrying requests to the server
+#[derive(Clone, Debug, PartialEq)]
+pub struct RetryOptions {
+    /// initial wait time before the first retry.
+    pub initial_interval: Duration,
+    /// randomization jitter that is used as a multiplier for the current retry interval
+    /// and is added or subtracted from the interval length.
+    pub randomization_factor: f64,
+    /// rate at which retry time should be increased, until it reaches max_interval.
+    pub multiplier: f64,
+    /// maximum amount of time to wait between retries.
+    pub max_interval: Duration,
+    /// maximum total amount of time requests should be retried for, if None is set then no limit
+    /// will be used.
+    pub max_elapsed_time: Option<Duration>,
+    /// maximum number of retry attempts.
+    pub max_retries: usize,
 }
 
-impl<SG> RetryClient<SG> {
-    /// Use the provided retry config with the provided client
-    pub fn new(client: SG, retry_config: RetryOptions) -> Self {
+impl Default for RetryOptions {
+    fn default() -> Self {
         Self {
-            client,
-            retry_config: Arc::new(retry_config),
+            initial_interval: Duration::from_millis(100), // 100 ms wait by default.
+            randomization_factor: 0.2,                    // +-20% jitter.
+            multiplier: 1.7, // each next retry delay will increase by 70%
+            max_interval: Duration::from_secs(5), // until it reaches 5 seconds.
+            max_elapsed_time: Some(Duration::from_secs(10)), // 10 seconds total allocated time for all retries.
+            max_retries: 10,
         }
     }
 }
 
-impl<SG> RetryClient<SG> {
-    /// Return the inner client type
-    pub fn get_client(&self) -> &SG {
-        &self.client
+impl RetryOptions {
+    pub(crate) const fn task_poll_retry_policy() -> Self {
+        Self {
+            initial_interval: Duration::from_millis(200),
+            randomization_factor: 0.2,
+            multiplier: 2.0,
+            max_interval: Duration::from_secs(10),
+            max_elapsed_time: None,
+            max_retries: 0,
+        }
     }
 
-    /// Return the inner client type mutably
-    pub fn get_client_mut(&mut self) -> &mut SG {
-        &mut self.client
+    pub(crate) const fn throttle_retry_policy() -> Self {
+        Self {
+            initial_interval: Duration::from_secs(1),
+            randomization_factor: 0.2,
+            multiplier: 2.0,
+            max_interval: Duration::from_secs(10),
+            max_elapsed_time: None,
+            max_retries: 0,
+        }
     }
 
-    /// Disable retry and return the inner client type
-    pub fn into_inner(self) -> SG {
-        self.client
+    /// A retry policy that never retires
+    pub const fn no_retries() -> Self {
+        Self {
+            initial_interval: Duration::from_secs(0),
+            randomization_factor: 0.0,
+            multiplier: 1.0,
+            max_interval: Duration::from_secs(0),
+            max_elapsed_time: None,
+            max_retries: 1,
+        }
     }
 
     pub(crate) fn get_call_info<R>(
         &self,
         call_name: &'static str,
-        request: Option<&Request<R>>,
+        request: Option<&tonic::Request<R>>,
     ) -> CallInfo {
         let mut call_type = CallType::Normal;
         let mut retry_short_circuit = None;
@@ -80,7 +121,7 @@ impl<SG> RetryClient<SG> {
         } else if call_type == CallType::TaskLongPoll {
             RetryOptions::task_poll_retry_policy()
         } else {
-            (*self.retry_config).clone()
+            self.clone()
         };
         CallInfo {
             call_type,
@@ -90,32 +131,38 @@ impl<SG> RetryClient<SG> {
         }
     }
 
-    pub(crate) fn make_future_retry<R, F, Fut>(
-        info: CallInfo,
-        factory: F,
-    ) -> FutureRetry<F, TonicErrorHandler<SystemClock>>
-    where
-        F: FnMut() -> Fut + Unpin,
-        Fut: Future<Output = Result<R>>,
-    {
-        FutureRetry::new(
-            factory,
-            TonicErrorHandler::new(info, RetryOptions::throttle_retry_policy()),
-        )
+    pub(crate) fn into_exp_backoff<C>(self, clock: C) -> exponential::ExponentialBackoff<C> {
+        exponential::ExponentialBackoff {
+            current_interval: self.initial_interval,
+            initial_interval: self.initial_interval,
+            randomization_factor: self.randomization_factor,
+            multiplier: self.multiplier,
+            max_interval: self.max_interval,
+            max_elapsed_time: self.max_elapsed_time,
+            clock,
+            start_time: Instant::now(),
+        }
     }
 }
 
-impl<SG> NamespacedClient for RetryClient<SG>
-where
-    SG: NamespacedClient,
-{
-    fn namespace(&self) -> String {
-        self.client.namespace()
+impl From<RetryOptions> for backoff::ExponentialBackoff {
+    fn from(c: RetryOptions) -> Self {
+        c.into_exp_backoff(SystemClock::default())
     }
+}
 
-    fn identity(&self) -> String {
-        self.client.identity()
-    }
+pub(crate) fn make_future_retry<R, F, Fut>(
+    info: CallInfo,
+    factory: F,
+) -> FutureRetry<F, TonicErrorHandler<SystemClock>>
+where
+    F: FnMut() -> Fut + Unpin,
+    Fut: Future<Output = Result<R, tonic::Status>>,
+{
+    FutureRetry::new(
+        factory,
+        TonicErrorHandler::new(info, RetryOptions::throttle_retry_policy()),
+    )
 }
 
 #[derive(Debug)]
@@ -525,12 +572,11 @@ mod tests {
     ) {
         // A bit odd, but we don't need a real client to test the retry client passes through the
         // correct retry config
-        let fake_retry = RetryClient::new((), TEST_RETRY_CONFIG);
         let mut req = req.into_request();
         req.extensions_mut().insert(IsWorkerTaskLongPoll);
         for i in 1..=50 {
             let mut err_handler = TonicErrorHandler::new(
-                fake_retry.get_call_info::<R>(call_name, Some(&req)),
+                TEST_RETRY_CONFIG.get_call_info::<R>(call_name, Some(&req)),
                 RetryOptions::throttle_retry_policy(),
             );
             let result = err_handler.handle(i, Status::new(Code::Unknown, "Ahh"));
@@ -557,13 +603,12 @@ mod tests {
         )]
         (call_name, req): (&'static str, R),
     ) {
-        let fake_retry = RetryClient::new((), TEST_RETRY_CONFIG);
         let mut req = req.into_request();
         req.extensions_mut().insert(IsWorkerTaskLongPoll);
         // For some reason we will get cancelled in these situations occasionally (always?) too
         for code in [Code::Cancelled, Code::DeadlineExceeded] {
             let mut err_handler = TonicErrorHandler::new(
-                fake_retry.get_call_info::<R>(call_name, Some(&req)),
+                TEST_RETRY_CONFIG.get_call_info::<R>(call_name, Some(&req)),
                 RetryOptions::throttle_retry_policy(),
             );
             for i in 1..=5 {

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

@@ -213,28 +213,48 @@ impl ClientWorkerSetImpl {
         Ok(())
     }
 
-    fn unregister(
-        &mut self,
-        worker_instance_key: Uuid,
-    ) -> Result<Arc<dyn ClientWorker + Send + Sync>, anyhow::Error> {
-        let worker = self
-            .all_workers
-            .remove(&worker_instance_key)
-            .ok_or_else(|| {
-                anyhow::anyhow!("Worker with worker_instance_key {worker_instance_key} not found")
-            })?;
+    /// Slot provider should be unregistered at the beginning of worker shutdown, in order to disable
+    /// eager workflow start.
+    fn unregister_slot_provider(&mut self, worker_instance_key: Uuid) -> Result<(), anyhow::Error> {
+        let worker = self.all_workers.get(&worker_instance_key).ok_or_else(|| {
+            anyhow::anyhow!("Worker not in all_workers during slot provider unregister")
+        })?;
 
         let slot_key = SlotKey::new(
             worker.namespace().to_string(),
             worker.task_queue().to_string(),
         );
-
         if let Some(slot_vec) = self.slot_providers.get_mut(&slot_key) {
             slot_vec.retain(|info| info.worker_id != worker_instance_key);
             if slot_vec.is_empty() {
                 self.slot_providers.remove(&slot_key);
             }
         }
+        Ok(())
+    }
+
+    fn finalize_unregister(
+        &mut self,
+        worker_instance_key: Uuid,
+    ) -> Result<Arc<dyn ClientWorker + Send + Sync>, anyhow::Error> {
+        if let Some(worker) = self.all_workers.get(&worker_instance_key)
+            && let Some(slot_vec) = self.slot_providers.get(&SlotKey::new(
+                worker.namespace().to_string(),
+                worker.task_queue().to_string(),
+            ))
+            && slot_vec
+                .iter()
+                .any(|info| info.worker_id == worker_instance_key)
+        {
+            return Err(anyhow::anyhow!(
+                "Worker still in slot_providers during finalize"
+            ));
+        }
+
+        let worker = self
+            .all_workers
+            .remove(&worker_instance_key)
+            .ok_or_else(|| anyhow::anyhow!("Worker not found in all_workers"))?;
 
         if let Some(w) = self.shared_worker.get_mut(worker.namespace()) {
             let (callback, is_empty) = w.unregister_callback(worker.worker_instance_key());
@@ -278,7 +298,7 @@ pub trait SharedNamespaceWorkerTrait {
 /// Enables local workers to make themselves visible to a shared client instance.
 ///
 /// For slot managing, there can only be one worker registered per
-/// namespace+queue_name+client, others will return an error.
+/// namespace+queue_name+connection, others will return an error.
 /// It also provides a convenient method to find compatible slots within the collection.
 pub struct ClientWorkerSet {
     worker_grouping_key: Uuid,
@@ -323,12 +343,24 @@ impl ClientWorkerSet {
             .register(worker, skip_client_worker_set_check)
     }
 
-    /// Unregisters a local worker, typically when that worker starts shutdown.
-    pub fn unregister_worker(
+    /// Disables Eager Workflow Start for this worker. This must be called before
+    /// `finalize_unregister`, otherwise `finalize_unregister` will return an err.
+    pub fn unregister_slot_provider(&self, worker_instance_key: Uuid) -> Result<(), anyhow::Error> {
+        self.worker_manager
+            .write()
+            .unregister_slot_provider(worker_instance_key)
+    }
+
+    /// Finalizes unregistering of worker from client. This must be called at the end of worker
+    /// shutdown in order to finalize shutdown for worker heartbeat properly. Must call after
+    /// `unregister_slot_provider`, otherwise an err will be returned.
+    pub fn finalize_unregister(
         &self,
         worker_instance_key: Uuid,
     ) -> Result<Arc<dyn ClientWorker + Send + Sync>, anyhow::Error> {
-        self.worker_manager.write().unregister(worker_instance_key)
+        self.worker_manager
+            .write()
+            .finalize_unregister(worker_instance_key)
     }
 
     /// Returns the worker grouping key, which is unique for each worker.
@@ -651,14 +683,14 @@ mod tests {
             let worker_instance_key = mock_provider.worker_instance_key();
 
             let result = manager.register_worker(Arc::new(mock_provider), false);
-            if result.is_ok() {
-                successful_registrations += 1;
-                worker_keys.push(worker_instance_key);
-            } else {
+            if let Err(err) = result {
                 // Should get error for overlapping worker task types
-                assert!(result.unwrap_err().to_string().contains(
+                assert!(err.to_string().contains(
                     "Registration of multiple workers with overlapping worker task types"
                 ));
+            } else {
+                successful_registrations += 1;
+                worker_keys.push(worker_instance_key);
             }
         }
 
@@ -666,9 +698,12 @@ mod tests {
         assert_eq!(3, manager.num_providers());
 
         let count = worker_keys.iter().fold(0, |count, key| {
-            manager.unregister_worker(*key).unwrap();
+            manager.unregister_slot_provider(*key).unwrap();
+            manager.finalize_unregister(*key).unwrap();
             // expect error since worker is already unregistered
-            let result = manager.unregister_worker(*key);
+            let result = manager.unregister_slot_provider(*key);
+            assert!(result.is_err());
+            let result = manager.finalize_unregister(*key);
             assert!(result.is_err());
             count + 1
         });
@@ -896,7 +931,10 @@ mod tests {
             assert_eq!(providers[1].build_id, Some("build-1".to_string()));
         }
 
-        manager.unregister_worker(worker2_instance_key).unwrap();
+        manager
+            .unregister_slot_provider(worker2_instance_key)
+            .unwrap();
+        manager.finalize_unregister(worker2_instance_key).unwrap();
 
         {
             let impl_ref = manager.worker_manager.read();
@@ -1016,7 +1054,10 @@ mod tests {
         drop(impl_ref);
 
         // Unregister first worker
-        manager.unregister_worker(worker_instance_key1).unwrap();
+        manager
+            .unregister_slot_provider(worker_instance_key1)
+            .unwrap();
+        manager.finalize_unregister(worker_instance_key1).unwrap();
 
         // After unregistering first worker: 1 slot provider, 1 heartbeat worker, shared worker still exists
         assert_eq!(1, manager.num_providers());
@@ -1036,7 +1077,10 @@ mod tests {
         drop(impl_ref);
 
         // Unregister second worker
-        manager.unregister_worker(worker_instance_key2).unwrap();
+        manager
+            .unregister_slot_provider(worker_instance_key2)
+            .unwrap();
+        manager.finalize_unregister(worker_instance_key2).unwrap();
 
         // After unregistering last worker: 0 slot providers, 0 heartbeat workers, shared worker is removed
         assert_eq!(0, manager.num_providers());
@@ -1447,8 +1491,11 @@ mod tests {
         );
 
         manager
-            .unregister_worker(wf_worker_key)
-            .expect("should unregister workflow worker");
+            .unregister_slot_provider(wf_worker_key)
+            .expect("should unregister slot provider for workflow worker");
+        manager
+            .finalize_unregister(wf_worker_key)
+            .expect("should finalize unregister for workflow worker");
 
         // Activity worker should still be registered
         assert_eq!(1, manager.num_providers());
@@ -1460,9 +1507,37 @@ mod tests {
         );
 
         manager
-            .unregister_worker(act_worker_key)
-            .expect("should unregister activity worker");
+            .unregister_slot_provider(act_worker_key)
+            .expect("should unregister slot provider for activity worker");
+        manager
+            .finalize_unregister(act_worker_key)
+            .expect("should finalize unregister for activity worker");
 
         assert_eq!(0, manager.num_providers());
     }
+
+    #[test]
+    fn worker_unregister_order() {
+        let manager = ClientWorkerSet::new();
+        let worker = new_mock_provider_with_heartbeat(
+            "namespace1".to_string(),
+            "queue1".to_string(),
+            true,
+            None,
+        );
+        let worker_instance_key = worker.worker_instance_key();
+        manager.register_worker(Arc::new(worker), false).unwrap();
+
+        let res = manager.finalize_unregister(worker_instance_key);
+        assert!(res.is_err());
+        let err_string = res.err().map(|e| e.to_string()).unwrap();
+        assert!(err_string.contains("Worker still in slot_providers during finalize"));
+
+        // previous incorrect call to finalize_unregister should not cause any state leaks when
+        // properly removed later
+        manager
+            .unregister_slot_provider(worker_instance_key)
+            .unwrap();
+        manager.finalize_unregister(worker_instance_key).unwrap();
+    }
 }