diff --git a/src/connection.rs b/src/connection.rs index 6e8c737e..d9ae544f 100644 --- a/src/connection.rs +++ b/src/connection.rs @@ -12,10 +12,7 @@ use std::{ }; use futures::StreamExt; -use tokio::{ - sync::mpsc, - time::{Duration, Instant}, -}; +use tokio::{sync::mpsc, time::Duration}; use tokio_util::sync::CancellationToken; use tracing::{info, info_span, warn}; @@ -46,6 +43,7 @@ impl Drop for ActiveConnection { pub fn active_connection_count() -> u64 { ACTIVE_CONNECTIONS.load(Ordering::Relaxed) } use crate::{ + fairness::Fairness, hooks::{ConnectionContext, ProtocolHooks}, push::{FrameLike, PushHandle, PushQueues}, response::{FrameStream, WireframeError}, @@ -66,7 +64,7 @@ enum Event { } /// Configuration controlling fairness when draining push queues. -#[derive(Clone, Copy)] +#[derive(Clone, Copy, Debug)] pub struct FairnessConfig { /// Number of consecutive high-priority frames to process before /// checking the low-priority queue. @@ -110,9 +108,7 @@ pub struct ConnectionActor { counter: Option, hooks: ProtocolHooks, ctx: ConnectionContext, - fairness: FairnessConfig, - high_counter: usize, - high_start: Option, + fairness: Fairness, connection_id: Option, peer_addr: Option, } @@ -170,9 +166,7 @@ where counter: Some(counter), hooks, ctx, - fairness: FairnessConfig::default(), - high_counter: 0, - high_start: None, + fairness: Fairness::new(FairnessConfig::default()), connection_id: None, peer_addr: None, }; @@ -188,7 +182,7 @@ where } /// Replace the fairness configuration. - pub fn set_fairness(&mut self, fairness: FairnessConfig) { self.fairness = fairness; } + pub fn set_fairness(&mut self, fairness: FairnessConfig) { self.fairness.set_config(fairness); } /// Set or replace the current streaming response. pub fn set_response(&mut self, stream: Option>) { self.response = stream; } @@ -307,7 +301,7 @@ where self.after_high(out, state); } else { Self::handle_closed_receiver(&mut self.high_rx, state); - self.reset_high_counter(); + self.fairness.reset(); } } @@ -369,12 +363,9 @@ where /// Update counters and opportunistically drain the low-priority queue. fn after_high(&mut self, out: &mut Vec, state: &mut ActorState) { - self.high_counter += 1; - if self.high_counter == 1 { - self.high_start = Some(Instant::now()); - } + self.fairness.after_high(); - if self.should_yield_to_low_priority() + if self.fairness.should_yield() && let Some(rx) = &mut self.low_rx { match rx.try_recv() { @@ -392,26 +383,8 @@ where } } - /// Determine if processing should yield to the low-priority queue. - fn should_yield_to_low_priority(&self) -> bool { - let threshold_hit = self.fairness.max_high_before_low > 0 - && self.high_counter >= self.fairness.max_high_before_low; - let time_hit = self - .fairness - .time_slice - .zip(self.high_start) - .is_some_and(|(slice, start)| start.elapsed() >= slice); - threshold_hit || time_hit - } - /// Reset counters after processing a low-priority frame. - fn after_low(&mut self) { self.reset_high_counter(); } - - /// Clear the burst counter and associated timestamp. - fn reset_high_counter(&mut self) { - self.high_counter = 0; - self.high_start = None; - } + fn after_low(&mut self) { self.fairness.after_low(); } /// Push a frame from the response stream into `out` or handle completion. /// diff --git a/src/fairness.rs b/src/fairness.rs new file mode 100644 index 00000000..8463b70f --- /dev/null +++ b/src/fairness.rs @@ -0,0 +1,106 @@ +//! Helpers tracking fairness counters for connection processing. +//! +//! This module encapsulates the logic for deciding when high-priority +//! processing should yield to low-priority traffic based on configured +//! thresholds and optional time slices. + +use tokio::time::Instant; + +use crate::connection::FairnessConfig; + +#[derive(Debug)] +pub(crate) struct Fairness { + config: FairnessConfig, + high_counter: usize, + high_start: Option, +} + +impl Fairness { + pub(crate) fn new(config: FairnessConfig) -> Self { + Self { + config, + high_counter: 0, + high_start: None, + } + } + + pub(crate) fn set_config(&mut self, config: FairnessConfig) { + self.config = config; + self.reset(); + } + + pub(crate) fn after_high(&mut self) { + self.high_counter += 1; + if self.high_counter == 1 { + self.high_start = Some(Instant::now()); + } + } + + pub(crate) fn should_yield(&self) -> bool { + let threshold_hit = self.config.max_high_before_low > 0 + && self.high_counter >= self.config.max_high_before_low; + let time_hit = self + .config + .time_slice + .zip(self.high_start) + .is_some_and(|(slice, start)| start.elapsed() >= slice); + threshold_hit || time_hit + } + + pub(crate) fn after_low(&mut self) { self.reset(); } + + pub(crate) fn reset(&mut self) { + self.high_counter = 0; + self.high_start = None; + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use tokio::time::{self, Duration}; + + use super::*; + + #[rstest] + #[tokio::test] + async fn yield_after_threshold() { + let cfg = FairnessConfig { + max_high_before_low: 2, + time_slice: None, + }; + let mut fairness = Fairness::new(cfg); + fairness.after_high(); + assert!(!fairness.should_yield()); + fairness.after_high(); + assert!(fairness.should_yield()); + } + + #[rstest] + #[tokio::test] + async fn after_low_resets_counter() { + let cfg = FairnessConfig { + max_high_before_low: 1, + time_slice: None, + }; + let mut fairness = Fairness::new(cfg); + fairness.after_high(); + assert!(fairness.should_yield()); + fairness.after_low(); + assert!(!fairness.should_yield()); + } + + #[rstest] + #[tokio::test] + async fn time_slice_triggers_yield() { + time::pause(); + let cfg = FairnessConfig { + max_high_before_low: 0, + time_slice: Some(Duration::from_millis(5)), + }; + let mut fairness = Fairness::new(cfg); + fairness.after_high(); + time::advance(Duration::from_millis(6)).await; + assert!(fairness.should_yield()); + } +} diff --git a/src/lib.rs b/src/lib.rs index 34fb32a1..e4c19ec8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,6 +9,7 @@ pub mod serializer; pub use serializer::{BincodeSerializer, Serializer}; pub mod connection; pub mod extractor; +mod fairness; pub mod frame; pub mod hooks; pub mod message;