@temporalio/core-bridge 0.16.4 → 0.18.0

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

@@ -1,40 +1,48 @@
 mod activities;
-mod config;
 mod dispatcher;
+mod wft_delivery;
 
-pub use crate::worker::config::{WorkerConfig, WorkerConfigBuilder};
+pub use temporal_sdk_core_api::worker::{WorkerConfig, WorkerConfigBuilder};
+
+pub(crate) use activities::{
+    ExecutingLAId, LocalActRequest, LocalActivityExecutionResult, LocalActivityResolution,
+    NewLocalAct,
+};
 pub(crate) use dispatcher::WorkerDispatcher;
 
 use crate::{
     errors::CompleteWfError,
-    machines::{EmptyWorkflowCommandErr, WFMachinesError},
     pollers::{
-        new_activity_task_buffer, new_workflow_task_buffer, BoxedActPoller, BoxedWFPoller,
-        GatewayRef, Poller, WorkflowTaskPoller,
+        new_activity_task_buffer, new_workflow_task_buffer, BoxedActPoller, BoxedWFPoller, Poller,
+        WorkflowTaskPoller,
     },
-    protosext::{legacy_query_failure, ValidPollWFTQResponse, WorkflowTaskCompletion},
-    task_token::TaskToken,
+    protosext::{legacy_query_failure, ValidPollWFTQResponse},
     telemetry::metrics::{
         activity_poller, workflow_poller, workflow_sticky_poller, MetricsContext,
     },
+    worker::{
+        activities::{DispatchOrTimeoutLA, LACompleteAction, LocalActivityManager},
+        wft_delivery::WFTSource,
+    },
     workflow::{
         workflow_tasks::{
             ActivationAction, FailedActivationOutcome, NewWfTaskOutcome,
             ServerCommandsWithWorkflowInfo, WorkflowTaskManager,
         },
-        WorkflowCachingPolicy,
+        EmptyWorkflowCommandErr, LocalResolution, WFMachinesError, WorkflowCachingPolicy,
     },
     ActivityHeartbeat, CompleteActivityError, PollActivityError, PollWfError,
 };
-use activities::WorkerActivityTasks;
+use activities::{LocalInFlightActInfo, WorkerActivityTasks};
 use futures::{Future, TryFutureExt};
 use std::{convert::TryInto, sync::Arc};
