Fixes two-slots-two-pollers

Author: Sushisource

core-api/src/worker.rs

@@ -389,6 +389,11 @@ impl SlotSupplierPermit {
     pub fn user_data_mut<T: Any + Send + Sync>(&mut self) -> Option<&mut T> {
         self.user_data.as_mut().and_then(|b| b.downcast_mut())
     }
+
+    /// Deconstruct this permit and return the inner data
+    pub fn into_user_data(self) -> Option<Box<dyn Any + Send + Sync>> {
+        self.user_data
+    }
 }
 
 #[derive(Debug, Copy, Clone, derive_more::Display, Eq, PartialEq)]

core/src/abstractions.rs

@@ -24,7 +24,7 @@ use tokio_util::sync::CancellationToken;
 /// as handling associated metrics tracking.
 #[derive(Clone)]
 pub(crate) struct MeteredPermitDealer<SK: SlotKind> {
-    supplier: Arc<dyn SlotSupplier<SlotKind = SK> + Send + Sync>,
+    pub(crate) supplier: Arc<dyn SlotSupplier<SlotKind = SK> + Send + Sync>,
     /// The number of permit owners who have acquired a permit, but are not yet meaningfully using
     /// that permit. This is useful for giving a more semantically accurate count of used task
     /// slots, since we typically wait for a permit first before polling, but that slot isn't used
@@ -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,
@@ -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

core/src/worker/workflow/wft_poller.rs

@@ -7,10 +7,13 @@ use crate::{
     worker::{client::WorkerClient, wft_poller_behavior},
 };
 use futures_util::{Stream, stream};
-use std::sync::{Arc, OnceLock};
+use std::sync::{
+    Arc, OnceLock,
+    atomic::{AtomicBool, Ordering},
+};
 use temporal_sdk_core_api::worker::{WorkerConfig, WorkflowSlotKind};
 use temporal_sdk_core_protos::temporal::api::workflowservice::v1::PollWorkflowTaskQueueResponse;
-use tokio::sync::watch;
+use tokio::sync::{Notify, watch};
 use tokio_util::sync::CancellationToken;
 
 pub(crate) fn make_wft_poller(
@@ -74,6 +77,8 @@ 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>>,
+    wait_for_first_nonsticky_poll: Notify,
+    have_done_first_poll: AtomicBool,
 }
 impl WFTPollerShared {
     pub(crate) fn new() -> Self {
@@ -82,6 +87,8 @@ impl WFTPollerShared {
             last_seen_sticky_backlog: (rx, tx),
             sticky_active: OnceLock::new(),
             non_sticky_active: OnceLock::new(),
+            wait_for_first_nonsticky_poll: Notify::new(),
+            have_done_first_poll: AtomicBool::new(false),
         }
     }
     pub(crate) fn set_sticky_active(&self, rx: watch::Receiver<usize>) {
@@ -93,15 +100,24 @@ impl WFTPollerShared {
     /// Makes either the sticky or non-sticky poller wait pre-permit-acquisition so that we can
     /// balance which kind of queue we poll appropriately.
     pub(crate) async fn wait_if_needed(&self, is_sticky: bool) {
+        // Sticky shouldn't start polling until after the first non-sticky poll has been allowed
+        if is_sticky {
+            self.wait_for_first_nonsticky_poll.notified().await;
+        }
+        info!(
+            "Waiting if needed. Sticky: {} / backlog {}",
+            is_sticky,
+            *self.last_seen_sticky_backlog.0.borrow()
+        );
         // If there's a sticky backlog, prioritize it.
         if !is_sticky {
             let backlog = *self.last_seen_sticky_backlog.0.borrow();
-            if backlog > 1 {
+            if backlog >= 1 {
                 let _ = self
                     .last_seen_sticky_backlog
                     .0
                     .clone()
-                    .wait_for(|v| *v <= 1)
+                    .wait_for(|v| *v < 1)
                     .await;
             }
         }
@@ -111,6 +127,12 @@ impl WFTPollerShared {
             if let Some((sticky_active, non_sticky_active)) =
                 self.sticky_active.get().zip(self.non_sticky_active.get())
             {
+                info!(
+                    "Balance (sticky {}), non-sticky {} sticky {}",
+                    is_sticky,
+                    *non_sticky_active.borrow(),
+                    *sticky_active.borrow()
+                );
                 if is_sticky {
                     let _ = sticky_active
                         .clone()
@@ -119,9 +141,13 @@ impl WFTPollerShared {
                 } else {
                     let _ = non_sticky_active
                         .clone()
-                        .wait_for(|v| *v <= *sticky_active.borrow())
+                        .wait_for(|v| {
+                            *v < *sticky_active.borrow()
+                                || !self.have_done_first_poll.load(Ordering::Relaxed)
+                        })
                         .await;
                 }
+                info!("Done balance (sticky {})", is_sticky);
             }
         }
     }

tests/integ_tests/polling_tests.rs

@@ -1,19 +1,21 @@
+use crate::integ_tests::activity_functions::echo;
 use assert_matches::assert_matches;
 use std::{sync::Arc, time::Duration};
 use temporal_client::{WfClientExt, WorkflowClientTrait, WorkflowOptions};
+use temporal_sdk::{ActivityOptions, WfContext};
 use temporal_sdk_core::{
     ClientOptionsBuilder, ephemeral_server::TemporalDevServerConfigBuilder, init_worker,
 };
-use temporal_sdk_core_api::Worker;
+use temporal_sdk_core_api::{Worker, worker::PollerBehavior};
 use temporal_sdk_core_protos::coresdk::{
-    IntoCompletion,
+    AsJsonPayloadExt, IntoCompletion,
     activity_task::activity_task as act_task,
     workflow_activation::{FireTimer, WorkflowActivationJob, workflow_activation_job},
     workflow_commands::{ActivityCancellationType, RequestCancelActivity, StartTimer},
     workflow_completion::WorkflowActivationCompletion,
 };
 use temporal_sdk_core_test_utils::{
-    WorkerTestHelpers, default_cached_download, drain_pollers_and_shutdown,
+    CoreWfStarter, WorkerTestHelpers, default_cached_download, drain_pollers_and_shutdown,
     init_core_and_create_wf, init_integ_telem, integ_worker_config, schedule_activity_cmd,
 };
 use tokio::time::timeout;
@@ -206,3 +208,41 @@ async fn switching_worker_client_changes_poll() {
     server1.shutdown().await.unwrap();
     server2.shutdown().await.unwrap();
 }
+
+#[tokio::test]
+async fn only_one_workflow_slot_and_two_pollers() {
+    let wf_name = "only_one_workflow_slot_and_two_pollers";
+    let mut starter = CoreWfStarter::new(wf_name);
+    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 {
+        for _ in 0..3 {
+            ctx.activity(ActivityOptions {
+                activity_type: "echo_activity".to_string(),
+                start_to_close_timeout: Some(Duration::from_secs(5)),
+                input: "hi!".as_json_payload().expect("serializes fine"),
+                ..Default::default()
+            })
+            .await;
+        }
+        Ok(().into())
+    });
+    worker.register_activity("echo_activity", echo);
+    worker
+        .submit_wf(
+            wf_name.to_string(),
+            wf_name.to_owned(),
+            vec![],
+            WorkflowOptions::default(),
+        )
+        .await
+        .unwrap();
+    // If we don't fail the workflow on nondeterminism, we'll get stuck here retrying the WFT
+    worker.run_until_done().await.unwrap();
+}