@temporalio/core-bridge 1.5.2 → 1.7.0

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

@@ -1,4 +1,5 @@
 use crate::{
+    abstractions::take_cell::TakeCell,
     worker::{activities::PendingActivityCancel, client::WorkerClient},
     TaskToken,
 };
@@ -17,7 +18,7 @@ use temporal_sdk_core_protos::{
 use tokio::{
     sync::{
         mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender},
-        Mutex, Notify,
+        Notify,
     },
     task::JoinHandle,
 };
@@ -26,12 +27,9 @@ use tokio_util::sync::CancellationToken;
 /// Used to supply new heartbeat events to the activity heartbeat manager, or to send a shutdown
 /// request.
 pub(crate) struct ActivityHeartbeatManager {
-    /// Cancellations that have been received when heartbeating are queued here and can be consumed
-    /// by [fetch_cancellations]
-    incoming_cancels: Mutex<UnboundedReceiver<PendingActivityCancel>>,
     shutdown_token: CancellationToken,
     /// Used during `shutdown` to await until all inflight requests are sent.
-    join_handle: Mutex<Option<JoinHandle<()>>>,
+    join_handle: TakeCell<JoinHandle<()>>,
     heartbeat_tx: UnboundedSender<HeartbeatAction>,
 }
 
@@ -74,15 +72,115 @@ pub enum ActivityHeartbeatError {
     /// to heartbeat.
     #[error("Unable to parse activity heartbeat timeout.")]
     InvalidHeartbeatTimeout,
-    /// Core is shutting down and thus new heartbeats are not accepted
-    #[error("New heartbeat requests are not accepted while shutting down")]
-    ShuttingDown,
 }
 
 /// Manages activity heartbeating for a worker. Allows sending new heartbeats or requesting and
 /// awaiting for the shutdown. When shutdown is requested, signal gets sent to all processors, which
 /// allows them to complete gracefully.
 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.
+    /// Returns the manager and a channel that buffers cancellation notifications to be sent to Lang.
+    pub(super) fn new(
+        client: Arc<dyn WorkerClient>,
+        cancels_tx: UnboundedSender<PendingActivityCancel>,
+    ) -> Self {
+        let (heartbeat_stream_state, heartbeat_tx_source, shutdown_token) =
+            HeartbeatStreamState::new();
+        let heartbeat_tx = heartbeat_tx_source.clone();
+
+        let join_handle = tokio::spawn(
+            // The stream of incoming heartbeats uses unfold to carry state across each item in the
+            // stream. The closure checks if, for any given activity, we should heartbeat or not
+            // depending on its delay and when we last issued a heartbeat for it.
+            futures::stream::unfold(heartbeat_stream_state, move |mut hb_states| {
+                async move {
+                    let hb = tokio::select! {
+                        biased;
+
+                        _ = hb_states.cancellation_token.cancelled() => {
+                            return None
+                        }
+                        hb = hb_states.incoming_hbs.recv() => match hb {
+                            None => return None,
+                            Some(hb) => hb,
+                        }
+                    };
+
+                    Some((
+                        match hb {
+                            HeartbeatAction::SendHeartbeat(hb) => hb_states.record(hb),
+                            HeartbeatAction::CompleteReport(tt) => hb_states.handle_report_completed(tt),
+                            HeartbeatAction::CompleteThrottle(tt) => hb_states.handle_throttle_completed(tt),
+                            HeartbeatAction::Evict{ token, on_complete } => hb_states.evict(token, on_complete),
+                        },
+                        hb_states,
+                    ))
+                }
+            })
+                // Filters out `None`s
+                .filter_map(|opt| async { opt })
+                .for_each_concurrent(None, move |action| {
+                    let heartbeat_tx = heartbeat_tx_source.clone();
+                    let sg = client.clone();
+                    let cancels_tx = cancels_tx.clone();
+                    async move {
+                        match action {
+                            HeartbeatExecutorAction::Sleep(tt, duration, cancellation_token) => {
+                                tokio::select! {
+                                _ = cancellation_token.cancelled() => (),
+                                _ = tokio::time::sleep(duration) => {
+                                    let _ = heartbeat_tx.send(HeartbeatAction::CompleteThrottle(tt));
+                                },
+                            };
+                            }
+                            HeartbeatExecutorAction::Report { task_token: tt, details } => {
+                                match sg
+                                    .record_activity_heartbeat(tt.clone(), details.into_payloads())
+                                    .await
+                                {
+                                    Ok(RecordActivityTaskHeartbeatResponse { cancel_requested }) => {
+                                        if cancel_requested {
+                                            cancels_tx
+                                                .send(PendingActivityCancel::new(
+                                                    tt.clone(),
+                                                    ActivityCancelReason::Cancelled,
+                                                ))
+                                                .expect(
+                                                    "Receive half of heartbeat cancels not blocked",
+                                                );
+                                        }
+                                    }
+                                    // Send cancels for any activity that learns its workflow already
+                                    // 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(),
+                                                ActivityCancelReason::NotFound,
+                                            ))
+                                            .expect("Receive half of heartbeat cancels not blocked");
+                                    }
+                                    Err(e) => {
+                                        warn!("Error when recording heartbeat: {:?}", e);
+                                    }
+                                };
+                                let _ = heartbeat_tx.send(HeartbeatAction::CompleteReport(tt));
+                            }
+                        }
+                    }
+                }),
+        );
+
+        Self {
+            join_handle: TakeCell::new(join_handle),
+            shutdown_token,
+            heartbeat_tx,
+        }
+    }
     /// Records a new heartbeat, the first call will result in an immediate call to the server,
     /// while rapid successive calls would accumulate for up to `delay` and then latest heartbeat
     /// details will be sent to the server.