+use temporal_client::{ServerGatewayApis, WorkflowTaskCompletion};
 use temporal_sdk_core_protos::{
     coresdk::{
-        activity_result::activity_result,
+        activity_result::activity_execution_result,
         activity_task::ActivityTask,
-        workflow_activation::WfActivation,
-        workflow_completion::{self, wf_activation_completion, WfActivationCompletion},
+        workflow_activation::{remove_from_cache::EvictionReason, WorkflowActivation},
+        workflow_completion::{self, workflow_activation_completion, WorkflowActivationCompletion},
     },
     temporal::api::{
         enums::v1::{TaskQueueKind, WorkflowTaskFailedCause},
@@ -42,39 +50,42 @@ use temporal_sdk_core_protos::{
         taskqueue::v1::{StickyExecutionAttributes, TaskQueue},
         workflowservice::v1::{PollActivityTaskQueueResponse, PollWorkflowTaskQueueResponse},
     },
+    TaskToken,
 };
-use tokio::sync::{watch, Mutex, Semaphore};
+use tokio::sync::{watch, Notify, Semaphore};
 use tonic::Code;
 use tracing_futures::Instrument;
 
 /// A worker polls on a certain task queue
 pub struct Worker {
     config: WorkerConfig,
-    server_gateway: Arc<GatewayRef>,
+    server_gateway: Arc<dyn ServerGatewayApis + Send + Sync>,
 
     /// Will be populated when this worker should poll on a sticky WFT queue
     sticky_name: Option<String>,
 
-    // TODO: Worth moving inside wf task mgr too?
     /// Buffers workflow task polling in the event we need to return a pending activation while
     /// a poll is ongoing. Sticky and nonsticky polling happens inside of it.
-    wf_task_poll_buffer: BoxedWFPoller,
-    /// Workflow task management
-    wft_manager: WorkflowTaskManager,
+    wf_task_source: WFTSource,
+    /// Workflow task management TODO: No pub
+    pub(crate) wft_manager: WorkflowTaskManager,
     /// Manages activity tasks for this worker/task queue
     at_task_mgr: Option<WorkerActivityTasks>,
+    /// Manages local activities
+    local_act_mgr: LocalActivityManager,
     /// Ensures we stay at or below this worker's maximum concurrent workflow limit
     workflows_semaphore: Semaphore,
     /// Used to wake blocked workflow task polling when there is some change to workflow activations
     /// that should cause us to restart the loop
-    pending_activations_notification_receiver: Mutex<watch::Receiver<bool>>,
-    /// Watched during shutdown to wait for all WFTs to complete
-    wfts_drained: watch::Receiver<bool>,
-    /// notifies when all WFTs have been drained after shutdown
-    wfts_drained_sender: watch::Sender<bool>,
+    pending_activations_notify: Arc<Notify>,
+    /// Watched during shutdown to wait for all WFTs to complete. Should be notified any time
+    /// a WFT is completed.
+    wfts_drained_notify: Arc<Notify>,
     /// Has shutdown been called?
     shutdown_requested: watch::Receiver<bool>,
     shutdown_sender: watch::Sender<bool>,
+    /// Will be called at the end of each activation completion
+    post_activate_hook: Option<Box<dyn Fn(&Self) + Send + Sync>>,
 
     metrics: MetricsContext,
 }
@@ -83,7 +94,7 @@ impl Worker {
     pub(crate) fn new(
         config: WorkerConfig,
         sticky_queue_name: Option<String>,
-        sg: Arc<GatewayRef>,
+        sg: Arc<dyn ServerGatewayApis + Send + Sync>,
         metrics: MetricsContext,
     ) -> Self {
         metrics.worker_registered();
@@ -96,7 +107,7 @@ impl Worker {
         let max_sticky_polls = config.max_sticky_polls();
         let wft_metrics = metrics.with_new_attrs([workflow_poller()]);
         let mut wf_task_poll_buffer = new_workflow_task_buffer(
-            sg.gw.clone(),
+            sg.clone(),
             config.task_queue.clone(),
             false,
             max_nonsticky_polls,
@@ -106,7 +117,7 @@ impl Worker {
         let sticky_queue_poller = sticky_queue_name.as_ref().map(|sqn| {
             let sticky_metrics = metrics.with_new_attrs([workflow_sticky_poller()]);
             let mut sp = new_workflow_task_buffer(
-                sg.gw.clone(),
+                sg.clone(),
                 sqn.clone(),
                 true,
                 max_sticky_polls,
@@ -119,7 +130,7 @@ impl Worker {
             None
         } else {
             let mut ap = new_activity_task_buffer(
-                sg.gw.clone(),
+                sg.clone(),
                 config.task_queue.clone(),
                 config.max_concurrent_at_polls,
                 config.max_concurrent_at_polls * 2,
@@ -148,7 +159,7 @@ impl Worker {
     pub(crate) fn new_with_pollers(
         config: WorkerConfig,
         sticky_queue_name: Option<String>,
-        sg: Arc<GatewayRef>,
+        sg: Arc<dyn ServerGatewayApis + Send + Sync>,
         wft_poller: BoxedWFPoller,
         act_poller: Option<BoxedActPoller>,
         metrics: MetricsContext,
@@ -160,63 +171,76 @@ impl Worker {
                 max_cached_workflows: config.max_cached_workflows,
             }
         };
-        let (pan_tx, pan_rx) = watch::channel(true);
-        let (wftd_tx, wftd_rx) = watch::channel(false);
+        let pa_notif = Arc::new(Notify::new());
+        let wfts_drained_notify = Arc::new(Notify::new());
         let (shut_tx, shut_rx) = watch::channel(false);
         Self {
             server_gateway: sg.clone(),
             sticky_name: sticky_queue_name,
-            wf_task_poll_buffer: wft_poller,
-            wft_manager: WorkflowTaskManager::new(pan_tx, cache_policy, metrics.clone()),
+            wf_task_source: WFTSource::new(wft_poller),
+            wft_manager: WorkflowTaskManager::new(pa_notif.clone(), cache_policy, metrics.clone()),
             at_task_mgr: act_poller.map(|ap| {
                 WorkerActivityTasks::new(
                     config.max_outstanding_activities,
                     ap,
-                    sg.gw.clone(),
+                    sg.clone(),
                     metrics.clone(),
                     config.max_heartbeat_throttle_interval,
                     config.default_heartbeat_throttle_interval,
                 )
             }),
+            local_act_mgr: LocalActivityManager::new(
+                config.max_outstanding_local_activities,
+                sg.get_options().namespace.clone(),
+            ),
             workflows_semaphore: Semaphore::new(config.max_outstanding_workflow_tasks),
             config,
             shutdown_requested: shut_rx,
             shutdown_sender: shut_tx,
-            wfts_drained: wftd_rx,
-            wfts_drained_sender: wftd_tx,
-            pending_activations_notification_receiver: Mutex::new(pan_rx),
+            post_activate_hook: None,
+            pending_activations_notify: pa_notif,
+            wfts_drained_notify,
             metrics,
         }
     }
 
-    /// Will shutdown the worker. Does not resolve until all outstanding workflow tasks have been
-    /// completed
-    pub(crate) async fn shutdown(&self) {
+    /// Begins the shutdown process, tells pollers they should stop. Is idempotent.
+    pub(crate) fn initiate_shutdown(&self) {
         let _ = self.shutdown_sender.send(true);
+        // First, we want to stop polling of both activity and workflow tasks
         if let Some(atm) = self.at_task_mgr.as_ref() {
             atm.notify_shutdown();
         }
-        self.wf_task_poll_buffer.notify_shutdown();
-        // Notify in case shutdown was requested while there were no more outstanding WFTs.
-        // This is required because the only other place where we notify wfts_drained is on
-        // activation completion and activation polling checks for wfts_drained.
-        self.maybe_notify_wtfs_drained();
-        // wait until all outstanding workflow tasks have been completed before shutting down
-        if !*self.wfts_drained.borrow() {
-            self.wfts_drained
-                .clone()
-                .changed()
-                .await
-                .expect("wfts_drained should not be dropped");
+        self.wf_task_source.stop_pollers();
+    }
+
+    /// Will shutdown the worker. Does not resolve until all outstanding workflow tasks have been
+    /// completed
+    pub(crate) async fn shutdown(&self) {
+        self.initiate_shutdown();
+        // Next we need to wait for all local activities to finish so no more workflow task
+        // heartbeats will be generated
+        self.local_act_mgr.shutdown_and_wait_all_finished().await;
+        // Then we need to wait for any tasks generated as a result of completing WFTs, which
+        // heartbeating generates
+        self.wf_task_source
+            .wait_for_tasks_from_complete_to_drain()
+            .await;
+        // wait until all outstanding workflow tasks have been completed
+        self.all_wfts_drained().await;
+        // Wait for activities to finish
+        if let Some(acts) = self.at_task_mgr.as_ref() {
+            acts.wait_all_finished().await;
         }
     }
 
     /// Finish shutting down by consuming the background pollers and freeing all resources
     pub(crate) async fn finalize_shutdown(self) {
-        self.wf_task_poll_buffer.shutdown_box().await;
-        if let Some(b) = self.at_task_mgr {
-            b.shutdown().await;
-        }
+        tokio::join!(self.wf_task_source.shutdown(), async {
+            if let Some(b) = self.at_task_mgr {
+                b.shutdown().await;
+            }
+        });
     }
 
     pub(crate) fn outstanding_workflow_tasks(&self) -> usize {
@@ -228,25 +252,36 @@ impl Worker {
         self.workflows_semaphore.available_permits()
     }
 
-    /// Wait until not at the outstanding activity limit, and then poll this worker's task queue for
-    /// new activities.
+    /// Get new activity tasks (may be local or nonlocal). Local activities are returned first
+    /// before polling the server if there are any.
     ///
     /// Returns `Ok(None)` in the event of a poll timeout or if the polling loop should otherwise
     /// be restarted
     pub(crate) async fn activity_poll(&self) -> Result<Option<ActivityTask>, PollActivityError> {
-        // No activity polling is allowed if this worker said it only handles local activities
-        let act_mgr = self
-            .at_task_mgr
-            .as_ref()
-            .ok_or_else(|| PollActivityError::NoWorkerForQueue(self.config.task_queue.clone()))?;
+        let act_mgr_poll = async {
+            if let Some(ref act_mgr) = self.at_task_mgr {
+                act_mgr.poll().await
+            } else {
+                let _ = self.shutdown_requested.clone().changed().await;
+                Err(PollActivityError::ShutDown)
+            }
+        };
 
         tokio::select! {
             biased;
 
-            r = act_mgr.poll() => r,
-            _ = self.shutdown_notifier() => {
-                Err(PollActivityError::ShutDown)
-            }
+            r = self.local_act_mgr.next_pending() => {
+                match r {
+                    Some(DispatchOrTimeoutLA::Dispatch(r)) => Ok(Some(r)),
+                    Some(DispatchOrTimeoutLA::Timeout { run_id, resolution, task }) => {
+                        self.notify_local_result(
+                            &run_id, LocalResolution::LocalActivity(resolution)).await;
+                        Ok(task)
+                    },
+                    None => Ok(None)
+                }
+            },
+            r = act_mgr_poll => r,
         }
     }
 
@@ -263,10 +298,35 @@ impl Worker {
     pub(crate) async fn complete_activity(
         &self,
         task_token: TaskToken,
-        status: activity_result::Status,
+        status: activity_execution_result::Status,
     ) -> Result<(), CompleteActivityError> {
+        if task_token.is_local_activity_task() {
+            let as_la_res: LocalActivityExecutionResult = status.try_into()?;
+            match self.local_act_mgr.complete(&task_token, &as_la_res) {
+                LACompleteAction::Report(info) => {
+                    self.complete_local_act(as_la_res, info, None).await
+                }
+                LACompleteAction::LangDoesTimerBackoff(backoff, info) => {
+                    // This la needs to write a failure marker, and then we will tell lang how
+                    // long of a timer to schedule to back off for. We do this because there are
+                    // no other situations where core generates "internal" commands so it is much
+                    // simpler for lang to reply with the timer / next LA command than to do it
+                    // internally. Plus, this backoff hack we'd like to eliminate eventually.
+                    self.complete_local_act(as_la_res, info, Some(backoff))
+                        .await
+                }
+                LACompleteAction::WillBeRetried => {
+                    // Nothing to do here
+                }
+                LACompleteAction::Untracked => {
+                    warn!("Tried to complete untracked local activity {}", task_token);
+                }
+            }
+            return Ok(());
+        }
+
         if let Some(atm) = &self.at_task_mgr {
-            atm.complete(task_token, status, self.server_gateway.gw.as_ref())
+            atm.complete(task_token, status, self.server_gateway.as_ref())
                 .await
         } else {
             error!(
@@ -276,8 +336,7 @@ impl Worker {
             Ok(())
         }
     }
-
-    pub(crate) async fn next_workflow_activation(&self) -> Result<WfActivation, PollWfError> {
+    pub(crate) async fn next_workflow_activation(&self) -> Result<WorkflowActivation, PollWfError> {
         // The poll needs to be in a loop because we can't guarantee tail call optimization in Rust
         // (simply) and we really, really need that for long-poll retries.
         loop {
@@ -297,8 +356,6 @@ impl Worker {
                     _ => continue,
                 }
             }
-            let mut pending_activations_notification =
-                self.pending_activations_notification_receiver.lock().await;
 
             let selected_f = tokio::select! {
                 biased;
@@ -307,7 +364,7 @@ impl Worker {
                 // the loop right away to provide any potential new pending activation.
                 // Continue here means that we unnecessarily add another permit to the poll buffer,
                 // this will go away when polling is done in the background.
-                _ = pending_activations_notification.changed() => continue,
+                _ = self.pending_activations_notify.notified() => continue,
                 r = self.workflow_poll_or_wfts_drained() => r,
             }?;
 
@@ -326,32 +383,47 @@ impl Worker {
 
     pub(crate) async fn complete_workflow_activation(
         &self,
-        completion: WfActivationCompletion,
+        completion: WorkflowActivationCompletion,
     ) -> Result<(), CompleteWfError> {
         let wfstatus = completion.status;
-        let did_complete_wft = match wfstatus {
-            Some(wf_activation_completion::Status::Successful(success)) => {
+        let report_outcome = match wfstatus {
+            Some(workflow_activation_completion::Status::Successful(success)) => {
                 self.wf_activation_success(&completion.run_id, success)
                     .await
             }
-            Some(wf_activation_completion::Status::Failed(failure)) => {
-                self.wf_activation_failed(&completion.run_id, failure).await
+
+            Some(workflow_activation_completion::Status::Failed(failure)) => {
+                self.wf_activation_failed(
+                    &completion.run_id,
+                    WorkflowTaskFailedCause::Unspecified,
+                    EvictionReason::LangFail,
+                    failure,
+                )
+                .await
+            }
+            None => {
+                return Err(CompleteWfError::MalformedWorkflowCompletion {
+                    reason: "Workflow completion had empty status field".to_owned(),
+                    completion: None,
+                })
             }
-            None => Err(CompleteWfError::MalformedWorkflowCompletion {
-                reason: "Workflow completion had empty status field".to_owned(),
-                completion: None,
-            }),
         }?;
-        self.after_workflow_activation(&completion.run_id, did_complete_wft);
-        Ok(())
-    }
 
-    fn maybe_notify_wtfs_drained(&self) {
-        if *self.shutdown_requested.borrow() && self.outstanding_workflow_tasks() == 0 {
-            self.wfts_drained_sender
-                .send(true)
-                .expect("wfts_drained sender shouldn't be dropped");
+        self.wft_manager
+            .after_wft_report(&completion.run_id, report_outcome.reported_to_server);
+        if report_outcome.reported_to_server || report_outcome.failed {
+            // If we failed the WFT but didn't report anything, we still want to release the WFT
+            // permit since the server will eventually time out the task and we've already evicted
+            // the run.
+            self.return_workflow_task_permit();
+        }
+        self.wfts_drained_notify.notify_waiters();
+
+        if let Some(h) = &self.post_activate_hook {
+            h(self);
         }
+
+        Ok(())
     }
 
     /// Tell the worker a workflow task has completed, for tracking max outstanding WFTs
@@ -359,37 +431,43 @@ impl Worker {
         self.workflows_semaphore.add_permits(1);
     }
 
-    pub(crate) fn request_wf_eviction(&self, run_id: &str, reason: impl Into<String>) {
-        self.wft_manager.request_eviction(run_id, reason);
+    pub(crate) fn request_wf_eviction(
+        &self,
+        run_id: &str,
+        message: impl Into<String>,
+        reason: EvictionReason,
+    ) {
+        self.wft_manager.request_eviction(run_id, message, reason);
+    }
+
+    /// Sets a function to be called at the end of each activation completion
+    pub(crate) fn set_post_activate_hook(
+        &mut self,
+        callback: impl Fn(&Self) + Send + Sync + 'static,
+    ) {
+        self.post_activate_hook = Some(Box::new(callback))
     }
 
     /// Resolves with WFT poll response or `PollWfError::ShutDown` if WFTs have been drained
     async fn workflow_poll_or_wfts_drained(
         &self,
     ) -> Result<Option<ValidPollWFTQResponse>, PollWfError> {
+        let mut shutdown_requested = self.shutdown_requested.clone();
         loop {
-            if *self.wfts_drained.borrow() {
-                debug!("Returning shutdown error");
-                return Err(PollWfError::ShutDown);
-            } else if *self.shutdown_requested.borrow() {
-                self.wfts_drained
-                    .clone()
-                    .changed()
-                    .await
-                    .expect("wfts_drained should not be dropped");
-            } else {
-                let mut shutdown_requested = self.shutdown_requested.clone();
-                tokio::select! {
-                    biased;
-
-                    r = self.workflow_poll()
-                        .map_err(Into::into) => match r {
-                         Err(PollWfError::ShutDown) => {},
-                        _ => return r,
-                    },
-                    _ = shutdown_requested.changed() => {},
-                }
-            };
+            tokio::select! {
+                biased;
+
+                r = self.workflow_poll().map_err(Into::into) => {
+                    if matches!(r, Err(PollWfError::ShutDown)) {
+                        // Don't actually return shutdown until workflow tasks are drained.
+                        // Outstanding tasks being completed will generate new pending activations
+                        // which will cause us to abort this function.
+                        self.all_wfts_drained().await;
+                    }
+                    return r
+                },
+                _ = shutdown_requested.changed() => {},
+            }
         }
     }
 
@@ -399,6 +477,17 @@ impl Worker {
     /// Returns `Ok(None)` in the event of a poll timeout, or if there was some gRPC error that
     /// callers can't do anything about.
     async fn workflow_poll(&self) -> Result<Option<ValidPollWFTQResponse>, PollWfError> {
+        // We can't say we're shut down if there are outstanding LAs, as they could end up WFT
+        // heartbeating which is a "new" workflow task that we need to accept and process as long as
+        // the LA is outstanding. Similarly, if we already have such tasks (from a WFT completion),
+        // then we must fetch them from the source before we can say workflow polling is shutdown.
+        if *self.shutdown_requested.borrow()
+            && !self.wf_task_source.has_tasks_from_complete()
+            && self.local_act_mgr.num_outstanding() == 0
+        {
+            return Err(PollWfError::ShutDown);
+        }
+
         let sem = self
             .workflows_semaphore
             .acquire()
@@ -406,8 +495,8 @@ impl Worker {
             .expect("outstanding workflow tasks semaphore not dropped");
 
         let res = self
-            .wf_task_poll_buffer
-            .poll()
+            .wf_task_source
+            .next_wft()
             .await
             .ok_or(PollWfError::ShutDown)??;
 
@@ -444,12 +533,12 @@ impl Worker {
     async fn apply_server_work(
         &self,
         work: ValidPollWFTQResponse,
-    ) -> Result<Option<WfActivation>, PollWfError> {
+    ) -> Result<Option<WorkflowActivation>, PollWfError> {
         let we = work.workflow_execution.clone();
         let tt = work.task_token.clone();
         let res = self
             .wft_manager
-            .apply_new_poll_resp(work, &self.server_gateway)
+            .apply_new_poll_resp(work, self.server_gateway.clone())
             .await;
         Ok(match res {
             NewWfTaskOutcome::IssueActivation(a) => {
@@ -462,10 +551,10 @@ impl Worker {
                 self.return_workflow_task_permit();
                 None
             }
-            NewWfTaskOutcome::Autocomplete => {
+            NewWfTaskOutcome::Autocomplete | NewWfTaskOutcome::LocalActsOutstanding => {
                 debug!(workflow_execution=?we,
-                       "No work for lang to perform after polling server. Sending autocomplete.");
-                self.complete_workflow_activation(WfActivationCompletion {
+                       "No new work for lang to perform after polling server");
+                self.complete_workflow_activation(WorkflowActivationCompletion {
                     task_queue: self.config.task_queue.clone(),
                     run_id: we.run_id,
                     status: Some(workflow_completion::Success::from_variants(vec![]).into()),
@@ -496,6 +585,7 @@ impl Worker {
                 self.request_wf_eviction(
                     &we.run_id,
                     format!("Error while applying poll response to workflow: {:?}", e),
+                    e.evict_reason(),
                 );
                 None
             }
@@ -509,7 +599,7 @@ impl Worker {
         &self,
         run_id: &str,
         success: workflow_completion::Success,
-    ) -> Result<bool, CompleteWfError> {
+    ) -> Result<WFTReportOutcome, CompleteWfError> {
         // Convert to wf commands
         let cmds = success
             .commands
@@ -523,13 +613,18 @@ impl Worker {
                 completion: None,
             })?;
 
-        match self.wft_manager.successful_activation(run_id, cmds).await {
+        match self
+            .wft_manager
+            .successful_activation(run_id, cmds, |acts| self.local_act_mgr.enqueue(acts))
+            .await
+        {
             Ok(Some(ServerCommandsWithWorkflowInfo {
                 task_token,
                 action:
                     ActivationAction::WftComplete {
                         commands,
                         query_responses,
+                        force_new_wft,
                     },
             })) => {
                 debug!("Sending commands to server: {:?}", &commands);
@@ -541,19 +636,33 @@ impl Worker {
                     commands,
                     query_responses,
                     sticky_attributes: None,
-                    return_new_workflow_task: false,
-                    force_create_new_workflow_task: false,
+                    return_new_workflow_task: force_new_wft,
+                    force_create_new_workflow_task: force_new_wft,
                 };
                 let sticky_attrs = self.get_sticky_attrs();
+                // Do not return new WFT if we would not cache, because returned new WFTs are always
+                // partial.
+                if sticky_attrs.is_none() {
+                    completion.return_new_workflow_task = false;
+                }
                 completion.sticky_attributes = sticky_attrs;
+
                 self.handle_wft_reporting_errs(run_id, || async {
-                    self.server_gateway
+                    let maybe_wft = self
+                        .server_gateway
                         .complete_workflow_task(completion)
                         .instrument(span!(tracing::Level::DEBUG, "Complete WFT call"))
-                        .await
+                        .await?;
+                    if let Some(wft) = maybe_wft.workflow_task {
+                        self.wf_task_source.add_wft_from_completion(wft);
+                    }
+                    Ok(())
                 })
                 .await?;
-                Ok(true)
+                Ok(WFTReportOutcome {
+                    reported_to_server: true,
+                    failed: false,
+                })
             }
             Ok(Some(ServerCommandsWithWorkflowInfo {
                 task_token,
@@ -563,9 +672,15 @@ impl Worker {
                 self.server_gateway
                     .respond_legacy_query(task_token, result)
                     .await?;
-                Ok(true)
+                Ok(WFTReportOutcome {
+                    reported_to_server: true,
+                    failed: false,
+                })
             }
-            Ok(None) => Ok(false),
+            Ok(None) => Ok(WFTReportOutcome {
+                reported_to_server: false,
+                failed: false,
+            }),
             Err(update_err) => {
                 // Automatically fail the workflow task in the event we couldn't update machines
                 let fail_cause = if matches!(&update_err.source, WFMachinesError::Nondeterminism(_))
@@ -574,30 +689,14 @@ impl Worker {
                 } else {
                     WorkflowTaskFailedCause::Unspecified
                 };
-
-                warn!(run_id, error=?update_err, "Failing workflow task");
-
-                if let Some(ref tt) = update_err.task_token {
-                    let wft_fail_str = format!("{:?}", update_err);
-                    self.handle_wft_reporting_errs(run_id, || async {
-                        self.server_gateway
-                            .fail_workflow_task(
-                                tt.clone(),
-                                fail_cause,
-                                Some(Failure::application_failure(wft_fail_str.clone(), false)),
-                            )
-                            .await
-                    })
-                    .await?;
-                    // We must evict the workflow since we've failed a WFT
-                    self.request_wf_eviction(
-                        run_id,
-                        format!("Workflow task failure: {}", wft_fail_str),
-                    );
-                    Ok(true)
-                } else {
-                    Ok(false)
-                }
+                let wft_fail_str = format!("{:?}", update_err);
+                self.wf_activation_failed(
+                    run_id,
+                    fail_cause,
+                    update_err.evict_reason(),
+                    Failure::application_failure(wft_fail_str.clone(), false).into(),
+                )
+                .await
             }
         }
     }
@@ -608,39 +707,45 @@ impl Worker {
     async fn wf_activation_failed(
         &self,
         run_id: &str,
+        cause: WorkflowTaskFailedCause,
+        reason: EvictionReason,
         failure: workflow_completion::Failure,
-    ) -> Result<bool, CompleteWfError> {
-        Ok(match self.wft_manager.failed_activation(run_id) {
-            FailedActivationOutcome::Report(tt) => {
-                self.handle_wft_reporting_errs(run_id, || async {
-                    self.server_gateway
-                        .fail_workflow_task(
-                            tt,
-                            WorkflowTaskFailedCause::Unspecified,
-                            failure.failure.map(Into::into),
-                        )
-                        .await
-                })
-                .await?;
-                true
-            }
-            FailedActivationOutcome::ReportLegacyQueryFailure(task_token) => {
-                self.server_gateway
-                    .respond_legacy_query(task_token, legacy_query_failure(failure))
+    ) -> Result<WFTReportOutcome, CompleteWfError> {
+        Ok(
+            match self.wft_manager.failed_activation(
+                run_id,
+                reason,
+                format!("Workflow activation completion failed: {:?}", failure),
+            ) {
+                FailedActivationOutcome::Report(tt) => {
+                    warn!(run_id, failure=?failure, "Failing workflow activation");
+                    self.handle_wft_reporting_errs(run_id, || async {
+                        self.server_gateway
+                            .fail_workflow_task(tt, cause, failure.failure.map(Into::into))
+                            .await
+                    })
                     .await?;
-                true
-            }
-            FailedActivationOutcome::NoReport => false,
-        })
-    }
-
-    fn after_workflow_activation(&self, run_id: &str, did_complete_wft: bool) {
-        self.wft_manager.after_wft_report(run_id);
-        if did_complete_wft {
-            self.return_workflow_task_permit();
-        }
-        self.wft_manager.on_activation_done(run_id);
-        self.maybe_notify_wtfs_drained();
+                    WFTReportOutcome {
+                        reported_to_server: true,
+                        failed: true,
+                    }
+                }
+                FailedActivationOutcome::ReportLegacyQueryFailure(task_token) => {
+                    warn!(run_id, failure=?failure, "Failing legacy query request");
+                    self.server_gateway
+                        .respond_legacy_query(task_token, legacy_query_failure(failure))
+                        .await?;
+                    WFTReportOutcome {
+                        reported_to_server: true,
+                        failed: true,
+                    }
+                }
+                FailedActivationOutcome::NoReport => WFTReportOutcome {
+                    reported_to_server: false,
+                    failed: true,
+                },
+            },
+        )
     }
 
     /// Handle server errors from either completing or failing a workflow task. Returns any errors
@@ -653,20 +758,20 @@ impl Worker {
     where
         Fut: Future<Output = Result<T, tonic::Status>>,
     {
-        let mut should_evict = false;
+        let mut should_evict = None;
         let res = match completer().await {
             Err(err) => {
                 match err.code() {
                     // Silence unhandled command errors since the lang SDK cannot do anything about
                     // them besides poll again, which it will do anyway.
                     tonic::Code::InvalidArgument if err.message() == "UnhandledCommand" => {
-                        warn!(error = %err, "Unhandled command response when completing");
-                        should_evict = true;
+                        warn!(error = %err, run_id, "Unhandled command response when completing");
+                        should_evict = Some(EvictionReason::UnhandledCommand);
                         Ok(())
                     }
                     tonic::Code::NotFound => {
-                        warn!(error = %err, "Task not found when completing");
-                        should_evict = true;
+                        warn!(error = %err, run_id, "Task not found when completing");
+                        should_evict = Some(EvictionReason::TaskNotFound);
                         Ok(())
                     }
                     _ => Err(err),
@@ -674,12 +779,46 @@ impl Worker {
             }
             _ => Ok(()),
         };
-        if should_evict {
-            self.request_wf_eviction(run_id, "Error reporting WFT to server");
+        if let Some(reason) = should_evict {
+            self.request_wf_eviction(run_id, "Error reporting WFT to server", reason);
         }
         res.map_err(Into::into)
     }
 
+    async fn complete_local_act(
+        &self,
+        la_res: LocalActivityExecutionResult,
+        info: LocalInFlightActInfo,
+        backoff: Option<prost_types::Duration>,
+    ) {
+        self.notify_local_result(
+            &info.la_info.workflow_exec_info.run_id,
+            LocalResolution::LocalActivity(LocalActivityResolution {
+                seq: info.la_info.schedule_cmd.seq,
+                result: la_res,
+                runtime: info.dispatch_time.elapsed(),
+                attempt: info.attempt,
+                backoff,
+                original_schedule_time: Some(info.la_info.schedule_time),
+            }),
+        )
+        .await
+    }
+
+    async fn notify_local_result(&self, run_id: &str, res: LocalResolution) {
+        if let Err(e) = self.wft_manager.notify_of_local_result(run_id, res).await {
+            error!(
+                "Problem with local resolution on run {}: {:?} -- will evict the workflow",
+                run_id, e
+            );
+            self.request_wf_eviction(
+                run_id,
+                "Issue while processing local resolution",
+                e.evict_reason(),
+            );
+        }
+    }
+
     /// Return the sticky execution attributes that should be used to complete workflow tasks
     /// for this worker (if any).
     fn get_sticky_attrs(&self) -> Option<StickyExecutionAttributes> {
@@ -696,102 +835,89 @@ impl Worker {
             })
     }
 
-    /// A future that resolves to true the shutdown flag has been set to true, false is simply
-    /// a signal that a poll loop should be restarted. Only meant to be called from polling funcs.
-    async fn shutdown_notifier(&self) {
-        if *self.shutdown_requested.borrow() {
-            return;
+    /// Resolves when there are no more outstanding WFTs
+    async fn all_wfts_drained(&self) {
+        while self.outstanding_workflow_tasks() != 0 {
+            self.wfts_drained_notify.notified().await;
         }
-        let _ = self.shutdown_requested.clone().changed().await;
     }
 }
 
-impl WorkerConfig {
-    fn max_nonsticky_polls(&self) -> usize {
-        ((self.max_concurrent_wft_polls as f32 * self.nonsticky_to_sticky_poll_ratio) as usize)
-            .max(1)
-    }
-    fn max_sticky_polls(&self) -> usize {
-        self.max_concurrent_wft_polls
-            .saturating_sub(self.max_nonsticky_polls())
-            .max(1)
-    }
+struct WFTReportOutcome {
+    reported_to_server: bool,
+    failed: bool,
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::{pollers::MockServerGatewayApis, test_help::fake_sg_opts};
+    use temporal_client::mocks::mock_gateway;
     use temporal_sdk_core_protos::temporal::api::workflowservice::v1::PollActivityTaskQueueResponse;
 
     #[tokio::test]
     async fn activity_timeouts_dont_eat_permits() {
-        let mut mock_gateway = MockServerGatewayApis::new();
+        let mut mock_gateway = mock_gateway();
         mock_gateway
             .expect_poll_activity_task()
             .returning(|_| Ok(PollActivityTaskQueueResponse::default()));
-        let gwref = GatewayRef::new(Arc::new(mock_gateway), fake_sg_opts());
 
         let cfg = WorkerConfigBuilder::default()
             .task_queue("whatever")
             .max_outstanding_activities(5_usize)
             .build()
             .unwrap();
-        let worker = Worker::new(cfg, None, Arc::new(gwref), Default::default());
+        let worker = Worker::new(cfg, None, Arc::new(mock_gateway), Default::default());
         assert_eq!(worker.activity_poll().await.unwrap(), None);
         assert_eq!(worker.at_task_mgr.unwrap().remaining_activity_capacity(), 5);
     }
 
     #[tokio::test]
     async fn workflow_timeouts_dont_eat_permits() {
-        let mut mock_gateway = MockServerGatewayApis::new();
+        let mut mock_gateway = mock_gateway();
         mock_gateway
             .expect_poll_workflow_task()
             .returning(|_, _| Ok(PollWorkflowTaskQueueResponse::default()));
-        let gwref = GatewayRef::new(Arc::new(mock_gateway), fake_sg_opts());
 
         let cfg = WorkerConfigBuilder::default()
             .task_queue("whatever")
             .max_outstanding_workflow_tasks(5_usize)
             .build()
             .unwrap();
-        let worker = Worker::new(cfg, None, Arc::new(gwref), Default::default());
+        let worker = Worker::new(cfg, None, Arc::new(mock_gateway), Default::default());
         assert_eq!(worker.workflow_poll().await.unwrap(), None);
         assert_eq!(worker.workflows_semaphore.available_permits(), 5);
     }
 
     #[tokio::test]
     async fn activity_errs_dont_eat_permits() {
-        let mut mock_gateway = MockServerGatewayApis::new();
+        let mut mock_gateway = mock_gateway();
         mock_gateway
             .expect_poll_activity_task()
             .returning(|_| Err(tonic::Status::internal("ahhh")));
-        let gwref = GatewayRef::new(Arc::new(mock_gateway), fake_sg_opts());
 
         let cfg = WorkerConfigBuilder::default()
             .task_queue("whatever")
             .max_outstanding_activities(5_usize)
             .build()
             .unwrap();
-        let worker = Worker::new(cfg, None, Arc::new(gwref), Default::default());
+        let worker = Worker::new(cfg, None, Arc::new(mock_gateway), Default::default());
         assert!(worker.activity_poll().await.is_err());
         assert_eq!(worker.at_task_mgr.unwrap().remaining_activity_capacity(), 5);
     }
 
     #[tokio::test]
     async fn workflow_errs_dont_eat_permits() {
-        let mut mock_gateway = MockServerGatewayApis::new();
+        let mut mock_gateway = mock_gateway();
         mock_gateway
             .expect_poll_workflow_task()
             .returning(|_, _| Err(tonic::Status::internal("ahhh")));
-        let gwref = GatewayRef::new(Arc::new(mock_gateway), fake_sg_opts());
 
         let cfg = WorkerConfigBuilder::default()
             .task_queue("whatever")
             .max_outstanding_workflow_tasks(5_usize)
             .build()
             .unwrap();
-        let worker = Worker::new(cfg, None, Arc::new(gwref), Default::default());
+        let worker = Worker::new(cfg, None, Arc::new(mock_gateway), Default::default());
         assert!(worker.workflow_poll().await.is_err());
         assert_eq!(worker.workflows_semaphore.available_permits(), 5);
     }