Poll concurrently at will (#545)

* Unleash the pollers
* Refactor things so semaphore can live inside poller
* Enforce config requirement of pollers no more than slots
* Make unusual panic I noticed non-fatal

Author: Sushisource

core-api/src/worker.rs

@@ -1,6 +1,9 @@
 use std::time::Duration;
 use tokio::sync::mpsc::UnboundedSender;
 
+const MAX_OUTSTANDING_WFT_DEFAULT: usize = 100;
+const MAX_CONCURRENT_WFT_POLLS_DEFAULT: usize = 5;
+
 /// Defines per-worker configuration options
 #[derive(Debug, Clone, derive_builder::Builder, serde::Serialize, serde::Deserialize)]
 #[builder(setter(into), build_fn(validate = "Self::validate"))]
@@ -31,7 +34,7 @@ pub struct WorkerConfig {
     /// "outstanding" until all activations it requires have been completed.
     ///
     /// Cannot be larger than `max_cached_workflows`.
-    #[builder(default = "100")]
+    #[builder(default = "MAX_OUTSTANDING_WFT_DEFAULT")]
     pub max_outstanding_workflow_tasks: usize,
     /// The maximum number of activity tasks that will ever be given to this worker concurrently
     #[builder(default = "100")]
@@ -43,7 +46,7 @@ pub struct WorkerConfig {
     /// 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.
-    #[builder(default = "5")]
+    #[builder(default = "MAX_CONCURRENT_WFT_POLLS_DEFAULT")]
     pub max_concurrent_wft_polls: usize,
     /// [WorkerConfig::max_concurrent_wft_polls] * this number = the number of max pollers that will
     /// be allowed for the nonsticky queue when sticky tasks are enabled. If both defaults are used,
@@ -161,6 +164,18 @@ impl WorkerConfigBuilder {
                 );
             }
         }
+        if self
+            .max_concurrent_wft_polls
+            .unwrap_or(MAX_CONCURRENT_WFT_POLLS_DEFAULT)
+            > self
+                .max_outstanding_workflow_tasks
+                .unwrap_or(MAX_OUTSTANDING_WFT_DEFAULT)
+        {
+            return Err(
+                "`max_concurrent_wft_polls` cannot exceed `max_outstanding_workflow_tasks`"
+                    .to_owned(),
+            );
+        }
         Ok(())
     }
 }

core/src/abstractions.rs

@@ -4,7 +4,6 @@ pub mod take_cell;
 
 use crate::MetricsContext;
 use derive_more::DebugCustom;