@@ -95,9 +193,6 @@ impl ActivityHeartbeatManager {
         hb: ActivityHeartbeat,
         throttle_interval: Duration,
     ) -> Result<(), ActivityHeartbeatError> {
-        if self.shutdown_token.is_cancelled() {
-            return Err(ActivityHeartbeatError::ShuttingDown);
-        }
         self.heartbeat_tx
             .send(HeartbeatAction::SendHeartbeat(ValidActivityHeartbeat {
                 task_token: TaskToken(hb.task_token),
@@ -121,19 +216,11 @@ impl ActivityHeartbeatManager {
         completed.notified().await;
     }
 
-    /// Returns a future that resolves any time there is a new activity cancel that must be
-    /// dispatched to lang
-    pub(super) async fn next_pending_cancel(&self) -> Option<PendingActivityCancel> {
-        self.incoming_cancels.lock().await.recv().await
-    }
-
-    // TODO: Can own self now!
     /// Initiates shutdown procedure by stopping lifecycle loop and awaiting for all in-flight
     /// heartbeat requests to be flushed to the server.
     pub(super) async fn shutdown(&self) {
         self.shutdown_token.cancel();
-        let mut handle = self.join_handle.lock().await;
-        if let Some(h) = handle.take() {
+        if let Some(h) = self.join_handle.take_once() {
             let handle_r = h.await;
             if let Err(e) = handle_r {
                 if !e.is_cancelled() {
@@ -301,110 +388,6 @@ 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(client: Arc<dyn WorkerClient>) -> Self {
-        let (heartbeat_stream_state, heartbeat_tx_source, shutdown_token) =
-            HeartbeatStreamState::new();
-        let (cancels_tx, cancels_rx) = unbounded_channel();
-        let heartbeat_tx = heartbeat_tx_source.clone();
-
-        let join_handle = tokio::spawn(
-            // The stream of incoming heartbeats uses unfold to carry state across each item in the
-            // stream. The closure checks if, for any given activity, we should heartbeat or not
-            // depending on its delay and when we last issued a heartbeat for it.
-            futures::stream::unfold(heartbeat_stream_state, move |mut hb_states| {
-                async move {
-                    let hb = tokio::select! {
-                        biased;
-
-                        _ = hb_states.cancellation_token.cancelled() => {
-                            return None
-                        }
-                        hb = hb_states.incoming_hbs.recv() => match hb {
-                            None => return None,
-                            Some(hb) => hb,
-                        }
-                    };
-
-                    Some((
-                        match hb {
-                            HeartbeatAction::SendHeartbeat(hb) => hb_states.record(hb),
-                            HeartbeatAction::CompleteReport(tt) => hb_states.handle_report_completed(tt),
-                            HeartbeatAction::CompleteThrottle(tt) => hb_states.handle_throttle_completed(tt),
-                            HeartbeatAction::Evict{ token, on_complete } => hb_states.evict(token, on_complete),
-                        },
-                        hb_states,
-                    ))
-                }
-            })
-            // Filters out `None`s
-            .filter_map(|opt| async { opt })
-            .for_each_concurrent(None, move |action| {
-                let heartbeat_tx = heartbeat_tx_source.clone();
-                let sg = client.clone();
-                let cancels_tx = cancels_tx.clone();
-                async move {
-                    match action {
-                        HeartbeatExecutorAction::Sleep(tt, duration, cancellation_token) => {
-                            tokio::select! {
-                                _ = cancellation_token.cancelled() => (),
-                                _ = tokio::time::sleep(duration) => {
-                                    let _ = heartbeat_tx.send(HeartbeatAction::CompleteThrottle(tt));
-                                },
-                            };
-                        }
-                        HeartbeatExecutorAction::Report { task_token: tt, details } => {
-                            match sg
-                                .record_activity_heartbeat(tt.clone(), details.into_payloads())
-                                .await
-                            {
-                                Ok(RecordActivityTaskHeartbeatResponse { cancel_requested }) => {
-                                    if cancel_requested {
-                                        cancels_tx
-                                            .send(PendingActivityCancel::new(
-                                                tt.clone(),
-                                                ActivityCancelReason::Cancelled,
-                                            ))
-                                            .expect(
-                                                "Receive half of heartbeat cancels not blocked",
-                                            );
-                                    }
-                                }
-                                // Send cancels for any activity that learns its workflow already
-                                // 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(),
-                                            ActivityCancelReason::NotFound,
-                                        ))
-                                        .expect("Receive half of heartbeat cancels not blocked");
-                                }
-                                Err(e) => {
-                                    warn!("Error when recording heartbeat: {:?}", e);
-                                }
-                            };
-                            let _ = heartbeat_tx.send(HeartbeatAction::CompleteReport(tt));
-                        }
-                    }
-                }
-            }),
-        );
-
-        Self {
-            incoming_cancels: Mutex::new(cancels_rx),
-            join_handle: Mutex::new(Some(join_handle)),
-            shutdown_token,
-            heartbeat_tx,
-        }
-    }
-}
-
 #[cfg(test)]
 mod test {
     use super::*;
@@ -425,7 +408,8 @@ mod test {
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
             .times(2);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client));
+        let (cancel_tx, _cancel_rx) = unbounded_channel();
+        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client), cancel_tx);
         let fake_task_token = vec![1, 2, 3];
         // Send 2 heartbeat requests for 20ms apart.
         // The first heartbeat should be sent right away, and
@@ -446,14 +430,14 @@ mod test {
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
             .times(3);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client));
+        let (cancel_tx, _cancel_rx) = unbounded_channel();
+        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client), cancel_tx);
         let fake_task_token = vec![1, 2, 3];
         // Heartbeats always get sent if recorded less frequently than the throttle interval
         for i in 0_u8..3 {
             record_heartbeat(&hm, fake_task_token.clone(), i, Duration::from_millis(10));
             sleep(Duration::from_millis(20)).await;
         }
