Just do the obvious thing with max slots

Author: Sushisource

core-api/src/worker.rs

@@ -521,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/pollers/poll_buffer.rs

@@ -724,7 +724,7 @@ mod tests {
             CancellationToken::new(),
             None::<fn(usize)>,
             WorkflowTaskOptions {
-                wft_poller_shared: Some(Arc::new(WFTPollerShared::new(false))),
+                wft_poller_shared: Some(Arc::new(WFTPollerShared::new(Some(10)))),
             },
         );
 

core/src/worker/workflow/wft_poller.rs

@@ -34,7 +34,7 @@ pub(crate) fn make_wft_poller(
     let poller_behavior = wft_poller_behavior(config, false);
     let wft_poller_shared = if sticky_queue_name.is_some() {
         Some(Arc::new(WFTPollerShared::new(
-            poller_behavior.is_autoscaling(),
+            wft_slots.available_permits(),
         )))
     } else {
         None
@@ -81,16 +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>>,
-    using_autoscaling: bool,
+    max_slots: Option<usize>,
 }
 impl WFTPollerShared {
-    pub(crate) fn new(using_autoscaling: bool) -> 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(),
-            using_autoscaling,
+            max_slots,
         }
     }
     pub(crate) fn set_sticky_active(&self, rx: watch::Receiver<usize>) {
@@ -112,12 +112,10 @@ impl WFTPollerShared {
                 .await;
         }
 
-        // Unless autoscaling, if there's no sticky backlog, balance poller counts. This logic
-        // allows the poller to proceed if it has the same or fewer pollers as it's opposite. There
-        // is a preference for the sticky poller when counts are equal. This does not mean we always
-        // have equal numbers of pollers, as later on the scaler will also prevent polling based on
-        // the scaling information provided independently by the sticky/nonsticky queues.
-        if !self.using_autoscaling && *self.last_seen_sticky_backlog.0.borrow() == 0 {
+        // 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())
             {
@@ -126,10 +124,19 @@ impl WFTPollerShared {
                 loop {
                     let num_sticky_active = *sticky_active.borrow_and_update();
                     let num_non_sticky_active = *non_sticky_active.borrow_and_update();
-                    if (is_sticky && num_sticky_active <= num_non_sticky_active)
-                        || (!is_sticky && (num_non_sticky_active < num_sticky_active))
-                    {
-                        break;
+                    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() => (),
@@ -139,6 +146,7 @@ impl WFTPollerShared {
             }
         }
     }
+
     pub(crate) fn record_sticky_backlog(&self, v: usize) {
         let _ = self.last_seen_sticky_backlog.1.send(v);
     }

tests/integ_tests/polling_tests.rs

@@ -212,16 +212,29 @@ async fn switching_worker_client_changes_poll() {
     server2.shutdown().await.unwrap();
 }
 
+#[rstest::rstest]
 #[tokio::test]
-async fn small_workflow_slots_and_pollers() {
+async fn small_workflow_slots_and_pollers(#[values(false, true)] use_autoscaling: bool) {
     let wf_name = "only_one_workflow_slot_and_two_pollers";
     let mut starter = CoreWfStarter::new(wf_name);
+    if use_autoscaling {
+        starter
+            .worker_config
+            .workflow_task_poller_behavior(PollerBehavior::Autoscaling {
+                minimum: 1,
+                maximum: 5,
+                initial: 1,
+            });
+    } else {
+        starter
+            .worker_config
+            .workflow_task_poller_behavior(PollerBehavior::SimpleMaximum(2));
+    }
     starter
         .worker_config
         .max_outstanding_workflow_tasks(2_usize)
         .max_outstanding_local_activities(1_usize)
         .activity_task_poller_behavior(PollerBehavior::SimpleMaximum(1))
-        .workflow_task_poller_behavior(PollerBehavior::SimpleMaximum(2))
         .max_outstanding_activities(1_usize);
     let mut worker = starter.worker().await;
     worker.register_wf(wf_name.to_owned(), |ctx: WfContext| async move {
@@ -246,9 +259,10 @@ async fn small_workflow_slots_and_pollers() {
         )
         .await
         .unwrap();
+    let wf2id = format!("{}-2", starter.get_task_queue());
     worker
         .submit_wf(
-            format!("{}-2", starter.get_task_queue()),
+            wf2id.clone(),
             wf_name.to_owned(),
             vec![],
             WorkflowOptions::default(),
@@ -264,4 +278,17 @@ async fn small_workflow_slots_and_pollers() {
         .iter()
         .any(|e| e.event_type() == EventType::WorkflowTaskTimedOut);
     assert!(!any_task_timeouts);
+    let history = starter
+        .get_client()
+        .await
+        .get_workflow_execution_history(wf2id, None, vec![])
+        .await
+        .unwrap()
+        .history
+        .unwrap();
+    let any_task_timeouts = history
+        .events
+        .iter()
+        .any(|e| e.event_type() == EventType::WorkflowTaskTimedOut);
+    assert!(!any_task_timeouts);
 }