-use futures::{stream, Stream, StreamExt};
 use std::{
     fmt::{Debug, Formatter},
     sync::{
@@ -47,6 +46,11 @@ impl MeteredSemaphore {
         self.sem.available_permits()
     }
 
+    #[cfg(test)]
+    pub fn unused_permits(&self) -> usize {
+        self.sem.available_permits() + self.unused_claimants.load(Ordering::Acquire)
+    }
+
     pub async fn acquire_owned(&self) -> Result<OwnedMeteredSemPermit, AcquireError> {
         let res = self.sem.clone().acquire_owned().await?;
         Ok(self.build_owned(res))
@@ -114,8 +118,8 @@ impl ClosableMeteredSemaphore {
 
 impl ClosableMeteredSemaphore {
     #[cfg(test)]
-    pub fn available_permits(&self) -> usize {
-        self.inner.available_permits()
+    pub fn unused_permits(&self) -> usize {
+        self.inner.unused_permits()
     }
 
     /// Request to close the semaphore and prevent new permits from being acquired.
@@ -231,35 +235,6 @@ impl UsedMeteredSemPermit {
     }
 }
 
-/// From the input stream, create a new stream which only pulls from the input stream when allowed.
-/// When allowed is determined by the passed in `proceeder` emitting an item. The input stream is
-/// only pulled from when that future resolves.
-///
-/// This is *almost* identical to `zip`, but does not terminate early if the input stream closes.
-/// The proceeder must allow the poll before the returned stream closes. If the proceeder terminates
-/// the overall stream will terminate.
-pub(crate) fn stream_when_allowed<S, AS>(
-    input: S,
-    proceeder: AS,
-) -> impl Stream<Item = (S::Item, AS::Item)>
-where
-    S: Stream + Send + 'static,
-    AS: Stream + Send + 'static,
-{
-    let stream = stream::unfold(
-        (proceeder.boxed(), input.boxed()),
-        |(mut proceeder, mut input)| async {
-            let v = proceeder.next().await;
-            if let Some(v) = v {
-                input.next().await.map(|i| ((i, v), (proceeder, input)))
-            } else {
-                None
-            }
-        },
-    );
-    stream
-}
-
 macro_rules! dbg_panic {
   ($($arg:tt)*) => {
       error!($($arg)*);
@@ -271,49 +246,6 @@ pub(crate) use dbg_panic;
 #[cfg(test)]
 mod tests {
     use super::*;
-    use futures::pin_mut;
-    use std::task::Poll;
-    use tokio::sync::mpsc::unbounded_channel;
-    use tokio_stream::wrappers::UnboundedReceiverStream;
-
-    #[test]
-    fn stream_when_allowed_works() {
-        let inputs = stream::iter([1, 2, 3]);
-        let (allow_tx, allow_rx) = unbounded_channel();
-        let when_allowed = stream_when_allowed(inputs, UnboundedReceiverStream::new(allow_rx));
-
-        let waker = futures::task::noop_waker_ref();
-        let mut cx = std::task::Context::from_waker(waker);
-        pin_mut!(when_allowed);
-
-        allow_tx.send(()).unwrap();
-        assert_eq!(
-            when_allowed.poll_next_unpin(&mut cx),
-            Poll::Ready(Some((1, ())))
-        );
-        // Now, it won't be ready
-        for _ in 1..10 {
-            assert_eq!(when_allowed.poll_next_unpin(&mut cx), Poll::Pending);
-        }
-        allow_tx.send(()).unwrap();
-        assert_eq!(
-            when_allowed.poll_next_unpin(&mut cx),
-            Poll::Ready(Some((2, ())))
-        );
-        for _ in 1..10 {
-            assert_eq!(when_allowed.poll_next_unpin(&mut cx), Poll::Pending);
-        }
-        allow_tx.send(()).unwrap();
-        assert_eq!(
-            when_allowed.poll_next_unpin(&mut cx),
-            Poll::Ready(Some((3, ())))
-        );
-        for _ in 1..10 {
-            assert_eq!(when_allowed.poll_next_unpin(&mut cx), Poll::Pending);
-        }
-        allow_tx.send(()).unwrap();
-        assert_eq!(when_allowed.poll_next_unpin(&mut cx), Poll::Ready(None));
-    }
 
     #[tokio::test]
     async fn closable_semaphore_permit_drop_returns_permit() {

core/src/core_tests/activity_tasks.rs

@@ -1058,26 +1058,34 @@ async fn graceful_shutdown(#[values(true, false)] at_max_outstanding: bool) {
     let _task_q = "q";
     let grace_period = Duration::from_millis(200);
     let mut tasks = three_tasks();
-    let mut mock_client = mock_workflow_client();
-    mock_client
-        .expect_poll_activity_task()
+    let mut mock_act_poller = mock_poller();
+    mock_act_poller
+        .expect_poll()
         .times(3)
-        .returning(move |_, _| Ok(tasks.pop_front().unwrap()));
+        .returning(move || Some(Ok(tasks.pop_front().unwrap())));
+    mock_act_poller
+        .expect_poll()
+        .times(1)
+        .returning(move || None);
     // They shall all be reported as failed
+    let mut mock_client = mock_workflow_client();
     mock_client
         .expect_fail_activity_task()
         .times(3)
         .returning(|_, _| Ok(Default::default()));
 
     let max_outstanding = if at_max_outstanding { 3_usize } else { 100 };
-    let worker = Worker::new_test(
-        test_worker_cfg()
+    let mw = MockWorkerInputs {
+        act_poller: Some(Box::from(mock_act_poller)),
+        config: test_worker_cfg()
             .graceful_shutdown_period(grace_period)
             .max_outstanding_activities(max_outstanding)
+            .max_concurrent_at_polls(1_usize) // Makes test logic simple
             .build()
             .unwrap(),
-        mock_client,
-    );
+        ..Default::default()
+    };
+    let worker = mock_worker(MocksHolder::from_mock_worker(mock_client, mw));
 
     let _1 = worker.poll_activity_task().await.unwrap();
 

core/src/core_tests/workers.rs

@@ -1,13 +1,13 @@
 use crate::{
     prost_dur,
     test_help::{
-        build_fake_worker, build_mock_pollers, canned_histories, mock_manual_poller, mock_worker,
-        MockPollCfg, MockWorkerInputs, MocksHolder, ResponseType, WorkerExt,
+        build_fake_worker, build_mock_pollers, canned_histories, mock_worker, MockPollCfg,
+        MockWorkerInputs, MocksHolder, ResponseType, WorkerExt,
     },
     worker::client::mocks::mock_workflow_client,
     PollActivityError, PollWfError,
 };
-use futures::FutureExt;
+use futures_util::{stream, stream::StreamExt};
 use std::{cell::RefCell, time::Duration};
 use temporal_sdk_core_api::Worker;
 use temporal_sdk_core_protos::{
@@ -16,7 +16,9 @@ use temporal_sdk_core_protos::{
         workflow_commands::{workflow_command, CompleteWorkflowExecution, StartTimer},
         workflow_completion::WorkflowActivationCompletion,
     },
-    temporal::api::workflowservice::v1::RespondWorkflowTaskCompletedResponse,
+    temporal::api::workflowservice::v1::{
+        PollWorkflowTaskQueueResponse, RespondWorkflowTaskCompletedResponse,
+    },
 };
 use temporal_sdk_core_test_utils::start_timer_cmd;
 use tokio::sync::{watch, Barrier};
@@ -87,20 +89,18 @@ async fn shutdown_worker_can_complete_pending_activation() {
 #[tokio::test]
 async fn worker_shutdown_during_poll_doesnt_deadlock() {
     let (tx, rx) = watch::channel(false);
-    let mut mock_poller = mock_manual_poller();
     let rx = rx.clone();
-    mock_poller.expect_poll().returning(move || {
-        let mut rx = rx.clone();
-        async move {
-            // Don't resolve polls until worker shuts down
-            rx.changed().await.unwrap();
-            // We don't want to return a real response here because it would get buffered and
-            // then we'd have real work to do to be able to finish shutdown.
-            Some(Ok(Default::default()))
-        }
-        .boxed()
+    let stream = stream::unfold(rx, |mut rx| async move {
+        // Don't resolve polls until worker shuts down
+        rx.changed().await.unwrap();
+        // We don't want to return a real response here because it would get buffered and
+        // then we'd have real work to do to be able to finish shutdown.
+        Some((
+            Ok(PollWorkflowTaskQueueResponse::default().try_into().unwrap()),
+            rx,
+        ))
     });
-    let mw = MockWorkerInputs::new_from_poller(Box::new(mock_poller));
+    let mw = MockWorkerInputs::new(stream.boxed());
     let mut mock_client = mock_workflow_client();
     mock_client
         .expect_complete_workflow_task()

core/src/core_tests/workflow_tasks.rs

@@ -18,7 +18,7 @@ use rstest::{fixture, rstest};
 use std::{
     collections::{HashMap, HashSet, VecDeque},
     sync::{
-        atomic::{AtomicU64, Ordering},
+        atomic::{AtomicU64, AtomicUsize, Ordering},
         Arc,
     },
     time::Duration,
@@ -60,6 +60,7 @@ use temporal_sdk_core_test_utils::{fanout_tasks, start_timer_cmd, WorkerTestHelp
 use tokio::{
     join,
     sync::{Barrier, Semaphore},
+    time,
 };
 
 #[fixture(hist_batches = &[])]
@@ -1535,8 +1536,6 @@ async fn failing_wft_doesnt_eat_permit_forever() {
             .unwrap();
     }
     assert_eq!(worker.outstanding_workflow_tasks().await, 0);
-    // 1 permit is in use because the next task is buffered and has re-used the permit
-    assert_eq!(worker.available_wft_permits().await, 1);
     // We should be "out of work" because the mock service thinks we didn't complete the last task,
     // which we didn't, because we don't spam failures. The real server would eventually time out
     // the task. Mock doesn't understand that, so the WFT permit is released because eventually a
@@ -1545,7 +1544,7 @@ async fn failing_wft_doesnt_eat_permit_forever() {
     outstanding_mock_tasks.unwrap().release_run(&run_id);
     let activation = worker.poll_workflow_activation().await.unwrap();
     // There should be no change in permits, since this just unbuffered the buffered task
-    assert_eq!(worker.available_wft_permits().await, 1);
+    assert_eq!(worker.available_wft_permits(), 1);
     worker
         .complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
             activation.run_id,
@@ -1554,7 +1553,7 @@ async fn failing_wft_doesnt_eat_permit_forever() {
         .await
         .unwrap();
     worker.shutdown().await;
-    assert_eq!(worker.available_wft_permits().await, 2);
+    assert_eq!(worker.available_wft_permits(), 2);
 }
 
 #[tokio::test]
@@ -1875,7 +1874,7 @@ async fn poll_faster_than_complete_wont_overflow_cache() {
 
     join!(blocking_poll, complete_evict);
     // p5 outstanding and final poll outstanding -- hence one permit available
-    assert_eq!(core.available_wft_permits().await, 1);
+    assert_eq!(core.available_wft_permits(), 1);
     assert_eq!(core.cached_workflows().await, 3);
 }
 
@@ -2588,3 +2587,72 @@ async fn history_length_with_fail_and_timeout(
         .unwrap();
     worker.run_until_done().await.unwrap();
 }
+
+#[tokio::test]
+async fn poller_wont_run_ahead_of_task_slots() {
+    let popped_tasks = Arc::new(AtomicUsize::new(0));
+    let ptc = popped_tasks.clone();
+    let mut bunch_of_first_tasks = (1..50).map(move |i| {
+        ptc.fetch_add(1, Ordering::Relaxed);
+        hist_to_poll_resp(
+            &canned_histories::single_timer(&format!("{i}")),
+            format!("wf-{i}"),
+            1.into(),
+        )
+        .resp
+    });
+    let mut mock_client = mock_workflow_client();
+    mock_client
+        .expect_poll_workflow_task()
+        .returning(move |_, _| Ok(bunch_of_first_tasks.next().unwrap()));
+    mock_client
+        .expect_complete_workflow_task()
+        .returning(|_| Ok(Default::default()));
+
+    let worker = Worker::new_test(
+        test_worker_cfg()
+            .max_cached_workflows(10_usize)
+            .max_outstanding_workflow_tasks(10_usize)
+            .max_concurrent_wft_polls(10_usize)
+            .no_remote_activities(true)
+            .build()
+            .unwrap(),
+        mock_client,
+    );
+
+    // Should be able to get up to 10 tasks
+    let mut tasks = vec![];
+    for _ in 0..10 {
+        tasks.push(worker.poll_workflow_activation().await.unwrap());
+    }
+
+    assert_eq!(worker.outstanding_workflow_tasks().await, 10);
+    assert_eq!(worker.available_wft_permits(), 0);
+    assert_eq!(worker.unused_wft_permits(), 0);
+
+    // This one should hang until we complete some tasks since we're at the limit
+    let hung_poll = async {
+        // This should end up getting shut down after the other routine finishes tasks
+        assert_matches!(
+            worker.poll_workflow_activation().await.unwrap_err(),
+            PollWfError::ShutDown
+        );
+    };
+    // Wait for a bit concurrently with above, verify no extra tasks got taken, shutdown
+    let ender = async {
+        time::sleep(Duration::from_millis(300)).await;
+        // initiate shutdown, then complete open tasks
+        worker.initiate_shutdown();
+        for t in tasks {
+            worker
+                .complete_workflow_activation(WorkflowActivationCompletion::empty(t.run_id))
+                .await
+                .unwrap();
+        }
+        worker.shutdown().await;
+    };
+    join!(hung_poll, ender);
+    // We shouldn't have got more than the 10 tasks from the poller -- verifying that the concurrent
+    // polling is not exceeding the task limit
+    assert_eq!(popped_tasks.load(Ordering::Relaxed), 10);
+}