@temporalio/core-bridge 0.19.2 → 0.20.2

package/sdk-core/core/src/abstractions.rs

@@ -0,0 +1,38 @@
+//! This module contains very generic helpers that can be used codebase-wide
+
+use crate::MetricsContext;
+use tokio::sync::{AcquireError, Semaphore, SemaphorePermit};
+
+/// Wraps a [Semaphore] with a function call that is fed the available permits any time a permit is
+/// acquired or restored through the provided methods
+pub(crate) struct MeteredSemaphore {
+    pub sem: Semaphore,
+    metrics_ctx: MetricsContext,
+    record_fn: fn(&MetricsContext, usize),
+}
+
+impl MeteredSemaphore {
+    pub fn new(
+        inital_permits: usize,
+        metrics_ctx: MetricsContext,
+        record_fn: fn(&MetricsContext, usize),
+    ) -> Self {
+        Self {
+            sem: Semaphore::new(inital_permits),
+            metrics_ctx,
+            record_fn,
+        }
+    }
+
+    pub async fn acquire(&self) -> Result<SemaphorePermit<'_>, AcquireError> {
+        let res = self.sem.acquire().await;
+        (self.record_fn)(&self.metrics_ctx, self.sem.available_permits());
+        res
+    }
+
+    /// Adds just one permit
+    pub fn add_permit(&self) {
+        self.sem.add_permits(1);
+        (self.record_fn)(&self.metrics_ctx, self.sem.available_permits());
+    }
+}

package/sdk-core/core/src/core_tests/activity_tasks.rs

@@ -1,29 +1,26 @@
 use crate::{
     job_assert,
     test_help::{
-        build_fake_core, canned_histories, gen_assert_and_reply, mock_core,
-        mock_core_with_opts_no_workers, mock_manual_poller, mock_poller, mock_poller_from_resps,
-        poll_and_reply, MockWorker, MocksHolder, TEST_Q,
+        build_fake_worker, canned_histories, gen_assert_and_reply, mock_manual_poller, mock_poller,
+        mock_worker, poll_and_reply, test_worker_cfg, MockWorker, MocksHolder,
     },
+    worker::client::mocks::{mock_manual_workflow_client, mock_workflow_client},
     workflow::WorkflowCachingPolicy::NonSticky,
-    ActivityHeartbeat, ActivityTask, Core, CoreInitOptionsBuilder, CoreSDK, WorkerConfigBuilder,
+    ActivityHeartbeat, Worker, WorkerConfigBuilder,
 };
 use futures::FutureExt;
 use std::{
     cell::RefCell,
     collections::{hash_map::Entry, HashMap, VecDeque},
     rc::Rc,
-    sync::{
-        atomic::{AtomicUsize, Ordering},
-        Arc,
-    },
+    sync::atomic::{AtomicUsize, Ordering},
     time::Duration,
 };
