rzmq 0.5.18

//! An adaptive, probabilistic throttle for managing I/O fairness in async actors.
//!
//! This module provides an `AdaptiveThrottle` that allows an actor to balance its
//! time between different workloads (e.g., incoming vs. outgoing data) without
//! explicit, complex scheduling logic in the actor itself.
//!
//! # Model
//!
//! The throttle uses a "DEBT/CREDIT" system, represented by a single `balance`
//! counter.
//! - **Ingress** (work coming in) adds to the balance, creating "debt."
//! - **Egress** (work going out) subtracts from the balance, paying it off.
//!
//! The throttle learns the normal operating balance for a given workload and only
//! intervenes when the system deviates significantly from this learned norm. The
//! intervention is a probabilistic `yield`, giving other tasks a chance to run.
//!
//! # Usage
//!
//! The API is designed around a simple, two-step cycle for each I/O operation:
//!
//! 1.  Call `throttle.begin_work()` before the operation to get a `ThrottleGuard`.
//!     This performs an anticipatory update to the throttle's state.
//! 2.  After the operation, call `guard.should_throttle().await`. This will
//!     run the probabilistic logic and potentially yield to the async runtime.
//!
//! ```rust,ignore
//! // Inside an actor's select! loop arm
//!
//! // 1. Announce the work and get the guard.
//! let guard = self.throttle.begin_work(Direction::Ingress);
//!
//! // 2. Do the I/O work. The outcome is not tracked by the throttle.
//! self.handle_incoming_message(message).await;
//!
//! // 3. Finalize the cycle, potentially yielding.
//! guard.should_throttle().await;
//! ```

pub mod strategies;
pub mod types;

use types::ThrottleStateView;
pub use types::{AdaptiveThrottleConfig, Direction};

use std::sync::atomic::{AtomicI32, AtomicU32, Ordering};
use std::sync::Arc;

use atomic_float::AtomicF64;
use rand::random;

use crate::throttle::types::Priority;

/// This struct holds both the mutable atomic state and the immutable config.
/// It is the single unit that will be shared via an `Arc` for cheap, thread-safe access.
#[derive(Debug)]
struct InternalSharedState {
  current_balance: AtomicI32,
  learned_balance: AtomicF64,
  consecutive_ingress: AtomicU32,
  consecutive_egress: AtomicU32,
  ops_since_nudge: AtomicU32,
  config: AdaptiveThrottleConfig,
}

// --- Main Throttle Struct ---

/// An adaptive, probabilistic throttle that can be shared across tasks.
///
/// It learns the natural I/O balance of a workload and gently yields CPU time
/// when the workload deviates significantly from its norm, preventing starvation
/// while maximizing throughput for bursty traffic.
#[derive(Debug, Clone)]
pub struct AdaptiveThrottle {
  shared: Arc<InternalSharedState>,
}

impl AdaptiveThrottle {
  /// Creates a new `AdaptiveThrottle` with the specified configuration.
  pub fn new(config: AdaptiveThrottleConfig) -> Self {
    let mut cfg = config.clone();
    // clamp learning rate
    if cfg.adaptive_learning_rate < 0.01 {
      cfg.adaptive_learning_rate = 0.01;
    }
    if cfg.adaptive_learning_rate > 0.2 {
      cfg.adaptive_learning_rate = 0.2;
    }
    let state = InternalSharedState {
      current_balance: AtomicI32::new(0),
      learned_balance: AtomicF64::new(0.0),
      consecutive_ingress: AtomicU32::new(0),
      consecutive_egress: AtomicU32::new(0),
      ops_since_nudge: AtomicU32::new(0),
      config: cfg,
    };
    Self {
      shared: Arc::new(state),
    }
  }

  /// Records the intent to perform an I/O operation and returns a `ThrottleGuard`.
  ///
  /// This method performs an "anticipatory" update to the throttle's internal
  /// state, reflecting the work that is about to happen. The returned guard
  /// must be used to complete the work cycle.
  pub fn begin_work(&self, dir: Direction) -> ThrottleGuard {
    self.begin_work_bulk(dir, 1)
  }

