Prevent overshooting / many idle pollers

Author: Sushisource

core/src/abstractions.rs

@@ -7,8 +7,9 @@ use std::{
     fmt::{Debug, Formatter},
     sync::{
         atomic::{AtomicBool, AtomicUsize, Ordering},
-        Arc,
+        Arc, OnceLock,
     },
+    time::Instant,
 };
 use temporal_sdk_core_api::worker::{
     SlotKind, SlotMarkUsedContext, SlotReleaseContext, SlotReservationContext, SlotSupplier,
@@ -122,6 +123,7 @@ where
         let supp_c = self.supplier.clone();
         let supp_c_c = self.supplier.clone();
         let mets = self.metrics_ctx.clone();
+        let mets2 = self.metrics_ctx.clone();
         let metric_rec =
             // When being called from the drop impl, the permit isn't actually dropped yet, so
             // account for that with the `add_one` parameter.
@@ -134,6 +136,15 @@ where
                 mets.task_slots_used((ep_rx_c.borrow().saturating_sub(unused) + extra) as u64);
             };
         let mrc = metric_rec.clone();
+        // TODO: Delete or flag these new metrics
+        let use_time = Arc::new(OnceLock::<Instant>::new());
+        let ut = use_time.clone();
+        let use_time_record = move || {
+            if let Some(ut) = ut.get() {
+                mets2.task_execution_latency(ut.elapsed());
+            }
+        };
+
         mrc(false);
 
         OwnedMeteredSemPermit {
@@ -144,11 +155,13 @@ where
             },
             use_fn: Box::new(move |info| {
                 supp_c.mark_slot_used(info);
+                let _ = use_time.set(Instant::now());
                 metric_rec(false)
             }),
             release_fn: Box::new(move |info| {
                 supp_c_c.release_slot(info);
                 ep_tx_c.send_modify(|ep| *ep -= 1);
+                use_time_record();
                 mrc(true)
             }),
         }

core/src/core_tests/activity_tasks.rs

