@temporalio/core-bridge 0.17.2 → 0.18.0

package/sdk-core/{src → core/src}/workflow/workflow_tasks/mod.rs

@@ -4,33 +4,46 @@ mod cache_manager;
 mod concurrency_manager;
 
 use crate::{
-    errors::{WorkflowMissingError, WorkflowUpdateError},
-    machines::{ProtoCommand, WFCommand, WFMachinesError},
     pending_activations::PendingActivations,
-    pollers::GatewayRef,
-    protosext::{ValidPollWFTQResponse, WfActivationExt},
-    task_token::TaskToken,
+    protosext::{ValidPollWFTQResponse, WorkflowActivationExt},
     telemetry::metrics::MetricsContext,
+    worker::{LocalActRequest, LocalActivityResolution},
     workflow::{
+        machines::WFMachinesError,
         workflow_tasks::{
             cache_manager::WorkflowCacheManager, concurrency_manager::WorkflowConcurrencyManager,
         },
-        HistoryPaginator, HistoryUpdate, WorkflowCachingPolicy, WorkflowManager, LEGACY_QUERY_ID,
+        HistoryPaginator, HistoryUpdate, LocalResolution, WFCommand, WorkflowCachingPolicy,
+        WorkflowManager, LEGACY_QUERY_ID,
     },
 };
 use crossbeam::queue::SegQueue;
 use futures::FutureExt;
 use parking_lot::Mutex;