  /// Records the intent to perform `weight` operations simultaneously,
  /// debiting the balance and EMA proportionally to the entire batch.
  pub fn begin_work_bulk(&self, dir: Direction, weight: u32) -> ThrottleGuard {
    if !self.shared.config.enabled {
      return ThrottleGuard {
        shared: None,
        direction: dir,
        should_throttle_fn: bypass_should_throttle,
      };
    }

    let delta = self.shared.config.credit_per_message * (weight as i32);
    let new_balance = match dir {
      Direction::Ingress => {
        self
          .shared
          .current_balance
          .fetch_add(delta, Ordering::Relaxed)
          + delta
      }
      Direction::Egress => {
        self
          .shared
          .current_balance
          .fetch_sub(delta, Ordering::Relaxed)
          - delta
      }
    };

    // Periodic EMA nudge
    let since = self.shared.ops_since_nudge.fetch_add(weight, Ordering::Relaxed) + weight;
    if since >= self.shared.config.nudge_interval_ops {
      let α = self.shared.config.adaptive_learning_rate;
      let old_learned = self.shared.learned_balance.load(Ordering::Relaxed);
      let updated = α * (new_balance as f64) + (1.0 - α) * old_learned;
      self
        .shared
        .learned_balance
        .store(updated, Ordering::Relaxed);
      self.shared.ops_since_nudge.store(0, Ordering::Relaxed);
    }

    // Track direction-specific streaks proportionally
    match dir {
      Direction::Ingress => {
        self
          .shared
          .consecutive_ingress
          .fetch_add(weight, Ordering::Relaxed);
        self.shared.consecutive_egress.store(0, Ordering::Relaxed);
      }
      Direction::Egress => {
        self
          .shared
          .consecutive_egress
          .fetch_add(weight, Ordering::Relaxed);
        self.shared.consecutive_ingress.store(0, Ordering::Relaxed);
      }
    }

    ThrottleGuard {
      shared: Some(self.shared.clone()),
      direction: dir,
      should_throttle_fn: active_should_throttle,
    }
  }
}

// --- Throttle Guard ---

/// A temporary guard representing an in-progress I/O operation.
/// Its purpose is to call the throttling logic when the operation is complete.
pub struct ThrottleGuard {
  shared: Option<Arc<InternalSharedState>>,
  direction: Direction,
  should_throttle_fn: fn(&ThrottleGuard) -> bool,
}

impl ThrottleGuard {
  pub fn get_current_balance(&self) -> i32 {
    self
      .shared
      .as_ref()
      .map_or(0, |s| s.current_balance.load(Ordering::Relaxed))
  }

  #[inline(always)]
  pub fn should_throttle(&self) -> bool {
    (self.should_throttle_fn)(self)
  }
}

fn bypass_should_throttle(_guard: &ThrottleGuard) -> bool {
  false
}