-use temporal_client::mocks::{fake_sg_opts, mock_gateway, mock_manual_gateway};
+use temporal_sdk_core_api::Worker as WorkerTrait;
 use temporal_sdk_core_protos::{
     coresdk::{
         activity_result::{activity_resolution, ActivityExecutionResult, ActivityResolution},
-        activity_task::activity_task,
+        activity_task::{activity_task, ActivityTask},
         workflow_activation::{workflow_activation_job, ResolveActivity, WorkflowActivationJob},
         workflow_commands::{
             ActivityCancellationType, CompleteWorkflowExecution, RequestCancelActivity,
@@ -38,11 +35,11 @@ use temporal_sdk_core_protos::{
     },
 };
 use temporal_sdk_core_test_utils::{fanout_tasks, start_timer_cmd};
-use tokio::{join, sync::Notify, time::sleep};
+use tokio::{join, time::sleep};
 
 #[tokio::test]
 async fn max_activities_respected() {
-    let task_q = "q";
+    let _task_q = "q";
     let mut tasks = VecDeque::from(vec![
         PollActivityTaskQueueResponse {
             task_token: vec![1],
@@ -60,47 +57,38 @@ async fn max_activities_respected() {
             ..Default::default()
         },
     ]);
-    let mut mock_gateway = mock_gateway();
-    mock_gateway
+    let mut mock_client = mock_workflow_client();
+    mock_client
         .expect_poll_activity_task()
         .times(3)
-        .returning(move |_| Ok(tasks.pop_front().unwrap()));
-    mock_gateway
+        .returning(move |_, _| Ok(tasks.pop_front().unwrap()));
+    mock_client
         .expect_complete_activity_task()
         .returning(|_, _| Ok(RespondActivityTaskCompletedResponse::default()));
 
-    let core = CoreSDK::new(
-        mock_gateway,
-        CoreInitOptionsBuilder::default()
-            .gateway_opts(fake_sg_opts())
-            .build()
-            .unwrap(),
-    );
-    core.register_worker(
-        WorkerConfigBuilder::default()
-            .task_queue(TEST_Q)
+    let worker = Worker::new_test(
+        test_worker_cfg()
             .max_outstanding_activities(2_usize)
             .build()
             .unwrap(),
-    )
-    .unwrap();
+        mock_client,
+    );
 
     // We allow two outstanding activities, therefore first two polls should return right away
-    let r1 = core.poll_activity_task(task_q).await.unwrap();
-    let _r2 = core.poll_activity_task(task_q).await.unwrap();
+    let r1 = worker.poll_activity_task().await.unwrap();
+    let _r2 = worker.poll_activity_task().await.unwrap();
     // Third should block until we complete one of the first two
     let last_finisher = AtomicUsize::new(0);
     tokio::join! {
         async {
-            core.complete_activity_task(ActivityTaskCompletion {
+            worker.complete_activity_task(ActivityTaskCompletion {
                 task_token: r1.task_token,
-                task_queue: TEST_Q.to_string(),
                 result: Some(ActivityExecutionResult::ok(vec![1].into()))
             }).await.unwrap();
             last_finisher.store(1, Ordering::SeqCst);
         },
         async {
-            core.poll_activity_task(task_q).await.unwrap();
+            worker.poll_activity_task().await.unwrap();
             last_finisher.store(2, Ordering::SeqCst);
         }
     };
@@ -110,19 +98,14 @@ async fn max_activities_respected() {
 
 #[tokio::test]
 async fn activity_not_found_returns_ok() {
-    let mut mock_gateway = mock_gateway();
+    let mut mock_client = mock_workflow_client();
     // Mock won't even be called, since we weren't tracking activity
-    mock_gateway.expect_complete_activity_task().times(0);
+    mock_client.expect_complete_activity_task().times(0);
 
-    let core = mock_core(MocksHolder::from_gateway_with_responses(
-        mock_gateway,
-        [],
-        [],
-    ));
+    let core = mock_worker(MocksHolder::from_client_with_responses(mock_client, [], []));
 
     core.complete_activity_task(ActivityTaskCompletion {
         task_token: vec![1],
-        task_queue: TEST_Q.to_string(),
         result: Some(ActivityExecutionResult::ok(vec![1].into())),
     })
     .await
@@ -132,8 +115,8 @@ async fn activity_not_found_returns_ok() {
 
 #[tokio::test]
 async fn heartbeats_report_cancels_only_once() {
-    let mut mock_gateway = mock_gateway();
-    mock_gateway
+    let mut mock_client = mock_workflow_client();
+    mock_client
         .expect_record_activity_heartbeat()
         .times(2)
         .returning(|_, _| {
@@ -141,17 +124,17 @@ async fn heartbeats_report_cancels_only_once() {
                 cancel_requested: true,
             })
         });
-    mock_gateway
+    mock_client
         .expect_complete_activity_task()
         .times(1)
         .returning(|_, _| Ok(RespondActivityTaskCompletedResponse::default()));
-    mock_gateway
+    mock_client
         .expect_cancel_activity_task()
         .times(1)
         .returning(|_, _| Ok(RespondActivityTaskCanceledResponse::default()));
 
-    let core = mock_core(MocksHolder::from_gateway_with_responses(
-        mock_gateway,
+    let core = mock_worker(MocksHolder::from_client_with_responses(
+        mock_client,
         [],
         [
             PollActivityTaskQueueResponse {
@@ -169,15 +152,14 @@ async fn heartbeats_report_cancels_only_once() {
         ],
     ));
 
-    let act = core.poll_activity_task(TEST_Q).await.unwrap();
+    let act = core.poll_activity_task().await.unwrap();
     core.record_activity_heartbeat(ActivityHeartbeat {
         task_token: act.task_token.clone(),
-        task_queue: TEST_Q.to_string(),
         details: vec![vec![1_u8, 2, 3].into()],
     });
     // We have to wait a beat for the heartbeat to be processed
     sleep(Duration::from_millis(10)).await;
-    let act = core.poll_activity_task(TEST_Q).await.unwrap();
+    let act = core.poll_activity_task().await.unwrap();
     assert_matches!(
         &act,
         ActivityTask {
@@ -192,7 +174,6 @@ async fn heartbeats_report_cancels_only_once() {
     sleep(Duration::from_millis(10)).await;
     core.record_activity_heartbeat(ActivityHeartbeat {
         task_token: act.task_token.clone(),
-        task_queue: TEST_Q.to_string(),
         details: vec![vec![1_u8, 2, 3].into()],
     });
     // Wait delay again to flush heartbeat
@@ -200,14 +181,14 @@ async fn heartbeats_report_cancels_only_once() {
     // Now complete it as cancelled
     core.complete_activity_task(ActivityTaskCompletion {
         task_token: act.task_token,
-        task_queue: TEST_Q.to_string(),
+
         result: Some(ActivityExecutionResult::cancel_from_details(None)),
     })
     .await
     .unwrap();
     // Since cancels always come before new tasks, if we get a new non-cancel task, we did not
     // double-issue cancels.
-    let act = core.poll_activity_task(TEST_Q).await.unwrap();
+    let act = core.poll_activity_task().await.unwrap();
     assert_matches!(
         &act,
         ActivityTask {
@@ -219,7 +200,7 @@ async fn heartbeats_report_cancels_only_once() {
     // Complete it so shutdown goes through
     core.complete_activity_task(ActivityTaskCompletion {
         task_token: act.task_token,
-        task_queue: TEST_Q.to_string(),
+
         result: Some(ActivityExecutionResult::ok(vec![1].into())),
     })
     .await
@@ -250,8 +231,8 @@ async fn activity_cancel_interrupts_poll() {
         .times(2)
         .returning(move || poll_resps.pop_front().unwrap());
 
-    let mut mock_gateway = mock_manual_gateway();
-    mock_gateway
+    let mut mock_client = mock_manual_workflow_client();
+    mock_client
         .expect_record_activity_heartbeat()
         .times(1)
         .returning(|_, _| {
@@ -262,37 +243,36 @@ async fn activity_cancel_interrupts_poll() {
             }
             .boxed()
         });
-    mock_gateway
+    mock_client
         .expect_complete_activity_task()
         .times(1)
         .returning(|_, _| async { Ok(RespondActivityTaskCompletedResponse::default()) }.boxed());
 
-    let mock_worker = MockWorker {
+    let mw = MockWorker {
         act_poller: Some(Box::from(mock_poller)),
         ..Default::default()
     };
-
-    let core = mock_core(MocksHolder::from_mock_workers(mock_gateway, [mock_worker]));
+    let core = mock_worker(MocksHolder::from_mock_worker(mock_client.into(), mw));
     let last_finisher = AtomicUsize::new(0);
     // Perform first poll to get the activity registered
-    let act = core.poll_activity_task(TEST_Q).await.unwrap();
+    let act = core.poll_activity_task().await.unwrap();
     // Poll should block until heartbeat is sent, issuing the cancel, and interrupting the poll
     tokio::join! {
         async {
             core.record_activity_heartbeat(ActivityHeartbeat {
                 task_token: act.task_token,
-                task_queue: TEST_Q.to_string(),
+
                 details: vec![vec![1_u8, 2, 3].into()],
             });
             last_finisher.store(1, Ordering::SeqCst);
         },
         async {
-            let act = core.poll_activity_task(TEST_Q).await.unwrap();
+            let act = core.poll_activity_task().await.unwrap();
             // Must complete this activity for shutdown to finish
             core.complete_activity_task(
                 ActivityTaskCompletion {
                     task_token: act.task_token,
-                    task_queue: TEST_Q.to_string(),
+
                     result: Some(ActivityExecutionResult::ok(vec![1].into())),
                 }
             ).await.unwrap();
@@ -306,7 +286,7 @@ async fn activity_cancel_interrupts_poll() {
 
 #[tokio::test]
 async fn activity_poll_timeout_retries() {
-    let mock_gateway = mock_gateway();
+    let mock_client = mock_workflow_client();
     let mut calls = 0;
     let mut mock_act_poller = mock_poller();
     mock_act_poller.expect_poll().times(3).returning(move || {
@@ -320,12 +300,12 @@ async fn activity_poll_timeout_retries() {
             }))
         }
     });
-    let mock_worker = MockWorker {
+    let mw = MockWorker {
         act_poller: Some(Box::from(mock_act_poller)),
         ..Default::default()
     };
-    let core = mock_core(MocksHolder::from_mock_workers(mock_gateway, [mock_worker]));
-    let r = core.poll_activity_task(TEST_Q).await.unwrap();
+    let core = mock_worker(MocksHolder::from_mock_worker(mock_client.into(), mw));
+    let r = core.poll_activity_task().await.unwrap();
     assert_matches!(r.task_token.as_slice(), b"hello!");
 }
 
@@ -335,7 +315,7 @@ async fn many_concurrent_heartbeat_cancels() {
     // them after a few successful heartbeats
     const CONCURRENCY_NUM: usize = 5;
 
-    let mut mock_gateway = mock_manual_gateway();
+    let mut mock_client = mock_manual_workflow_client();
     let mut poll_resps = VecDeque::from(
         (0..CONCURRENCY_NUM)
             .map(|i| {
@@ -361,13 +341,13 @@ async fn many_concurrent_heartbeat_cancels() {
         .boxed(),
     );
     let mut calls_map = HashMap::<_, i32>::new();
-    mock_gateway
+    mock_client
         .expect_poll_activity_task()
-        .returning(move |_| poll_resps.pop_front().unwrap());
-    mock_gateway
+        .returning(move |_, _| poll_resps.pop_front().unwrap());
+    mock_client
         .expect_cancel_activity_task()
         .returning(move |_, _| async move { Ok(Default::default()) }.boxed());
-    mock_gateway
+    mock_client
         .expect_record_activity_heartbeat()
         .returning(move |tt, _| {
             let calls = match calls_map.entry(tt) {
@@ -391,30 +371,27 @@ async fn many_concurrent_heartbeat_cancels() {
             .boxed()
         });
 
-    let core = &mock_core_with_opts_no_workers(mock_gateway, CoreInitOptionsBuilder::default());
-    core.register_worker(
-        WorkerConfigBuilder::default()
-            .task_queue(TEST_Q)
+    let worker = &Worker::new_test(
+        test_worker_cfg()
             .max_outstanding_activities(CONCURRENCY_NUM)
             // Only 1 poll at a time to avoid over-polling and running out of responses
             .max_concurrent_at_polls(1_usize)
             .build()
             .unwrap(),
-    )
-    .unwrap();
+        mock_client,
+    );
 
     // Poll all activities first so they are registered
     for _ in 0..CONCURRENCY_NUM {
-        core.poll_activity_task(TEST_Q).await.unwrap();
+        worker.poll_activity_task().await.unwrap();
     }
 
     // Spawn "activities"
     fanout_tasks(CONCURRENCY_NUM, |i| async move {
         let task_token = i.to_be_bytes().to_vec();
         for _ in 0..12 {
-            core.record_activity_heartbeat(ActivityHeartbeat {
+            worker.record_activity_heartbeat(ActivityHeartbeat {
                 task_token: task_token.clone(),
-                task_queue: TEST_Q.to_string(),
                 details: vec![],
             });
             sleep(Duration::from_millis(50)).await;
@@ -424,7 +401,7 @@ async fn many_concurrent_heartbeat_cancels() {
 
     // Read all the cancellations and reply to them concurrently
     fanout_tasks(CONCURRENCY_NUM, |_| async move {
-        let r = core.poll_activity_task(TEST_Q).await.unwrap();
+        let r = worker.poll_activity_task().await.unwrap();
         assert_matches!(
             r,
             ActivityTask {
@@ -432,23 +409,23 @@ async fn many_concurrent_heartbeat_cancels() {
                 ..
             }
         );
-        core.complete_activity_task(ActivityTaskCompletion {
-            task_token: r.task_token.clone(),
-            task_queue: TEST_Q.to_string(),
-            result: Some(ActivityExecutionResult::cancel_from_details(None)),
-        })
-        .await
-        .unwrap();
+        worker
+            .complete_activity_task(ActivityTaskCompletion {
+                task_token: r.task_token.clone(),
+                result: Some(ActivityExecutionResult::cancel_from_details(None)),
+            })
+            .await
+            .unwrap();
     })
     .await;
 
-    core.shutdown().await;
+    worker.shutdown().await;
 }
 
 #[tokio::test]
 async fn activity_timeout_no_double_resolve() {
     let t = canned_histories::activity_double_resolve_repro();
-    let core = build_fake_core("fake_wf_id", t, &[3]);
+    let core = build_fake_worker("fake_wf_id", t, &[3]);
     let activity_id = 1;
 
     poll_and_reply(
@@ -497,84 +474,10 @@ async fn activity_timeout_no_double_resolve() {
     core.shutdown().await;
 }
 
-#[tokio::test]
-async fn only_returns_cancels_for_desired_queue() {
-    let mut mock_gateway = mock_gateway();
-    let seen_cancel = Arc::new(Notify::new());
-    let sc = seen_cancel.clone();
-    mock_gateway
-        .expect_record_activity_heartbeat()
-        .times(1)
-        .returning(move |_, _| {
-            // Mark the activity as needing cancel when heartbeated
-            sc.notify_one();
-            Ok(RecordActivityTaskHeartbeatResponse {
-                cancel_requested: true,
-            })
-        });
-    mock_gateway
-        .expect_cancel_activity_task()
-        .times(1)
-        .returning(|_, _| Ok(RespondActivityTaskCanceledResponse::default()));
-
-    let mut w1 = MockWorker::for_queue("q1");
-    w1.act_poller = Some(mock_poller_from_resps([PollActivityTaskQueueResponse {
-        task_token: vec![1],
-        activity_id: "act1".to_string(),
-        heartbeat_timeout: Some(Duration::from_millis(1).into()),
-        ..Default::default()
-    }]));
-    let mut mock_act_poller = mock_poller();
-    mock_act_poller
-        .expect_poll()
-        .returning(|| Some(Ok(Default::default())));
-    let mut w2 = MockWorker::for_queue("q2");
-    w2.act_poller = Some(Box::from(mock_act_poller));
-
-    let core = mock_core(MocksHolder::from_mock_workers(mock_gateway, [w1, w2]));
-    // First poll should get the activity
-    core.poll_activity_task("q1").await.unwrap();
-    // Now record a heartbeat which will get the cancel response and mark the act as cancelled
-    core.record_activity_heartbeat(ActivityHeartbeat {
-        task_token: vec![1],
-        task_queue: "q1".to_string(),
-        details: vec![],
-    });
-    // Worker two's poll should never resolve, since it's just getting pretend long poll timeouts
-    let q1poll = async {
-        // Wait for cancel to propagate
-        seen_cancel.notified().await;
-        core.poll_activity_task("q1").await
-    };
-    let q1_res = tokio::select! {
-        _ = core.poll_activity_task("q2") => panic!("q2 poll resolved!"),
-        r = q1poll => {
-            r.unwrap()
-        }
-    };
-    assert_matches!(
-        &q1_res,
-        ActivityTask {
-            variant: Some(activity_task::Variant::Cancel(_)),
-            ..
-        }
-    );
-    // act needs to complete to finalize shutdown
-    core.complete_activity_task(ActivityTaskCompletion {
-        task_token: q1_res.task_token,
-        task_queue: "q1".to_string(),
-        result: Some(ActivityExecutionResult::cancel_from_details(None)),
-    })
-    .await
-    .unwrap();
-
-    core.shutdown().await;
-}
-
 #[tokio::test]
 async fn can_heartbeat_acts_during_shutdown() {
-    let mut mock_gateway = mock_gateway();
-    mock_gateway
+    let mut mock_client = mock_workflow_client();
+    mock_client
         .expect_record_activity_heartbeat()
         .times(1)
         .returning(|_, _| {
@@ -582,13 +485,13 @@ async fn can_heartbeat_acts_during_shutdown() {
                 cancel_requested: false,
             })
         });
-    mock_gateway
+    mock_client
         .expect_complete_activity_task()
         .times(1)
         .returning(|_, _| Ok(RespondActivityTaskCompletedResponse::default()));
 
-    let core = mock_core(MocksHolder::from_gateway_with_responses(
-        mock_gateway,
+    let core = mock_worker(MocksHolder::from_client_with_responses(
+        mock_client,
         [],
         [PollActivityTaskQueueResponse {
             task_token: vec![1],
@@ -598,7 +501,7 @@ async fn can_heartbeat_acts_during_shutdown() {
         }],
     ));
 
-    let act = core.poll_activity_task(TEST_Q).await.unwrap();
+    let act = core.poll_activity_task().await.unwrap();
     let complete_order = RefCell::new(vec![]);
     // Start shutdown before completing the activity
     let shutdown_fut = async {
@@ -608,12 +511,12 @@ async fn can_heartbeat_acts_during_shutdown() {
     let complete_fut = async {
         core.record_activity_heartbeat(ActivityHeartbeat {
             task_token: act.task_token.clone(),
-            task_queue: TEST_Q.to_string(),
+
             details: vec![vec![1_u8, 2, 3].into()],
         });
         core.complete_activity_task(ActivityTaskCompletion {
             task_token: act.task_token,
-            task_queue: TEST_Q.to_string(),
+
             result: Some(ActivityExecutionResult::ok(vec![1].into())),
         })
         .await
@@ -629,10 +532,10 @@ async fn can_heartbeat_acts_during_shutdown() {
 #[tokio::test]
 async fn complete_act_with_fail_flushes_heartbeat() {
     let last_hb = 50;
-    let mut mock_gateway = mock_gateway();
+    let mut mock_client = mock_workflow_client();
     let last_seen_payload = Rc::new(RefCell::new(None));
     let lsp = last_seen_payload.clone();
-    mock_gateway
+    mock_client
         .expect_record_activity_heartbeat()
         // Two times b/c we always record the first heartbeat, and we'll flush the last
         .times(2)
@@ -642,13 +545,13 @@ async fn complete_act_with_fail_flushes_heartbeat() {
                 cancel_requested: false,
             })
         });
-    mock_gateway
+    mock_client
         .expect_fail_activity_task()
         .times(1)
         .returning(|_, _| Ok(RespondActivityTaskFailedResponse::default()));
 
-    let core = mock_core(MocksHolder::from_gateway_with_responses(
-        mock_gateway,
+    let core = mock_worker(MocksHolder::from_client_with_responses(
+        mock_client,
         [],
         [PollActivityTaskQueueResponse {
             task_token: vec![1],
@@ -658,18 +561,16 @@ async fn complete_act_with_fail_flushes_heartbeat() {
         }],
     ));
 
-    let act = core.poll_activity_task(TEST_Q).await.unwrap();
+    let act = core.poll_activity_task().await.unwrap();
     // Record a bunch of heartbeats
     for i in 1..=last_hb {
         core.record_activity_heartbeat(ActivityHeartbeat {
             task_token: act.task_token.clone(),
-            task_queue: TEST_Q.to_string(),
             details: vec![vec![i].into()],
         });
     }
     core.complete_activity_task(ActivityTaskCompletion {
         task_token: act.task_token.clone(),
-        task_queue: TEST_Q.to_string(),
         result: Some(ActivityExecutionResult::fail("Ahh".into())),
     })
     .await
@@ -680,3 +581,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_client = mock_workflow_client();
+    mock_client
+        .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_test(cfg, mock_client);
+    worker.poll_activity_task().await.unwrap();
+}