-use std::{fmt::Debug, time::Instant};
-use temporal_sdk_core_protos::coresdk::{
-    workflow_activation::{
-        create_evict_activation, create_query_activation, wf_activation_job, QueryWorkflow,
-        WfActivation,
+use std::{
+    fmt::Debug,
+    ops::Add,
+    sync::Arc,
+    time::{Duration, Instant},
+};
+use temporal_client::ServerGatewayApis;
+use temporal_sdk_core_protos::{
+    coresdk::{
+        workflow_activation::{
+            create_query_activation, remove_from_cache::EvictionReason, workflow_activation_job,
+            QueryWorkflow, WorkflowActivation,
+        },
+        workflow_commands::QueryResult,
+        FromPayloadsExt,
     },
-    workflow_commands::QueryResult,
-    FromPayloadsExt,
+    temporal::api::command::v1::Command as ProtoCommand,
+    TaskToken,
 };
-use tokio::sync::watch;
+use tokio::{sync::Notify, time::timeout_at};
+
+/// What percentage of a WFT timeout we are willing to wait before sending a WFT heartbeat when
+/// necessary.
+const WFT_HEARTBEAT_TIMEOUT_FRACTION: f32 = 0.8;
 
 /// Centralizes concerns related to applying new workflow tasks and reporting the activations they
 /// produce.
@@ -43,10 +56,13 @@ pub struct WorkflowTaskManager {
     /// when cancelling an activity in try-cancel/abandon mode), or for other reasons such as a
     /// requested eviction. They queue here.
     pending_activations: PendingActivations,
+    /// Holds activations which are purely query activations needed to respond to legacy queries.
+    /// Activations may only be added here for runs which do not have other pending activations.
+    pending_legacy_queries: SegQueue<WorkflowActivation>,
     /// Holds poll wft responses from the server that need to be applied
     ready_buffered_wft: SegQueue<ValidPollWFTQResponse>,
     /// Used to wake blocked workflow task polling
-    pending_activations_notifier: watch::Sender<bool>,
+    pending_activations_notifier: Arc<Notify>,
     /// Lock guarded cache manager, which is the authority for limit-based workflow machine eviction
     /// from the cache.
     // TODO: Also should be moved inside concurrency manager, but there is some complexity around
@@ -98,7 +114,7 @@ pub struct WorkflowTaskInfo {
 #[derive(Debug, derive_more::From)]
 pub(crate) enum NewWfTaskOutcome {
     /// A new activation for the workflow should be issued to lang
-    IssueActivation(WfActivation),
+    IssueActivation(WorkflowActivation),
     /// The poll loop should be restarted, there is nothing to do
     TaskBuffered,
     /// The workflow task should be auto-completed with an empty command list, as it must be replied
@@ -108,6 +124,8 @@ pub(crate) enum NewWfTaskOutcome {
     CacheMiss,
     /// The workflow task ran into problems while being applied and we must now evict the workflow
     Evict(WorkflowUpdateError),
+    /// No action should be taken. Possibly we are waiting for local activities to complete
+    LocalActsOutstanding,
 }
 
 #[derive(Debug)]
@@ -118,17 +136,18 @@ pub enum FailedActivationOutcome {
 }
 
 #[derive(Debug)]
-pub struct ServerCommandsWithWorkflowInfo {
+pub(crate) struct ServerCommandsWithWorkflowInfo {
     pub task_token: TaskToken,
     pub action: ActivationAction,
 }
 
 #[derive(Debug)]
-pub enum ActivationAction {
+pub(crate) enum ActivationAction {
     /// We should respond that the workflow task is complete
     WftComplete {
         commands: Vec<ProtoCommand>,
         query_responses: Vec<QueryResult>,
+        force_new_wft: bool,
     },
     /// We should respond to a legacy query request
     RespondLegacyQuery { result: QueryResult },
@@ -149,13 +168,14 @@ macro_rules! machine_mut {
 
 impl WorkflowTaskManager {
     pub(crate) fn new(
-        pending_activations_notifier: watch::Sender<bool>,
+        pending_activations_notifier: Arc<Notify>,
         eviction_policy: WorkflowCachingPolicy,
         metrics: MetricsContext,
     ) -> Self {
         Self {
             workflow_machines: WorkflowConcurrencyManager::new(),
             pending_activations: Default::default(),
+            pending_legacy_queries: Default::default(),
             ready_buffered_wft: Default::default(),
             pending_activations_notifier,
             cache_manager: Mutex::new(WorkflowCacheManager::new(eviction_policy, metrics.clone())),
@@ -163,7 +183,11 @@ impl WorkflowTaskManager {
         }
     }
 
-    pub(crate) fn next_pending_activation(&self) -> Option<WfActivation> {
+    pub(crate) fn next_pending_activation(&self) -> Option<WorkflowActivation> {
+        // Dispatch pending legacy queries first
+        if let leg_q @ Some(_) = self.pending_legacy_queries.pop() {
+            return leg_q;
+        }
         // It is important that we do not issue pending activations for any workflows which already
         // have an outstanding activation. If we did, it can result in races where an in-progress
         // completion may appear to be the last in a task (no more pending activations) because
@@ -172,37 +196,69 @@ impl WorkflowTaskManager {
         let maybe_act = self
             .pending_activations
             .pop_first_matching(|rid| self.workflow_machines.get_activation(rid).is_none());
-        if let Some(act) = maybe_act.as_ref() {
-            if let Err(WorkflowMissingError { run_id }) = self.insert_outstanding_activation(act) {
-                self.request_eviction(&run_id, "Pending activation present for missing run");
+        if let Some(pending_info) = maybe_act {
+            if let Ok(act) = self
+                .workflow_machines
+                .access_sync(&pending_info.run_id, |wfm| wfm.machines.get_wf_activation())
+                .and_then(|mut act| {
+                    if let Some(reason) = pending_info.needs_eviction {
+                        act.append_evict_job(reason);
+                    }
+                    self.insert_outstanding_activation(&act)?;
+                    Ok(act)
+                })
+            {
+                self.cache_manager.lock().touch(&act.run_id);
+                Some(act)
+            } else {
+                self.request_eviction(
+                    &pending_info.run_id,
+                    "Tried to apply pending activation for missing run",
+                    EvictionReason::Fatal,
+                );
                 // Continue trying to return a valid pending activation
-                return self.next_pending_activation();
+                self.next_pending_activation()
             }
-            self.cache_manager.lock().touch(&act.run_id);
+        } else {
+            None
         }
-        maybe_act
     }
 
-    pub fn next_buffered_poll(&self) -> Option<ValidPollWFTQResponse> {
+    pub(crate) fn next_buffered_poll(&self) -> Option<ValidPollWFTQResponse> {
         self.ready_buffered_wft.pop()
     }
 
-    pub fn outstanding_wft(&self) -> usize {
+    pub(crate) fn outstanding_wft(&self) -> usize {
         self.workflow_machines.outstanding_wft()
     }
 
+    /// Returns the event id of the most recently processed event for the provided run id.
+    pub(crate) fn most_recently_processed_event(
+        &self,
+        run_id: &str,
+    ) -> Result<i64, WorkflowMissingError> {
+        self.workflow_machines
+            .access_sync(run_id, |wfm| wfm.machines.last_processed_event)
+    }
+
     /// Request a workflow eviction. This will queue up an activation to evict the workflow from
     /// the lang side. Workflow will not *actually* be evicted until lang replies to that activation
     ///
     /// Returns, if found, the number of attempts on the current workflow task
-    pub fn request_eviction(&self, run_id: &str, reason: impl Into<String>) -> Option<u32> {
+    pub(crate) fn request_eviction(
+        &self,
+        run_id: &str,
+        message: impl Into<String>,
+        reason: EvictionReason,
+    ) -> Option<u32> {
         if self.workflow_machines.exists(run_id) {
             if !self.activation_has_eviction(run_id) {
-                debug!(%run_id, "Eviction requested");
+                let message = message.into();
+                debug!(%run_id, %message, "Eviction requested");
                 // Queue up an eviction activation
                 self.pending_activations
-                    .push(create_evict_activation(run_id.to_string(), reason.into()));
-                let _ = self.pending_activations_notifier.send(true);
+                    .notify_needs_eviction(run_id, message, reason);
+                self.pending_activations_notifier.notify_waiters();
             }
             self.workflow_machines
                 .get_task(run_id)
@@ -235,15 +291,14 @@ impl WorkflowTaskManager {
     }
 
     /// Given a validated poll response from the server, prepare an activation (if there is one) to
-    /// be sent to lang. If applying the response to the workflow's state does not produce a new
-    /// activation, `None` is returned.
+    /// be sent to lang.
     ///
     /// The new activation is immediately considered to be an outstanding workflow task - so it is
     /// expected that new activations will be dispatched to lang right away.
     pub(crate) async fn apply_new_poll_resp(
         &self,
         work: ValidPollWFTQResponse,
-        gateway: &GatewayRef,
+        gateway: Arc<dyn ServerGatewayApis + Send + Sync>,
     ) -> NewWfTaskOutcome {
         let mut work = if let Some(w) = self.workflow_machines.buffer_resp_if_outstanding_work(work)
         {
@@ -285,7 +340,7 @@ impl WorkflowTaskManager {
                 debug!("Dispatching legacy query {:?}", &lq);
                 next_activation
                     .jobs
-                    .push(wf_activation_job::Variant::QueryWorkflow(lq).into());
+                    .push(workflow_activation_job::Variant::QueryWorkflow(lq).into());
             }
             None
         } else {
@@ -304,7 +359,20 @@ impl WorkflowTaskManager {
             .expect("Workflow machines must exist, we just created/updated them");
 
         if next_activation.jobs.is_empty() {
-            NewWfTaskOutcome::Autocomplete
+            let outstanding_las = self
+                .workflow_machines
+                .access_sync(&next_activation.run_id, |wfm| {
+                    wfm.machines.outstanding_local_activity_count()
+                })
+                .expect("Workflow machines must exist, we just created/updated them");
+            if outstanding_las > 0 {
+                // If there are outstanding local activities, we don't want to autocomplete the
+                // workflow task. We want to give them a chance to complete. If they take longer
+                // than the WFT timeout, we will force a new WFT just before the timeout.
+                NewWfTaskOutcome::LocalActsOutstanding
+            } else {
+                NewWfTaskOutcome::Autocomplete
+            }
         } else {
             if let Err(wme) = self.insert_outstanding_activation(&next_activation) {
                 return NewWfTaskOutcome::Evict(wme.into());
@@ -319,15 +387,20 @@ impl WorkflowTaskManager {
         &self,
         run_id: &str,
         mut commands: Vec<WFCommand>,
+        local_activity_request_sink: impl FnOnce(Vec<LocalActRequest>) -> Vec<LocalActivityResolution>,
     ) -> Result<Option<ServerCommandsWithWorkflowInfo>, WorkflowUpdateError> {
         // No-command replies to evictions can simply skip everything
         if commands.is_empty() && self.activation_has_eviction(run_id) {
             return Ok(None);
         }
 
-        let (task_token, is_leg_query_task) =
+        let (task_token, is_leg_query_task, start_time) =
             if let Some(entry) = self.workflow_machines.get_task(run_id) {
-                (entry.info.task_token.clone(), entry.legacy_query.is_some())
+                (
+                    entry.info.task_token.clone(),
+                    entry.legacy_query.is_some(),
+                    entry.start_time,
+                )
             } else {
                 if !self.activation_has_eviction(run_id) {
                     // Don't bother warning if this was an eviction, since it's normal to issue
@@ -340,7 +413,7 @@ impl WorkflowTaskManager {
                 return Ok(None);
             };
 
-        // If the only command in the activation is a legacy query response, that means we need
+        // If the only command from the activation is a legacy query response, that means we need
         // to respond differently than a typical activation.
         let ret = if matches!(&commands.as_slice(),
                     &[WFCommand::QueryResponse(qr)] if qr.query_id == LEGACY_QUERY_ID)
@@ -363,7 +436,11 @@ impl WorkflowTaskManager {
                     if let WFCommand::QueryResponse(qr) = commands.remove(i) {
                         if qr.query_id == LEGACY_QUERY_ID {
                             return Err(WorkflowUpdateError {
-                                source: WFMachinesError::Fatal("Legacy query activation response included other commands, this is not allowed and constitutes an error in the lang SDK".to_string()),
+                                source: WFMachinesError::Fatal(
+                                    "Legacy query activation response included other commands, \
+                                    this is not allowed and constitutes an error in the lang SDK"
+                                        .to_string(),
+                                ),
                                 run_id: run_id.to_string(),
                             });
                         }
@@ -374,25 +451,57 @@ impl WorkflowTaskManager {
                 }
             }
 
-            // Send commands from lang into the machines
-            machine_mut!(self, run_id, |wfm: &mut WorkflowManager| {
-                wfm.push_commands(commands).boxed()
-            })?;
-            // Check if the workflow run needs another activation and queue it up if there is one
-            // by pushing it into the pending activations list
-            let next_activation = machine_mut!(self, run_id, move |mgr: &mut WorkflowManager| mgr
-                .get_next_activation()
-                .boxed())?;
-            if !next_activation.jobs.is_empty() {
-                self.pending_activations.push(next_activation);
-                let _ = self.pending_activations_notifier.send(true);
+            let (are_pending, server_cmds, local_activities, wft_timeout) = machine_mut!(
+                self,
+                run_id,
+                |wfm: &mut WorkflowManager| {
+                    async move {
+                        // Send commands from lang into the machines then check if the workflow run
+                        // needs another activation and mark it if so
+                        wfm.push_commands(commands).await?;
+                        let are_pending = wfm.apply_next_task_if_ready().await?;
+                        // We want to fetch the outgoing commands only after a next WFT may have
+                        // been applied, as outgoing server commands may be affected.
+                        let outgoing_cmds = wfm.get_server_commands();
+                        let new_local_acts = wfm.drain_queued_local_activities();
+
+                        let wft_timeout: Duration = wfm
+                            .machines
+                            .started_attrs()
+                            .and_then(|attrs| attrs.workflow_task_timeout.clone())
+                            .and_then(|tt| tt.try_into().ok())
+                            .ok_or_else(|| {
+                                WFMachinesError::Fatal(
+                                    "Workflow's start attribs were missing a well formed task timeout"
+                                        .to_string(),
+                                )
+                            })?;
+
+                        Ok((are_pending, outgoing_cmds, new_local_acts, wft_timeout))
+                    }
+                    .boxed()
+                }
+            )?;
+
+            if are_pending {
+                self.needs_activation(run_id);
             }
-            // We want to fetch the outgoing commands only after any new activation has been queued,
-            // as doing so may have altered the outgoing commands.
-            let server_cmds = machine_mut!(self, run_id, |wfm: &mut WorkflowManager| {
-                async move { Ok(wfm.get_server_commands()) }.boxed()
-            })?;
+            let immediate_resolutions = local_activity_request_sink(local_activities);
+            for resolution in immediate_resolutions {
+                self.notify_of_local_result(run_id, LocalResolution::LocalActivity(resolution))
+                    .await?;
+            }
+
+            // The heartbeat deadline is 80% of the WFT timeout
+            let wft_heartbeat_deadline =
+                start_time.add(wft_timeout.mul_f32(WFT_HEARTBEAT_TIMEOUT_FRACTION));
+            // Wait on local activities to resolve if there are any, or for the WFT timeout to
+            // be about to expire, in which case we will need to send a WFT heartbeat.
+            let must_heartbeat = self
+                .wait_for_local_acts_or_heartbeat(run_id, wft_heartbeat_deadline)
+                .await;
             let is_query_playback = is_leg_query_task && query_responses.is_empty();
+
             // We only actually want to send commands back to the server if there are no more
             // pending activations and we are caught up on replay. We don't want to complete a wft
             // if we already saw the final event in the workflow, or if we are playing back for the
@@ -404,6 +513,8 @@ impl WorkflowTaskManager {
                 Some(ServerCommandsWithWorkflowInfo {
                     task_token,
                     action: ActivationAction::WftComplete {
+                        // TODO: Don't force if also sending complete execution cmd
+                        force_new_wft: must_heartbeat,
                         commands: server_cmds.commands,
                         query_responses,
                     },
@@ -416,6 +527,7 @@ impl WorkflowTaskManager {
                     action: ActivationAction::WftComplete {
                         commands: vec![],
                         query_responses,
+                        force_new_wft: false,
                     },
                 })
             }
@@ -425,7 +537,12 @@ impl WorkflowTaskManager {
 
     /// Record that an activation failed, returns enum that indicates if failure should be reported
     /// to the server
-    pub(crate) fn failed_activation(&self, run_id: &str) -> FailedActivationOutcome {
+    pub(crate) fn failed_activation(
+        &self,
+        run_id: &str,
+        reason: EvictionReason,
+        failstr: String,
+    ) -> FailedActivationOutcome {
         let tt = if let Some(tt) = self
             .workflow_machines
             .get_task(run_id)
@@ -450,7 +567,7 @@ impl WorkflowTaskManager {
         } else {
             // Blow up any cached data associated with the workflow
             let should_report = self
-                .request_eviction(run_id, "Activation failed")
+                .request_eviction(run_id, failstr, reason)
                 .map_or(true, |attempt| attempt <= 1);
             if should_report {
                 FailedActivationOutcome::Report(tt)
@@ -467,8 +584,8 @@ impl WorkflowTaskManager {
     async fn instantiate_or_update_workflow(
         &self,
         poll_wf_resp: ValidPollWFTQResponse,
-        gateway: &GatewayRef,
-    ) -> Result<(WorkflowTaskInfo, WfActivation), WorkflowUpdateError> {
+        gateway: Arc<dyn ServerGatewayApis + Send + Sync>,
+    ) -> Result<(WorkflowTaskInfo, WorkflowActivation), WorkflowUpdateError> {
         let run_id = poll_wf_resp.workflow_execution.run_id.clone();
 
         let wft_info = WorkflowTaskInfo {
@@ -486,12 +603,12 @@ impl WorkflowTaskManager {
                         poll_wf_resp.workflow_execution.workflow_id.clone(),
                         poll_wf_resp.workflow_execution.run_id.clone(),
                         poll_wf_resp.next_page_token,
-                        gateway.gw.clone(),
+                        gateway.clone(),
                     ),
                     poll_wf_resp.previous_started_event_id,
                 ),
                 &poll_wf_resp.workflow_execution.workflow_id,
-                &gateway.options.namespace,
+                &gateway.get_options().namespace,
                 &poll_wf_resp.workflow_type,
                 &self.metrics,
             )
@@ -503,7 +620,7 @@ impl WorkflowTaskManager {
                     let query_jobs = poll_wf_resp
                         .query_requests
                         .into_iter()
-                        .map(|q| wf_activation_job::Variant::QueryWorkflow(q).into());
+                        .map(|q| workflow_activation_job::Variant::QueryWorkflow(q).into());
                     activation.jobs.extend(query_jobs);
                 }
 
@@ -513,13 +630,13 @@ impl WorkflowTaskManager {
         }
     }
 
-    /// Called after every WFT completion or failure, updates outstanding task status & issues
-    /// evictions if required. It is important this is called *after* reporting a successful WFT
-    /// to server, as some replies (task not found) may require an eviction, which could be avoided
-    /// if this is called too early.
+    /// Called after every workflow activation completion or failure, updates outstanding task
+    /// status & issues evictions if required. It is important this is called *after* potentially
+    /// reporting a successful WFT to server, as some replies (task not found) may require an
+    /// eviction, which could be avoided if this is called too early.
     ///
-    /// Returns true if WFT is complete
-    pub(crate) fn after_wft_report(&self, run_id: &str) -> bool {
+    /// Returns true if WFT was marked completed internally
+    pub(crate) fn after_wft_report(&self, run_id: &str, did_complete_wft: bool) -> bool {
         let mut just_evicted = false;
 
         if let Some(OutstandingActivation::Normal {
@@ -529,9 +646,10 @@ impl WorkflowTaskManager {
             self.evict_run(run_id);
             just_evicted = true;
         };
+
         // Workflows with no more pending activations (IE: They have completed a WFT) must be
         // removed from the outstanding tasks map
-        if !self.pending_activations.has_pending(run_id) {
+        let retme = if !self.pending_activations.has_pending(run_id) {
             if !just_evicted {
                 // Check if there was a legacy query which must be fulfilled, and if there is create
                 // a new pending activation for it.
@@ -542,8 +660,8 @@ impl WorkflowTaskManager {
                 {
                     if let Some(query) = ot.legacy_query.take() {
                         let na = create_query_activation(run_id.to_string(), [query]);
-                        self.pending_activations.push(na);
-                        let _ = self.pending_activations_notifier.send(true);
+                        self.pending_legacy_queries.push(na);
+                        self.pending_activations_notifier.notify_waiters();
                         return false;
                     }
                 }
@@ -551,7 +669,11 @@ impl WorkflowTaskManager {
                 // Evict run id if cache is full. Non-sticky will always evict.
                 let maybe_evicted = self.cache_manager.lock().insert(run_id);
                 if let Some(evicted_run_id) = maybe_evicted {
-                    self.request_eviction(&evicted_run_id, "Workflow cache full");
+                    self.request_eviction(
+                        &evicted_run_id,
+                        "Workflow cache full",
+                        EvictionReason::CacheFull,
+                    );
                 }
 
                 // If there was a buffered poll response from the server, it is now ready to
@@ -563,9 +685,14 @@ impl WorkflowTaskManager {
 
             // The evict may or may not have already done this, but even when we aren't evicting
             // we want to clear the outstanding workflow task since it's now complete.
-            return self.workflow_machines.complete_wft(run_id).is_some();
-        }
-        false
+            self.workflow_machines
+                .complete_wft(run_id, did_complete_wft)
+                .is_some()
+        } else {
+            false
+        };
+        self.on_activation_done(run_id);
+        retme
     }
 
     /// Must be called after *every* activation is replied to, regardless of whether or not we
@@ -573,20 +700,41 @@ impl WorkflowTaskManager {
     /// every activation we issue to lang has exactly one reply.
     ///
     /// Any subsequent action that needs to be taken will be created as a new activation
-    pub(crate) fn on_activation_done(&self, run_id: &str) {
+    fn on_activation_done(&self, run_id: &str) {
         if self.workflow_machines.delete_activation(run_id).is_some() {
-            let _ = self.pending_activations_notifier.send(true);
+            self.pending_activations_notifier.notify_waiters();
         }
         // It's possible the activation is already removed due to completing an eviction
     }
 
+    /// Let a workflow know that something we've been waiting locally on has resolved, like a local
+    /// activity or side effect
+    #[instrument(level = "debug", skip(self, resolved))]
+    pub(crate) async fn notify_of_local_result(
+        &self,
+        run_id: &str,
+        resolved: LocalResolution,
+    ) -> Result<(), WorkflowUpdateError> {
+        self.workflow_machines
+            .access_sync(run_id, |wfm: &mut WorkflowManager| {
+                wfm.notify_of_local_result(resolved)
+            })?
+            .map_err(|wfme| WorkflowUpdateError {
+                source: wfme,
+                run_id: run_id.to_string(),
+            })?;
+
+        self.needs_activation(run_id);
+        Ok(())
+    }
+
     fn make_buffered_poll_ready(&self, buffd: ValidPollWFTQResponse) {
         self.ready_buffered_wft.push(buffd);
     }
 
     fn insert_outstanding_activation(
         &self,
-        act: &WfActivation,
+        act: &WorkflowActivation,
     ) -> Result<(), WorkflowMissingError> {
         let act_type = if act.is_legacy_query() {
             OutstandingActivation::LegacyQuery
@@ -619,4 +767,72 @@ impl WorkflowTaskManager {
             .map(OutstandingActivation::has_eviction)
             .unwrap_or_default()
     }
+
+    fn needs_activation(&self, run_id: &str) {
+        self.pending_activations.notify_needs_activation(run_id);
+        self.pending_activations_notifier.notify_waiters();
+    }
+
+    /// Wait for either all local activities to resolve, or for 80% of the WFT timeout, in which
+    /// case we will "heartbeat" by completing the WFT, even if there are no commands to send.
+    ///
+    /// Returns true if we must heartbeat
+    async fn wait_for_local_acts_or_heartbeat(
+        &self,
+        run_id: &str,
+        wft_heartbeat_deadline: Instant,
+    ) -> bool {
+        loop {
+            let la_count = self
+                .workflow_machines
+                .access_sync(run_id, |wfm| {
+                    wfm.machines.outstanding_local_activity_count()
+                })
+                .expect("Workflow cannot go missing while we are waiting on LAs");
+            if la_count == 0 {
+                return false;
+            } else if Instant::now() >= wft_heartbeat_deadline {
+                // We must heartbeat b/c there are still pending local activities
+                return true;
+            }
+            // Since an LA resolution always results in a new pending activation, we can wait on
+            // notifications of that to re-check if they're all resolved.
+            let _ = timeout_at(
+                wft_heartbeat_deadline.into(),
+                self.pending_activations_notifier.notified(),
+            )
+            .await;
+        }
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct WorkflowUpdateError {
+    /// Underlying workflow error
+    pub source: WFMachinesError,
+    /// The run id of the erring workflow
+    #[allow(dead_code)] // Useful in debug output
+    pub run_id: String,
+}
+
+impl WorkflowUpdateError {
+    pub fn evict_reason(&self) -> EvictionReason {
+        self.source.evict_reason()
+    }
+}
+
+impl From<WorkflowMissingError> for WorkflowUpdateError {
+    fn from(wme: WorkflowMissingError) -> Self {
+        Self {
+            source: WFMachinesError::Fatal("Workflow machines missing".to_string()),
+            run_id: wme.run_id,
+        }
+    }
+}
+
+/// The workflow machines were expected to be in the cache but were not
+#[derive(Debug)]
+pub(crate) struct WorkflowMissingError {
+    /// The run id of the erring workflow
+    pub run_id: String,
 }