@temporalio/core-bridge 0.19.2 → 0.20.0

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

@@ -1,4 +1,7 @@
-use crate::{pollers::ServerGatewayApis, worker::activities::PendingActivityCancel, TaskToken};
+use crate::{
+    worker::{activities::PendingActivityCancel, client::WorkerClientBag},
+    TaskToken,
+};
 use futures::StreamExt;
 use std::{
     collections::{hash_map::Entry, HashMap},
@@ -56,7 +59,6 @@ enum HeartbeatExecutorAction {
     Report {
         task_token: TaskToken,
         details: Vec<common::Payload>,
-        on_reported: Option<Arc<Notify>>,
     },
 }
 
@@ -139,6 +141,8 @@ impl ActivityHeartbeatManager {
 struct ActivityHeartbeatState {
     /// If None and throttle interval is over, untrack this task token
     last_recorded_details: Option<Vec<common::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,
     throttle_interval: Duration,
     throttled_cancellation_token: Option<CancellationToken>,
@@ -161,6 +165,7 @@ impl ActivityHeartbeatState {
 #[derive(Debug)]
 struct HeartbeatStreamState {
     tt_to_state: HashMap<TaskToken, ActivityHeartbeatState>,
+    tt_needs_flush: HashMap<TaskToken, Arc<Notify>>,
     incoming_hbs: UnboundedReceiver<HeartbeatAction>,
     /// Token that can be used to cancel the entire stream.
     /// Requests to the server are not cancelled with this token.
@@ -175,6 +180,7 @@ impl HeartbeatStreamState {
             Self {
                 cancellation_token: cancellation_token.clone(),
                 tt_to_state: Default::default(),
+                tt_needs_flush: Default::default(),
                 incoming_hbs,
             },
             heartbeat_tx,
@@ -193,13 +199,13 @@ impl HeartbeatStreamState {
                     // None is used to mark that after throttling we can stop tracking this task
                     // token.
                     last_recorded_details: None,
+                    is_record_in_flight: true,
                     throttled_cancellation_token: None,
                 };
                 e.insert(state);
                 Some(HeartbeatExecutorAction::Report {
                     task_token: hb.task_token,
                     details: hb.details,
-                    on_reported: None,
                 })
             }
             Entry::Occupied(mut o) => {
@@ -212,7 +218,11 @@ impl HeartbeatStreamState {
 
     /// Heartbeat report to server completed
     fn handle_report_completed(&mut self, tt: TaskToken) -> Option<HeartbeatExecutorAction> {
+        if let Some(not) = self.tt_needs_flush.remove(&tt) {
+            not.notify_one();
+        }
         if let Some(st) = self.tt_to_state.get_mut(&tt) {
+            st.is_record_in_flight = false;
             let cancellation_token = self.cancellation_token.child_token();
             st.throttled_cancellation_token = Some(cancellation_token.clone());
             // Always sleep for simplicity even if the duration is 0
@@ -236,10 +246,10 @@ impl HeartbeatStreamState {
                     // Reset the cancellation token and schedule another report
                     state.throttled_cancellation_token = None;
                     state.last_send_requested = Instant::now();
+                    state.is_record_in_flight = true;
                     Some(HeartbeatExecutorAction::Report {
                         task_token: tt,
                         details,
-                        on_reported: None,
                     })
                 } else {
                     // Nothing to report, forget this task token
@@ -265,11 +275,14 @@ impl HeartbeatStreamState {
                 let _ = cancel_tok.cancel();
             }
             if let Some(last_deets) = state.last_recorded_details {
+                self.tt_needs_flush.insert(tt.clone(), on_complete);
                 return Some(HeartbeatExecutorAction::Report {
                     task_token: tt,
                     details: last_deets,
-                    on_reported: Some(on_complete),
                 });
+            } else if state.is_record_in_flight {
+                self.tt_needs_flush.insert(tt, on_complete);
+                return None;
             }
         }
         // Since there's nothing to flush immediately report back that eviction is finished
@@ -281,7 +294,7 @@ impl HeartbeatStreamState {
 impl ActivityHeartbeatManager {
     /// Creates a new instance of an activity heartbeat manager and returns a handle to the user,
     /// which allows to send new heartbeats and initiate the shutdown.
-    pub fn new(sg: Arc<impl ServerGatewayApis + Send + Sync + 'static + ?Sized>) -> Self {
+    pub fn new(client: Arc<WorkerClientBag>) -> Self {
         let (heartbeat_stream_state, heartbeat_tx_source, shutdown_token) =
             HeartbeatStreamState::new();
         let (cancels_tx, cancels_rx) = unbounded_channel();
@@ -320,7 +333,7 @@ impl ActivityHeartbeatManager {
             .filter_map(|opt| async { opt })
             .for_each_concurrent(None, move |action| {
                 let heartbeat_tx = heartbeat_tx_source.clone();
-                let sg = sg.clone();
+                let sg = client.clone();
                 let cancels_tx = cancels_tx.clone();
                 async move {
                     match action {
@@ -332,7 +345,7 @@ impl ActivityHeartbeatManager {
                                 },
                             };
                         }
-                        HeartbeatExecutorAction::Report { task_token: tt, details, on_reported } => {
+                        HeartbeatExecutorAction::Report { task_token: tt, details } => {
                             match sg
                                 .record_activity_heartbeat(tt.clone(), details.into_payloads())
                                 .await
@@ -353,6 +366,8 @@ impl ActivityHeartbeatManager {
                                 // finished (which is one thing not found implies - other reasons
                                 // would seem equally valid).
                                 Err(s) if s.code() == tonic::Code::NotFound => {
+                                    debug!(task_token = %tt,
+                                           "Activity not found when recording heartbeat");
                                     cancels_tx
                                         .send(PendingActivityCancel::new(
                                             tt.clone(),
@@ -364,9 +379,6 @@ impl ActivityHeartbeatManager {
                                     warn!("Error when recording heartbeat: {:?}", e);
                                 }
                             };
-                            if let Some(onrep) = on_reported {
-                                onrep.notify_one();
-                            }
                             let _ = heartbeat_tx.send(HeartbeatAction::CompleteReport(tt));
                         }
                     }
@@ -386,10 +398,9 @@ impl ActivityHeartbeatManager {
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::test_help::TEST_Q;
-    use std::time::Duration;
 
-    use temporal_client::mocks::mock_gateway;
+    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,
@@ -400,12 +411,12 @@ mod test {
     /// every 1/2 of the heartbeat timeout.
     #[tokio::test]
     async fn process_heartbeats_and_shutdown() {
-        let mut mock_gateway = mock_gateway();
-        mock_gateway
+        let mut mock_client = mock_workflow_client();
+        mock_client
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
             .times(2);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_gateway));
+        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client.into()));
         let fake_task_token = vec![1, 2, 3];
         // Send 2 heartbeat requests for 20ms apart.
         // The first heartbeat should be sent right away, and
@@ -421,12 +432,12 @@ mod test {
 
     #[tokio::test]
     async fn send_heartbeats_less_frequently_than_throttle_interval() {
-        let mut mock_gateway = mock_gateway();
-        mock_gateway
+        let mut mock_client = mock_workflow_client();
+        mock_client
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
             .times(3);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_gateway));
+        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client.into()));
         let fake_task_token = vec![1, 2, 3];
         // Heartbeats always get sent if recorded less frequently than the throttle intreval
         for i in 0_u8..3 {
@@ -441,12 +452,12 @@ mod test {
     /// interactions with the server - one immediately and one after 500ms after the throttle_interval.
     #[tokio::test]
     async fn process_tight_loop_and_shutdown() {
-        let mut mock_gateway = mock_gateway();
-        mock_gateway
+        let mut mock_client = mock_workflow_client();
+        mock_client
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
             .times(1);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_gateway));
+        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client.into()));
         let fake_task_token = vec![1, 2, 3];
         // Send a whole bunch of heartbeats very fast. We should still only send one total.
         for i in 0_u8..50 {
@@ -460,12 +471,12 @@ mod test {
     /// This test reports one heartbeat and waits for the throttle_interval to elapse before sending another
     #[tokio::test]
     async fn report_heartbeat_after_timeout() {
-        let mut mock_gateway = mock_gateway();
-        mock_gateway
+        let mut mock_client = mock_workflow_client();
+        mock_client
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
             .times(2);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_gateway));
+        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client.into()));
         let fake_task_token = vec![1, 2, 3];
         record_heartbeat(&hm, fake_task_token.clone(), 0, Duration::from_millis(100));
         sleep(Duration::from_millis(500)).await;
@@ -477,12 +488,12 @@ mod test {
 
     #[tokio::test]
     async fn evict_works() {
-        let mut mock_gateway = mock_gateway();
-        mock_gateway
+        let mut mock_client = mock_workflow_client();
+        mock_client
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
             .times(2);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_gateway));
+        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client.into()));
         let fake_task_token = vec![1, 2, 3];
         record_heartbeat(&hm, fake_task_token.clone(), 0, Duration::from_millis(100));
         // Let it propagate
@@ -495,20 +506,33 @@ mod test {
         hm.shutdown().await;
     }
 
+    #[tokio::test]
+    async fn evict_immediate_after_record() {
+        let mut mock_client = mock_workflow_client();
+        mock_client
+            .expect_record_activity_heartbeat()
+            .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
+            .times(1);
+        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client.into()));
+        let fake_task_token = vec![1, 2, 3];
+        record_heartbeat(&hm, fake_task_token.clone(), 0, Duration::from_millis(100));
+        hm.evict(fake_task_token.clone().into()).await;
+        hm.shutdown().await;
+    }
+
     /// Recording new heartbeats after shutdown is not allowed, and will result in error.
     #[tokio::test]
     async fn record_after_shutdown() {
-        let mut mock_gateway = mock_gateway();
-        mock_gateway
+        let mut mock_client = mock_workflow_client();
+        mock_client
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
             .times(0);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_gateway));
+        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client.into()));
         hm.shutdown().await;
         match hm.record(
             ActivityHeartbeat {
                 task_token: vec![1, 2, 3],
-                task_queue: TEST_Q.to_string(),
                 details: vec![Payload {
                     // payload doesn't matter in this case, as it shouldn't get sent anyways.
                     ..Default::default()
@@ -534,7 +558,6 @@ mod test {
         hm.record(
             ActivityHeartbeat {
                 task_token,
-                task_queue: TEST_Q.to_string(),
                 details: vec![Payload {
                     metadata: Default::default(),
                     data: vec![payload_data],

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

@@ -1,4 +1,7 @@
-use crate::{protosext::ValidScheduleLA, retry_logic::RetryPolicyExt, TaskToken};
+use crate::{
+    abstractions::MeteredSemaphore, protosext::ValidScheduleLA, retry_logic::RetryPolicyExt,
+    MetricsContext, TaskToken,
+};
 use parking_lot::Mutex;
 use std::{
     collections::HashMap,
@@ -16,7 +19,7 @@ use temporal_sdk_core_protos::{
 use tokio::{
     sync::{
         mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender},
-        Notify, Semaphore,
+        Notify,
     },
     task::JoinHandle,
     time::sleep,
@@ -103,7 +106,7 @@ pub(crate) struct LocalActivityManager {
     /// Just so we can provide activity tasks the same namespace as the worker
     namespace: String,
     /// Constrains number of currently executing local activities
-    semaphore: Semaphore,
+    semaphore: MeteredSemaphore,
     /// Sink for new activity execution requests
     act_req_tx: UnboundedSender<NewOrRetry>,
     /// Cancels need a different queue since they should be taken first, and don't take a permit
@@ -135,13 +138,21 @@ impl LAMData {
 }
 
 impl LocalActivityManager {
-    pub(crate) fn new(max_concurrent: usize, namespace: String) -> Self {
+    pub(crate) fn new(
+        max_concurrent: usize,
+        namespace: String,
+        metrics_context: MetricsContext,
+    ) -> Self {
         let (act_req_tx, act_req_rx) = unbounded_channel();
         let (cancels_req_tx, cancels_req_rx) = unbounded_channel();
         let shutdown_complete_tok = CancellationToken::new();
         Self {
             namespace,
-            semaphore: Semaphore::new(max_concurrent),
+            semaphore: MeteredSemaphore::new(
+                max_concurrent,
+                metrics_context,
+                MetricsContext::available_task_slots,
+            ),
             act_req_tx,
             cancels_req_tx,
             complete_notify: Notify::new(),
@@ -161,6 +172,15 @@ impl LocalActivityManager {
         }
     }
 
+    #[cfg(test)]
+    fn test(max_concurrent: usize) -> Self {
+        Self::new(
+            max_concurrent,
+            "fake_ns".to_string(),
+            MetricsContext::default(),
+        )
+    }
+
     pub(crate) fn num_outstanding(&self) -> usize {
         self.dat.lock().outstanding_activity_tasks.len()
     }
@@ -358,7 +378,7 @@ impl LocalActivityManager {
                 workflow_execution: Some(new_la.workflow_exec_info),
                 activity_id: sa.activity_id,
                 activity_type: sa.activity_type,
-                header_fields: sa.header_fields,
+                header_fields: sa.headers,
                 input: sa.arguments,
                 heartbeat_details: vec![],
                 scheduled_time: Some(new_la.schedule_time.into()),
@@ -387,7 +407,7 @@ impl LocalActivityManager {
                 seq_num: info.la_info.schedule_cmd.seq,
             };
             dlock.id_to_tt.remove(&exec_id);
-            self.semaphore.add_permits(1);
+            self.semaphore.add_permit();
 
             match status {
                 LocalActivityExecutionResult::Completed(_)
@@ -507,7 +527,7 @@ struct RcvChans {
 }
 
 impl RcvChans {
-    async fn next(&mut self, new_sem: &Semaphore) -> Option<NewOrCancel> {
+    async fn next(&mut self, new_sem: &MeteredSemaphore) -> Option<NewOrCancel> {
         tokio::select! {
             cancel = async { self.cancels_req_rx.recv().await } => {
                 Some(NewOrCancel::Cancel(cancel.expect("Send halves of LA manager are not dropped")))
@@ -633,7 +653,7 @@ mod tests {
 
     #[tokio::test]
     async fn max_concurrent_respected() {
-        let lam = LocalActivityManager::new(1, "whatever".to_string());
+        let lam = LocalActivityManager::test(1);
         lam.enqueue((1..=50).map(|i| {
             NewLocalAct {
                 schedule_cmd: ValidScheduleLA {
@@ -673,7 +693,7 @@ mod tests {
 
     #[tokio::test]
     async fn no_work_doesnt_deadlock_with_complete() {
-        let lam = LocalActivityManager::new(5, "whatever".to_string());
+        let lam = LocalActivityManager::test(5);
         lam.enqueue([NewLocalAct {
             schedule_cmd: ValidScheduleLA {
                 seq: 1,
@@ -705,7 +725,7 @@ mod tests {
 
     #[tokio::test]
     async fn can_cancel_in_flight() {
-        let lam = LocalActivityManager::new(5, "whatever".to_string());
+        let lam = LocalActivityManager::test(5);
         lam.enqueue([NewLocalAct {
             schedule_cmd: ValidScheduleLA {
                 seq: 1,
@@ -732,7 +752,7 @@ mod tests {
 
     #[tokio::test]
     async fn respects_timer_backoff_threshold() {
-        let lam = LocalActivityManager::new(1, "whatever".to_string());
+        let lam = LocalActivityManager::test(1);
         lam.enqueue([NewLocalAct {
             schedule_cmd: ValidScheduleLA {
                 seq: 1,
@@ -767,7 +787,7 @@ mod tests {
 
     #[tokio::test]
     async fn can_cancel_during_local_backoff() {
-        let lam = LocalActivityManager::new(1, "whatever".to_string());
+        let lam = LocalActivityManager::test(1);
         lam.enqueue([NewLocalAct {
             schedule_cmd: ValidScheduleLA {
                 seq: 1,
@@ -816,7 +836,7 @@ mod tests {
 
     #[tokio::test]
     async fn local_backoff_clears_handle_map_when_started() {
-        let lam = LocalActivityManager::new(1, "whatever".to_string());
+        let lam = LocalActivityManager::test(1);
         lam.enqueue([NewLocalAct {
             schedule_cmd: ValidScheduleLA {
                 seq: 1,
@@ -852,7 +872,7 @@ mod tests {
 
     #[tokio::test]
     async fn sched_to_start_timeout() {
-        let lam = LocalActivityManager::new(1, "whatever".to_string());
+        let lam = LocalActivityManager::test(1);
         let timeout = Duration::from_millis(100);
         lam.enqueue([NewLocalAct {
             schedule_cmd: ValidScheduleLA {
@@ -893,7 +913,7 @@ mod tests {
     #[case::start(false)]
     #[tokio::test]
     async fn local_x_to_close_timeout(#[case] is_schedule: bool) {
-        let lam = LocalActivityManager::new(1, "whatever".to_string());
+        let lam = LocalActivityManager::test(1);
         let timeout = Duration::from_millis(100);
         let close_timeouts = if is_schedule {
             LACloseTimeouts::ScheduleOnly(timeout)
@@ -937,7 +957,7 @@ mod tests {
 
     #[tokio::test]
     async fn idempotency_enforced() {
-        let lam = LocalActivityManager::new(10, "whatever".to_string());
+        let lam = LocalActivityManager::test(10);
         let new_la = NewLocalAct {
             schedule_cmd: ValidScheduleLA {
                 seq: 1,

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

@@ -8,10 +8,14 @@ pub(crate) use local_activities::{
 };
 
 use crate::{
+    abstractions::MeteredSemaphore,
     pollers::BoxedActPoller,
-    telemetry::metrics::{activity_type, workflow_type, MetricsContext},
-    worker::activities::activity_heartbeat_manager::ActivityHeartbeatError,
-    CompleteActivityError, PollActivityError, ServerGatewayApis, TaskToken,
+    telemetry::metrics::{activity_type, activity_worker_type, workflow_type, MetricsContext},
+    worker::{
+        activities::activity_heartbeat_manager::ActivityHeartbeatError,
+        client::{WorkerClient, WorkerClientBag},
+    },
+    CompleteActivityError, PollActivityError, TaskToken,
 };
 use activity_heartbeat_manager::ActivityHeartbeatManager;
 use dashmap::DashMap;
@@ -31,7 +35,7 @@ use temporal_sdk_core_protos::{
         workflowservice::v1::PollActivityTaskQueueResponse,
     },
 };
-use tokio::sync::{Notify, Semaphore};
+use tokio::sync::Notify;
 
 #[derive(Debug, derive_more::Constructor)]
 struct PendingActivityCancel {
@@ -88,7 +92,7 @@ pub(crate) struct WorkerActivityTasks {
     /// ongoing.
     poller: BoxedActPoller,
     /// Ensures we stay at or below this worker's maximum concurrent activity limit
-    activities_semaphore: Semaphore,
+    activities_semaphore: MeteredSemaphore,
     /// Wakes every time an activity is removed from the outstanding map
     complete_notify: Notify,
 
@@ -102,16 +106,20 @@ impl WorkerActivityTasks {
     pub(crate) fn new(
         max_activity_tasks: usize,
         poller: BoxedActPoller,
-        sg: Arc<impl ServerGatewayApis + Send + Sync + 'static + ?Sized>,
+        client: Arc<WorkerClientBag>,
         metrics: MetricsContext,
         max_heartbeat_throttle_interval: Duration,
         default_heartbeat_throttle_interval: Duration,
     ) -> Self {
         Self {
-            heartbeat_manager: ActivityHeartbeatManager::new(sg),
+            heartbeat_manager: ActivityHeartbeatManager::new(client),
             outstanding_activity_tasks: Default::default(),
             poller,
-            activities_semaphore: Semaphore::new(max_activity_tasks),
+            activities_semaphore: MeteredSemaphore::new(
+                max_activity_tasks,
+                metrics.with_new_attrs([activity_worker_type()]),
+                MetricsContext::available_task_slots,
+            ),
             complete_notify: Notify::new(),
             metrics,
             max_heartbeat_throttle_interval,
@@ -196,7 +204,7 @@ impl WorkerActivityTasks {
         &self,
         task_token: TaskToken,
         status: aer::Status,
-        gateway: &(dyn ServerGatewayApis + Send + Sync),
+        client: &dyn WorkerClient,
     ) -> Result<(), CompleteActivityError> {
         if let Some((_, act_info)) = self.outstanding_activity_tasks.remove(&task_token) {
             let act_metrics = self.metrics.with_new_attrs([
@@ -204,7 +212,7 @@ 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_permits(1);
+            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.
@@ -214,13 +222,13 @@ impl WorkerActivityTasks {
             if !known_not_found {
                 let maybe_net_err = match status {
                     aer::Status::WillCompleteAsync(_) => None,
-                    aer::Status::Completed(ar::Success { result }) => gateway
+                    aer::Status::Completed(ar::Success { result }) => client
                         .complete_activity_task(task_token.clone(), result.map(Into::into))
                         .await
                         .err(),
                     aer::Status::Failed(ar::Failure { failure }) => {
                         act_metrics.act_execution_failed();
-                        gateway
+                        client
                             .fail_activity_task(task_token.clone(), failure.map(Into::into))
                             .await
                             .err()
@@ -238,7 +246,7 @@ impl WorkerActivityTasks {
                                 "Expected activity cancelled status with CanceledFailureInfo");
                             None
                         };
-                        gateway
+                        client
                             .cancel_activity_task(task_token.clone(), details.map(Into::into))
                             .await
                             .err()
@@ -316,8 +324,7 @@ impl WorkerActivityTasks {
                 }
                 Ok(Some(ActivityTask::cancel_from_ids(task_token.0, reason)))
             } else {
-                warn!(task_token = ?task_token,
-                              "Unknown activity task when issuing cancel");
+                debug!(task_token = ?task_token, "Unknown activity task when issuing cancel");
                 // If we can't find the activity here, it's already been completed,
                 // in which case issuing a cancel again is pointless.
                 Ok(None)
@@ -331,6 +338,6 @@ impl WorkerActivityTasks {
 
     #[cfg(test)]
     pub(crate) fn remaining_activity_capacity(&self) -> usize {
-        self.activities_semaphore.available_permits()
+        self.activities_semaphore.sem.available_permits()
     }
 }

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

@@ -0,0 +1,86 @@
+use super::*;
+use futures::Future;
+
+#[cfg(test)]
+/// Create a mock client primed with basic necessary expectations
+pub(crate) fn mock_workflow_client() -> MockWorkerClient {
+    MockWorkerClient::new()
+}
+
+/// Create a mock manual client primed with basic necessary expectations
+pub(crate) fn mock_manual_workflow_client() -> MockManualWorkerClient {
+    MockManualWorkerClient::new()
+}
+
+// Need a version of the mock that can return futures so we can return potentially pending
+// results. This is really annoying b/c of the async trait stuff. Need
+// https://github.com/asomers/mockall/issues/189 to be fixed for it to go away.
+mockall::mock! {
+    pub ManualWorkerClient {}
+    impl WorkerClient for ManualWorkerClient {
+        fn poll_workflow_task<'a, 'b>(&'a self, task_queue: String, is_sticky: bool)
+            -> impl Future<Output = Result<PollWorkflowTaskQueueResponse>> + Send + 'b
+            where 'a: 'b, Self: 'b;
+
+        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;
+
+        fn complete_workflow_task<'a, 'b>(
+            &self,
+            request: WorkflowTaskCompletion,
+        ) -> impl Future<Output = Result<RespondWorkflowTaskCompletedResponse>> + Send + 'b
+            where 'a: 'b, Self: 'b;
+
+        fn complete_activity_task<'a, 'b>(
+            &self,
+            task_token: TaskToken,
+            result: Option<Payloads>,
+        ) -> impl Future<Output = Result<RespondActivityTaskCompletedResponse>> + Send + 'b
+            where 'a: 'b, Self: 'b;
+
+        fn cancel_activity_task<'a, 'b>(
+            &self,
+            task_token: TaskToken,
+            details: Option<Payloads>,
+        ) -> impl Future<Output = Result<RespondActivityTaskCanceledResponse>> + Send + 'b
+            where 'a: 'b, Self: 'b;
+
+        fn fail_activity_task<'a, 'b>(
+            &self,
+            task_token: TaskToken,
+            failure: Option<Failure>,
+        ) -> impl Future<Output = Result<RespondActivityTaskFailedResponse>> + Send + 'b
+            where 'a: 'b, Self: 'b;
+
+        fn fail_workflow_task<'a, 'b>(
+            &self,
+            task_token: TaskToken,
+            cause: WorkflowTaskFailedCause,
+            failure: Option<Failure>,
+        ) -> impl Future<Output = Result<RespondWorkflowTaskFailedResponse>> + Send + 'b
+            where 'a: 'b, Self: 'b;
+
+        fn record_activity_heartbeat<'a, 'b>(
+           &self,
+           task_token: TaskToken,
+           details: Option<Payloads>,
+        ) -> impl Future<Output = Result<RecordActivityTaskHeartbeatResponse>> + Send + 'b
+            where 'a: 'b, Self: 'b;
+
+        fn get_workflow_execution_history<'a, 'b>(
+            &self,
+            workflow_id: String,
+            run_id: Option<String>,
+            page_token: Vec<u8>
+        ) -> impl Future<Output = Result<GetWorkflowExecutionHistoryResponse>> + Send + 'b
+            where 'a: 'b, Self: 'b;
+
+        fn respond_legacy_query<'a, 'b>(
+            &self,
+            task_token: TaskToken,
+            query_result: QueryResult,
+        ) -> impl Future<Output = Result<RespondQueryTaskCompletedResponse>> + Send + 'b
+            where 'a: 'b, Self: 'b;
+    }
+}