@temporalio/core-bridge 0.23.0 → 1.0.0

package/sdk-core/core/src/worker/activities/activity_heartbeat_manager.rs

@@ -9,8 +9,10 @@ use std::{
     time::{self, Duration, Instant},
 };
 use temporal_sdk_core_protos::{
-    coresdk::{activity_task::ActivityCancelReason, common, ActivityHeartbeat, IntoPayloadsExt},
-    temporal::api::workflowservice::v1::RecordActivityTaskHeartbeatResponse,
+    coresdk::{activity_task::ActivityCancelReason, ActivityHeartbeat, IntoPayloadsExt},
+    temporal::api::{
+        common::v1::Payload, workflowservice::v1::RecordActivityTaskHeartbeatResponse,
+    },
 };
 use tokio::{
     sync::{
@@ -47,7 +49,7 @@ enum HeartbeatAction {
 #[derive(Debug)]
 pub struct ValidActivityHeartbeat {
     pub task_token: TaskToken,
-    pub details: Vec<common::Payload>,
+    pub details: Vec<Payload>,
     pub throttle_interval: time::Duration,
 }
 
@@ -58,7 +60,7 @@ enum HeartbeatExecutorAction {
     /// Report heartbeat to the server
     Report {
         task_token: TaskToken,
-        details: Vec<common::Payload>,
+        details: Vec<Payload>,
     },
 }
 
@@ -132,7 +134,15 @@ impl ActivityHeartbeatManager {
         let _ = self.shutdown_token.cancel();
         let mut handle = self.join_handle.lock().await;
         if let Some(h) = handle.take() {
-            h.await.expect("shutdown should exit cleanly");
+            let handle_r = h.await;
+            if let Err(e) = handle_r {
+                if !e.is_cancelled() {
+                    error!(
+                        "Unexpected error joining heartbeating tasks during shutdown: {:?}",
+                        e
+                    )
+                }
+            }
         }
     }
 }
@@ -140,7 +150,7 @@ impl ActivityHeartbeatManager {
 #[derive(Debug)]
 struct ActivityHeartbeatState {
     /// If None and throttle interval is over, untrack this task token
-    last_recorded_details: Option<Vec<common::Payload>>,
+    last_recorded_details: Option<Vec<Payload>>,
     /// True if we've queued up a request to record against server, but it hasn't yet completed
     is_record_in_flight: bool,
     last_send_requested: Instant,
@@ -401,9 +411,8 @@ mod test {
 
     use crate::worker::client::mocks::mock_workflow_client;
     use std::time::Duration;
-    use temporal_sdk_core_protos::{
-        coresdk::common::Payload,
-        temporal::api::workflowservice::v1::RecordActivityTaskHeartbeatResponse,
+    use temporal_sdk_core_protos::temporal::api::{
+        common::v1::Payload, workflowservice::v1::RecordActivityTaskHeartbeatResponse,
     };
     use tokio::time::sleep;
 

package/sdk-core/core/src/worker/activities/local_activities.rs

@@ -1,5 +1,7 @@
 use crate::{
-    abstractions::MeteredSemaphore, protosext::ValidScheduleLA, retry_logic::RetryPolicyExt,
+    abstractions::{MeteredSemaphore, OwnedMeteredSemPermit},
+    protosext::ValidScheduleLA,
+    retry_logic::RetryPolicyExt,
     MetricsContext, TaskToken,
 };
 use parking_lot::Mutex;
@@ -12,9 +14,8 @@ use temporal_sdk_core_protos::{
     coresdk::{
         activity_result::{Cancellation, Failure as ActFail, Success},
         activity_task::{activity_task, ActivityCancelReason, ActivityTask, Cancel, Start},
-        common::WorkflowExecution,
     },
-    temporal::api::enums::v1::TimeoutType,
+    temporal::api::{common::v1::WorkflowExecution, enums::v1::TimeoutType},
 };
 use tokio::{
     sync::{
@@ -44,6 +45,7 @@ pub(crate) struct LocalInFlightActInfo {
     pub la_info: NewLocalAct,
     pub dispatch_time: Instant,
     pub attempt: u32,
+    _permit: OwnedMeteredSemPermit,
 }
 
 #[derive(Debug, Clone)]
@@ -181,6 +183,7 @@ impl LocalActivityManager {
         )
     }
 
+    #[cfg(test)]
     pub(crate) fn num_outstanding(&self) -> usize {
         self.dat.lock().outstanding_activity_tasks.len()
     }
@@ -192,11 +195,11 @@ impl LocalActivityManager {
 
     pub(crate) fn enqueue(
         &self,
-        reqs: impl IntoIterator<Item = LocalActRequest> + Debug,
+        reqs: impl IntoIterator<Item = LocalActRequest>,
     ) -> Vec<LocalActivityResolution> {
-        debug!("Queuing local activities: {:?}", &reqs);
         let mut immediate_resolutions = vec![];
         for req in reqs {
+            debug!(local_activity = ?req, "Queuing local activity");
             match req {
                 LocalActRequest::New(act) => {
                     let id = ExecutingLAId {
@@ -267,7 +270,7 @@ impl LocalActivityManager {
     /// Returns the next pending local-activity related action, or None if shutdown has initiated
     /// and there are no more remaining actions to take.
     pub(crate) async fn next_pending(&self) -> Option<DispatchOrTimeoutLA> {
-        let new_or_retry = match self.rcvs.lock().await.next(&self.semaphore).await? {
+        let (new_or_retry, permit) = match self.rcvs.lock().await.next(&self.semaphore).await? {
             NewOrCancel::Cancel(c) => {
                 return match c {
                     CancelOrTimeout::Cancel(c) => Some(DispatchOrTimeoutLA::Dispatch(c)),
@@ -308,7 +311,7 @@ impl LocalActivityManager {
                     }
                 };
             }
-            NewOrCancel::New(n) => n,
+            NewOrCancel::New(n, perm) => (n, perm),
         };
 
         // It is important that there are no await points after receiving from the channel, as
@@ -363,6 +366,7 @@ impl LocalActivityManager {
                 la_info: orig,
                 dispatch_time: Instant::now(),
                 attempt,
+                _permit: permit,
             },
         );
         if let Some(to) = dat.timeout_tasks.get_mut(&id) {
@@ -407,7 +411,6 @@ impl LocalActivityManager {
                 seq_num: info.la_info.schedule_cmd.seq,
             };
             dlock.id_to_tt.remove(&exec_id);
-            self.semaphore.add_permit();
 
             match status {
                 LocalActivityExecutionResult::Completed(_)
@@ -513,7 +516,7 @@ enum CancelOrTimeout {
 }
 
 enum NewOrCancel {
-    New(NewOrRetry),
+    New(NewOrRetry, OwnedMeteredSemPermit),
     Cancel(CancelOrTimeout),
 }
 
@@ -531,13 +534,13 @@ impl RcvChans {
             cancel = async { self.cancels_req_rx.recv().await } => {
                 Some(NewOrCancel::Cancel(cancel.expect("Send halves of LA manager are not dropped")))
             }
-            maybe_new_or_retry = async {
+            (maybe_new_or_retry, perm) = async {
                 // Wait for a permit to take a task and forget it. Permits are removed until a
                 // completion.
-                new_sem.acquire().await.expect("is never closed").forget();
-                self.act_req_rx.recv().await
+                let perm = new_sem.acquire_owned().await.expect("is never closed");
+                (self.act_req_rx.recv().await, perm)
             } => Some(NewOrCancel::New(
-                maybe_new_or_retry.expect("Send halves of LA manager are not dropped")
+                maybe_new_or_retry.expect("Send halves of LA manager are not dropped"), perm
             )),
             _ = self.shutdown.cancelled() => None
         }
@@ -636,9 +639,9 @@ impl Drop for TimeoutBag {
 mod tests {
     use super::*;
     use crate::protosext::LACloseTimeouts;
-    use temporal_sdk_core_protos::{
-        coresdk::common::RetryPolicy,
-        temporal::api::failure::v1::{failure::FailureInfo, ApplicationFailureInfo, Failure},
+    use temporal_sdk_core_protos::temporal::api::{
+        common::v1::RetryPolicy,
+        failure::v1::{failure::FailureInfo, ApplicationFailureInfo, Failure},
     };
     use tokio::{sync::mpsc::error::TryRecvError, task::yield_now};
 

package/sdk-core/core/src/worker/activities.rs

@@ -8,7 +8,7 @@ pub(crate) use local_activities::{
 };
 
 use crate::{
-    abstractions::MeteredSemaphore,
+    abstractions::{MeteredSemaphore, OwnedMeteredSemPermit},
     pollers::BoxedActPoller,
     telemetry::metrics::{activity_type, activity_worker_type, workflow_type, MetricsContext},
     worker::{
@@ -19,6 +19,12 @@ use crate::{
 };
 use activity_heartbeat_manager::ActivityHeartbeatManager;
 use dashmap::DashMap;
+use governor::{
+    clock::DefaultClock,
+    middleware::NoOpMiddleware,
+    state::{InMemoryState, NotKeyed},
+    Quota, RateLimiter,
+};
 use std::{
     convert::TryInto,
     sync::Arc,
@@ -52,7 +58,6 @@ struct InFlightActInfo {
 }
 
 /// Augments [InFlightActInfo] with details specific to remote activities
-#[derive(Debug)]
 struct RemoteInFlightActInfo {
     pub base: InFlightActInfo,
     /// Used to calculate aggregation delay between activity heartbeats.
@@ -63,12 +68,15 @@ struct RemoteInFlightActInfo {
     /// we have learned from heartbeating and issued a cancel task, in which case we may simply
     /// discard the reply.
     pub known_not_found: bool,
+    /// The permit from the max concurrent semaphore
+    _permit: OwnedMeteredSemPermit,
 }
 impl RemoteInFlightActInfo {
     fn new(
         activity_type: String,
         workflow_type: String,
         heartbeat_timeout: Option<prost_types::Duration>,
+        permit: OwnedMeteredSemPermit,
     ) -> Self {
         Self {
             base: InFlightActInfo {
@@ -79,10 +87,26 @@ impl RemoteInFlightActInfo {
             heartbeat_timeout,
             issued_cancel_to_lang: false,
             known_not_found: false,
+            _permit: permit,
         }
     }
 }
 
+struct NonPollActBuffer {
+    tx: async_channel::Sender<PermittedTqResp>,
+    rx: async_channel::Receiver<PermittedTqResp>,
+}
+impl NonPollActBuffer {
+    pub fn new() -> Self {
+        let (tx, rx) = async_channel::unbounded();
+        Self { tx, rx }
+    }
+
+    pub async fn next(&self) -> PermittedTqResp {
+        self.rx.recv().await.expect("Send half cannot be dropped")
+    }
+}
+
 pub(crate) struct WorkerActivityTasks {
     /// Centralizes management of heartbeat issuing / throttling
     heartbeat_manager: ActivityHeartbeatManager,
@@ -91,8 +115,13 @@ pub(crate) struct WorkerActivityTasks {
     /// Buffers activity task polling in the event we need to return a cancellation while a poll is
     /// ongoing.
     poller: BoxedActPoller,
+    /// Holds activity tasks we have received by non-polling means. EX: In direct response to
+    /// workflow task completion.
+    non_poll_tasks: NonPollActBuffer,
     /// Ensures we stay at or below this worker's maximum concurrent activity limit
-    activities_semaphore: MeteredSemaphore,
+    activities_semaphore: Arc<MeteredSemaphore>,
+    /// Enables per-worker rate-limiting of activity tasks
+    ratelimiter: Option<RateLimiter<NotKeyed, InMemoryState, DefaultClock, NoOpMiddleware>>,
     /// Wakes every time an activity is removed from the outstanding map
     complete_notify: Notify,
 
@@ -105,6 +134,7 @@ pub(crate) struct WorkerActivityTasks {
 impl WorkerActivityTasks {
     pub(crate) fn new(
         max_activity_tasks: usize,
+        max_worker_act_per_sec: Option<f64>,
         poller: BoxedActPoller,
         client: Arc<WorkerClientBag>,
         metrics: MetricsContext,
@@ -115,11 +145,15 @@ impl WorkerActivityTasks {
             heartbeat_manager: ActivityHeartbeatManager::new(client),
             outstanding_activity_tasks: Default::default(),
             poller,
-            activities_semaphore: MeteredSemaphore::new(
+            non_poll_tasks: NonPollActBuffer::new(),
+            activities_semaphore: Arc::new(MeteredSemaphore::new(
                 max_activity_tasks,
                 metrics.with_new_attrs([activity_worker_type()]),
                 MetricsContext::available_task_slots,
-            ),
+            )),
+            ratelimiter: max_worker_act_per_sec.and_then(|ps| {
+                Quota::with_period(Duration::from_secs_f64(ps.recip())).map(RateLimiter::direct)
+            }),
             complete_notify: Notify::new(),
             metrics,
             max_heartbeat_throttle_interval,
@@ -151,12 +185,15 @@ impl WorkerActivityTasks {
             // Acquire and subsequently forget a permit for an outstanding activity. When they are
             // completed, we must add a new permit to the semaphore, since holding the permit the
             // entire time lang does work would be a challenge.
-            let sem = self
+            let perm = self
                 .activities_semaphore
-                .acquire()
+                .acquire_owned()
                 .await
                 .expect("outstanding activity semaphore not closed");
-            (self.poller.poll().await, sem)
+            if let Some(ref rl) = self.ratelimiter {
+                rl.until_ready().await;
+            }
+            (self.poller.poll().await, perm)
         };
 
         tokio::select! {
@@ -165,7 +202,10 @@ impl WorkerActivityTasks {
             cancel_task = self.next_pending_cancel_task() => {
                 cancel_task
             }
-            (work, sem) = poll_with_semaphore => {
+            task = self.non_poll_tasks.next() => {
+                Ok(Some(self.about_to_issue_task(task)))
+            }
+            (work, permit) = poll_with_semaphore => {
                 match work {
                     Some(Ok(work)) => {
                         if work == PollActivityTaskQueueResponse::default() {
@@ -173,23 +213,10 @@ impl WorkerActivityTasks {
                             self.metrics.act_poll_timeout();
                             return Ok(None)
                         }
-
-                        if let Some(dur) = work.sched_to_start() {
-                            self.metrics
-                                .act_sched_to_start_latency(dur);
-                        }
-
-                        self.outstanding_activity_tasks.insert(
-                            work.task_token.clone().into(),
-                            RemoteInFlightActInfo::new(
-                                work.activity_type.clone().unwrap_or_default().name,
-                                work.workflow_type.clone().unwrap_or_default().name,
-                                work.heartbeat_timeout.clone()
-                            ),
-                        );
-                        // Only permanently take a permit in the event the poll finished properly
-                        sem.forget();
-                        Ok(Some(ActivityTask::start_from_poll_resp(work)))
+                        let work = self.about_to_issue_task(PermittedTqResp {
+                           resp: work, permit
+                        });
+                        Ok(Some(work))
                     }
                     None => {
                         Err(PollActivityError::ShutDown)
@@ -212,10 +239,9 @@ impl WorkerActivityTasks {
                 workflow_type(act_info.base.workflow_type.clone()),
             ]);
             act_metrics.act_execution_latency(act_info.base.start_time.elapsed());
-            self.activities_semaphore.add_permit();
-            self.heartbeat_manager.evict(task_token.clone()).await;
             let known_not_found = act_info.known_not_found;
             drop(act_info); // TODO: Get rid of dashmap. If we hold ref across await, bad stuff.
+            self.heartbeat_manager.evict(task_token.clone()).await;
             self.complete_notify.notify_waiters();
 
             // No need to report activities which we already know the server doesn't care about
@@ -304,6 +330,14 @@ impl WorkerActivityTasks {
         self.heartbeat_manager.record(details, throttle_interval)
     }
 
+    /// Returns a handle that the workflows management side can use to interact with this manager
+    pub(crate) fn get_handle_for_workflows(&self) -> ActivitiesFromWFTsHandle {
+        ActivitiesFromWFTsHandle {
+            sem: self.activities_semaphore.clone(),
+            tx: self.non_poll_tasks.tx.clone(),
+        }
+    }
+
     async fn next_pending_cancel_task(&self) -> Result<Option<ActivityTask>, PollActivityError> {
         let next_pc = self.heartbeat_manager.next_pending_cancel().await;
         // Issue cancellations for anything we noticed was cancelled during heartbeating
@@ -336,8 +370,101 @@ impl WorkerActivityTasks {
         }
     }
 
+    /// Called when there is a new act task about to be bubbled up out of the manager
+    fn about_to_issue_task(&self, task: PermittedTqResp) -> ActivityTask {
+        if let Some(dur) = task.resp.sched_to_start() {
+            self.metrics.act_sched_to_start_latency(dur);
+        };
+
+        self.outstanding_activity_tasks.insert(
+            task.resp.task_token.clone().into(),
+            RemoteInFlightActInfo::new(
+                task.resp.activity_type.clone().unwrap_or_default().name,
+                task.resp.workflow_type.clone().unwrap_or_default().name,
+                task.resp.heartbeat_timeout.clone(),
+                task.permit,
+            ),
+        );
+
+        ActivityTask::start_from_poll_resp(task.resp)
+    }
+
     #[cfg(test)]
     pub(crate) fn remaining_activity_capacity(&self) -> usize {
-        self.activities_semaphore.sem.available_permits()
+        self.activities_semaphore.available_permits()
+    }
+}
+
+/// Provides facilities for the workflow side of things to interact with the activity manager.
+/// Allows for the handling of activities returned by WFT completions.
+pub(crate) struct ActivitiesFromWFTsHandle {
+    sem: Arc<MeteredSemaphore>,
+    tx: async_channel::Sender<PermittedTqResp>,
+}
+
+impl ActivitiesFromWFTsHandle {
+    /// Returns a handle that can be used to reserve an activity slot. EX: When requesting eager
+    /// dispatch of an activity to this worker upon workflow task completion
+    pub(crate) fn reserve_slot(&self) -> Option<OwnedMeteredSemPermit> {
+        self.sem.try_acquire_owned().ok()
+    }
+
+    /// Queue new activity tasks for dispatch received from non-polling sources (ex: eager returns
+    /// from WFT completion)
+    pub(crate) fn add_tasks(&self, tasks: impl IntoIterator<Item = PermittedTqResp>) {
+        for t in tasks.into_iter() {
+            self.tx.try_send(t).expect("Receive half cannot be dropped");
+        }
+    }
+}
+
+pub(crate) struct PermittedTqResp {
+    pub permit: OwnedMeteredSemPermit,
+    pub resp: PollActivityTaskQueueResponse,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{
+        test_help::mock_poller_from_resps, worker::client::mocks::mock_manual_workflow_client,
+    };
+
+    #[tokio::test]
+    async fn per_worker_ratelimit() {
+        let poller = mock_poller_from_resps([
+            PollActivityTaskQueueResponse {
+                task_token: vec![1],
+                activity_id: "act1".to_string(),
+                ..Default::default()
+            }
+            .into(),
+            PollActivityTaskQueueResponse {
+                task_token: vec![2],
+                activity_id: "act2".to_string(),
+                ..Default::default()
+            }
+            .into(),
+        ]);
+        let client = WorkerClientBag::new(
+            Box::new(mock_manual_workflow_client()),
+            "fake_namespace".to_string(),
+        );
+        let atm = WorkerActivityTasks::new(
+            10,
+            Some(2.0),
+            poller,
+            Arc::new(client),
+            MetricsContext::default(),
+            Duration::from_secs(1),
+            Duration::from_secs(1),
+        );
+        let start = Instant::now();
+        atm.poll().await.unwrap().unwrap();
+        atm.poll().await.unwrap().unwrap();
+        // At least half a second will have elapsed since we only allow 2 tasks per second.
+        // With no ratelimit, even on a slow CI server with lots of load, this would typically take
+        // low single digit ms or less.
+        assert!(start.elapsed() > Duration::from_secs_f64(0.5));
     }
 }

package/sdk-core/core/src/worker/client.rs

@@ -118,6 +118,7 @@ pub(crate) trait WorkerClient: Sync + Send {
 #[async_trait::async_trait]
 impl<'a, T> WorkerClient for T
 where
+    // TODO: This should be workflow service... no reason to marry worker trait to sdk client trait
     T: Borrow<dyn WorkflowClientTrait + 'a + Send + Sync> + Send + Sync,
 {
     async fn poll_workflow_task(