Fix bad balacing between sticky & nonsticky WFT polling with low slot/poller counts (#925)

Author: Sushisource

core-api/src/worker.rs

@@ -40,8 +40,8 @@ pub struct WorkerConfig {
     #[builder(setter(into = false, strip_option), default)]
     pub tuner: Option<Arc<dyn WorkerTuner + Send + Sync>>,
     /// Maximum number of concurrent poll workflow task requests we will perform at a time on this
-    /// worker's task queue. See also [WorkerConfig::nonsticky_to_sticky_poll_ratio]. Must be at
-    /// least 1.
+    /// worker's task queue. See also [WorkerConfig::nonsticky_to_sticky_poll_ratio].
+    /// If using SimpleMaximum, Must be at least 2 when `max_cached_workflows` > 0, or is an error.
     #[builder(default = "PollerBehavior::SimpleMaximum(5)")]
     pub workflow_task_poller_behavior: PollerBehavior,
     /// Only applies when using [PollerBehavior::SimpleMaximum]
@@ -135,7 +135,8 @@ pub struct WorkerConfig {
 
     /// The maximum allowed number of workflow tasks that will ever be given to this worker at one
     /// time. Note that one workflow task may require multiple activations - so the WFT counts as
-    /// "outstanding" until all activations it requires have been completed.
+    /// "outstanding" until all activations it requires have been completed. Must be at least 2 if
+    /// `max_cached_workflows` is > 0, or is an error.
     ///
     /// Mutually exclusive with `tuner`
     #[builder(setter(into, strip_option), default)]
@@ -224,6 +225,41 @@ impl WorkerConfigBuilder {
             }
         }
 
+        if matches!(self.max_outstanding_workflow_tasks.as_ref(), Some(Some(v)) if *v == 0) {
+            return Err("`max_outstanding_workflow_tasks` must be > 0".to_owned());
+        }
+        if matches!(self.max_outstanding_activities.as_ref(), Some(Some(v)) if *v == 0) {
+            return Err("`max_outstanding_activities` must be > 0".to_owned());
+        }
+        if matches!(self.max_outstanding_local_activities.as_ref(), Some(Some(v)) if *v == 0) {
+            return Err("`max_outstanding_local_activities` must be > 0".to_owned());
+        }
+        if matches!(self.max_outstanding_nexus_tasks.as_ref(), Some(Some(v)) if *v == 0) {
+            return Err("`max_outstanding_nexus_tasks` must be > 0".to_owned());
+        }
+
+        if let Some(cache) = self.max_cached_workflows.as_ref() {
+            if *cache > 0 {
+                if let Some(Some(max_wft)) = self.max_outstanding_workflow_tasks.as_ref() {
+                    if *max_wft < 2 {
+                        return Err(
+                            "`max_cached_workflows` > 0 requires `max_outstanding_workflow_tasks` >= 2"
+                                .to_owned(),
+                        );
+                    }
+                }
+                if let Some(b) = self.workflow_task_poller_behavior.as_ref() {
+                    if matches!(b, PollerBehavior::SimpleMaximum(u) if *u < 2) {
+                        return Err(
+                            "`max_cached_workflows` > 0 requires `workflow_task_poller_behavior` to be at least 2"
+                                .to_owned(),
+                        );
+                    }
+                    b.validate()?
+                }
+            }
+        }
+
         if self.tuner.is_some()
             && (self.max_outstanding_workflow_tasks.is_some()
                 || self.max_outstanding_activities.is_some()
@@ -264,7 +300,9 @@ impl WorkerConfigBuilder {
 /// This trait allows users to customize the performance characteristics of workers dynamically.
 /// For more, see the docstrings of the traits in the return types of its functions.
 pub trait WorkerTuner {
-    /// Return a [SlotSupplier] for workflow tasks
+    /// Return a [SlotSupplier] for workflow tasks. Note that workflow task slot suppliers must be
+    /// willing to hand out a minimum of one non-sticky slot and one sticky slot if workflow caching
+    /// is enabled, otherwise the worker may fail to process new tasks.
     fn workflow_task_slot_supplier(
         &self,
     ) -> Arc<dyn SlotSupplier<SlotKind = WorkflowSlotKind> + Send + Sync>;
@@ -483,7 +521,7 @@ pub enum PollerBehavior {
     /// requires a slot to be available before beginning polling.
     Autoscaling {
         /// At least this many poll calls will always be attempted (assuming slots are available).
-        /// Cannot be less than two for workflow tasks, or one for other tasks.
+        /// Cannot be zero.
         minimum: usize,
         /// At most this many poll calls will ever be open at once. Must be >= `minimum`.
         maximum: usize,

core/src/abstractions.rs

@@ -393,6 +393,35 @@ macro_rules! dbg_panic {
 }
 pub(crate) use dbg_panic;
 
+pub(crate) struct ActiveCounter<F: Fn(usize)>(watch::Sender<usize>, Option<Arc<F>>);
+impl<F> ActiveCounter<F>
+where
+    F: Fn(usize),
+{
+    pub(crate) fn new(a: watch::Sender<usize>, change_fn: Option<Arc<F>>) -> Self {
+        a.send_modify(|v| {
+            *v += 1;
+            if let Some(cfn) = change_fn.as_ref() {
+                cfn(*v);
+            }
+        });
+        Self(a, change_fn)
+    }
+}
+impl<F> Drop for ActiveCounter<F>
+where
+    F: Fn(usize),
+{
+    fn drop(&mut self) {
+        self.0.send_modify(|v| {
+            *v -= 1;
+            if let Some(cfn) = self.1.as_ref() {
+                cfn(*v)
+            };
+        });
+    }
+}
+
 #[cfg(test)]
 pub(crate) mod tests {
     use super::*;

core/src/pollers/poll_buffer.rs

@@ -1,5 +1,5 @@
 use crate::{
-    abstractions::{MeteredPermitDealer, OwnedMeteredSemPermit, dbg_panic},
+    abstractions::{ActiveCounter, MeteredPermitDealer, OwnedMeteredSemPermit, dbg_panic},
     pollers::{self, Poller},
     worker::{
         WFTPollerShared,
@@ -53,7 +53,7 @@ pub(crate) struct LongPollBuffer<T, SK: SlotKind> {
 }
 
 pub(crate) struct WorkflowTaskOptions {
-    pub(crate) wft_poller_shared: Arc<WFTPollerShared>,
+    pub(crate) wft_poller_shared: Option<Arc<WFTPollerShared>>,
 }
 
 pub(crate) struct ActivityTaskOptions {
@@ -77,27 +77,27 @@ impl LongPollBuffer<PollWorkflowTaskQueueResponse, WorkflowSlotKind> {
     ) -> Self {
         let is_sticky = sticky_queue.is_some();
         let poll_scaler = PollScaler::new(poller_behavior, num_pollers_handler, shutdown.clone());
-        if is_sticky {
-            options
-                .wft_poller_shared
-                .set_sticky_active(poll_scaler.active_rx.clone());
-        } else {
-            options
-                .wft_poller_shared
-                .set_non_sticky_active(poll_scaler.active_rx.clone());
-        };
-        let shared = options.wft_poller_shared.clone();
-        let pre_permit_delay = Some(move || {
-            let shared = shared.clone();
-            async move {
-                shared.wait_if_needed(is_sticky).await;
+        if let Some(wftps) = options.wft_poller_shared.as_ref() {
+            if is_sticky {
+                wftps.set_sticky_active(poll_scaler.active_rx.clone());
+            } else {
+                wftps.set_non_sticky_active(poll_scaler.active_rx.clone());
+            };
+        }
+        let pre_permit_delay = options.wft_poller_shared.clone().map(|wftps| {
+            move || {
+                let shared = wftps.clone();
+                async move {
+                    shared.wait_if_needed(is_sticky).await;
+                }
             }
         });
-        let post_poll_fn = Some(move |t: &PollWorkflowTaskQueueResponse| {
-            if is_sticky {
-                options
-                    .wft_poller_shared
-                    .record_sticky_backlog(t.backlog_count_hint as usize)
+
+        let post_poll_fn = options.wft_poller_shared.clone().map(|wftps| {
+            move |t: &PollWorkflowTaskQueueResponse| {
+                if is_sticky {
+                    wftps.record_sticky_backlog(t.backlog_count_hint as usize)
+                }
             }
         });
         let no_retry = if matches!(poller_behavior, PollerBehavior::Autoscaling { .. }) {
@@ -391,35 +391,6 @@ async fn handle_task_panic(t: JoinHandle<()>) {
     }
 }
 
-struct ActiveCounter<F: Fn(usize)>(watch::Sender<usize>, Option<Arc<F>>);
-impl<F> ActiveCounter<F>
-where
-    F: Fn(usize),
-{
-    fn new(a: watch::Sender<usize>, change_fn: Option<Arc<F>>) -> Self {
-        a.send_modify(|v| {
-            *v += 1;
-            if let Some(cfn) = change_fn.as_ref() {
-                cfn(*v);
-            }
-        });
-        Self(a, change_fn)
-    }
-}
-impl<F> Drop for ActiveCounter<F>
-where
-    F: Fn(usize),
-{
-    fn drop(&mut self) {
-        self.0.send_modify(|v| {
-            *v -= 1;
-            if let Some(cfn) = self.1.as_ref() {
-                cfn(*v)
-            };
-        });
-    }
-}
-
 /// The PollScaler is responsible for managing the number of pollers based on the current load.
 ///
 /// It does so by receiving suggestions from the server about whether to scale up or down. It will
@@ -753,7 +724,7 @@ mod tests {
             CancellationToken::new(),
             None::<fn(usize)>,
             WorkflowTaskOptions {
-                wft_poller_shared: Arc::new(WFTPollerShared::new()),
+                wft_poller_shared: Some(Arc::new(WFTPollerShared::new(Some(10)))),
             },
         );
 

core/src/worker/mod.rs

@@ -867,9 +867,10 @@ fn wft_poller_behavior(config: &WorkerConfig, is_sticky: bool) -> PollerBehavior
                 config.nonsticky_to_sticky_poll_ratio,
             ))
         } else {
-            PollerBehavior::SimpleMaximum(m.saturating_sub(
-                calc_max_nonsticky(m, config.nonsticky_to_sticky_poll_ratio).max(1),
-            ))
+            PollerBehavior::SimpleMaximum(
+                m.saturating_sub(calc_max_nonsticky(m, config.nonsticky_to_sticky_poll_ratio))
+                    .max(1),
+            )
         }
     } else {
         config.workflow_task_poller_behavior

core/src/worker/workflow/wft_poller.rs

@@ -31,12 +31,19 @@ pub(crate) fn make_wft_poller(
 > + Sized
 + 'static {
     let wft_metrics = metrics.with_new_attrs([workflow_poller()]);
-    let wft_poller_shared = Arc::new(WFTPollerShared::new());
+    let poller_behavior = wft_poller_behavior(config, false);
+    let wft_poller_shared = if sticky_queue_name.is_some() {
+        Some(Arc::new(WFTPollerShared::new(
+            wft_slots.available_permits(),
+        )))
+    } else {
+        None
+    };
     let wf_task_poll_buffer = LongPollBuffer::new_workflow_task(
         client.clone(),
         config.task_queue.clone(),
         None,
-        wft_poller_behavior(config, false),
+        poller_behavior,
         wft_slots.clone(),
         shutdown_token.child_token(),
         Some(move |np| {
@@ -74,14 +81,16 @@ pub(crate) struct WFTPollerShared {
     last_seen_sticky_backlog: (watch::Receiver<usize>, watch::Sender<usize>),
     sticky_active: OnceLock<watch::Receiver<usize>>,
     non_sticky_active: OnceLock<watch::Receiver<usize>>,
+    max_slots: Option<usize>,
 }
 impl WFTPollerShared {
-    pub(crate) fn new() -> Self {
+    pub(crate) fn new(max_slots: Option<usize>) -> Self {
         let (tx, rx) = watch::channel(0);
         Self {
             last_seen_sticky_backlog: (rx, tx),
             sticky_active: OnceLock::new(),
             non_sticky_active: OnceLock::new(),
+            max_slots,
         }
     }
     pub(crate) fn set_sticky_active(&self, rx: watch::Receiver<usize>) {
@@ -95,36 +104,49 @@ impl WFTPollerShared {
     pub(crate) async fn wait_if_needed(&self, is_sticky: bool) {
         // If there's a sticky backlog, prioritize it.
         if !is_sticky {
-            let backlog = *self.last_seen_sticky_backlog.0.borrow();
-            if backlog > 1 {
-                let _ = self
-                    .last_seen_sticky_backlog
-                    .0
-                    .clone()
-                    .wait_for(|v| *v <= 1)
-                    .await;
-            }
+            let _ = self
+                .last_seen_sticky_backlog
+                .0
+                .clone()
+                .wait_for(|v| *v == 0)
+                .await;
         }
 
-        // If there's no meaningful sticky backlog, balance poller counts
-        if *self.last_seen_sticky_backlog.0.borrow() <= 1 {
+        // We need to make sure there's at least one poller of both kinds available. So, we check
+        // that we won't end up using every available permit with one kind of poller. In practice
+        // this is only ever likely to be an issue with very small numbers of slots.
+        if let Some(max_slots) = self.max_slots {
             if let Some((sticky_active, non_sticky_active)) =
                 self.sticky_active.get().zip(self.non_sticky_active.get())
             {
-                if is_sticky {
-                    let _ = sticky_active
-                        .clone()
-                        .wait_for(|v| *v <= *non_sticky_active.borrow())
-                        .await;
-                } else {
-                    let _ = non_sticky_active
-                        .clone()
-                        .wait_for(|v| *v <= *sticky_active.borrow())
-                        .await;
+                let mut sticky_active = sticky_active.clone();
+                let mut non_sticky_active = non_sticky_active.clone();
+                loop {
+                    let num_sticky_active = *sticky_active.borrow_and_update();
+                    let num_non_sticky_active = *non_sticky_active.borrow_and_update();
+                    let both_are_zero = num_sticky_active == 0 && num_non_sticky_active == 0;
+                    if both_are_zero {
+                        if !is_sticky {
+                            break;
+                        }
+                    } else {
+                        let would_exceed_max_slots =
+                            (num_sticky_active + num_non_sticky_active + 1) >= max_slots;
+                        let must_wait = would_exceed_max_slots
+                            && (num_sticky_active == 0 || num_non_sticky_active == 0);
+                        if !must_wait {
+                            break;
+                        }
+                    }
+                    tokio::select! {
+                        _ = sticky_active.changed() => (),
+                        _ = non_sticky_active.changed() => (),
+                    }
                 }
             }
         }
     }
+
     pub(crate) fn record_sticky_backlog(&self, v: usize) {
         let _ = self.last_seen_sticky_backlog.1.send(v);
     }

tests/integ_tests/metrics_tests.rs

@@ -154,7 +154,7 @@ async fn one_slot_worker_reports_available_slot() {
         // Need to use two for WFTs because there are a minimum of 2 pollers b/c of sticky polling
         .max_outstanding_workflow_tasks(2_usize)
         .max_outstanding_nexus_tasks(1_usize)
-        .workflow_task_poller_behavior(PollerBehavior::SimpleMaximum(1_usize))
+        .workflow_task_poller_behavior(PollerBehavior::SimpleMaximum(2_usize))
         .build()
         .unwrap();