@@ -91,7 +91,7 @@ async fn max_activities_respected() {
     mock_client
         .expect_poll_activity_task()
         .times(3)
-        .returning(move |_, _, _| Ok(tasks.pop_front().unwrap()));
+        .returning(move |_, _| Ok(tasks.pop_front().unwrap()));
     mock_client
         .expect_complete_activity_task()
         .returning(|_, _| Ok(RespondActivityTaskCompletedResponse::default()));
@@ -376,7 +376,7 @@ async fn many_concurrent_heartbeat_cancels() {
     let mut calls_map = HashMap::<_, i32>::new();
     mock_client
         .expect_poll_activity_task()
-        .returning(move |_, _, _| poll_resps.pop_front().unwrap());
+        .returning(move |_, _| poll_resps.pop_front().unwrap());
     mock_client
         .expect_cancel_activity_task()
         .returning(move |_, _| async move { Ok(Default::default()) }.boxed());
@@ -618,8 +618,8 @@ async fn max_tq_acts_set_passed_to_poll_properly() {
     let mut mock_client = mock_workflow_client();
     mock_client
         .expect_poll_activity_task()
-        .returning(move |_, tps, _| {
-            assert_eq!(tps, Some(rate));
+        .returning(move |_, ao| {
+            assert_eq!(ao.max_tasks_per_sec, Some(rate));
             Ok(PollActivityTaskQueueResponse {
                 task_token: vec![1],
                 ..Default::default()
@@ -1027,7 +1027,7 @@ async fn cant_complete_activity_with_unset_result_payload() {
     let mut mock_client = mock_workflow_client();
     mock_client
         .expect_poll_activity_task()
-        .returning(move |_, _, _| {
+        .returning(move |_, _| {
             Ok(PollActivityTaskQueueResponse {
                 task_token: vec![1],
                 ..Default::default()

core/src/core_tests/mod.rs

@@ -53,7 +53,7 @@ async fn shutdown_interrupts_both_polls() {
     mock_client
         .expect_poll_activity_task()
         .times(1)
-        .returning(move |_, _, _| {
+        .returning(move |_, _| {
             async move {
                 BARR.wait().await;
                 sleep(Duration::from_secs(1)).await;

core/src/core_tests/workers.rs

@@ -267,7 +267,7 @@ async fn worker_does_not_panic_on_retry_exhaustion_of_nonfatal_net_err() {
 async fn worker_can_shutdown_after_never_polling_ok(#[values(true, false)] poll_workflow: bool) {
     let mut mock = mock_workflow_client();
     mock.expect_poll_activity_task()
-        .returning(|_, _, _| Err(tonic::Status::permission_denied("you shall not pass")));
+        .returning(|_, _| Err(tonic::Status::permission_denied("you shall not pass")));
     if poll_workflow {
         mock.expect_poll_workflow_task()
             .returning(|_, _| Err(tonic::Status::permission_denied("you shall not pass")));

core/src/pollers/poll_buffer.rs

@@ -1,7 +1,7 @@
 use crate::{
     abstractions::{dbg_panic, MeteredPermitDealer, OwnedMeteredSemPermit},
     pollers::{self, Poller},
-    worker::client::WorkerClient,
+    worker::client::{PollActivityOptions, PollOptions, PollWorkflowOptions, WorkerClient},
 };
 use futures_util::{future::BoxFuture, FutureExt, StreamExt};
 use governor::{Quota, RateLimiter};
@@ -20,7 +20,7 @@ use temporal_sdk_core_api::worker::{
     ActivitySlotKind, NexusSlotKind, PollerBehavior, SlotKind, WorkflowSlotKind,
 };
 use temporal_sdk_core_protos::temporal::api::{
-    taskqueue::v1::{PollerScalingDecision, TaskQueue},
+    taskqueue::v1::PollerScalingDecision,
     workflowservice::v1::{
         PollActivityTaskQueueResponse, PollNexusTaskQueueResponse, PollWorkflowTaskQueueResponse,
     },
@@ -55,7 +55,7 @@ where
     SK: SlotKind + 'static,
 {
     pub(crate) fn new<FT, DelayFut>(
-        poll_fn: impl Fn() -> FT + Send + Sync + 'static,
+        poll_fn: impl Fn(Option<Duration>) -> FT + Send + Sync + 'static,
         permit_dealer: MeteredPermitDealer<SK>,
         poller_behavior: PollerBehavior,
         shutdown: CancellationToken,
@@ -119,14 +119,23 @@ where
                     let pf = pf.clone();
                     let tx = tx.clone();
                     let report_handle = poll_scaler.get_report_handle();
+                    // Reduce poll timeout if we're frequently getting tasks, to avoid having many
+                    // outstanding long polls for a full minute after a burst subsides.
+                    let timeout_override =
+                        if report_handle.ingested_this_period.load(Ordering::Relaxed) > 2 {
+                            Some(Duration::from_secs(10))
+                        } else {
+                            None
+                        };
                     let poll_task = tokio::spawn(async move {
                         let r = tokio::select! {
-                            r = pf() => r,
+                            r = pf(timeout_override) => r,
                             _ = shutdown.cancelled() => return,
                         };
                         drop(active_guard);
-                        report_handle.poll_result(&r);
-                        let _ = tx.send(r.map(|r| (r, permit)));
+                        if report_handle.poll_result(&r) {
+                            let _ = tx.send(r.map(|r| (r, permit)));
+                        }
                     });
                     let _ = spawned_tx.send(poll_task);
                 }
@@ -248,6 +257,21 @@ where
             target: AtomicUsize::new(target),
             ever_saw_scaling_decision: AtomicBool::default(),
             behavior,
+            ingested_this_period: Default::default(),
+            ingested_last_period: Default::default(),
+            scale_up_allowed: AtomicBool::new(true),
+        });
+        let rhc = report_handle.clone();
+        // TODO: Track task
+        tokio::task::spawn(async move {
+            let mut interval = tokio::time::interval(Duration::from_millis(100));
+            loop {
+                interval.tick().await;
+                let ingested = rhc.ingested_this_period.swap(0, Ordering::Relaxed);
+                let ingested_last = rhc.ingested_last_period.swap(ingested, Ordering::Relaxed);
+                rhc.scale_up_allowed
+                    .store(ingested_last >= ingested, Ordering::Relaxed);
+            }
         });
         Self {
             report_handle,
@@ -286,15 +310,23 @@ struct PollScalerReportHandle {
     target: AtomicUsize,
     ever_saw_scaling_decision: AtomicBool,
     behavior: PollerBehavior,
+
+    ingested_this_period: AtomicUsize,
+    ingested_last_period: AtomicUsize,
+    scale_up_allowed: AtomicBool,
 }
 
 impl PollScalerReportHandle {
-    fn poll_result(&self, res: &Result<impl TaskPollerResult, tonic::Status>) {
+    /// Returns true if the response should be passed on, false if it should be swallowed
+    fn poll_result(&self, res: &Result<impl TaskPollerResult, tonic::Status>) -> bool {
         match res {
             Ok(res) => {
                 if let PollerBehavior::SimpleMaximum(_) = self.behavior {
                     // We don't do auto-scaling with the simple max
-                    return;
+                    return true;
+                }
+                if !res.is_empty() {
+                    self.ingested_this_period.fetch_add(1, Ordering::Relaxed);
                 }
                 if let Some(scaling_decision) = res.scaling_decision() {
                     match scaling_decision.poll_request_delta_suggestion.cmp(&0) {
@@ -304,10 +336,15 @@ impl PollScalerReportHandle {
                                 .poll_request_delta_suggestion
                                 .unsigned_abs() as usize,
                         ),
-                        cmp::Ordering::Greater => self.change_target(
-                            usize::saturating_add,
-                            scaling_decision.poll_request_delta_suggestion as usize,
-                        ),
+                        cmp::Ordering::Greater => {
+                            // Only allow scale up if in the last period it increased ingestion rate
+                            if self.scale_up_allowed.load(Ordering::Relaxed) {
+                                self.change_target(
+                                    usize::saturating_add,
+                                    scaling_decision.poll_request_delta_suggestion as usize,
+                                )
+                            }
+                        }
                         cmp::Ordering::Equal => {}
                     }
                     self.ever_saw_scaling_decision
@@ -334,9 +371,11 @@ impl PollScalerReportHandle {
                         // reclaim this poller
                         self.change_target(usize::saturating_sub, 1);
                     }
+                    return false;
                 }
             }
         }
+        true
     }
 
     #[inline]
@@ -405,24 +444,37 @@ pub(crate) type PollWorkflowTaskBuffer =
     LongPollBuffer<PollWorkflowTaskQueueResponse, WorkflowSlotKind>;
 pub(crate) fn new_workflow_task_buffer(
     client: Arc<dyn WorkerClient>,
-    task_queue: TaskQueue,
+    task_queue: String,
+    sticky_queue: Option<String>,
     poller_behavior: PollerBehavior,
     permit_dealer: MeteredPermitDealer<WorkflowSlotKind>,
     shutdown: CancellationToken,
     num_pollers_handler: Option<impl Fn(usize) + Send + Sync + 'static>,
 ) -> PollWorkflowTaskBuffer {
-    let no_poll_retry = if matches!(poller_behavior, PollerBehavior::Autoscaling { .. }) {
+    let no_retry = if matches!(poller_behavior, PollerBehavior::Autoscaling { .. }) {
         Some(NoRetryOnMatching {
             predicate: poll_scaling_error_matcher,
         })
     } else {
         None
     };
     LongPollBuffer::new(
-        move || {
+        move |timeout_override| {
             let client = client.clone();
             let task_queue = task_queue.clone();
-            async move { client.poll_workflow_task(task_queue, no_poll_retry).await }
+            let sticky_queue_name = sticky_queue.clone();
+            async move {
+                client
+                    .poll_workflow_task(
+                        PollOptions {
+                            task_queue,
+                            no_retry,
+                            timeout_override,
+                        },
+                        PollWorkflowOptions { sticky_queue_name },
+                    )
+                    .await
+            }
         },
         permit_dealer,
         poller_behavior,
@@ -449,20 +501,29 @@ pub(crate) fn new_activity_task_buffer(
         Quota::with_period(Duration::from_secs_f64(ps.recip()))
             .map(|q| Arc::new(RateLimiter::direct(q)))
     });
-    let no_poll_retry = if matches!(poller_behavior, PollerBehavior::Autoscaling { .. }) {
+    let no_retry = if matches!(poller_behavior, PollerBehavior::Autoscaling { .. }) {
         Some(NoRetryOnMatching {
             predicate: poll_scaling_error_matcher,
         })
     } else {
         None
     };
     LongPollBuffer::new(
-        move || {
+        move |timeout_override| {
             let client = client.clone();
             let task_queue = task_queue.clone();
             async move {
                 client
-                    .poll_activity_task(task_queue, max_tps, no_poll_retry)
+                    .poll_activity_task(
+                        PollOptions {
+                            task_queue,
+                            no_retry,
+                            timeout_override,
+                        },
+                        PollActivityOptions {
+                            max_tasks_per_sec: max_tps,
+                        },
+                    )
                     .await
             }
         },
@@ -488,18 +549,26 @@ pub(crate) fn new_nexus_task_buffer(
     shutdown: CancellationToken,
     num_pollers_handler: Option<impl Fn(usize) + Send + Sync + 'static>,
 ) -> PollNexusTaskBuffer {
-    let no_poll_retry = if matches!(poller_behavior, PollerBehavior::Autoscaling { .. }) {
+    let no_retry = if matches!(poller_behavior, PollerBehavior::Autoscaling { .. }) {
         Some(NoRetryOnMatching {
             predicate: poll_scaling_error_matcher,
         })
     } else {
         None
     };
     LongPollBuffer::new(
-        move || {
+        move |timeout_override| {
             let client = client.clone();
             let task_queue = task_queue.clone();
-            async move { client.poll_nexus_task(task_queue, no_poll_retry).await }
+            async move {
+                client
+                    .poll_nexus_task(PollOptions {
+                        task_queue,
+                        no_retry,
+                        timeout_override,
+                    })
+                    .await
+            }
         },
         semaphore,
         poller_behavior,
@@ -592,7 +661,6 @@ mod tests {
     };
     use futures_util::FutureExt;
     use std::time::Duration;
-    use temporal_sdk_core_protos::temporal::api::enums::v1::TaskQueueKind;
     use tokio::select;
 
     #[tokio::test]
@@ -611,11 +679,8 @@ mod tests {
 
         let pb = new_workflow_task_buffer(
             Arc::new(mock_client),
-            TaskQueue {
-                name: "sometq".to_string(),
-                kind: TaskQueueKind::Normal as i32,
-                normal_name: "".to_string(),
-            },
+            "sometq".to_string(),
+            None,
             PollerBehavior::SimpleMaximum(1),
             fixed_size_permit_dealer(10),
             CancellationToken::new(),

core/src/telemetry/metrics.rs

@@ -57,6 +57,7 @@ struct Instruments {
     num_pollers: Arc<dyn Gauge>,
     task_slots_available: Arc<dyn Gauge>,
     task_slots_used: Arc<dyn Gauge>,
+    task_execution_latency: Arc<dyn HistogramDuration>,
     sticky_cache_hit: Arc<dyn Counter>,
     sticky_cache_miss: Arc<dyn Counter>,
     sticky_cache_size: Arc<dyn Gauge>,
@@ -280,6 +281,13 @@ impl MetricsContext {
         self.instruments.task_slots_used.record(num, &self.kvs)
     }
 
+    /// Record time it took to execute a task
+    pub(crate) fn task_execution_latency(&self, dur: Duration) {
+        self.instruments
+            .task_execution_latency
+            .record(dur, &self.kvs);
+    }
+
     /// Record current number of pollers. Context should include poller type / task queue tag.
     pub(crate) fn record_num_pollers(&self, num: usize) {
         self.instruments.num_pollers.record(num as u64, &self.kvs);
@@ -470,6 +478,11 @@ impl Instruments {
                 description: "Current number of used slots per task type".into(),
                 unit: "".into(),
             }),
+            task_execution_latency: meter.histogram_duration(MetricParameters {
+                name: "task_execution_latency".into(),
+                description: "Histogram of task execution latencies".into(),
+                unit: "".into(),
+            }),
             sticky_cache_hit: meter.counter(MetricParameters {
                 name: "sticky_cache_hit".into(),
                 description: "Count of times the workflow cache was used for a new workflow task"