@temporalio/core-bridge 1.6.0 → 1.7.1

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));
+                }
+            }
+        }
+    })
+}

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

@@ -1,5 +1,5 @@
 use crate::{
-    abstractions::{dbg_panic, MeteredSemaphore, OwnedMeteredSemPermit},
+    abstractions::{dbg_panic, MeteredSemaphore, OwnedMeteredSemPermit, UsedMeteredSemPermit},
     protosext::ValidScheduleLA,
     retry_logic::RetryPolicyExt,
     worker::workflow::HeartbeatTimeoutMsg,
@@ -20,7 +20,11 @@ use temporal_sdk_core_protos::{
         activity_result::{Cancellation, Failure as ActFail, Success},
         activity_task::{activity_task, ActivityCancelReason, ActivityTask, Cancel, Start},
     },
-    temporal::api::{common::v1::WorkflowExecution, enums::v1::TimeoutType},
+    temporal::api::{
+        common::v1::WorkflowExecution,
+        enums::v1::TimeoutType,
+        failure::v1::{failure, Failure as APIFailure, TimeoutFailureInfo},
+    },
 };
 use tokio::{
     sync::{
@@ -51,7 +55,7 @@ pub(crate) struct LocalInFlightActInfo {
     pub la_info: NewLocalAct,
     pub dispatch_time: Instant,
     pub attempt: u32,
-    _permit: OwnedMeteredSemPermit,
+    _permit: UsedMeteredSemPermit,
 }
 
 #[derive(Debug, Clone)]
@@ -70,7 +74,18 @@ impl LocalActivityExecutionResult {
         Self::Cancelled(Cancellation::from_details(None))
     }
     pub(crate) fn timeout(tt: TimeoutType) -> Self {
-        Self::TimedOut(ActFail::timeout(tt))
+        Self::TimedOut(ActFail {
+            failure: Some(APIFailure {
+                message: "Activity timed out".to_string(),
+                failure_info: Some(failure::FailureInfo::TimeoutFailureInfo(
+                    TimeoutFailureInfo {
+                        timeout_type: tt as i32,
+                        last_heartbeat_details: None,
+                    },
+                )),
+                ..Default::default()
+            }),
+        })
     }
 }
 
@@ -110,12 +125,17 @@ impl Debug for NewLocalAct {
 pub(crate) enum LocalActRequest {
     New(NewLocalAct),
     Cancel(ExecutingLAId),
+    #[from(ignore)]
     CancelAllInRun(String),
     StartHeartbeatTimeout {
         send_on_elapse: HeartbeatTimeoutMsg,
         deadline: Instant,
         abort_reg: AbortRegistration,
     },
+    /// Tell the LA manager that a workflow task was responded to (completed or failed) for a
+    /// certain run id
+    #[from(ignore)]
+    IndicateWorkflowTaskCompleted(String),
 }
 
 #[derive(Debug, Clone, Eq, PartialEq, Hash)]
@@ -153,6 +173,10 @@ struct LocalActivityInfo {
     /// Tasks / info about timeouts associated with this LA. May be empty for very brief periods
     /// while the LA id has been generated, but it has not yet been scheduled.
     timeout_bag: Option<TimeoutBag>,
+    /// True once the first workflow task this LA started in has elapsed
+    first_wft_has_ended: bool,
+    /// Attempts at executing this LA during the current WFT
+    attempts_in_wft: usize,
 }
 
 struct LAMData {
@@ -270,6 +294,8 @@ impl LocalActivityManager {
                                 task_token: tt,
                                 backing_off_task: None,
                                 timeout_bag: None,
+                                first_wft_has_ended: false,
+                                attempts_in_wft: 0,
                             });
 
                             // Set up timeouts for the new activity
@@ -324,6 +350,17 @@ impl LocalActivityManager {
                         }
                     }
                 }
+                LocalActRequest::IndicateWorkflowTaskCompleted(run_id) => {
+                    let mut dlock = self.dat.lock();
+                    let las_for_run = dlock
+                        .la_info
+                        .iter_mut()
+                        .filter(|(id, _)| id.run_id == run_id);
+                    for (_, lainf) in las_for_run {
+                        lainf.first_wft_has_ended = true;
+                        lainf.attempts_in_wft = 0;
+                    }
+                }
             }
         }
         immediate_resolutions
