temporalio 0.0.1 → 0.1.0

data/bridge/sdk-core/core/src/worker/activities/local_activities.rs

@@ -1,13 +1,18 @@
 use crate::{
-    abstractions::{MeteredSemaphore, OwnedMeteredSemPermit},
+    abstractions::{dbg_panic, MeteredSemaphore, OwnedMeteredSemPermit, UsedMeteredSemPermit},
     protosext::ValidScheduleLA,
     retry_logic::RetryPolicyExt,
+    worker::workflow::HeartbeatTimeoutMsg,
     MetricsContext, TaskToken,
 };
-use parking_lot::Mutex;
+use futures::{stream::BoxStream, Stream};
+use futures_util::{future, future::AbortRegistration, stream, StreamExt};
+use parking_lot::{Mutex, MutexGuard};
 use std::{
-    collections::HashMap,
+    collections::{hash_map::Entry, HashMap},
     fmt::{Debug, Formatter},
+    pin::Pin,
+    task::{Context, Poll},
     time::{Duration, Instant, SystemTime},
 };
 use temporal_sdk_core_protos::{
@@ -25,6 +30,7 @@ use tokio::{
     task::JoinHandle,
     time::sleep,
 };
+use tokio_stream::wrappers::UnboundedReceiverStream;
 use tokio_util::sync::CancellationToken;
 
 #[allow(clippy::large_enum_variant)] // Timeouts are relatively rare
@@ -45,10 +51,14 @@ pub(crate) struct LocalInFlightActInfo {
     pub la_info: NewLocalAct,
     pub dispatch_time: Instant,
     pub attempt: u32,
-    _permit: OwnedMeteredSemPermit,
+    _permit: UsedMeteredSemPermit,
 }
 
 #[derive(Debug, Clone)]
+#[cfg_attr(
+    feature = "save_wf_inputs",
+    derive(serde::Serialize, serde::Deserialize)
+)]
 pub(crate) enum LocalActivityExecutionResult {
     Completed(Success),
     Failed(ActFail),
@@ -65,6 +75,10 @@ impl LocalActivityExecutionResult {
 }
 
 #[derive(Debug, Clone)]
+#[cfg_attr(
+    feature = "save_wf_inputs",
+    derive(serde::Serialize, serde::Deserialize)
+)]
 pub(crate) struct LocalActivityResolution {
     pub seq: u32,
     pub result: LocalActivityExecutionResult,
@@ -96,6 +110,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)]
@@ -107,28 +132,43 @@ pub(crate) struct ExecutingLAId {
 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: 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
     cancels_req_tx: UnboundedSender<CancelOrTimeout>,
+    /// For the emission of heartbeat timeouts, back into the workflow machines. This channel
+    /// needs to come in from above us, because we cannot rely on callers getting the next
+    /// activation as a way to deliver heartbeats.
+    heartbeat_timeout_tx: UnboundedSender<HeartbeatTimeoutMsg>,
     /// Wakes every time a complete is processed
     complete_notify: Notify,
+    /// Set once workflows have finished shutting down, and thus we know we will no longer receive
+    /// any requests to spawn new LAs
+    workflows_have_shut_down: CancellationToken,
 
     rcvs: tokio::sync::Mutex<RcvChans>,
     shutdown_complete_tok: CancellationToken,
     dat: Mutex<LAMData>,
 }
 
+struct LocalActivityInfo {
+    task_token: TaskToken,
+    /// Tasks for the current backoff until the next retry, if any.
+    backing_off_task: Option<JoinHandle<()>>,
+    /// 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 {
+    /// Maps local activity identifiers to information about them
+    la_info: HashMap<ExecutingLAId, LocalActivityInfo>,
     /// Activities that have been issued to lang but not yet completed
     outstanding_activity_tasks: HashMap<TaskToken, LocalInFlightActInfo>,
-    id_to_tt: HashMap<ExecutingLAId, TaskToken>,
-    /// Tasks for activities which are currently backing off. May be used to cancel retrying them.
-    backing_off_tasks: HashMap<ExecutingLAId, JoinHandle<()>>,
-    /// Tasks for timing out activities which are currently in the queue or dispatched.
-    timeout_tasks: HashMap<ExecutingLAId, TimeoutBag>,
     next_tt_num: u32,
 }
 