fn active_should_throttle(guard: &ThrottleGuard) -> bool {
  let shared = guard.shared.as_ref().unwrap();
  let cfg = &shared.config;

  // Hard cap per-direction
  let cons = match guard.direction {
    Direction::Ingress => shared.consecutive_ingress.load(Ordering::Relaxed),
    Direction::Egress => shared.consecutive_egress.load(Ordering::Relaxed),
  };
  if cons >= cfg.yield_after_n_consecutive {
    shared.consecutive_ingress.store(0, Ordering::Relaxed);
    shared.consecutive_egress.store(0, Ordering::Relaxed);
    return true;
  }

  // Probabilistic
  let state = ThrottleStateView {
    current_balance: shared.current_balance.load(Ordering::Relaxed),
    learned_balance: shared.learned_balance.load(Ordering::Relaxed),
    config: cfg,
  };
  let mut p = (cfg.strategy)(&state);

  // If the current work direction is NOT the prioritized one, and we are
  // imbalanced in a way that favors the *prioritized* direction, we
  // dramatically increase the probability of yielding to give the
  // priority work a chance to run.
  let is_priority_work = match cfg.priority {
    Priority::Egress => guard.direction == Direction::Egress,
    Priority::Ingress => guard.direction == Direction::Ingress,
    Priority::None => true, // With no priority, all work is "priority" work.
  };

  if !is_priority_work {
    let is_imbalanced_towards_priority = match cfg.priority {
      // We are doing Ingress (non-priority), and the balance is high (debt),
      // which means we need to do more Egress (priority) work.
      Priority::Egress => state.current_balance > state.learned_balance as i32,
      // We are doing Egress (non-priority), and the balance is low (credit),
      // which means we need to do more Ingress (priority) work.
      Priority::Ingress => state.current_balance < state.learned_balance as i32,
      Priority::None => false,
    };

    if is_imbalanced_towards_priority {
      p *= cfg.priority_boost_factor;
    }
  }

  // piecewise: beyond 2× max_imbalance, always yield
  let dev = (state.current_balance as f64 - state.learned_balance).abs();
  let hard_cut = (cfg.max_imbalance as f64) * 2.0;
  if dev >= hard_cut {
    p = 1.0;
  }

  if random::<f64>() < p {
    shared.consecutive_ingress.store(0, Ordering::Relaxed);
    shared.consecutive_egress.store(0, Ordering::Relaxed);
    return true;
  }

  false
}

#[cfg(test)]
mod throttle_math_tests {
  use super::*;
  use super::types::{Direction, Priority, ThrottleStateView};
  use std::sync::atomic::Ordering;

  #[test]
  fn test_ema_convergence() {
    let mut config = AdaptiveThrottleConfig::default();
    config.adaptive_learning_rate = 0.1;
    config.nudge_interval_ops = 5;
    config.credit_per_message = 2;
    config.enabled = true;

    let throttle = AdaptiveThrottle::new(config);

    // After 5 Ingress ops: balance = 5 * 2 = 10.
    // Nudge fires at op 5: updated = 0.1 * 10 + 0.9 * 0.0 = 1.0.
    for _ in 0..5 {
      let _ = throttle.begin_work(Direction::Ingress);
    }

    let learned = throttle.shared.learned_balance.load(Ordering::Relaxed);
    assert!(
      (learned - 1.0).abs() < f64::EPSILON,
      "Learned balance did not converge. Expected 1.0, got {}",
      learned
    );
  }

  #[test]
  fn test_priority_boost_logic() {
    let mut config = AdaptiveThrottleConfig::default();
    config.priority = Priority::Egress;
    config.priority_boost_factor = 3.0;
    config.enabled = true;

    let throttle = AdaptiveThrottle::new(config);

    // Simulate high Ingress debt: deviation must exceed healthy_balance_width (default 100)
    // for the strategy to return a non-zero probability.
    throttle.shared.current_balance.store(150, Ordering::SeqCst);
    throttle.shared.learned_balance.store(0.0, Ordering::SeqCst);

    let state = ThrottleStateView {
      current_balance: 150,
      learned_balance: 0.0,
      config: &throttle.shared.config,
    };

    let base_p = (throttle.shared.config.strategy)(&state);
    assert!(base_p > 0.0, "Expected non-zero base probability for imbalanced state");

    // With Egress as priority and balance > learned, doing Ingress (non-priority)
    // should apply a 3x boost. Verify the multiplier matches config.
    let is_imbalanced_towards_priority = state.current_balance > state.learned_balance as i32;
    assert!(is_imbalanced_towards_priority, "Expected imbalance: balance 150 > learned 0");

    let boosted_p = base_p * throttle.shared.config.priority_boost_factor;
    assert!(
      (boosted_p - base_p * 3.0).abs() < 1e-10,
      "Boost factor mismatch: expected {}, got {}",
      base_p * 3.0,
      boosted_p
    );
  }
}