-        // sleep again to let heartbeats be flushed
         hm.shutdown().await;
     }
 
@@ -466,7 +450,8 @@ mod test {
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
             .times(1);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client));
+        let (cancel_tx, _cancel_rx) = unbounded_channel();
+        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client), cancel_tx);
         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 {
@@ -485,7 +470,8 @@ mod test {
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
             .times(2);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client));
+        let (cancel_tx, _cancel_rx) = unbounded_channel();
+        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client), cancel_tx);
         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;
@@ -502,7 +488,8 @@ mod test {
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
             .times(2);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client));
+        let (cancel_tx, _cancel_rx) = unbounded_channel();
+        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client), cancel_tx);
         let fake_task_token = vec![1, 2, 3];
         record_heartbeat(&hm, fake_task_token.clone(), 0, Duration::from_millis(100));
         // Let it propagate
@@ -522,42 +509,14 @@ mod test {
             .expect_record_activity_heartbeat()
             .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
             .times(1);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client));
+        let (cancel_tx, _cancel_rx) = unbounded_channel();
+        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client), cancel_tx);
         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_client = mock_workflow_client();
-        mock_client
-            .expect_record_activity_heartbeat()
-            .returning(|_, _| Ok(RecordActivityTaskHeartbeatResponse::default()))
-            .times(0);
-        let hm = ActivityHeartbeatManager::new(Arc::new(mock_client));
-        hm.shutdown().await;
-        match hm.record(
-            ActivityHeartbeat {
-                task_token: vec![1, 2, 3],
-                details: vec![Payload {
-                    // payload doesn't matter in this case, as it shouldn't get sent anyways.
-                    ..Default::default()
-                }],
-            },
-            Duration::from_millis(1000),
-        ) {
-            Ok(_) => {
-                unreachable!("heartbeat should not be recorded after the shutdown");
-            }
-            Err(e) => {
-                matches!(e, ActivityHeartbeatError::ShuttingDown);
-            }
-        }
-    }
-
     fn record_heartbeat(
         hm: &ActivityHeartbeatManager,
         task_token: Vec<u8>,

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

@@ -0,0 +1,89 @@
+use crate::abstractions::MeteredSemaphore;
+use crate::worker::activities::PermittedTqResp;
+use crate::{pollers::BoxedActPoller, MetricsContext};
+use futures::{stream, Stream};
+use governor::clock::DefaultClock;
+use governor::middleware::NoOpMiddleware;
+use governor::state::{InMemoryState, NotKeyed};
+use governor::RateLimiter;
+use std::sync::Arc;
+use temporal_sdk_core_protos::temporal::api::workflowservice::v1::PollActivityTaskQueueResponse;
+use tokio::select;
+use tokio_util::sync::CancellationToken;
+
+struct StreamState {
+    poller: BoxedActPoller,
+    semaphore: Arc<MeteredSemaphore>,
+    rate_limiter: Option<RateLimiter<NotKeyed, InMemoryState, DefaultClock, NoOpMiddleware>>,
+    metrics: MetricsContext,
+    shutdown_token: CancellationToken,
+    poller_was_shutdown: bool,
+}
+
+pub(crate) fn new_activity_task_poller(
+    poller: BoxedActPoller,
+    semaphore: Arc<MeteredSemaphore>,
+    rate_limiter: Option<RateLimiter<NotKeyed, InMemoryState, DefaultClock, NoOpMiddleware>>,
+    metrics: MetricsContext,
+    shutdown_token: CancellationToken,
+) -> impl Stream<Item = Result<PermittedTqResp, tonic::Status>> {
+    let state = StreamState {
+        poller,
+        semaphore,
+        rate_limiter,
+        metrics,
+        shutdown_token,
+        poller_was_shutdown: false,
+    };
+    stream::unfold(state, |mut state| async move {
+        loop {
+            let poll = async {
+                let permit = state
+                    .semaphore
+                    .acquire_owned()
+                    .await
+                    .expect("outstanding activity semaphore not closed");
+                if !state.poller_was_shutdown {
+                    if let Some(ref rl) = state.rate_limiter {
+                        rl.until_ready().await;
+                    }
+                }
+                loop {
+                    return match state.poller.poll().await {
+                        Some(Ok(resp)) => {
+                            if resp == PollActivityTaskQueueResponse::default() {
+                                // We get the default proto in the event that the long poll times out.
+                                debug!("Poll activity task timeout");
+                                state.metrics.act_poll_timeout();
+                                continue;
+                            }
+                            Some(Ok(PermittedTqResp { permit, resp }))
+                        }
+                        Some(Err(e)) => {
+                            warn!(error=?e, "Error while polling for activity tasks");
+                            Some(Err(e))
+                        }
+                        // If poller returns None, it's dead, thus we also return None to terminate this
+                        // stream.
+                        None => None,
+                    };
+                }
+            };
+            if state.poller_was_shutdown {
+                return poll.await.map(|res| (res, state));
+            }
+            select! {
+                biased;
+
+                _ = state.shutdown_token.cancelled() => {
+                    state.poller.notify_shutdown();
+                    state.poller_was_shutdown = true;
+                    continue;
+                }
+                res = poll => {
+                    return res.map(|res| (res, state));
+                }
+            }
+        }
+    })
+}