@temporalio/core-bridge 1.4.4 → 1.5.1

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

@@ -101,32 +101,49 @@ impl<SG> RetryClient<SG> {
         F: FnMut() -> Fut + Unpin,
         Fut: Future<Output = Result<R>>,
     {
-        FutureRetry::new(factory, TonicErrorHandler::new(rtc, call_name))
+        FutureRetry::new(
+            factory,
+            TonicErrorHandler::new(rtc, RetryConfig::throttle_retry_policy(), call_name),
+        )
     }
 }
 
 #[derive(Debug)]
 pub(crate) struct TonicErrorHandler<C: Clock> {
     backoff: ExponentialBackoff<C>,
+    throttle_backoff: ExponentialBackoff<C>,
     max_retries: usize,
     call_type: CallType,
     call_name: &'static str,
 }
 impl TonicErrorHandler<SystemClock> {
-    fn new(cfg: RetryConfig, call_name: &'static str) -> Self {
-        Self::new_with_clock(cfg, call_name, SystemClock::default())
+    fn new(cfg: RetryConfig, throttle_cfg: RetryConfig, call_name: &'static str) -> Self {
+        Self::new_with_clock(
+            cfg,
+            throttle_cfg,
+            call_name,
+            SystemClock::default(),
+            SystemClock::default(),
+        )
     }
 }
 impl<C> TonicErrorHandler<C>
 where
     C: Clock,
 {
-    fn new_with_clock(cfg: RetryConfig, call_name: &'static str, clock: C) -> Self {
+    fn new_with_clock(
+        cfg: RetryConfig,
+        throttle_cfg: RetryConfig,
+        call_name: &'static str,
+        clock: C,
+        throttle_clock: C,
+    ) -> Self {
         Self {
             max_retries: cfg.max_retries,
             call_type: CallType::from_call_name(call_name),
             call_name,
             backoff: cfg.into_exp_backoff(clock),
+            throttle_backoff: throttle_cfg.into_exp_backoff(throttle_clock),
         }
     }
     const fn should_log_retry_warning(&self, cur_attempt: usize) -> bool {
@@ -184,10 +201,12 @@ where
             match self.backoff.next_backoff() {
                 None => RetryPolicy::ForwardError(e), // None is returned when we've ran out of time
                 Some(backoff) => {
-                    if cfg!(test) {
-                        // Allow unit tests to do lots of retries quickly. This does *not* apply
-                        // during integration testing, importantly.
-                        RetryPolicy::WaitRetry(Duration::from_millis(1))
+                    // We treat ResourceExhausted as a special case and backoff more
+                    // so we don't overload the server
+                    if e.code() == Code::ResourceExhausted {
+                        let extended_backoff =
+                            backoff.max(self.throttle_backoff.next_backoff().unwrap_or_default());
+                        RetryPolicy::WaitRetry(extended_backoff)
                     } else {
                         RetryPolicy::WaitRetry(backoff)
                     }
@@ -529,6 +548,16 @@ mod tests {
     use std::{ops::Add, time::Instant};
     use tonic::Status;
 
+    /// Predefined retry configs with low durations to make unit tests faster
+    const TEST_RETRY_CONFIG: RetryConfig = RetryConfig {
+        initial_interval: Duration::from_millis(1),
+        randomization_factor: 0.0,
+        multiplier: 1.1,
+        max_interval: Duration::from_millis(2),
+        max_elapsed_time: None,
+        max_retries: 10,
+    };
+
     #[tokio::test]
     async fn non_retryable_errors() {
         for code in [
@@ -547,7 +576,7 @@ mod tests {
                 .expect_cancel_activity_task()
                 .returning(move |_, _| Err(Status::new(code, "non-retryable failure")))
                 .times(1);
-            let retry_client = RetryClient::new(mock_client, Default::default());
+            let retry_client = RetryClient::new(mock_client, TEST_RETRY_CONFIG);
             let result = retry_client
                 .cancel_activity_task(vec![1].into(), None)
                 .await;
@@ -575,12 +604,13 @@ mod tests {
             Code::Unimplemented,
         ] {
             for call_name in [POLL_WORKFLOW_METH_NAME, POLL_ACTIVITY_METH_NAME] {
-                let retry_cfg = RetryConfig::poll_retry_policy();
                 let mut err_handler = TonicErrorHandler {
-                    max_retries: retry_cfg.max_retries,
+                    max_retries: TEST_RETRY_CONFIG.max_retries,
                     call_type: CallType::LongPoll,
                     call_name,
-                    backoff: retry_cfg.into_exp_backoff(FixedClock(Instant::now())),
+                    backoff: TEST_RETRY_CONFIG.into_exp_backoff(FixedClock(Instant::now())),
+                    throttle_backoff: TEST_RETRY_CONFIG
+                        .into_exp_backoff(FixedClock(Instant::now())),
                 };
                 let result = err_handler.handle(1, Status::new(code, "Ahh"));
                 assert_matches!(result, RetryPolicy::WaitRetry(_));
@@ -599,12 +629,13 @@ mod tests {
     async fn long_poll_retryable_errors_never_fatal() {
         for code in RETRYABLE_ERROR_CODES {
             for call_name in [POLL_WORKFLOW_METH_NAME, POLL_ACTIVITY_METH_NAME] {
-                let retry_cfg = RetryConfig::poll_retry_policy();
                 let mut err_handler = TonicErrorHandler {
-                    max_retries: retry_cfg.max_retries,
+                    max_retries: TEST_RETRY_CONFIG.max_retries,
                     call_type: CallType::LongPoll,
                     call_name,
-                    backoff: retry_cfg.into_exp_backoff(FixedClock(Instant::now())),
+                    backoff: TEST_RETRY_CONFIG.into_exp_backoff(FixedClock(Instant::now())),
+                    throttle_backoff: TEST_RETRY_CONFIG
+                        .into_exp_backoff(FixedClock(Instant::now())),
                 };
                 let result = err_handler.handle(1, Status::new(code, "Ahh"));
                 assert_matches!(result, RetryPolicy::WaitRetry(_));
@@ -621,7 +652,12 @@ mod tests {
 
     #[tokio::test]
     async fn retryable_errors() {
-        for code in RETRYABLE_ERROR_CODES {
+        // Take out retry exhausted since it gets a special policy which would make this take ages
+        for code in RETRYABLE_ERROR_CODES
+            .iter()
+            .copied()
+            .filter(|p| p != &Code::ResourceExhausted)
+        {
             let mut mock_client = MockWorkflowClientTrait::new();
             mock_client
                 .expect_cancel_activity_task()
@@ -632,7 +668,7 @@ mod tests {
                 .returning(|_, _| Ok(Default::default()))
                 .times(1);
 
-            let retry_client = RetryClient::new(mock_client, Default::default());
+            let retry_client = RetryClient::new(mock_client, TEST_RETRY_CONFIG);
             let result = retry_client
                 .cancel_activity_task(vec![1].into(), None)
                 .await;
@@ -641,15 +677,53 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn retry_resource_exhausted() {
+        let mut err_handler = TonicErrorHandler {
+            max_retries: TEST_RETRY_CONFIG.max_retries,
+            call_type: CallType::Normal,
+            call_name: POLL_WORKFLOW_METH_NAME,
+            backoff: TEST_RETRY_CONFIG.into_exp_backoff(FixedClock(Instant::now())),
+            throttle_backoff: RetryConfig {
+                initial_interval: Duration::from_millis(2),
+                randomization_factor: 0.0,
+                multiplier: 4.0,
+                max_interval: Duration::from_millis(10),
+                max_elapsed_time: None,
+                max_retries: 10,
+            }
+            .into_exp_backoff(FixedClock(Instant::now())),
+        };
+        let result = err_handler.handle(1, Status::new(Code::ResourceExhausted, "leave me alone"));
+        match result {
+            RetryPolicy::WaitRetry(duration) => assert_eq!(duration, Duration::from_millis(2)),
+            _ => panic!(),
+        }
+        err_handler.backoff.clock.0 = err_handler.backoff.clock.0.add(Duration::from_millis(10));
+        err_handler.throttle_backoff.clock.0 = err_handler
+            .throttle_backoff
+            .clock
+            .0
+            .add(Duration::from_millis(10));
+        let result = err_handler.handle(2, Status::new(Code::ResourceExhausted, "leave me alone"));
+        match result {
+            RetryPolicy::WaitRetry(duration) => assert_eq!(duration, Duration::from_millis(8)),
+            _ => panic!(),
+        }
+    }
+
     #[tokio::test]
     async fn long_poll_retries_forever() {
         // 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((), Default::default());
+        let fake_retry = RetryClient::new((), TEST_RETRY_CONFIG);
         for i in 1..=50 {
             for call in [POLL_WORKFLOW_METH_NAME, POLL_ACTIVITY_METH_NAME] {
-                let mut err_handler =
-                    TonicErrorHandler::new(fake_retry.get_retry_config(call), call);
+                let mut err_handler = TonicErrorHandler::new(
+                    fake_retry.get_retry_config(call),
+                    fake_retry.get_retry_config(call),
+                    call,
+                );
                 let result = err_handler.handle(i, Status::new(Code::Unknown, "Ahh"));
                 assert_matches!(result, RetryPolicy::WaitRetry(_));
             }
@@ -658,12 +732,15 @@ mod tests {
 
     #[tokio::test]
     async fn long_poll_retries_deadline_exceeded() {
-        let fake_retry = RetryClient::new((), Default::default());
+        let fake_retry = RetryClient::new((), TEST_RETRY_CONFIG);
         // For some reason we will get cancelled in these situations occasionally (always?) too
         for code in [Code::Cancelled, Code::DeadlineExceeded] {
             for call in [POLL_WORKFLOW_METH_NAME, POLL_ACTIVITY_METH_NAME] {
-                let mut err_handler =
-                    TonicErrorHandler::new(fake_retry.get_retry_config(call), call);
+                let mut err_handler = TonicErrorHandler::new(
+                    fake_retry.get_retry_config(call),
+                    fake_retry.get_retry_config(call),
+                    call,
+                );
                 for i in 1..=5 {
                     let result = err_handler.handle(i, Status::new(code, "retryable failure"));
                     assert_matches!(result, RetryPolicy::WaitRetry(_));

package/sdk-core/core/Cargo.toml

@@ -20,7 +20,7 @@ async-trait = "0.1"
 base64 = "0.13"
 crossbeam = "0.8"
 dashmap = "5.0"
-derive_builder = "0.11"
+derive_builder = "0.12"
 derive_more = "0.99"
 enum_dispatch = "0.3"
 flate2 = "1.0"
@@ -34,7 +34,7 @@ lazy_static = "1.4"
 log = "0.4"
 lru = "0.8"
 mockall = "0.11"
-nix = "0.25"
+nix = "0.26"
 once_cell = "1.5"
 opentelemetry = { version = "0.18", features = ["rt-tokio"] }
 opentelemetry-otlp = { version = "0.11", features = ["tokio", "metrics"] }
@@ -45,7 +45,7 @@ prost = "0.11"
 prost-types = "0.11"
 rand = "0.8.3"
 reqwest = { version = "0.11", features = ["json", "stream", "rustls-tls", "tokio-rustls"], default-features = false }
-ringbuf = "0.2"
+ringbuf = "0.3"
 serde = "1.0"
 serde_json = "1.0"
 siphasher = "0.3"
@@ -59,7 +59,7 @@ tonic = { version = "0.8", features = ["tls", "tls-roots"] }
 tracing = { version = "0.1", features = ["log-always"] }
 tracing-futures = "0.2"
 tracing-opentelemetry = "0.18"
-tracing-subscriber = { version = "0.3", features = ["parking_lot", "env-filter"] }
+tracing-subscriber = { version = "0.3", features = ["parking_lot", "env-filter", "registry"] }
 url = "2.2"
 uuid = { version = "1.1", features = ["v4"] }
 zip = "0.6"
@@ -86,7 +86,7 @@ version = "0.1"
 assert_matches = "1.4"
 bimap = "0.6.1"
 criterion = "0.4"
-rstest = "0.15"
+rstest = "0.16"
 temporal-sdk-core-test-utils = { path = "../test-utils" }
 temporal-sdk = { path = "../sdk" }
 

package/sdk-core/core/benches/workflow_replay.rs

@@ -1,29 +1,30 @@
 use criterion::{criterion_group, criterion_main, Criterion};
 use futures::StreamExt;
 use std::time::Duration;
-use temporal_sdk::{WfContext, Worker, WorkflowFunction};
-use temporal_sdk_core::{telemetry_init, TelemetryOptionsBuilder};
+use temporal_sdk::{WfContext, WorkflowFunction};
+use temporal_sdk_core::replay::HistoryForReplay;
 use temporal_sdk_core_protos::DEFAULT_WORKFLOW_TYPE;
-use temporal_sdk_core_test_utils::{canned_histories, init_core_replay_preloaded};
+use temporal_sdk_core_test_utils::{canned_histories, replay_sdk_worker};
 
 pub fn criterion_benchmark(c: &mut Criterion) {
     let tokio_runtime = tokio::runtime::Builder::new_current_thread()
         .enable_time()
         .build()
         .unwrap();
-    telemetry_init(&TelemetryOptionsBuilder::default().build().unwrap()).unwrap();
     let _g = tokio_runtime.enter();
 
     let num_timers = 10;
     let t = canned_histories::long_sequential_timers(num_timers as usize);
-    let hist = t.get_full_history_info().unwrap().into();
+    let hist = HistoryForReplay::new(
+        t.get_full_history_info().unwrap().into(),
+        "whatever".to_string(),
+    );
 
     c.bench_function("Small history replay", |b| {
         b.iter(|| {
             tokio_runtime.block_on(async {
                 let func = timers_wf(num_timers);
-                let (worker, _) = init_core_replay_preloaded("replay_bench", &hist);
-                let mut worker = Worker::new_from_core(worker, "replay_bench".to_string());
+                let mut worker = replay_sdk_worker([hist.clone()]);
                 worker.register_wf(DEFAULT_WORKFLOW_TYPE, func);
                 worker.run().await.unwrap();
             })
@@ -32,14 +33,16 @@ pub fn criterion_benchmark(c: &mut Criterion) {
 
     let num_tasks = 50;
     let t = canned_histories::lots_of_big_signals(num_tasks);
-    let hist = t.get_full_history_info().unwrap().into();
+    let hist = HistoryForReplay::new(
+        t.get_full_history_info().unwrap().into(),
+        "whatever".to_string(),
+    );
 
     c.bench_function("Large payloads history replay", |b| {
         b.iter(|| {
             tokio_runtime.block_on(async {
                 let func = big_signals_wf(num_tasks);
-                let (worker, _) = init_core_replay_preloaded("large_hist_bench", &hist);
-                let mut worker = Worker::new_from_core(worker, "large_hist_bench".to_string());
+                let mut worker = replay_sdk_worker([hist.clone()]);
                 worker.register_wf(DEFAULT_WORKFLOW_TYPE, func);
                 worker.run().await.unwrap();
             })

package/sdk-core/core/src/abstractions.rs

@@ -4,10 +4,9 @@ use crate::MetricsContext;
 use futures::{stream, Stream, StreamExt};
 use std::{
     fmt::{Debug, Formatter},
-    future::Future,
     sync::Arc,
 };
-use tokio::sync::{AcquireError, Notify, OwnedSemaphorePermit, Semaphore, TryAcquireError};
+use tokio::sync::{AcquireError, OwnedSemaphorePermit, Semaphore, TryAcquireError};
 
 /// Wraps a [Semaphore] with a function call that is fed the available permits any time a permit is
 /// acquired or restored through the provided methods
@@ -82,24 +81,29 @@ impl Debug for OwnedMeteredSemPermit {
 }
 
 /// From the input stream, create a new stream which only pulls from the input stream when allowed.
-/// When allowed is determined by the passed in `proceeder` which must return a future every time
-/// it's called. The input stream is only pulled from when that future resolves.
-pub(crate) fn stream_when_allowed<S, F, FF>(
+/// When allowed is determined by the passed in `proceeder` emitting an item. The input stream is
+/// only pulled from when that future resolves.
+///
+/// This is *almost* identical to `zip`, but does not terminate early if the input stream closes.
+/// The proceeder must allow the poll before the returned stream closes. If the proceeder terminates
+/// the overall stream will terminate.
+pub(crate) fn stream_when_allowed<S, AS>(
     input: S,
-    proceeder: FF,
-) -> impl Stream<Item = (S::Item, F::Output)>
+    proceeder: AS,
+) -> impl Stream<Item = (S::Item, AS::Item)>
 where
     S: Stream + Send + 'static,
-    F: Future,
-    FF: FnMut() -> F,
+    AS: Stream + Send + 'static,
 {
-    let acceptable_notify = Arc::new(Notify::new());
-    acceptable_notify.notify_one();
     let stream = stream::unfold(
-        (proceeder, input.boxed()),
+        (proceeder.boxed(), input.boxed()),
         |(mut proceeder, mut input)| async {
-            let v = proceeder().await;
-            input.next().await.map(|i| ((i, v), (proceeder, input)))
+            let v = proceeder.next().await;
+            if let Some(v) = v {
+                input.next().await.map(|i| ((i, v), (proceeder, input)))
+            } else {
+                None
+            }
         },
     );
     stream
@@ -108,7 +112,7 @@ where
 macro_rules! dbg_panic {
   ($($arg:tt)*) => {
       error!($($arg)*);
-      debug_assert!(true, $($arg)*);
+      debug_assert!(false, $($arg)*);
   };
 }
 pub(crate) use dbg_panic;
@@ -117,19 +121,15 @@ pub(crate) use dbg_panic;
 mod tests {
     use super::*;
     use futures::pin_mut;
-    use std::{cell::RefCell, task::Poll};
+    use std::task::Poll;
     use tokio::sync::mpsc::unbounded_channel;
+    use tokio_stream::wrappers::UnboundedReceiverStream;
 
-    // This is fine. Test only / guaranteed to happen serially.
-    #[allow(clippy::await_holding_refcell_ref)]
     #[test]
     fn stream_when_allowed_works() {
         let inputs = stream::iter([1, 2, 3]);
         let (allow_tx, allow_rx) = unbounded_channel();
-        let allow_rx = RefCell::new(allow_rx);
-        let when_allowed = stream_when_allowed(inputs, || async {
-            allow_rx.borrow_mut().recv().await.unwrap()
-        });
+        let when_allowed = stream_when_allowed(inputs, UnboundedReceiverStream::new(allow_rx));
 
         let waker = futures::task::noop_waker_ref();
         let mut cx = std::task::Context::from_waker(waker);

package/sdk-core/core/src/core_tests/activity_tasks.rs

@@ -438,7 +438,7 @@ async fn many_concurrent_heartbeat_cancels() {
 #[tokio::test]
 async fn activity_timeout_no_double_resolve() {
     let t = canned_histories::activity_double_resolve_repro();
-    let core = build_fake_worker("fake_wf_id", t, &[3]);
+    let core = build_fake_worker("fake_wf_id", t, [3]);
     let activity_id = 1;
 
     poll_and_reply(