From 8fbad716a06158308311d428b4240c69ccc26560 Mon Sep 17 00:00:00 2001 From: James Watkins-Harvey Date: Thu, 12 Feb 2026 06:47:23 -0500 Subject: [PATCH 1/4] Poller autoscaler: add backoff on ResourceExhausted errors --- crates/sdk-core/Cargo.toml | 1 + crates/sdk-core/src/pollers/poll_buffer.rs | 142 ++++++++++++++++++++- 2 files changed, 136 insertions(+), 7 deletions(-) diff --git a/crates/sdk-core/Cargo.toml b/crates/sdk-core/Cargo.toml index c2131d9f9..ed2b9bede 100644 --- a/crates/sdk-core/Cargo.toml +++ b/crates/sdk-core/Cargo.toml @@ -33,6 +33,7 @@ antithesis_assertions = ["dep:antithesis_sdk"] anyhow = "1.0" antithesis_sdk = { version = "0.2.1", optional = true, default-features = false, features = ["full"] } assert_matches = { version = "1.5", optional = true } +backoff = "0.4" bimap = { version = "0.6.3", optional = true } async-trait = "0.1" bon = { workspace = true } diff --git a/crates/sdk-core/src/pollers/poll_buffer.rs b/crates/sdk-core/src/pollers/poll_buffer.rs index 80843e175..8791bd988 100644 --- a/crates/sdk-core/src/pollers/poll_buffer.rs +++ b/crates/sdk-core/src/pollers/poll_buffer.rs @@ -6,6 +6,7 @@ use crate::{ client::{PollActivityOptions, PollOptions, PollWorkflowOptions, WorkerClient}, }, }; +use backoff::{SystemClock, backoff::Backoff, exponential::ExponentialBackoff}; use crossbeam_utils::atomic::AtomicCell; use futures_util::{FutureExt, StreamExt, future::BoxFuture}; use governor::{Quota, RateLimiter}; @@ -366,13 +367,18 @@ where r = pf(timeout_override) => r, _ = poll_interruptor => return, }; - drop(active_guard); if let Ok(r) = &r && let Some(ppf) = post_pf.as_ref() { ppf(r); } - if report_handle.poll_result(&r) { + let (should_forward, backoff_duration) = report_handle.poll_result(&r); + if let Some(duration) = backoff_duration { + // Apply backoff BEFORE dropping active_guard to prevent next poll from starting + tokio::time::sleep(duration).await; + } + drop(active_guard); + if should_forward { let _ = tx.send(r.map(|r| (r, permit))); } }); @@ -486,6 +492,19 @@ where ingested_last_period: Default::default(), scale_up_allowed: AtomicBool::new(true), last_successful_poll_time, + // Use same backoff config as gRPC client's throttle_backoff for ResourceExhausted + // (1s initial, 10s max, 2x multiplier, 0.2 randomization, unlimited retries) + resource_exhausted_backoff: std::sync::Mutex::new(ExponentialBackoff { + // Copied from RetryOptions::throttle_retry_policy() + current_interval: Duration::from_secs(1), + initial_interval: Duration::from_secs(1), + randomization_factor: 0.2, + multiplier: 2.0, + max_interval: Duration::from_secs(10), + max_elapsed_time: None, + clock: SystemClock::default(), + start_time: std::time::Instant::now(), + }), }); let rhc = report_handle.clone(); let ingestor_task = if behavior.is_autoscaling() { @@ -549,18 +568,31 @@ struct PollScalerReportHandle { ingested_last_period: AtomicUsize, scale_up_allowed: AtomicBool, last_successful_poll_time: Arc>>, + + resource_exhausted_backoff: std::sync::Mutex>, } impl PollScalerReportHandle { - /// Returns true if the response should be passed on, false if it should be swallowed - fn poll_result(&self, res: &Result) -> bool { + /// Returns (should_forward, backoff_duration) + /// - should_forward: true if the response should be passed on, false if it should be swallowed + /// - backoff_duration: Some(duration) if we should sleep before the next poll + fn poll_result( + &self, + res: &Result, + ) -> (bool, Option) { match res { Ok(res) => { self.last_successful_poll_time .store(Some(SystemTime::now())); + + // Reset backoff on successful poll + if let Ok(mut backoff) = self.resource_exhausted_backoff.lock() { + backoff.reset(); + } + if let PollerBehavior::SimpleMaximum(_) = self.behavior { // We don't do auto-scaling with the simple max - return true; + return (true, None); } if !res.is_empty() { self.ingested_this_period.fetch_add(1, Ordering::Relaxed); @@ -602,6 +634,14 @@ impl PollScalerReportHandle { if e.code() == Code::ResourceExhausted { // Scale down significantly for resource exhaustion self.change_target(usize::saturating_div, 2); + + let backoff_duration = self + .resource_exhausted_backoff + .lock() + .unwrap() + .next_backoff(); + + return (false, backoff_duration); } else { // Other codes that would normally have made us back off briefly can // reclaim this poller @@ -612,13 +652,14 @@ impl PollScalerReportHandle { // logic. IE: We don't want to fail callers because we said we wanted to know // about ResourceExhausted errors, but we haven't seen a scaling decision yet, // so we're not reacting to errors, only propagating them. - return !e + let should_forward = !e .metadata() .contains_key(ERROR_RETURNED_DUE_TO_SHORT_CIRCUIT); + return (should_forward, None); } } } - true + (true, None) } #[inline] @@ -972,4 +1013,91 @@ mod tests { *v.borrow_mut() = None; }); } + + #[tokio::test] + async fn autoscaler_applies_backoff_on_resource_exhausted() { + use temporalio_common::protos::temporal::api::taskqueue::v1::PollerScalingDecision; + + let call_count = Arc::new(AtomicUsize::new(0)); + let call_count_clone = call_count.clone(); + let first_poll_done = Arc::new(AtomicBool::new(false)); + let first_poll_done_clone = first_poll_done.clone(); + + let mut mock_client = mock_manual_worker_client(); + mock_client + .expect_poll_workflow_task() + .returning(move |_, _| { + call_count_clone.fetch_add(1, Ordering::SeqCst); + let first_done = first_poll_done_clone.clone(); + async move { + // First poll: return empty response with scaling decision + // This sets ever_saw_scaling_decision to true + if !first_done.swap(true, Ordering::SeqCst) { + Ok(PollWorkflowTaskQueueResponse { + task_token: vec![], // Empty poll + poller_scaling_decision: Some(PollerScalingDecision { + // Aggressively scale up to max + poll_request_delta_suggestion: 100, + ..Default::default() + }), + ..Default::default() + }) + } else { + // All subsequent polls: return ResourceExhausted immediately + // This simulates the namespace RPS limit being exceeded + Err(tonic::Status::new( + Code::ResourceExhausted, + "namespace rate limit exceeded", + )) + } + } + .boxed() + }); + + let pb = Arc::new(LongPollBuffer::new_workflow_task( + Arc::new(mock_client), + "sometq".to_string(), + None, + PollerBehavior::Autoscaling { + minimum: 5, + maximum: 100, + initial: 10, + }, + fixed_size_permit_dealer(10), + CancellationToken::new(), + None::, + WorkflowTaskOptions { + wft_poller_shared: Some(Arc::new(WFTPollerShared::new(Some(10)))), + }, + Arc::new(AtomicCell::new(None)), + )); + + // Trigger the first poll to initialize and get the scaling decision + let pb_clone = pb.clone(); + tokio::spawn(async move { + let _ = pb_clone.poll().await; + }); + + // Wait for the first poll to complete + tokio::time::sleep(Duration::from_millis(20)).await; + + // Let the hot loop run for 100ms while we continue to attempt consuming + let start = std::time::Instant::now(); + tokio::time::sleep(Duration::from_millis(100)).await; + let elapsed = start.elapsed(); + let hot_loop_calls = call_count.load(Ordering::SeqCst); + + // Without backoff, this was producing ~6300 polls in ~100ms on my machine. + // With exponential backoff, I'm getting exactly 10 (initial poller count). + assert!( + hot_loop_calls == 10, + "Expected proper backoff with == 10 polls in 100ms, but got {} polls.", + hot_loop_calls + ); + + Arc::try_unwrap(pb) + .unwrap_or_else(|_| panic!("Failed to unwrap Arc")) + .shutdown() + .await; + } } From b58255a82daea34d499679bd83320d7553557603 Mon Sep 17 00:00:00 2001 From: James Watkins-Harvey Date: Thu, 12 Feb 2026 08:38:42 -0500 Subject: [PATCH 2/4] lint --- crates/sdk-core/src/pollers/poll_buffer.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/crates/sdk-core/src/pollers/poll_buffer.rs b/crates/sdk-core/src/pollers/poll_buffer.rs index 8791bd988..1908ffd9e 100644 --- a/crates/sdk-core/src/pollers/poll_buffer.rs +++ b/crates/sdk-core/src/pollers/poll_buffer.rs @@ -1038,7 +1038,6 @@ mod tests { poller_scaling_decision: Some(PollerScalingDecision { // Aggressively scale up to max poll_request_delta_suggestion: 100, - ..Default::default() }), ..Default::default() }) @@ -1084,7 +1083,6 @@ mod tests { // Let the hot loop run for 100ms while we continue to attempt consuming let start = std::time::Instant::now(); tokio::time::sleep(Duration::from_millis(100)).await; - let elapsed = start.elapsed(); let hot_loop_calls = call_count.load(Ordering::SeqCst); // Without backoff, this was producing ~6300 polls in ~100ms on my machine. From 0f75b7719f5c2180edf073c76d7283e24c97b691 Mon Sep 17 00:00:00 2001 From: James Watkins-Harvey Date: Thu, 12 Feb 2026 08:44:01 -0500 Subject: [PATCH 3/4] lint --- crates/sdk-core/src/pollers/poll_buffer.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/crates/sdk-core/src/pollers/poll_buffer.rs b/crates/sdk-core/src/pollers/poll_buffer.rs index 1908ffd9e..496bd7a39 100644 --- a/crates/sdk-core/src/pollers/poll_buffer.rs +++ b/crates/sdk-core/src/pollers/poll_buffer.rs @@ -1081,7 +1081,6 @@ mod tests { tokio::time::sleep(Duration::from_millis(20)).await; // Let the hot loop run for 100ms while we continue to attempt consuming - let start = std::time::Instant::now(); tokio::time::sleep(Duration::from_millis(100)).await; let hot_loop_calls = call_count.load(Ordering::SeqCst); From e7b3798686f00bf0177c399d1511cde5301384c0 Mon Sep 17 00:00:00 2001 From: James Watkins-Harvey Date: Thu, 12 Feb 2026 12:34:14 -0500 Subject: [PATCH 4/4] Use parking for mutex --- crates/sdk-core/src/pollers/poll_buffer.rs | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/crates/sdk-core/src/pollers/poll_buffer.rs b/crates/sdk-core/src/pollers/poll_buffer.rs index 496bd7a39..bc9ce942d 100644 --- a/crates/sdk-core/src/pollers/poll_buffer.rs +++ b/crates/sdk-core/src/pollers/poll_buffer.rs @@ -494,7 +494,7 @@ where last_successful_poll_time, // Use same backoff config as gRPC client's throttle_backoff for ResourceExhausted // (1s initial, 10s max, 2x multiplier, 0.2 randomization, unlimited retries) - resource_exhausted_backoff: std::sync::Mutex::new(ExponentialBackoff { + resource_exhausted_backoff: parking_lot::Mutex::new(ExponentialBackoff { // Copied from RetryOptions::throttle_retry_policy() current_interval: Duration::from_secs(1), initial_interval: Duration::from_secs(1), @@ -569,7 +569,7 @@ struct PollScalerReportHandle { scale_up_allowed: AtomicBool, last_successful_poll_time: Arc>>, - resource_exhausted_backoff: std::sync::Mutex>, + resource_exhausted_backoff: parking_lot::Mutex>, } impl PollScalerReportHandle { @@ -586,9 +586,7 @@ impl PollScalerReportHandle { .store(Some(SystemTime::now())); // Reset backoff on successful poll - if let Ok(mut backoff) = self.resource_exhausted_backoff.lock() { - backoff.reset(); - } + self.resource_exhausted_backoff.lock().reset(); if let PollerBehavior::SimpleMaximum(_) = self.behavior { // We don't do auto-scaling with the simple max @@ -635,11 +633,8 @@ impl PollScalerReportHandle { // Scale down significantly for resource exhaustion self.change_target(usize::saturating_div, 2); - let backoff_duration = self - .resource_exhausted_backoff - .lock() - .unwrap() - .next_backoff(); + let backoff_duration = + self.resource_exhausted_backoff.lock().next_backoff(); return (false, backoff_duration); } else {