@temporalio/core-bridge 0.15.0 → 0.17.0

package/sdk-core/src/machines/child_workflow_state_machine.rs

@@ -54,6 +54,12 @@ fsm! {
     Started --(ChildWorkflowExecutionTimedOut(i32), shared on_child_workflow_execution_timed_out) --> TimedOut;
     Started --(ChildWorkflowExecutionCancelled, shared on_child_workflow_execution_cancelled) --> Cancelled;
     Started --(ChildWorkflowExecutionTerminated, shared on_child_workflow_execution_terminated) --> Terminated;
+
+    // Ignore any spurious cancellations after resolution
+    Cancelled --(Cancel) --> Cancelled;
+    Failed --(Cancel) --> Failed;
+    TimedOut --(Cancel) --> TimedOut;
+    Completed --(Cancel) --> Completed;
 }
 
 pub struct ChildWorkflowExecutionStartedEvent {
@@ -288,7 +294,7 @@ pub(super) struct Terminated {}
 #[derive(Default, Clone)]
 pub(super) struct TimedOut {}
 
-#[derive(Default, Clone)]
+#[derive(Default, Clone, Debug)]
 pub(super) struct SharedState {
     initiated_event_id: i64,
     started_event_id: i64,
@@ -630,6 +636,7 @@ mod test {
     };
     use anyhow::anyhow;
     use rstest::{fixture, rstest};
+    use std::mem::discriminant;
     use temporal_sdk_core_protos::coresdk::{
         child_workflow::child_workflow_result,
         workflow_activation::resolve_child_workflow_execution_start::Status as StartStatus,
@@ -643,11 +650,11 @@ mod test {
     }
 
     impl Expectation {
-        fn try_from_u8(x: u8) -> Option<Self> {
+        const fn try_from_u8(x: u8) -> Option<Self> {
             Some(match x {
-                0 => Expectation::Success,
-                1 => Expectation::Failure,
-                2 => Expectation::StartFailure,
+                0 => Self::Success,
+                1 => Self::Failure,
+                2 => Self::StartFailure,
                 _ => return None,
             })
         }
@@ -693,12 +700,7 @@ mod test {
         };
         match (
             expectation,
-            start_res
-                .as_started()
-                .unwrap()
-                .result(&mut ctx)
-                .await
-                .status,
+            start_res.into_started().unwrap().result().await.status,
         ) {
             (Expectation::Success, Some(child_workflow_result::Status::Completed(_))) => {
                 Ok(().into())
@@ -803,4 +805,22 @@ mod test {
         );
         wfm.shutdown().await.unwrap();
     }
+
+    #[test]
+    fn cancels_ignored_terminal() {
+        for state in [
+            ChildWorkflowMachineState::Cancelled(Cancelled {}),
+            Failed {}.into(),
+            TimedOut {}.into(),
+            Completed {}.into(),
+        ] {
+            let mut s = ChildWorkflowMachine {
+                state: state.clone(),
+                shared_state: Default::default(),
+            };
+            let cmds = s.cancel().unwrap();
+            assert_eq!(cmds.len(), 0);
+            assert_eq!(discriminant(&state), discriminant(&s.state));
+        }
+    }
 }

package/sdk-core/src/machines/complete_workflow_state_machine.rs

@@ -117,7 +117,7 @@ pub(super) struct CompleteWorkflowCommandRecorded {}
 
 impl From<CompleteWorkflowCommandCreated> for CompleteWorkflowCommandRecorded {
     fn from(_: CompleteWorkflowCommandCreated) -> Self {
-        Default::default()
+        Self::default()
     }
 }
 

package/sdk-core/src/machines/continue_as_new_workflow_state_machine.rs

@@ -64,7 +64,7 @@ impl Created {
 
 impl From<ContinueAsNewWorkflowCommandCreated> for ContinueAsNewWorkflowCommandRecorded {
     fn from(_: ContinueAsNewWorkflowCommandCreated) -> Self {
-        Default::default()
+        Self::default()
     }
 }
 

package/sdk-core/src/machines/mod.rs

@@ -77,38 +77,34 @@ impl TryFrom<WorkflowCommand> for WFCommand {
 
     fn try_from(c: WorkflowCommand) -> Result<Self, Self::Error> {
         match c.variant.ok_or(EmptyWorkflowCommandErr)? {
-            workflow_command::Variant::StartTimer(s) => Ok(WFCommand::AddTimer(s)),
-            workflow_command::Variant::CancelTimer(s) => Ok(WFCommand::CancelTimer(s)),
-            workflow_command::Variant::ScheduleActivity(s) => Ok(WFCommand::AddActivity(s)),
+            workflow_command::Variant::StartTimer(s) => Ok(Self::AddTimer(s)),
+            workflow_command::Variant::CancelTimer(s) => Ok(Self::CancelTimer(s)),
+            workflow_command::Variant::ScheduleActivity(s) => Ok(Self::AddActivity(s)),
             workflow_command::Variant::RequestCancelActivity(s) => {
-                Ok(WFCommand::RequestCancelActivity(s))
+                Ok(Self::RequestCancelActivity(s))
             }
             workflow_command::Variant::CompleteWorkflowExecution(c) => {
-                Ok(WFCommand::CompleteWorkflow(c))
+                Ok(Self::CompleteWorkflow(c))
             }
-            workflow_command::Variant::FailWorkflowExecution(s) => Ok(WFCommand::FailWorkflow(s)),
-            workflow_command::Variant::RespondToQuery(s) => Ok(WFCommand::QueryResponse(s)),
+            workflow_command::Variant::FailWorkflowExecution(s) => Ok(Self::FailWorkflow(s)),
+            workflow_command::Variant::RespondToQuery(s) => Ok(Self::QueryResponse(s)),
             workflow_command::Variant::ContinueAsNewWorkflowExecution(s) => {
-                Ok(WFCommand::ContinueAsNew(s))
+                Ok(Self::ContinueAsNew(s))
             }
-            workflow_command::Variant::CancelWorkflowExecution(s) => {
-                Ok(WFCommand::CancelWorkflow(s))
-            }
-            workflow_command::Variant::SetPatchMarker(s) => Ok(WFCommand::SetPatchMarker(s)),
+            workflow_command::Variant::CancelWorkflowExecution(s) => Ok(Self::CancelWorkflow(s)),
+            workflow_command::Variant::SetPatchMarker(s) => Ok(Self::SetPatchMarker(s)),
             workflow_command::Variant::StartChildWorkflowExecution(s) => {
-                Ok(WFCommand::AddChildWorkflow(s))
+                Ok(Self::AddChildWorkflow(s))
             }
             workflow_command::Variant::RequestCancelExternalWorkflowExecution(s) => {
-                Ok(WFCommand::RequestCancelExternalWorkflow(s))
+                Ok(Self::RequestCancelExternalWorkflow(s))
             }
             workflow_command::Variant::SignalExternalWorkflowExecution(s) => {
-                Ok(WFCommand::SignalExternalWorkflow(s))
-            }
-            workflow_command::Variant::CancelSignalWorkflow(s) => {
-                Ok(WFCommand::CancelSignalWorkflow(s))
+                Ok(Self::SignalExternalWorkflow(s))
             }
+            workflow_command::Variant::CancelSignalWorkflow(s) => Ok(Self::CancelSignalWorkflow(s)),
             workflow_command::Variant::CancelUnstartedChildWorkflowExecution(s) => {
-                Ok(WFCommand::CancelUnstartedChild(s))
+                Ok(Self::CancelUnstartedChild(s))
             }
         }
     }
@@ -170,9 +166,7 @@ where
         + Clone
         + Send
         + 'static,
-    <SM as StateMachine>::Event: TryFrom<HistoryEvent>,
-    <SM as StateMachine>::Event: TryFrom<CommandType>,
-    <SM as StateMachine>::Event: Display,
+    <SM as StateMachine>::Event: TryFrom<HistoryEvent> + TryFrom<CommandType> + Display,
     WFMachinesError: From<<<SM as StateMachine>::Event as TryFrom<HistoryEvent>>::Error>,
     <SM as StateMachine>::Command: Debug + Display,
     <SM as StateMachine>::State: Display,
@@ -248,7 +242,8 @@ where
         let res = self.cancel();
         res.map_err(|e| match e {
             MachineError::InvalidTransition => WFMachinesError::Fatal(format!(
-                "Invalid transition while attempting to cancel in {}",
+                "Invalid transition while attempting to cancel {} in {}",
+                self.name(),
                 self.state(),
             )),
             MachineError::Underlying(e) => e.into(),

package/sdk-core/src/machines/patch_state_machine.rs

@@ -107,7 +107,7 @@ impl PatchMachine {
                 .into(),
             ),
         };
-        let mut machine = PatchMachine {
+        let mut machine = Self {
             state: initial_state,
             shared_state: state,
         };
@@ -210,7 +210,7 @@ impl TryFrom<HistoryEvent> for PatchMachineEvents {
 
     fn try_from(e: HistoryEvent) -> Result<Self, Self::Error> {
         match e.get_changed_marker_details() {
-            Some((id, _)) => Ok(PatchMachineEvents::MarkerRecorded(id)),
+            Some((id, _)) => Ok(Self::MarkerRecorded(id)),
             _ => Err(WFMachinesError::Nondeterminism(format!(
                 "Change machine cannot handle this event: {}",
                 e
@@ -422,7 +422,9 @@ mod tests {
             commands[0].command_type,
             CommandType::ScheduleActivityTask as i32
         );
-        let act = if !replaying {
+        let act = if replaying {
+            wfm.get_next_activation().await
+        } else {
             // Feed more history
             wfm.new_history(
                 patch_marker_single_activity(marker_type)
@@ -431,8 +433,6 @@ mod tests {
                     .into(),
             )
             .await
-        } else {
-            wfm.get_next_activation().await
         };
 
         if marker_type == MarkerType::Deprecated {
@@ -514,7 +514,9 @@ mod tests {
             if activity_id == expected_activity_id
         );
 
-        let act = if !replaying {
+        let act = if replaying {
+            wfm.get_next_activation().await
+        } else {
             // Feed more history. Since we are *not* replaying, we *always* "have" the change
             // and the history should have the has-change timer. v3 of course always has the change
             // regardless.
@@ -525,8 +527,6 @@ mod tests {
                     .into(),
             )
             .await
-        } else {
-            wfm.get_next_activation().await
         };
 
         let act = act.unwrap();

package/sdk-core/src/machines/signal_external_state_machine.rs

@@ -42,6 +42,10 @@ fsm! {
     SignalExternalCommandRecorded
       --(SignalExternalWorkflowExecutionFailed(SignalExternalWorkflowExecutionFailedCause),
          on_signal_external_workflow_execution_failed) --> Failed;
+
+    // Ignore any spurious cancellations after resolution
+    Cancelled --(Cancel) --> Cancelled;
+    Signaled --(Cancel) --> Signaled;
 }
 
 #[derive(Default, Clone)]
@@ -262,7 +266,7 @@ impl Cancellable for SignalExternalMachine {
                         ..Default::default()
                     }),
                 }
-                .into()]
+                .into()];
             }
             Some(_) => panic!("Signal external machine cancel produced unexpected result"),
             None => (),
@@ -285,6 +289,7 @@ mod tests {
         test_help::TestHistoryBuilder,
         workflow::managed_wf::ManagedWFFunc,
     };
+    use std::mem::discriminant;
     use temporal_sdk_core_protos::coresdk::workflow_activation::{
         wf_activation_job, WfActivationJob,
     };
@@ -382,4 +387,20 @@ mod tests {
         );
         wfm.shutdown().await.unwrap();
     }
+
+    #[test]
+    fn cancels_ignored_terminal() {
+        for state in [
+            SignalExternalMachineState::Cancelled(Cancelled {}),
+            Signaled {}.into(),
+        ] {
+            let mut s = SignalExternalMachine {
+                state: state.clone(),
+                shared_state: Default::default(),
+            };
+            let cmds = s.cancel().unwrap();
+            assert_eq!(cmds.len(), 0);
+            assert_eq!(discriminant(&state), discriminant(&s.state));
+        }
+    }
 }

package/sdk-core/src/machines/timer_state_machine.rs

@@ -39,6 +39,10 @@ fsm! {
         --(CommandCancelTimer, on_command_cancel_timer) --> CancelTimerCommandSent;
 
     CancelTimerCommandSent --(TimerCanceled) --> Canceled;
+
+    // Ignore any spurious cancellations after resolution
+    Canceled --(Cancel) --> Canceled;
+    Fired --(Cancel) --> Fired;
 }
 
 #[derive(Debug, derive_more::Display)]
@@ -195,13 +199,13 @@ impl StartCommandRecorded {
         dat: SharedState,
         attrs: TimerFiredEventAttributes,
     ) -> TimerMachineTransition<Fired> {
-        if dat.attrs.seq.to_string() != attrs.timer_id {
+        if dat.attrs.seq.to_string() == attrs.timer_id {
+            TransitionResult::ok(vec![TimerMachineCommand::Complete], Fired::default())
+        } else {
             TransitionResult::Err(WFMachinesError::Fatal(format!(
                 "Timer fired event did not have expected timer id {}, it was {}!",
                 dat.attrs.seq, attrs.timer_id
             )))
-        } else {
-            TransitionResult::ok(vec![TimerMachineCommand::Complete], Fired::default())
         }
     }
 
@@ -275,6 +279,7 @@ mod test {
         workflow::managed_wf::ManagedWFFunc,
     };
     use rstest::{fixture, rstest};
+    use std::mem::discriminant;
     use std::time::Duration;
 
     #[fixture]
@@ -420,4 +425,17 @@ mod test {
         assert_eq!(commands.len(), 0);
         wfm.shutdown().await.unwrap();
     }
+
+    #[test]
+    fn cancels_ignored_terminal() {
+        for state in [TimerMachineState::Canceled(Canceled {}), Fired {}.into()] {
+            let mut s = TimerMachine {
+                state: state.clone(),
+                shared_state: Default::default(),
+            };
+            let cmds = s.cancel().unwrap();
+            assert_eq!(cmds.len(), 0);
+            assert_eq!(discriminant(&state), discriminant(&s.state));
+        }
+    }
 }

package/sdk-core/src/machines/transition_coverage.rs

@@ -79,7 +79,7 @@ mod machine_coverage_report {
         workflow_task_state_machine::WorkflowTaskMachine,
     };
     use rustfsm::StateMachine;
-    use std::{fs::File, io::Write, ops::Deref};
+    use std::{fs::File, io::Write};
 
     // This "test" needs to exist so that we have a way to join the spawned thread. Otherwise
     // it'll just get abandoned.
@@ -89,7 +89,7 @@ mod machine_coverage_report {
     #[ignore]
     fn reporter() {
         // Make sure thread handle exists
-        let _ = COVERAGE_SENDER.deref();
+        let _ = &*COVERAGE_SENDER;
         // Join it
         THREAD_HANDLE
             .lock()
@@ -124,7 +124,7 @@ mod machine_coverage_report {
                 m @ "WorkflowTaskMachine" => cover_transitions(m, &mut wf_task, coverage),
                 m @ "FailWorkflowMachine" => cover_transitions(m, &mut fail_wf, coverage),
                 m @ "ContinueAsNewWorkflowMachine" => {
-                    cover_transitions(m, &mut cont_as_new, coverage)
+                    cover_transitions(m, &mut cont_as_new, coverage);
                 }
                 m @ "CancelWorkflowMachine" => cover_transitions(m, &mut cancel_wf, coverage),
                 m @ "PatchMachine" => cover_transitions(m, &mut version, coverage),

package/sdk-core/src/machines/workflow_machines.rs

@@ -150,7 +150,7 @@ where
     T: Into<wf_activation_job::Variant>,
 {
     fn from(v: T) -> Self {
-        MachineResponse::PushWFJob(v.into())
+        Self::PushWFJob(v.into())
     }
 }
 
@@ -171,7 +171,7 @@ pub(crate) enum WFMachinesError {
 
 impl From<TimestampOutOfSystemRangeError> for WFMachinesError {
     fn from(_: TimestampOutOfSystemRangeError) -> Self {
-        WFMachinesError::Fatal("Could not decode timestamp".to_string())
+        Self::Fatal("Could not decode timestamp".to_string())
     }
 }
 
@@ -210,7 +210,7 @@ impl WorkflowMachines {
     }
 
     /// Returns true if workflow has seen a terminal command
-    pub(crate) fn workflow_is_finished(&self) -> bool {
+    pub(crate) const fn workflow_is_finished(&self) -> bool {
         self.workflow_end_time.is_some()
     }
 
@@ -486,7 +486,7 @@ impl WorkflowMachines {
     }
 
     fn set_current_time(&mut self, time: SystemTime) -> SystemTime {
-        if self.current_wf_time.map(|t| t < time).unwrap_or(true) {
+        if self.current_wf_time.map_or(true, |t| t < time) {
             self.current_wf_time = Some(time);
         }
         self.current_wf_time
@@ -503,7 +503,7 @@ impl WorkflowMachines {
         let results = self.drive_me.fetch_workflow_iteration_output().await;
         let jobs = self.handle_driven_results(results)?;
         let has_new_lang_jobs = !jobs.is_empty();
-        for job in jobs.into_iter() {
+        for job in jobs {
             self.drive_me.send_job(job);
         }
         self.prepare_commands()?;
@@ -626,15 +626,15 @@ impl WorkflowMachines {
                     .machine_mut(c.machine)
                     .handle_command(c.command.command_type())?;
                 self.process_machine_responses(c.machine, machine_responses)?;
+                self.commands.push_back(c);
             }
-            self.commands.push_back(c);
         }
         debug!(commands = %self.commands.display(), "prepared commands");
         Ok(())
     }
 
     /// After a machine handles either an event or a command, it produces [MachineResponses] which
-    /// this function uses to drive sending jobs to lang, trigging new workflow tasks, etc.
+    /// this function uses to drive sending jobs to lang, triggering new workflow tasks, etc.
     fn process_machine_responses(
         &mut self,
         sm: MachineKey,
@@ -695,7 +695,7 @@ impl WorkflowMachines {
                     self.current_wf_task_commands.push_back(timer);
                 }
                 WFCommand::CancelTimer(attrs) => {
-                    jobs.extend(self.process_cancellation(CommandID::Timer(attrs.seq))?)
+                    jobs.extend(self.process_cancellation(CommandID::Timer(attrs.seq))?);
                 }
                 WFCommand::AddActivity(attrs) => {
                     let seq = attrs.seq;
@@ -705,7 +705,7 @@ impl WorkflowMachines {
                     self.current_wf_task_commands.push_back(activity);
                 }
                 WFCommand::RequestCancelActivity(attrs) => {
-                    jobs.extend(self.process_cancellation(CommandID::Activity(attrs.seq))?)
+                    jobs.extend(self.process_cancellation(CommandID::Activity(attrs.seq))?);
                 }
                 WFCommand::CompleteWorkflow(attrs) => {
                     self.metrics.wf_completed();
@@ -815,7 +815,7 @@ impl WorkflowMachines {
                     self.current_wf_task_commands.push_back(sigm);
                 }
                 WFCommand::CancelSignalWorkflow(attrs) => {
-                    jobs.extend(self.process_cancellation(CommandID::SignalExternal(attrs.seq))?)
+                    jobs.extend(self.process_cancellation(CommandID::SignalExternal(attrs.seq))?);
                 }
                 WFCommand::QueryResponse(_) => {
                     // Nothing to do here, queries are handled above the machine level
@@ -840,7 +840,7 @@ impl WorkflowMachines {
                     self.current_wf_task_commands.push_back(CommandAndMachine {
                         command: c,
                         machine: m_key,
-                    })
+                    });
                 }
                 MachineResponse::PushWFJob(j) => {
                     jobs.push(j);

package/sdk-core/src/pending_activations.rs

@@ -68,15 +68,14 @@ impl PendingActivations {
         let mut inner = self.inner.write();
         let mut key_queue = inner.queue.iter().copied();
         let maybe_key = key_queue.position(|k| {
-            if let Some(activation) = inner.activations.get(k) {
-                predicate(&activation.run_id)
-            } else {
-                false
-            }
+            inner
+                .activations
+                .get(k)
+                .map_or(false, |activation| predicate(&activation.run_id))
         });
 
         let maybe_key = maybe_key.map(|pos| inner.queue.remove(pos).unwrap());
-        if let Some(key) = maybe_key {
+        maybe_key.and_then(|key| {
             if let Some(pa) = inner.activations.remove(key) {
                 inner.by_run_id.remove(&pa.run_id);
                 Some(pa)
@@ -87,9 +86,7 @@ impl PendingActivations {
                 drop(inner); // Will deadlock when we recurse w/o this
                 self.pop()
             }
-        } else {
-            None
-        }
+        })
     }
 
     pub fn pop(&self) -> Option<WfActivation> {
@@ -119,17 +116,19 @@ fn merge_joblists(
         .as_mut_slice()
         .sort_by(evictions_always_last_compare);
     // Drop any duplicate evictions
-    let truncate_len = if let Some(last_non_evict_job) = existing_list.iter().rev().position(|j| {
-        !matches!(
-            j.variant,
-            Some(wf_activation_job::Variant::RemoveFromCache(_))
-        )
-    }) {
-        existing_list.len() - last_non_evict_job + 1
-    } else {
-        1
-    };
-    existing_list.truncate(truncate_len)
+    let truncate_len = existing_list
+        .iter()
+        .rev()
+        .position(|j| {
+            !matches!(
+                j.variant,
+                Some(wf_activation_job::Variant::RemoveFromCache(_))
+            )
+        })
+        .map_or(1, |last_non_evict_job| {
+            existing_list.len() - last_non_evict_job + 1
+        });
+    existing_list.truncate(truncate_len);
 }
 
 fn evictions_always_last_compare(a: &WfActivationJob, b: &WfActivationJob) -> Ordering {

package/sdk-core/src/pollers/gateway.rs

@@ -165,7 +165,7 @@ impl Default for RetryConfig {
 }
 
 impl RetryConfig {
-    pub(crate) fn poll_retry_policy() -> Self {
+    pub(crate) const fn poll_retry_policy() -> Self {
         Self {
             initial_interval: Duration::from_millis(200),
             randomization_factor: 0.2,
@@ -179,7 +179,7 @@ impl RetryConfig {
 
 impl From<RetryConfig> for ExponentialBackoff {
     fn from(c: RetryConfig) -> Self {
-        ExponentialBackoff {
+        Self {
             current_interval: c.initial_interval,
             initial_interval: c.initial_interval,
             randomization_factor: c.randomization_factor,
@@ -269,7 +269,7 @@ impl Interceptor for ServiceCallInterceptor {
                 .parse()
                 .unwrap_or_else(|_| MetadataValue::from_static("")),
         );
-        for (k, v) in self.opts.static_headers.iter() {
+        for (k, v) in &self.opts.static_headers {
             if let (Ok(k), Ok(v)) = (MetadataKey::from_str(k), MetadataValue::from_str(v)) {
                 metadata.insert(k, v);
             }
@@ -355,8 +355,11 @@ pub trait ServerGatewayApis {
 
     /// Fetch new workflow tasks from the provided queue. Should block indefinitely if there is no
     /// work.
-    async fn poll_workflow_task(&self, task_queue: String)
-        -> Result<PollWorkflowTaskQueueResponse>;
+    async fn poll_workflow_task(
+        &self,
+        task_queue: String,
+        is_sticky: bool,
+    ) -> Result<PollWorkflowTaskQueueResponse>;
 
     /// Fetch new activity tasks from the provided queue. Should block indefinitely if there is no
     /// work.
@@ -516,12 +519,17 @@ impl ServerGatewayApis for ServerGateway {
     async fn poll_workflow_task(
         &self,
         task_queue: String,
+        is_sticky: bool,
     ) -> Result<PollWorkflowTaskQueueResponse> {
         let mut request = PollWorkflowTaskQueueRequest {
             namespace: self.opts.namespace.clone(),
             task_queue: Some(TaskQueue {
                 name: task_queue,
-                kind: TaskQueueKind::Unspecified as i32,
+                kind: if is_sticky {
+                    TaskQueueKind::Sticky
+                } else {
+                    TaskQueueKind::Normal
+                } as i32,
             }),
             identity: self.opts.identity.clone(),
             binary_checksum: self.opts.worker_binary_id.clone(),
@@ -894,7 +902,7 @@ mockall::mock! {
         ) -> impl Future<Output = Result<StartWorkflowExecutionResponse>> + Send + 'b
             where 'a: 'b, Self: 'b;
 
-        fn poll_workflow_task<'a, 'b>(&'a self, task_queue: String)
+        fn poll_workflow_task<'a, 'b>(&'a self, task_queue: String, is_sticky: bool)
             -> impl Future<Output = Result<PollWorkflowTaskQueueResponse>> + Send + 'b
             where 'a: 'b, Self: 'b;
 

package/sdk-core/src/pollers/poll_buffer.rs

@@ -152,7 +152,7 @@ where
 
     async fn shutdown_box(self: Box<Self>) {
         let this = *self;
-        this.shutdown().await
+        this.shutdown().await;
     }
 }
 
@@ -192,7 +192,7 @@ impl Poller<PollWorkflowTaskQueueResponse> for WorkflowTaskPoller {
 
     async fn shutdown_box(self: Box<Self>) {
         let this = *self;
-        this.shutdown().await
+        this.shutdown().await;
     }
 }
 
@@ -200,6 +200,7 @@ pub type PollWorkflowTaskBuffer = LongPollBuffer<PollWorkflowTaskQueueResponse>;
 pub fn new_workflow_task_buffer(
     sg: Arc<impl ServerGatewayApis + Send + Sync + 'static + ?Sized>,
     task_queue: String,
+    is_sticky: bool,
     concurrent_pollers: usize,
     buffer_size: usize,
 ) -> PollWorkflowTaskBuffer {
@@ -207,7 +208,7 @@ pub fn new_workflow_task_buffer(
         move || {
             let sg = sg.clone();
             let task_queue = task_queue.clone();
-            async move { sg.poll_workflow_task(task_queue).await }
+            async move { sg.poll_workflow_task(task_queue, is_sticky).await }
         },
         concurrent_pollers,
         buffer_size,
@@ -246,7 +247,7 @@ mod tests {
         mock_gateway
             .expect_poll_workflow_task()
             .times(2)
-            .returning(move |_| {
+            .returning(move |_, _| {
                 async {
                     tokio::time::sleep(Duration::from_millis(100)).await;
                     Ok(Default::default())
@@ -255,7 +256,7 @@ mod tests {
             });
         let mock_gateway = Arc::new(mock_gateway);
 
-        let pb = new_workflow_task_buffer(mock_gateway, "someq".to_string(), 1, 1);
+        let pb = new_workflow_task_buffer(mock_gateway, "someq".to_string(), false, 1, 1);
 
         // Poll a bunch of times, "interrupting" it each time, we should only actually have polled
         // once since the poll takes a while