diff --git a/crates/sdk-core/src/pollers/mod.rs b/crates/sdk-core/src/pollers/mod.rs index 0a1e2f616..e6f555762 100644 --- a/crates/sdk-core/src/pollers/mod.rs +++ b/crates/sdk-core/src/pollers/mod.rs @@ -1,8 +1,8 @@ mod poll_buffer; +mod sticky; -pub(crate) use poll_buffer::{ - ActivityTaskOptions, LongPollBuffer, WorkflowTaskOptions, WorkflowTaskPoller, -}; +pub(crate) use poll_buffer::{ActivityTaskOptions, LongPollBuffer, WorkflowTaskOptions}; +pub(crate) use sticky::{StickyPollerShared, StickyTaskPoller, sticky_poller_behavior}; pub use temporalio_client::{Client, ClientOptions, ClientTlsOptions, RetryOptions, TlsOptions}; use crate::{ diff --git a/crates/sdk-core/src/pollers/poll_buffer.rs b/crates/sdk-core/src/pollers/poll_buffer.rs index b35249bc1..84a93cc2d 100644 --- a/crates/sdk-core/src/pollers/poll_buffer.rs +++ b/crates/sdk-core/src/pollers/poll_buffer.rs @@ -1,9 +1,9 @@ use crate::{ abstractions::{ActiveCounter, MeteredPermitDealer, OwnedMeteredSemPermit, dbg_panic}, - pollers::{self, Poller}, + pollers::{self, Poller, StickyPollerShared}, worker::{ ActivitySlotKind, NamespaceCapabilities, NexusSlotKind, PollerBehavior, SlotKind, - WFTPollerShared, WorkflowSlotKind, + WorkflowSlotKind, client::{ PollActivityOptions, PollNexusOptions, PollOptions, PollWorkflowOptions, WorkerClient, }, @@ -103,7 +103,7 @@ pub(crate) struct LongPollBuffer { } pub(crate) struct WorkflowTaskOptions { - pub(crate) wft_poller_shared: Option>, + pub(crate) wft_poller_shared: Option>, } pub(crate) struct ActivityTaskOptions { @@ -749,59 +749,6 @@ impl PollScalerReportHandle { } } -/// A poller capable of polling on a sticky and a nonsticky queue simultaneously for workflow tasks. -#[derive(derive_more::Constructor)] -pub(crate) struct WorkflowTaskPoller { - normal_poller: PollWorkflowTaskBuffer, - sticky_poller: Option, -} -type PollWorkflowTaskBuffer = LongPollBuffer; - -#[async_trait::async_trait] -impl - Poller<( - PollWorkflowTaskQueueResponse, - OwnedMeteredSemPermit, - )> for WorkflowTaskPoller -{ - async fn poll( - &self, - ) -> Option< - pollers::Result<( - PollWorkflowTaskQueueResponse, - OwnedMeteredSemPermit, - )>, - > { - if let Some(sq) = self.sticky_poller.as_ref() { - tokio::select! { - r = self.normal_poller.poll() => r, - r = sq.poll() => r, - } - } else { - self.normal_poller.poll().await - } - } - - fn notify_shutdown(&self) { - self.normal_poller.notify_shutdown(); - if let Some(sq) = self.sticky_poller.as_ref() { - sq.notify_shutdown(); - } - } - - async fn shutdown(mut self) { - self.normal_poller.shutdown().await; - if let Some(sq) = self.sticky_poller { - sq.shutdown().await; - } - } - - async fn shutdown_box(self: Box) { - let this = *self; - this.shutdown().await; - } -} - /// Returns true for errors that the poller scaler wants to see fn poll_scaling_error_matcher(err: &tonic::Status) -> bool { matches!( @@ -947,7 +894,7 @@ mod tests { CancellationToken::new(), None::, WorkflowTaskOptions { - wft_poller_shared: Some(Arc::new(WFTPollerShared::new(Some(10)))), + wft_poller_shared: Some(Arc::new(StickyPollerShared::new(Some(10)))), }, Arc::new(AtomicCell::new(None)), Arc::new(NamespaceCapabilities::default()), @@ -1004,7 +951,7 @@ mod tests { CancellationToken::new(), None::, WorkflowTaskOptions { - wft_poller_shared: Some(Arc::new(WFTPollerShared::new(Some(1)))), + wft_poller_shared: Some(Arc::new(StickyPollerShared::new(Some(1)))), }, Arc::new(AtomicCell::new(None)), Arc::new(NamespaceCapabilities::default()), @@ -1081,7 +1028,7 @@ mod tests { shutdown_token.clone(), None::, WorkflowTaskOptions { - wft_poller_shared: Some(Arc::new(WFTPollerShared::new(Some(10)))), + wft_poller_shared: Some(Arc::new(StickyPollerShared::new(Some(10)))), }, Arc::new(AtomicCell::new(None)), Arc::new(NamespaceCapabilities::default()), @@ -1190,7 +1137,7 @@ mod tests { CancellationToken::new(), None::, WorkflowTaskOptions { - wft_poller_shared: Some(Arc::new(WFTPollerShared::new(Some(INITIAL_POLLERS)))), + wft_poller_shared: Some(Arc::new(StickyPollerShared::new(Some(INITIAL_POLLERS)))), }, Arc::new(AtomicCell::new(None)), Arc::new(NamespaceCapabilities::default()), diff --git a/crates/sdk-core/src/pollers/sticky.rs b/crates/sdk-core/src/pollers/sticky.rs new file mode 100644 index 000000000..332dfbe64 --- /dev/null +++ b/crates/sdk-core/src/pollers/sticky.rs @@ -0,0 +1,226 @@ +use super::{LongPollBuffer, Poller}; +use crate::{ + WorkerConfig, + abstractions::OwnedMeteredSemPermit, + worker::{PollerBehavior, SlotKind}, +}; +use std::{fmt::Debug, sync::OnceLock}; +use tokio::sync::watch; + +/// Given an already-resolved task poller `behavior` (see +/// [crate::worker::resolve_effective_behavior]), applies the sticky/non-sticky split when it is a +/// [PollerBehavior::SimpleMaximum]; other behaviors are returned unchanged. +pub(crate) fn sticky_poller_behavior( + behavior: PollerBehavior, + config: &WorkerConfig, + is_sticky: bool, +) -> PollerBehavior { + fn calc_max_nonsticky(max_polls: usize, ratio: f32) -> usize { + ((max_polls as f32 * ratio) as usize).max(1) + } + + if let PollerBehavior::SimpleMaximum(m) = behavior { + if !is_sticky { + PollerBehavior::SimpleMaximum(calc_max_nonsticky( + m, + config.nonsticky_to_sticky_poll_ratio, + )) + } else { + PollerBehavior::SimpleMaximum( + m.saturating_sub(calc_max_nonsticky(m, config.nonsticky_to_sticky_poll_ratio)) + .max(1), + ) + } + } else { + behavior + } +} + +/// A poller capable of polling on a sticky and a nonsticky queue simultaneously. +#[derive(derive_more::Constructor)] +pub(crate) struct StickyTaskPoller { + normal_poller: LongPollBuffer, + sticky_poller: Option>, +} + +#[async_trait::async_trait] +impl Poller<(T, OwnedMeteredSemPermit)> for StickyTaskPoller +where + T: Send + Sync + Debug + 'static, + SK: SlotKind + 'static, +{ + async fn poll(&self) -> Option)>> { + if let Some(sq) = self.sticky_poller.as_ref() { + tokio::select! { + r = self.normal_poller.poll() => r, + r = sq.poll() => r, + } + } else { + self.normal_poller.poll().await + } + } + + fn notify_shutdown(&self) { + self.normal_poller.notify_shutdown(); + if let Some(sq) = self.sticky_poller.as_ref() { + sq.notify_shutdown(); + } + } + + async fn shutdown(mut self) { + self.normal_poller.shutdown().await; + if let Some(sq) = self.sticky_poller { + sq.shutdown().await; + } + } + + async fn shutdown_box(self: Box) { + let this = *self; + this.shutdown().await; + } +} + +/// Info that needs to be shared across the sticky and non-sticky pollers to prioritize sticky +/// when appropriate +pub(crate) struct StickyPollerShared { + last_seen_sticky_backlog: (watch::Receiver, watch::Sender), + sticky_active: OnceLock>, + sticky_target: OnceLock>, + non_sticky_active: OnceLock>, + max_slots: Option, +} +impl StickyPollerShared { + pub(crate) fn new(max_slots: Option) -> Self { + let (tx, rx) = watch::channel(0); + Self { + last_seen_sticky_backlog: (rx, tx), + sticky_active: OnceLock::new(), + sticky_target: OnceLock::new(), + non_sticky_active: OnceLock::new(), + max_slots, + } + } + pub(crate) fn set_sticky_active( + &self, + active_rx: watch::Receiver, + target_rx: watch::Receiver, + ) { + let _ = self.sticky_active.set(active_rx); + let _ = self.sticky_target.set(target_rx); + } + pub(crate) fn set_non_sticky_active(&self, active_rx: watch::Receiver) { + let _ = self.non_sticky_active.set(active_rx); + } + /// Makes either the sticky or non-sticky poller wait pre-permit-acquisition so that we can + /// balance which kind of queue we poll appropriately. + pub(crate) async fn wait_if_needed(&self, is_sticky: bool) { + // We need to make sure there's at least one poller of both kinds available. So, we check + // that we won't end up using every available permit with one kind of poller. In practice + // this is only ever likely to be an issue with very small numbers of slots. + if let Some(max_slots) = self.max_slots + && let Some(sticky_active) = self.sticky_active.get() + && let Some(sticky_target) = self.sticky_target.get() + && let Some(non_sticky_active) = self.non_sticky_active.get() + { + let mut sticky_active = sticky_active.clone(); + let mut sticky_target = sticky_target.clone(); + let mut non_sticky_active = non_sticky_active.clone(); + let mut sticky_backlog = self.last_seen_sticky_backlog.0.clone(); + + loop { + let num_sticky_active = *sticky_active.borrow_and_update(); + let num_sticky_target = *sticky_target.borrow_and_update(); + let num_non_sticky_active = *non_sticky_active.borrow_and_update(); + let num_sticky_backlog = *sticky_backlog.borrow_and_update(); + let sticky_should_be_chosen = num_sticky_backlog > 1 + && num_sticky_backlog > num_sticky_active + && num_sticky_active < num_sticky_target; + + let allow = || { + if !is_sticky { + // There should always be at least one non-sticky poller. + if num_non_sticky_active == 0 { + return true; + } + + // Do not allow an additional non-sticky poller to prevent starting a first sticky poller. + if num_sticky_active == 0 && num_non_sticky_active + 1 >= max_slots { + return false; + } + + // If there's a meaningful sticky backlog, prioritize sticky. + if sticky_should_be_chosen { + return false; + } + } else { + // There should always be at least one sticky poller. + if num_sticky_active == 0 { + return true; + } + + // Do not reserve a slot that the sticky scaler cannot use. + if num_sticky_active >= num_sticky_target { + return false; + } + + // Do not allow an additional sticky poller to prevent starting a first non-sticky poller. + if num_non_sticky_active == 0 && num_sticky_active + 1 >= max_slots { + return false; + } + + // If there's a meaningful sticky backlog, prioritize sticky. + if sticky_should_be_chosen { + return true; + } + } + + // Just balance the two poller types. + // FIXME: Do we need anything more here, to ensure proper balancing? + if num_sticky_active + num_non_sticky_active < max_slots { + return true; + } + + false + }; + + if allow() { + return; + } + + tokio::select! { + _ = sticky_active.changed() => (), + _ = sticky_target.changed() => (), + _ = non_sticky_active.changed() => (), + _ = sticky_backlog.changed() => (), + } + } + } + } + + pub(crate) fn record_sticky_backlog(&self, v: usize) { + let _ = self.last_seen_sticky_backlog.1.send(v); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_help::test_worker_cfg; + + #[test] + fn max_polls_calculated_properly() { + let cfg = { + let mut cfg = test_worker_cfg().build().unwrap(); + cfg.workflow_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(5_usize)); + cfg + }; + assert_eq!( + sticky_poller_behavior(PollerBehavior::SimpleMaximum(5), &cfg, false), + PollerBehavior::SimpleMaximum(1) + ); + assert_eq!( + sticky_poller_behavior(PollerBehavior::SimpleMaximum(5), &cfg, true), + PollerBehavior::SimpleMaximum(4) + ); + } +} diff --git a/crates/sdk-core/src/worker/mod.rs b/crates/sdk-core/src/worker/mod.rs index 6bb9b82b7..46c2c85a1 100644 --- a/crates/sdk-core/src/worker/mod.rs +++ b/crates/sdk-core/src/worker/mod.rs @@ -38,7 +38,6 @@ pub(crate) use activities::{ ExecutingLAId, LocalActRequest, LocalActivityExecutionResult, LocalActivityResolution, NewLocalAct, }; -pub(crate) use wft_poller::WFTPollerShared; #[allow(unreachable_pub)] // re-exported in test_help::integ_helpers pub use workflow::LEGACY_QUERY_ID; @@ -58,8 +57,7 @@ use crate::{ heartbeat::SharedNamespaceWorker, nexus::NexusManager, workflow::{ - LAReqSink, LocalResolution, WorkflowBasics, Workflows, wft_poller, - wft_poller::make_wft_poller, + LAReqSink, LocalResolution, WorkflowBasics, Workflows, wft_poller::make_wft_poller, }, }, }; @@ -2506,35 +2504,6 @@ pub(crate) enum TaskPollers { }, } -/// Given an already-resolved workflow task poller `behavior` (see [resolve_effective_behavior]), -/// applies the sticky/non-sticky split when it is a [PollerBehavior::SimpleMaximum]; other -/// behaviors are returned unchanged. -pub(crate) fn wft_poller_behavior( - behavior: PollerBehavior, - config: &WorkerConfig, - is_sticky: bool, -) -> PollerBehavior { - fn calc_max_nonsticky(max_polls: usize, ratio: f32) -> usize { - ((max_polls as f32 * ratio) as usize).max(1) - } - - if let PollerBehavior::SimpleMaximum(m) = behavior { - if !is_sticky { - PollerBehavior::SimpleMaximum(calc_max_nonsticky( - m, - config.nonsticky_to_sticky_poll_ratio, - )) - } else { - PollerBehavior::SimpleMaximum( - m.saturating_sub(calc_max_nonsticky(m, config.nonsticky_to_sticky_poll_ratio)) - .max(1), - ) - } - } else { - behavior - } -} - fn make_slots_info( dealer: &MeteredPermitDealer, slots_available: Arc, @@ -2636,23 +2605,6 @@ mod tests { ); } - #[test] - fn max_polls_calculated_properly() { - let cfg = { - let mut cfg = test_worker_cfg().build().unwrap(); - cfg.workflow_task_poller_behavior = Some(PollerBehavior::SimpleMaximum(5_usize)); - cfg - }; - assert_eq!( - wft_poller_behavior(PollerBehavior::SimpleMaximum(5), &cfg, false), - PollerBehavior::SimpleMaximum(1) - ); - assert_eq!( - wft_poller_behavior(PollerBehavior::SimpleMaximum(5), &cfg, true), - PollerBehavior::SimpleMaximum(4) - ); - } - #[test] fn max_polls_zero_is_err() { assert!( diff --git a/crates/sdk-core/src/worker/tuner/resource_based.rs b/crates/sdk-core/src/worker/tuner/resource_based.rs index 16cfdcb68..d0a984b03 100644 --- a/crates/sdk-core/src/worker/tuner/resource_based.rs +++ b/crates/sdk-core/src/worker/tuner/resource_based.rs @@ -165,13 +165,41 @@ pub(crate) struct ResourceBasedSlotsForType { last_slot_issued_tx: watch::Sender, last_slot_issued_rx: watch::Receiver, - // Only used for workflow slots - count of issued non-sticky slots + sticky_slots: StickySlotAccounting, + + _slot_kind: PhantomData, +} + +#[derive(Default)] +struct StickySlotAccounting { issued_nonsticky: AtomicUsize, - // Only used for workflow slots - count of issued sticky slots issued_sticky: AtomicUsize, +} - _slot_kind: PhantomData, +impl StickySlotAccounting { + // Always be willing to hand out at least 1 slot for sticky and 1 for non-sticky to + // avoid getting stuck. + fn needs_first_slot(&self, is_sticky: bool) -> bool { + self.counter(is_sticky).load(Ordering::Relaxed) == 0 + } + + fn record_issue(&self, is_sticky: bool) { + self.counter(is_sticky).fetch_add(1, Ordering::Relaxed); + } + + fn record_release(&self, is_sticky: bool) { + self.counter(is_sticky).fetch_sub(1, Ordering::Relaxed); + } + + fn counter(&self, is_sticky: bool) -> &AtomicUsize { + if is_sticky { + &self.issued_sticky + } else { + &self.issued_nonsticky + } + } } + /// Allows for the full customization of the PID options for a resource based tuner #[derive(Clone, Debug, bon::Builder)] #[non_exhaustive] @@ -327,11 +355,7 @@ where fn release_slot(&self, ctx: &dyn SlotReleaseContext) { // Really could use specialization here if let Some(SlotInfo::Workflow(info)) = ctx.info().map(|i| i.downcast()) { - if info.is_sticky { - self.issued_sticky.fetch_sub(1, Ordering::Relaxed); - } else { - self.issued_nonsticky.fetch_sub(1, Ordering::Relaxed); - } + self.sticky_slots.record_release(info.is_sticky); } } @@ -352,14 +376,11 @@ where last_slot_issued_tx: tx, last_slot_issued_rx: rx, inner, - issued_nonsticky: Default::default(), - issued_sticky: Default::default(), + sticky_slots: Default::default(), _slot_kind: PhantomData, } } - // Always be willing to hand out at least 1 slot for sticky and 1 for non-sticky to - // avoid getting stuck. fn issue_if_below_minimums( &self, ctx: &dyn SlotReservationContext, @@ -368,8 +389,7 @@ where return Some(self.issue_slot(ctx)); } if SK::kind() == SlotKindType::Workflow - && (ctx.is_sticky() && self.issued_sticky.load(Ordering::Relaxed) == 0 - || !ctx.is_sticky() && self.issued_nonsticky.load(Ordering::Relaxed) == 0) + && self.sticky_slots.needs_first_slot(ctx.is_sticky()) { return Some(self.issue_slot(ctx)); } @@ -381,11 +401,7 @@ where // Always be willing to hand out at least 1 slot for sticky and 1 for non-sticky to avoid // getting stuck. if SK::kind() == SlotKindType::Workflow { - if ctx.is_sticky() { - self.issued_sticky.fetch_add(1, Ordering::Relaxed); - } else { - self.issued_nonsticky.fetch_add(1, Ordering::Relaxed); - } + self.sticky_slots.record_issue(ctx.is_sticky()); } let _ = self.last_slot_issued_tx.send(Instant::now()); SlotSupplierPermit::default() diff --git a/crates/sdk-core/src/worker/workflow/wft_poller.rs b/crates/sdk-core/src/worker/workflow/wft_poller.rs index ee6dee2db..8737e8159 100644 --- a/crates/sdk-core/src/worker/workflow/wft_poller.rs +++ b/crates/sdk-core/src/worker/workflow/wft_poller.rs @@ -1,19 +1,18 @@ use crate::{ MetricsContext, WorkerConfig, abstractions::{MeteredPermitDealer, OwnedMeteredSemPermit}, - pollers::{BoxedWFPoller, LongPollBuffer, Poller, WorkflowTaskOptions, WorkflowTaskPoller}, + pollers::{ + BoxedWFPoller, LongPollBuffer, Poller, StickyPollerShared, StickyTaskPoller, + WorkflowTaskOptions, sticky_poller_behavior, + }, protosext::ValidPollWFTQResponse, telemetry::metrics::{workflow_poller, workflow_sticky_poller}, - worker::{NamespaceCapabilities, WorkflowSlotKind, client::WorkerClient, wft_poller_behavior}, + worker::{NamespaceCapabilities, WorkflowSlotKind, client::WorkerClient}, }; use crossbeam_utils::atomic::AtomicCell; use futures_util::{Stream, stream}; -use std::{ - sync::{Arc, OnceLock}, - time::SystemTime, -}; +use std::{sync::Arc, time::SystemTime}; use temporalio_common::protos::temporal::api::workflowservice::v1::PollWorkflowTaskQueueResponse; -use tokio::sync::watch; use tokio_util::sync::CancellationToken; #[allow(clippy::too_many_arguments)] @@ -49,7 +48,7 @@ pub(crate) fn make_wft_poller( .into_iter() .flatten() .min(); - Some(Arc::new(WFTPollerShared::new(balance_limit))) + Some(Arc::new(StickyPollerShared::new(balance_limit))) } else { None }; @@ -57,7 +56,7 @@ pub(crate) fn make_wft_poller( client.clone(), config.task_queue.clone(), None, - wft_poller_behavior(effective_behavior, config, false), + sticky_poller_behavior(effective_behavior, config, false), wft_slots.clone(), shutdown_token.child_token(), Some(move |np| { @@ -75,7 +74,7 @@ pub(crate) fn make_wft_poller( client.clone(), config.task_queue.clone(), Some(sqn.clone()), - wft_poller_behavior(effective_behavior, config, true), + sticky_poller_behavior(effective_behavior, config, true), wft_slots.clone().into_sticky(), shutdown_token.child_token(), Some(move |np| { @@ -86,135 +85,13 @@ pub(crate) fn make_wft_poller( capabilities, ) }); - let wf_task_poll_buffer = Box::new(WorkflowTaskPoller::new( + let wf_task_poll_buffer = Box::new(StickyTaskPoller::new( wf_task_poll_buffer, sticky_queue_poller, )); new_wft_poller(wf_task_poll_buffer, metrics.clone()) } -/// Info that needs to be shared across the sticky and non-sticky wft pollers to prioritize sticky -/// when appropriate -pub(crate) struct WFTPollerShared { - last_seen_sticky_backlog: (watch::Receiver, watch::Sender), - sticky_active: OnceLock>, - sticky_target: OnceLock>, - non_sticky_active: OnceLock>, - max_slots: Option, -} -impl WFTPollerShared { - pub(crate) fn new(max_slots: Option) -> Self { - let (tx, rx) = watch::channel(0); - Self { - last_seen_sticky_backlog: (rx, tx), - sticky_active: OnceLock::new(), - sticky_target: OnceLock::new(), - non_sticky_active: OnceLock::new(), - max_slots, - } - } - pub(crate) fn set_sticky_active( - &self, - active_rx: watch::Receiver, - target_rx: watch::Receiver, - ) { - let _ = self.sticky_active.set(active_rx); - let _ = self.sticky_target.set(target_rx); - } - pub(crate) fn set_non_sticky_active(&self, active_rx: watch::Receiver) { - let _ = self.non_sticky_active.set(active_rx); - } - /// Makes either the sticky or non-sticky poller wait pre-permit-acquisition so that we can - /// balance which kind of queue we poll appropriately. - pub(crate) async fn wait_if_needed(&self, is_sticky: bool) { - // We need to make sure there's at least one poller of both kinds available. So, we check - // that we won't end up using every available permit with one kind of poller. In practice - // this is only ever likely to be an issue with very small numbers of slots. - if let Some(max_slots) = self.max_slots - && let Some(sticky_active) = self.sticky_active.get() - && let Some(sticky_target) = self.sticky_target.get() - && let Some(non_sticky_active) = self.non_sticky_active.get() - { - let mut sticky_active = sticky_active.clone(); - let mut sticky_target = sticky_target.clone(); - let mut non_sticky_active = non_sticky_active.clone(); - let mut sticky_backlog = self.last_seen_sticky_backlog.0.clone(); - - loop { - let num_sticky_active = *sticky_active.borrow_and_update(); - let num_sticky_target = *sticky_target.borrow_and_update(); - let num_non_sticky_active = *non_sticky_active.borrow_and_update(); - let num_sticky_backlog = *sticky_backlog.borrow_and_update(); - let sticky_should_be_chosen = num_sticky_backlog > 1 - && num_sticky_backlog > num_sticky_active - && num_sticky_active < num_sticky_target; - - let allow = || { - if !is_sticky { - // There should always be at least one non-sticky poller. - if num_non_sticky_active == 0 { - return true; - } - - // Do not allow an additional non-sticky poller to prevent starting a first sticky poller. - if num_sticky_active == 0 && num_non_sticky_active + 1 >= max_slots { - return false; - } - - // If there's a meaningful sticky backlog, prioritize sticky. - if sticky_should_be_chosen { - return false; - } - } else { - // There should always be at least one sticky poller. - if num_sticky_active == 0 { - return true; - } - - // Do not reserve a slot that the sticky scaler cannot use. - if num_sticky_active >= num_sticky_target { - return false; - } - - // Do not allow an additional sticky poller to prevent starting a first non-sticky poller. - if num_non_sticky_active == 0 && num_sticky_active + 1 >= max_slots { - return false; - } - - // If there's a meaningful sticky backlog, prioritize sticky. - if sticky_should_be_chosen { - return true; - } - } - - // Just balance the two poller types. - // FIXME: Do we need anything more here, to ensure proper balancing? - if num_sticky_active + num_non_sticky_active < max_slots { - return true; - } - - false - }; - - if allow() { - return; - } - - tokio::select! { - _ = sticky_active.changed() => (), - _ = sticky_target.changed() => (), - _ = non_sticky_active.changed() => (), - _ = sticky_backlog.changed() => (), - } - } - } - } - - pub(crate) fn record_sticky_backlog(&self, v: usize) { - let _ = self.last_seen_sticky_backlog.1.send(v); - } -} - fn new_wft_poller( poller: BoxedWFPoller, metrics: MetricsContext, @@ -302,6 +179,7 @@ mod tests { task::{Context, Poll, Wake}, }; use temporalio_common::protos::temporal::api::enums::v1::EventType; + use tokio::sync::watch; struct WakeCounter(AtomicUsize); @@ -313,7 +191,7 @@ mod tests { #[tokio::test] async fn sticky_priority_respects_poller_target() { - let shared = WFTPollerShared::new(Some(4)); + let shared = StickyPollerShared::new(Some(4)); let (_sticky_tx, sticky_rx) = watch::channel(2); let (_non_sticky_tx, non_sticky_rx) = watch::channel(1); let (sticky_target_tx, sticky_target_rx) = watch::channel(3); @@ -330,7 +208,7 @@ mod tests { #[tokio::test] async fn target_change_wakes_balancer() { - let shared = WFTPollerShared::new(Some(4)); + let shared = StickyPollerShared::new(Some(4)); let (_sticky_tx, sticky_rx) = watch::channel(2); let (_non_sticky_tx, non_sticky_rx) = watch::channel(1); let (sticky_target_tx, sticky_target_rx) = watch::channel(3);