Enable setting activity task queue throttle (#283)

Author: Sushisource

.buildkite/docker/Dockerfile

@@ -1,4 +1,4 @@
-FROM rust:1.58
+FROM rust:1.59
 
 RUN rustup component add rustfmt && \
 	rustup component add clippy

client/src/lib.rs

@@ -62,6 +62,7 @@ use uuid::Uuid;
 
 #[cfg(any(feature = "mocks", test))]
 use futures::Future;
+use temporal_sdk_core_protos::temporal::api::taskqueue::v1::TaskQueueMetadata;
 
 static LONG_POLL_METHOD_NAMES: [&str; 2] = ["PollWorkflowTaskQueue", "PollActivityTaskQueue"];
 /// The server times out polls after 60 seconds. Set our timeout to be slightly beyond that.
@@ -532,8 +533,11 @@ pub trait ServerGatewayApis {
 
     /// Fetch new activity tasks from the provided queue. Should block indefinitely if there is no
     /// work.
-    async fn poll_activity_task(&self, task_queue: String)
-        -> Result<PollActivityTaskQueueResponse>;
+    async fn poll_activity_task(
+        &self,
+        task_queue: String,
+        max_tasks_per_sec: Option<f64>,
+    ) -> Result<PollActivityTaskQueueResponse>;
 
     /// Notifies the server that workflow tasks for a given workflow should be sent to the normal
     /// non-sticky task queue. This normally happens when workflow has been evicted from the cache.
@@ -731,6 +735,7 @@ impl ServerGatewayApis for ServerGateway {
     async fn poll_activity_task(
         &self,
         task_queue: String,
+        max_tasks_per_sec: Option<f64>,
     ) -> Result<PollActivityTaskQueueResponse> {
         let request = PollActivityTaskQueueRequest {
             namespace: self.namespace.clone(),
@@ -739,7 +744,9 @@ impl ServerGatewayApis for ServerGateway {
                 kind: TaskQueueKind::Normal as i32,
             }),
             identity: self.opts.identity.clone(),
-            task_queue_metadata: None,
+            task_queue_metadata: max_tasks_per_sec.map(|tps| TaskQueueMetadata {
+                max_tasks_per_second: Some(tps),
+            }),
         };
 
         Ok(self

client/src/mocks.rs

@@ -54,7 +54,7 @@ mockall::mock! {
             -> impl Future<Output = Result<PollWorkflowTaskQueueResponse>> + Send + 'b
             where 'a: 'b, Self: 'b;
 
-        fn poll_activity_task<'a, 'b>(&self, task_queue: String)
+        fn poll_activity_task<'a, 'b>(&self, task_queue: String, max_tasks_per_sec: Option<f64>)
             -> impl Future<Output = Result<PollActivityTaskQueueResponse>> + Send + 'b
             where 'a: 'b, Self: 'b;
 

client/src/retry.rs

@@ -229,8 +229,14 @@ where
     async fn poll_activity_task(
         &self,
         task_queue: String,
+        max_tasks_per_sec: Option<f64>,
     ) -> Result<PollActivityTaskQueueResponse> {
-        retry_call!(self, poll_activity_task, task_queue.clone())
+        retry_call!(
+            self,
+            poll_activity_task,
+            task_queue.clone(),
+            max_tasks_per_sec
+        )
     }
 
     async fn reset_sticky_task_queue(
@@ -483,14 +489,16 @@ mod tests {
                 .times(1);
             mock_gateway
                 .expect_poll_activity_task()
-                .returning(move |_| Err(Status::new(code, "non-retryable failure")))
+                .returning(move |_, _| Err(Status::new(code, "non-retryable failure")))
                 .times(1);
             let retry_gateway = RetryGateway::new(mock_gateway, Default::default());
             let result = retry_gateway
                 .poll_workflow_task("tq".to_string(), false)
                 .await;
             assert!(result.is_err());
-            let result = retry_gateway.poll_activity_task("tq".to_string()).await;
+            let result = retry_gateway
+                .poll_activity_task("tq".to_string(), None)
+                .await;
             assert!(result.is_err());
         }
     }
@@ -530,11 +538,11 @@ mod tests {
             .times(1);
         mock_gateway
             .expect_poll_activity_task()
-            .returning(move |_| Err(Status::new(Code::Unknown, "retryable failure")))
+            .returning(move |_, _| Err(Status::new(Code::Unknown, "retryable failure")))
             .times(50);
         mock_gateway
             .expect_poll_activity_task()
-            .returning(|_| Ok(Default::default()))
+            .returning(|_, _| Ok(Default::default()))
             .times(1);
 
         let retry_gateway = RetryGateway::new(mock_gateway, Default::default());
@@ -543,7 +551,9 @@ mod tests {
             .poll_workflow_task("tq".to_string(), false)
             .await;
         assert!(result.is_ok());
-        let result = retry_gateway.poll_activity_task("tq".to_string()).await;
+        let result = retry_gateway
+            .poll_activity_task("tq".to_string(), None)
+            .await;
         assert!(result.is_ok());
     }
 
@@ -562,11 +572,11 @@ mod tests {
                 .times(1);
             mock_gateway
                 .expect_poll_activity_task()
-                .returning(move |_| Err(Status::new(code, "retryable failure")))
+                .returning(move |_, _| Err(Status::new(code, "retryable failure")))
                 .times(5);
             mock_gateway
                 .expect_poll_activity_task()
-                .returning(|_| Ok(Default::default()))
+                .returning(|_, _| Ok(Default::default()))
                 .times(1);
 
             let retry_gateway = RetryGateway::new(mock_gateway, Default::default());
@@ -575,7 +585,9 @@ mod tests {
                 .poll_workflow_task("tq".to_string(), false)
                 .await;
             assert!(result.is_ok());
-            let result = retry_gateway.poll_activity_task("tq".to_string()).await;
+            let result = retry_gateway
+                .poll_activity_task("tq".to_string(), None)
+                .await;
             assert!(result.is_ok());
         }
     }

core-api/src/worker.rs

@@ -4,7 +4,6 @@ use std::time::Duration;
 #[derive(Debug, Clone, derive_builder::Builder)]
 #[builder(setter(into), build_fn(validate = "Self::validate"))]
 #[non_exhaustive]
-// TODO: per-second queue limits
 pub struct WorkerConfig {
     /// The Temporal service namespace this worker is bound to
     pub namespace: String,
@@ -65,6 +64,13 @@ pub struct WorkerConfig {
     /// `heartbeat_timeout * 0.8`.
     #[builder(default = "Duration::from_secs(30)")]
     pub default_heartbeat_throttle_interval: Duration,
+
+    /// Sets the maximum number of activities per second the task queue will dispatch, controlled
+    /// server-side. Note that this only takes effect upon an activity poll request. If multiple
+    /// workers on the same queue have different values set, they will thrash with the last poller
+    /// winning.
+    #[builder(setter(strip_option), default)]
+    pub max_task_queue_activities_per_second: Option<f64>,
 }
 
 impl WorkerConfig {

core/src/core_tests/activity_tasks.rs

@@ -5,9 +5,10 @@ use crate::{
         mock_worker, poll_and_reply, test_worker_cfg, MockWorker, MocksHolder,
     },
     workflow::WorkflowCachingPolicy::NonSticky,
-    ActivityHeartbeat,
+    ActivityHeartbeat, MetricsContext, Worker, WorkerConfigBuilder,
 };
 use futures::FutureExt;
+use std::sync::Arc;
 use std::{
     cell::RefCell,
     collections::{hash_map::Entry, HashMap, VecDeque},
@@ -16,7 +17,7 @@ use std::{
     time::Duration,
 };
 use temporal_client::mocks::{mock_gateway, mock_manual_gateway};
-use temporal_sdk_core_api::Worker;
+use temporal_sdk_core_api::Worker as WorkerTrait;
 use temporal_sdk_core_protos::{
     coresdk::{
         activity_result::{activity_resolution, ActivityExecutionResult, ActivityResolution},
@@ -61,7 +62,7 @@ async fn max_activities_respected() {
     mock_gateway
         .expect_poll_activity_task()
         .times(3)
-        .returning(move |_| Ok(tasks.pop_front().unwrap()));
+        .returning(move |_, _| Ok(tasks.pop_front().unwrap()));
     mock_gateway
         .expect_complete_activity_task()
         .returning(|_, _| Ok(RespondActivityTaskCompletedResponse::default()));
@@ -347,7 +348,7 @@ async fn many_concurrent_heartbeat_cancels() {
     let mut calls_map = HashMap::<_, i32>::new();
     mock_gateway
         .expect_poll_activity_task()
-        .returning(move |_| poll_resps.pop_front().unwrap());
+        .returning(move |_, _| poll_resps.pop_front().unwrap());
     mock_gateway
         .expect_cancel_activity_task()
         .returning(move |_, _| async move { Ok(Default::default()) }.boxed());
@@ -585,3 +586,28 @@ async fn complete_act_with_fail_flushes_heartbeat() {
     let last_seen_payload = &last_seen_payload.take().unwrap().payloads[0];
     assert_eq!(last_seen_payload.data, &[last_hb]);
 }
+
+#[tokio::test]
+async fn max_tq_acts_set_passed_to_poll_properly() {
+    let rate = 9.28;
+    let mut mock_gateway = mock_gateway();
+    mock_gateway
+        .expect_poll_activity_task()
+        .returning(move |_, tps| {
+            assert_eq!(tps, Some(rate));
+            Ok(PollActivityTaskQueueResponse {
+                task_token: vec![1],
+                ..Default::default()
+            })
+        });
+
+    let cfg = WorkerConfigBuilder::default()
+        .namespace("enchi")
+        .task_queue("cat")
+        .max_concurrent_at_polls(1_usize)
+        .max_task_queue_activities_per_second(rate)
+        .build()
+        .unwrap();
+    let worker = Worker::new(cfg, None, Arc::new(mock_gateway), MetricsContext::default());
+    worker.poll_activity_task().await.unwrap();
+}

core/src/core_tests/mod.rs

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

core/src/pollers/poll_buffer.rs

@@ -221,12 +221,13 @@ pub fn new_activity_task_buffer(
     task_queue: String,
     concurrent_pollers: usize,
     buffer_size: usize,
+    max_tps: Option<f64>,
 ) -> PollActivityTaskBuffer {
     LongPollBuffer::new(
         move || {
             let sg = sg.clone();
             let task_queue = task_queue.clone();
-            async move { sg.poll_activity_task(task_queue).await }
+            async move { sg.poll_activity_task(task_queue, max_tps).await }
         },
         concurrent_pollers,
         buffer_size,

core/src/protosext/mod.rs

@@ -361,7 +361,7 @@ impl ValidScheduleLA {
         };
         let retry_policy = v
             .retry_policy
-            .ok_or(anyhow!("Retry policy must be defined!"))?;
+            .ok_or_else(|| anyhow!("Retry policy must be defined!"))?;
         let local_retry_threshold = v
             .local_retry_threshold
             .clone()

core/src/worker/mod.rs

@@ -217,6 +217,7 @@ impl Worker {
                 config.task_queue.clone(),
                 config.max_concurrent_at_polls,
                 config.max_concurrent_at_polls * 2,
+                config.max_task_queue_activities_per_second,
             );
             let act_metrics = metrics.with_new_attrs([activity_poller()]);
             ap.set_num_pollers_handler(move |np| act_metrics.record_num_pollers(np));
@@ -343,7 +344,7 @@ impl Worker {
     ///
     /// Returns `Ok(None)` in the event of a poll timeout or if the polling loop should otherwise
     /// be restarted
-    pub(crate) async fn activity_poll(&self) -> Result<Option<ActivityTask>, PollActivityError> {
+    async fn activity_poll(&self) -> Result<Option<ActivityTask>, PollActivityError> {
         let act_mgr_poll = async {
             if let Some(ref act_mgr) = self.at_task_mgr {
                 act_mgr.poll().await
@@ -948,7 +949,7 @@ mod tests {
         let mut mock_gateway = mock_gateway();
         mock_gateway
             .expect_poll_activity_task()
-            .returning(|_| Ok(PollActivityTaskQueueResponse::default()));
+            .returning(|_, _| Ok(PollActivityTaskQueueResponse::default()));
 
         let cfg = test_worker_cfg()
             .max_outstanding_activities(5_usize)
@@ -980,7 +981,7 @@ mod tests {
         let mut mock_gateway = mock_gateway();
         mock_gateway
             .expect_poll_activity_task()
-            .returning(|_| Err(tonic::Status::internal("ahhh")));
+            .returning(|_, _| Err(tonic::Status::internal("ahhh")));
 
         let cfg = test_worker_cfg()
             .max_outstanding_activities(5_usize)

sdk-core-protos/src/history_info.rs

@@ -42,9 +42,9 @@ impl HistoryInfo {
                 attrs
                     .workflow_type
                     .as_ref()
-                    .ok_or(anyhow!(
-                        "No workflow type defined in execution started attributes"
-                    ))?
+                    .ok_or_else(|| {
+                        anyhow!("No workflow type defined in execution started attributes")
+                    })?
                     .name
                     .clone()
             }