@@ -143,43 +183,47 @@ impl LocalActivityManager {
     pub(crate) fn new(
         max_concurrent: usize,
         namespace: String,
+        heartbeat_timeout_tx: UnboundedSender<HeartbeatTimeoutMsg>,
         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();
+        let semaphore = MeteredSemaphore::new(
+            max_concurrent,
+            metrics_context,
+            MetricsContext::available_task_slots,
+        );
         Self {
             namespace,
-            semaphore: MeteredSemaphore::new(
-                max_concurrent,
-                metrics_context,
-                MetricsContext::available_task_slots,
-            ),
+            rcvs: tokio::sync::Mutex::new(RcvChans::new(
+                act_req_rx,
+                semaphore,
+                cancels_req_rx,
+                shutdown_complete_tok.clone(),
+            )),
             act_req_tx,
             cancels_req_tx,
+            heartbeat_timeout_tx,
             complete_notify: Notify::new(),
-            rcvs: tokio::sync::Mutex::new(RcvChans {
-                act_req_rx,
-                cancels_req_rx,
-                shutdown: shutdown_complete_tok.clone(),
-            }),
             shutdown_complete_tok,
             dat: Mutex::new(LAMData {
                 outstanding_activity_tasks: Default::default(),
-                id_to_tt: Default::default(),
-                backing_off_tasks: Default::default(),
-                timeout_tasks: Default::default(),
+                la_info: Default::default(),
                 next_tt_num: 0,
             }),
+            workflows_have_shut_down: Default::default(),
         }
     }
 
     #[cfg(test)]
     fn test(max_concurrent: usize) -> Self {
+        let (hb_tx, _hb_rx) = unbounded_channel();
         Self::new(
             max_concurrent,
             "fake_ns".to_string(),
-            MetricsContext::default(),
+            hb_tx,
+            MetricsContext::no_op(),
         )
     }
 
@@ -190,76 +234,116 @@ impl LocalActivityManager {
 
     #[cfg(test)]
     fn num_in_backoff(&self) -> usize {
-        self.dat.lock().backing_off_tasks.len()
+        self.dat
+            .lock()
+            .la_info
+            .values()
+            .filter(|lai| lai.backing_off_task.is_some())
+            .count()
     }
 
     pub(crate) fn enqueue(
         &self,
         reqs: impl IntoIterator<Item = LocalActRequest>,
     ) -> Vec<LocalActivityResolution> {
+        if self.workflows_have_shut_down.is_cancelled() {
+            dbg_panic!("Tried to enqueue local activity after workflows were shut down");
+            return vec![];
+        }
         let mut immediate_resolutions = vec![];
         for req in reqs {
-            debug!(local_activity = ?req, "Queuing local activity");
             match req {
                 LocalActRequest::New(act) => {
+                    debug!(local_activity=?act, "Queuing local activity");
                     let id = ExecutingLAId {
                         run_id: act.workflow_exec_info.run_id.clone(),
                         seq_num: act.schedule_cmd.seq,
                     };
                     let mut dlock = self.dat.lock();
-                    if dlock.id_to_tt.contains_key(&id) {
-                        // Do not queue local activities which are in fact already executing.
-                        // This can happen during evictions.
-                        debug!("Tried to queue already-executing local activity {:?}", &id);
-                        continue;
-                    }
-                    // Pre-generate and insert the task token now, before we may or may not dispatch
-                    // the activity, so we can enforce idempotency. Prevents two identical LAs
-                    // ending up in the queue at once.
                     let tt = dlock.gen_next_token();
-                    dlock.id_to_tt.insert(id.clone(), tt);
-
-                    // Set up timeouts for the new activity
-                    match TimeoutBag::new(&act, self.cancels_req_tx.clone()) {
-                        Ok(tb) => {
-                            dlock.timeout_tasks.insert(id, tb);
-
-                            self.act_req_tx
-                                .send(NewOrRetry::New(act))
-                                .expect("Receive half of LA request channel cannot be dropped");
+                    match dlock.la_info.entry(id) {
+                        Entry::Occupied(o) => {
+                            // Do not queue local activities which are in fact already executing.
+                            // This can happen during evictions.
+                            debug!(
+                                "Tried to queue already-executing local activity {:?}",
+                                o.key()
+                            );
+                            continue;
+                        }
+                        Entry::Vacant(ve) => {
+                            // Insert the task token now, before we may or may not dispatch the
+                            // activity, so we can enforce idempotency. Prevents two identical LAs
+                            // ending up in the queue at once.
+                            let lai = ve.insert(LocalActivityInfo {
+                                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
+                            match TimeoutBag::new(&act, self.cancels_req_tx.clone()) {
+                                Ok(tb) => {
+                                    lai.timeout_bag = Some(tb);
+
+                                    self.act_req_tx.send(NewOrRetry::New(act)).expect(
+                                        "Receive half of LA request channel cannot be dropped",
+                                    );
+                                }
+                                Err(res) => immediate_resolutions.push(res),
+                            }
                         }
-                        Err(res) => immediate_resolutions.push(res),
                     }
                 }
+                LocalActRequest::StartHeartbeatTimeout {
+                    send_on_elapse,
+                    deadline,
+                    abort_reg,
+                } => {
+                    let chan = self.heartbeat_timeout_tx.clone();
+                    tokio::spawn(future::Abortable::new(
+                        async move {
+                            tokio::time::sleep_until(deadline.into()).await;
+                            let _ = chan.send(send_on_elapse);
+                        },
+                        abort_reg,
+                    ));
+                }
                 LocalActRequest::Cancel(id) => {
+                    debug!(id=?id, "Cancelling local activity");
                     let mut dlock = self.dat.lock();
-
-                    // First check if this ID is currently backing off, if so abort the backoff
-                    // task
-                    if let Some(t) = dlock.backing_off_tasks.remove(&id) {
-                        t.abort();
-                        immediate_resolutions.push(LocalActivityResolution {
-                            seq: id.seq_num,
-                            result: LocalActivityExecutionResult::Cancelled(
-                                Cancellation::from_details(None),
-                            ),
-                            runtime: Duration::from_secs(0),
-                            attempt: 0,
-                            backoff: None,
-                            original_schedule_time: None,
-                        });
-                        continue;
+                    if let Some(lai) = dlock.la_info.get_mut(&id) {
+                        if let Some(immediate_res) = self.cancel_one_la(id.seq_num, lai) {
+                            immediate_resolutions.push(immediate_res);
+                        }
                     }
-
-                    if let Some(tt) = dlock.id_to_tt.get(&id) {
-                        self.cancels_req_tx
-                            .send(CancelOrTimeout::Cancel(ActivityTask {
-                                task_token: tt.0.clone(),
-                                variant: Some(activity_task::Variant::Cancel(Cancel {
-                                    reason: ActivityCancelReason::Cancelled as i32,
-                                })),
-                            }))
-                            .expect("Receive half of LA cancel channel cannot be dropped");
+                }
+                LocalActRequest::CancelAllInRun(run_id) => {
+                    debug!(run_id=%run_id, "Cancelling all local activities for run");
+                    let mut dlock = self.dat.lock();
+                    // Even if we've got 100k+ LAs this should only take a ms or two. Not worth
+                    // adding another map to keep in sync.
+                    let las_for_run = dlock
+                        .la_info
+                        .iter_mut()
+                        .filter(|(id, _)| id.run_id == run_id);
+                    for (laid, lainf) in las_for_run {
+                        if let Some(immediate_res) = self.cancel_one_la(laid.seq_num, lainf) {
+                            immediate_resolutions.push(immediate_res);
+                        }
+                    }
+                }
+                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;
                     }
                 }
             }
@@ -270,7 +354,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, permit) = match self.rcvs.lock().await.next(&self.semaphore).await? {
+        let (new_or_retry, permit) = match self.rcvs.lock().await.next().await? {
             NewOrCancel::Cancel(c) => {
                 return match c {
                     CancelOrTimeout::Cancel(c) => Some(DispatchOrTimeoutLA::Dispatch(c)),
@@ -283,12 +367,13 @@ impl LocalActivityManager {
                             let tt = self
                                 .dat
                                 .lock()
-                                .id_to_tt
+                                .la_info
                                 .get(&ExecutingLAId {
                                     run_id: run_id.clone(),
                                     seq_num: resolution.seq,
                                 })
-                                .map(Clone::clone);
+                                .as_ref()
+                                .map(|lai| lai.task_token.clone());
                             if let Some(task_token) = tt {
                                 self.complete(&task_token, &resolution.result);
                                 Some(ActivityTask {
@@ -323,18 +408,21 @@ impl LocalActivityManager {
             }
             NewOrRetry::Retry { in_flight, attempt } => (in_flight, attempt),
         };
-        let orig = new_la.clone();
+        let la_info_for_in_flight_map = new_la.clone();
         let id = ExecutingLAId {
             run_id: new_la.workflow_exec_info.run_id.clone(),
             seq_num: new_la.schedule_cmd.seq,
         };
+        let orig_sched_time = new_la.schedule_cmd.original_schedule_time;
         let sa = new_la.schedule_cmd;
 
         let mut dat = self.dat.lock();
         // If this request originated from a local backoff task, clear the entry for it. We
         // don't await the handle because we know it must already be done, and there's no
         // meaningful value.
-        dat.backing_off_tasks.remove(&id);
+        dat.la_info
+            .get_mut(&id)
+            .map(|lai| lai.backing_off_task.take());
 
         // If this task sat in the queue for too long, return a timeout for it instead
         if let Some(s2s) = sa.schedule_to_start_timeout.as_ref() {
@@ -348,30 +436,27 @@ impl LocalActivityManager {
                         runtime: sat_for,
                         attempt,
                         backoff: None,
-                        original_schedule_time: Some(new_la.schedule_time),
+                        original_schedule_time: orig_sched_time,
                     },
                     task: None,
                 });
             }
         }
 
-        let tt = dat
-            .id_to_tt
-            .get(&id)
-            .expect("Task token must exist")
-            .clone();
+        let la_info = dat.la_info.get_mut(&id).expect("Activity must exist");
+        let tt = la_info.task_token.clone();
+        if let Some(to) = la_info.timeout_bag.as_mut() {
+            to.mark_started();
+        }
         dat.outstanding_activity_tasks.insert(
             tt.clone(),
             LocalInFlightActInfo {
-                la_info: orig,
+                la_info: la_info_for_in_flight_map,
                 dispatch_time: Instant::now(),
                 attempt,
-                _permit: permit,
+                _permit: permit.into_used(),
             },
         );
-        if let Some(to) = dat.timeout_tasks.get_mut(&id) {
-            to.mark_started();
-        }
 
         let (schedule_to_close, start_to_close) = sa.close_timeouts.into_sched_and_start();
         Some(DispatchOrTimeoutLA::Dispatch(ActivityTask {
@@ -406,11 +491,23 @@ impl LocalActivityManager {
     ) -> LACompleteAction {
         let mut dlock = self.dat.lock();
         if let Some(info) = dlock.outstanding_activity_tasks.remove(task_token) {
+            if self.workflows_have_shut_down.is_cancelled() {
+                // If workflows are already shut down, the results of all this don't matter.
+                // Just say we're done if there's nothing outstanding any more.
+                self.set_shutdown_complete_if_ready(&mut dlock);
+            }
+
             let exec_id = ExecutingLAId {
                 run_id: info.la_info.workflow_exec_info.run_id.clone(),
                 seq_num: info.la_info.schedule_cmd.seq,
             };
-            dlock.id_to_tt.remove(&exec_id);
+            let maybe_old_lai = dlock.la_info.remove(&exec_id);
+            if let Some(ref oldlai) = maybe_old_lai {
+                if let Some(ref bot) = oldlai.backing_off_task {
+                    dbg_panic!("Just-resolved LA should not have backoff task");
+                    bot.abort();
+                }
+            }
 
             match status {
                 LocalActivityExecutionResult::Completed(_)
@@ -446,8 +543,6 @@ impl LocalActivityManager {
                         }
                         // Immediately create a new task token for the to-be-retried LA
                         let tt = dlock.gen_next_token();
-                        dlock.id_to_tt.insert(exec_id.clone(), tt);
-
                         // Send the retry request after waiting the backoff duration
                         let send_chan = self.act_req_tx.clone();
                         let jh = tokio::spawn(async move {
@@ -460,7 +555,22 @@ impl LocalActivityManager {
                                 })
                                 .expect("Receive half of LA request channel cannot be dropped");
                         });
-                        dlock.backing_off_tasks.insert(exec_id, jh);
+                        dlock.la_info.insert(
+                            exec_id,
+                            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),
+                            },
+                        );
 
                         LACompleteAction::WillBeRetried
                     } else {
@@ -473,11 +583,70 @@ impl LocalActivityManager {
         }
     }
 
-    pub(crate) async fn shutdown_and_wait_all_finished(&self) {
-        while !self.dat.lock().outstanding_activity_tasks.is_empty() {
+    pub(crate) fn workflows_have_shutdown(&self) {
+        self.workflows_have_shut_down.cancel();
+        self.set_shutdown_complete_if_ready(&mut self.dat.lock());
+    }
+
+    pub(crate) async fn wait_all_outstanding_tasks_finished(&self) {
+        while !self.set_shutdown_complete_if_ready(&mut self.dat.lock()) {
             self.complete_notify.notified().await;
         }
-        self.shutdown_complete_tok.cancel();
+    }
+
+    /// 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 {
+            self.shutdown_complete_tok.cancel();
+        }
+        nothing_outstanding
+    }
+
+    fn cancel_one_la(
+        &self,
+        seq: u32,
+        lai: &mut LocalActivityInfo,
+    ) -> Option<LocalActivityResolution> {
+        // First check if this ID is currently backing off, if so abort the backoff
+        // task
+        if let Some(t) = lai.backing_off_task.take() {
+            t.abort();
+            return Some(LocalActivityResolution {
+                seq,
+                result: LocalActivityExecutionResult::Cancelled(Cancellation::from_details(None)),
+                runtime: Duration::from_secs(0),
+                attempt: 0,
+                backoff: None,
+                original_schedule_time: None,
+            });
+        }
+
+        self.cancels_req_tx
+            .send(CancelOrTimeout::Cancel(ActivityTask {
+                task_token: lai.task_token.0.clone(),
+                variant: Some(activity_task::Variant::Cancel(Cancel {
+                    reason: ActivityCancelReason::Cancelled as i32,
+                })),
+            }))
+            .expect("Receive half of LA cancel channel cannot be dropped");
+        None
     }
 }
 
@@ -515,37 +684,51 @@ enum CancelOrTimeout {
     },
 }
 
+#[allow(clippy::large_enum_variant)]
 enum NewOrCancel {
     New(NewOrRetry, OwnedMeteredSemPermit),
     Cancel(CancelOrTimeout),
 }
 
+#[pin_project::pin_project]
 struct RcvChans {
-    /// Activities that need to be executed by lang
-    act_req_rx: UnboundedReceiver<NewOrRetry>,
-    /// Cancels to send to lang or apply internally
-    cancels_req_rx: UnboundedReceiver<CancelOrTimeout>,
-    shutdown: CancellationToken,
+    #[pin]
+    inner: BoxStream<'static, NewOrCancel>,
 }
 
 impl RcvChans {
-    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")))
-            }
-            (maybe_new_or_retry, perm) = async {
-                // Wait for a permit to take a task and forget it. Permits are removed until a
-                // completion.
-                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"), perm
-            )),
-            _ = self.shutdown.cancelled() => None
+    fn new(
+        new_reqs: UnboundedReceiver<NewOrRetry>,
+        new_sem: MeteredSemaphore,
+        cancels: UnboundedReceiver<CancelOrTimeout>,
+        shutdown_completed: CancellationToken,
+    ) -> Self {
+        let cancel_stream = UnboundedReceiverStream::new(cancels).map(NewOrCancel::Cancel);
+        let new_stream = UnboundedReceiverStream::new(new_reqs)
+            // Get a permit for each new activity request
+            .zip(stream::unfold(new_sem, |new_sem| async move {
+                let permit = new_sem
+                    .acquire_owned()
+                    .await
+                    .expect("Local activity semaphore is never closed");
+                Some((permit, new_sem))
+            }))
+            .map(|(req, permit)| NewOrCancel::New(req, permit));
+        Self {
+            inner: tokio_stream::StreamExt::merge(cancel_stream, new_stream)
+                .take_until(async move { shutdown_completed.cancelled().await })
+                .boxed(),
         }
     }
 }
+impl Stream for RcvChans {
+    type Item = NewOrCancel;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.project();
+        this.inner.poll_next(cx)
+    }
+}
 
 struct TimeoutBag {
     sched_to_close_handle: JoinHandle<()>,
@@ -566,17 +749,21 @@ impl TimeoutBag {
         let (schedule_to_close, start_to_close) =
             new_la.schedule_cmd.close_timeouts.into_sched_and_start();
 
+        let sched_time = new_la
+            .schedule_cmd
+            .original_schedule_time
+            .unwrap_or(new_la.schedule_time);
         let resolution = LocalActivityResolution {
             seq: new_la.schedule_cmd.seq,
             result: LocalActivityExecutionResult::timeout(TimeoutType::ScheduleToClose),
             runtime: Default::default(),
             attempt: new_la.schedule_cmd.attempt,
             backoff: None,
-            original_schedule_time: Some(new_la.schedule_time),
+            original_schedule_time: new_la.schedule_cmd.original_schedule_time,
         };
         // Remove any time already elapsed since the scheduling time
         let schedule_to_close = schedule_to_close
-            .map(|s2c| s2c.saturating_sub(new_la.schedule_time.elapsed().unwrap_or_default()));
+            .map(|s2c| s2c.saturating_sub(sched_time.elapsed().unwrap_or_default()));
         if let Some(ref s2c) = schedule_to_close {
             if s2c.is_zero() {
                 return Err(resolution);
@@ -639,18 +826,19 @@ impl Drop for TimeoutBag {
 mod tests {
     use super::*;
     use crate::{prost_dur, protosext::LACloseTimeouts};
+    use futures_util::FutureExt;
     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};
+    use tokio::task::yield_now;
 
     impl DispatchOrTimeoutLA {
         fn unwrap(self) -> ActivityTask {
             match self {
                 DispatchOrTimeoutLA::Dispatch(t) => t,
-                DispatchOrTimeoutLA::Timeout { .. } => {
-                    panic!("Timeout returned when expected a task")
+                _ => {
+                    panic!("Non-dispatched action returned")
                 }
             }
         }
@@ -1025,18 +1213,66 @@ mod tests {
         lam.next_pending().await.unwrap().unwrap();
         assert_eq!(lam.num_outstanding(), 1);
         // There should be nothing else in the queue
-        assert_eq!(
-            lam.rcvs.lock().await.act_req_rx.try_recv().unwrap_err(),
-            TryRecvError::Empty
-        );
+        assert!(lam.rcvs.lock().await.next().now_or_never().is_none());
 
         // Verify that if we now enqueue the same act again, after the task is outstanding, we still
         // don't add it.
         lam.enqueue([new_la.into()]);
         assert_eq!(lam.num_outstanding(), 1);
-        assert_eq!(
-            lam.rcvs.lock().await.act_req_rx.try_recv().unwrap_err(),
-            TryRecvError::Empty
-        );
+        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);
+        }
     }
 }