@@ -432,7 +469,7 @@ impl LocalActivityManager {
                 la_info: la_info_for_in_flight_map,
                 dispatch_time: Instant::now(),
                 attempt,
-                _permit: permit,
+                _permit: permit.into_used(),
             },
         );
 
@@ -538,6 +575,14 @@ impl LocalActivityManager {
                             LocalActivityInfo {
                                 task_token: tt,
                                 backing_off_task: Some(jh),
+                                first_wft_has_ended: maybe_old_lai
+                                    .as_ref()
+                                    .map(|old| old.first_wft_has_ended)
+                                    .unwrap_or_default(),
+                                attempts_in_wft: maybe_old_lai
+                                    .as_ref()
+                                    .map(|old| old.attempts_in_wft + 1)
+                                    .unwrap_or(1),
                                 timeout_bag: maybe_old_lai.and_then(|old| old.timeout_bag),
                             },
                         );
@@ -564,9 +609,26 @@ impl LocalActivityManager {
         }
     }
 
+    /// Try to close the activity stream as soon as worker shutdown is initiated. This is required
+    /// for activity-only workers where since workflows are not polled and the activity poller might
+    /// get "stuck".
+    pub(crate) fn shutdown_initiated(&self) {
+        self.set_shutdown_complete_if_ready(&mut self.dat.lock());
+    }
+
+    pub(crate) fn get_nonfirst_attempt_count(&self, for_run_id: &str) -> usize {
+        let dlock = self.dat.lock();
+        dlock
+            .la_info
+            .iter()
+            .filter(|(id, info)| id.run_id == for_run_id && info.first_wft_has_ended)
+            .map(|(_, info)| info.attempts_in_wft)
+            .sum()
+    }
+
     fn set_shutdown_complete_if_ready(&self, dlock: &mut MutexGuard<LAMData>) -> bool {
         let nothing_outstanding = dlock.outstanding_activity_tasks.is_empty();
-        if nothing_outstanding {
+        if nothing_outstanding && self.workflows_have_shut_down.is_cancelled() {
             self.shutdown_complete_tok.cancel();
         }
         nothing_outstanding
@@ -1174,4 +1236,58 @@ mod tests {
         assert_eq!(lam.num_outstanding(), 1);
         assert!(lam.rcvs.lock().await.next().now_or_never().is_none());
     }
+
+    #[tokio::test]
+    async fn nonfirst_la_attempt_count_is_accurate() {
+        let run_id = "run_id";
+        let lam = LocalActivityManager::test(10);
+        let new_la = NewLocalAct {
+            schedule_cmd: ValidScheduleLA {
+                seq: 1,
+                activity_id: 1.to_string(),
+                retry_policy: RetryPolicy {
+                    initial_interval: Some(prost_dur!(from_millis(1))),
+                    backoff_coefficient: 1.0,
+                    ..Default::default()
+                },
+                local_retry_threshold: Duration::from_secs(500),
+                ..Default::default()
+            },
+            workflow_type: "".to_string(),
+            workflow_exec_info: WorkflowExecution {
+                workflow_id: "".to_string(),
+                run_id: run_id.to_string(),
+            },
+            schedule_time: SystemTime::now(),
+        };
+        lam.enqueue([new_la.clone().into()]);
+        let spinfail = || async {
+            for _ in 1..=10 {
+                let next = lam.next_pending().await.unwrap().unwrap();
+                let tt = TaskToken(next.task_token);
+                lam.complete(
+                    &tt,
+                    &LocalActivityExecutionResult::Failed(Default::default()),
+                );
+            }
+        };
+
+        // Fail a bunch of times
+        spinfail().await;
+        // Nonfirst attempt count should still be zero
+        let count = lam.get_nonfirst_attempt_count(run_id);
+        assert_eq!(count, 0);
+
+        for _ in 1..=2 {
+            // This should work over multiple WFTs
+            // say the first wft was completed
+            lam.enqueue([LocalActRequest::IndicateWorkflowTaskCompleted(
+                run_id.to_string(),
+            )]);
+            // Do some more attempts
+            spinfail().await;
+            let count = lam.get_nonfirst_attempt_count(run_id);
+            assert_eq!(count, 10);
+        }
+    }
 }