use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};
mod controller;
pub use controller::{
BrokerPhase, BrokerReport, Model, OpeningBracketOutcome, OpeningBracketReport,
ProducerCapReason, Rejection, RunningBroker,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(deny_unknown_fields, default)]
pub struct EngagementPolicy {
pub min_threads_per_stream: usize,
}
impl Default for EngagementPolicy {
fn default() -> Self {
Self {
min_threads_per_stream: 8,
}
}
}
impl EngagementPolicy {
pub const fn engages(&self, budget: usize, streams: usize) -> bool {
streams != 0 && budget >= self.min_threads_per_stream.saturating_mul(streams)
}
pub const fn always() -> Self {
Self {
min_threads_per_stream: 0,
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Work {
pub busy_nanos: u64,
pub items: u64,
}
impl Work {
fn delta(self, prev: Work) -> Work {
Work {
busy_nanos: self.busy_nanos.saturating_sub(prev.busy_nanos),
items: self.items.saturating_sub(prev.items),
}
}
}
pub trait Consumer: Send + Sync {
fn set_threads(&self, n: usize) -> Result<(), ResizeError>;
fn live_threads(&self) -> usize;
fn work(&self) -> Work;
fn set_measurement_mode(&self, _mode: ProducerMeasurementMode) {}
}
pub trait Producer: Send + Sync {
fn set_limit(&self, n: usize) -> Result<(), ResizeError>;
fn limit(&self) -> usize;
fn active_slots(&self) -> usize;
fn buffered_items(&self) -> Option<u64> {
None
}
fn auxiliary_threads(&self) -> usize {
0
}
fn set_measurement_mode(&self, _mode: ProducerMeasurementMode) {}
fn set_monitoring_interval(&self, _interval: Duration) {}
fn measurement_stats(&self) -> Option<ProducerMeasurementStats> {
None
}
fn finish_measurement(&self) -> Option<ProducerMeasurementStats> {
self.measurement_stats()
}
fn pressure(&self) -> ProducerPressure;
fn work(&self) -> Work;
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "snake_case")]
pub enum ProducerMeasurementMode {
#[default]
Calibration,
Monitoring,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "snake_case")]
pub enum SteadyStatePolicy {
#[default]
Responsive,
FreezeAfterConvergence,
FreezeAfterFullCalibration,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum OpeningPolicy {
#[default]
Fixed,
Bracket(OpeningBracketConfig),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct OpeningBracketConfig {
pub max_points: usize,
pub horizon: Duration,
pub total_budget: Duration,
}
impl Default for OpeningBracketConfig {
fn default() -> Self {
Self {
max_points: 3,
horizon: Duration::from_millis(200),
total_budget: Duration::from_secs(4),
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize)]
pub struct ProducerMeasurementStats {
pub busy_nanos: u64,
pub accounted_busy_nanos: Option<u64>,
pub completed_worker_cpu_nanos: Option<u64>,
pub completed_auxiliary_cpu_nanos: Option<u64>,
pub cpu_accounting_failures: Option<usize>,
pub sampler_cpu_nanos: Option<u64>,
pub sampler_cpu_accounting_failures: usize,
pub calibration_samples: u64,
pub monitoring_samples: u64,
pub mode_changes: u64,
pub final_mode: ProducerMeasurementMode,
pub observation_nanos: u64,
pub calibration_interval_micros: u64,
pub monitoring_interval_micros: u64,
}
#[doc(hidden)]
pub struct ThreadCpuTimer(Option<cpu_time::ThreadTime>);
impl ThreadCpuTimer {
pub fn start() -> Self {
Self(cpu_time::ThreadTime::try_now().ok())
}
pub fn elapsed(&self) -> Option<Duration> {
self.0.as_ref()?.try_elapsed().ok()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProducerPressure {
Starved,
Satisfied,
SourceBound,
Inelastic,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BrokerConfigError(&'static str);
impl std::fmt::Display for BrokerConfigError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "invalid thread-broker configuration: {}", self.0)
}
}
impl std::error::Error for BrokerConfigError {}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ResizeError(Box<str>);
impl ResizeError {
pub fn new(message: impl Into<String>) -> Self {
Self(message.into().into_boxed_str())
}
}
impl std::fmt::Display for ResizeError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
impl std::error::Error for ResizeError {}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ResizeSide {
Consumer,
Producer,
}
#[derive(Debug)]
pub enum BrokerErrorKind {
ResizeRefused {
side: ResizeSide,
requested: usize,
source: ResizeError,
},
ResizeTimedOut {
side: ResizeSide,
target: usize,
observed: usize,
timeout: Duration,
},
WorkCountersRegressed {
side: ResizeSide,
previous: Work,
observed: Work,
},
ThreadSpawn(std::io::Error),
ThreadPanicked,
}
#[derive(Debug)]
pub struct BrokerError {
kind: BrokerErrorKind,
report: Option<Box<BrokerReport>>,
}
impl BrokerError {
pub(crate) fn new(kind: BrokerErrorKind) -> Self {
Self { kind, report: None }
}
pub(crate) fn with_report(mut self, report: BrokerReport) -> Self {
self.report = Some(Box::new(report));
self
}
pub fn kind(&self) -> &BrokerErrorKind {
&self.kind
}
pub fn report(&self) -> Option<&BrokerReport> {
self.report.as_deref()
}
}
impl std::fmt::Display for BrokerError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match &self.kind {
BrokerErrorKind::ResizeRefused {
side,
requested,
source,
} => write!(
f,
"thread broker {side:?} resize to {requested} was refused: {source}"
),
BrokerErrorKind::ResizeTimedOut {
side,
target,
observed,
timeout,
} => write!(
f,
"thread broker timed out after {timeout:?} waiting for {side:?} to shrink to {target} (observed {observed})"
),
BrokerErrorKind::WorkCountersRegressed {
side,
previous,
observed,
} => write!(
f,
"thread broker observed decreasing {side:?} work counters: {previous:?} -> {observed:?}"
),
BrokerErrorKind::ThreadSpawn(source) => {
write!(f, "could not spawn thread-broker sampler: {source}")
}
BrokerErrorKind::ThreadPanicked => f.write_str("thread-broker sampler panicked"),
}
}
}
impl std::error::Error for BrokerError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match &self.kind {
BrokerErrorKind::ResizeRefused { source, .. } => Some(source),
BrokerErrorKind::ThreadSpawn(source) => Some(source),
_ => None,
}
}
}
#[derive(Debug, Clone, Copy)]
pub struct BrokerConfig {
pub sample_interval: Duration,
pub steady_probe_interval: Option<Duration>,
pub warmup: Duration,
pub smoothing_windows: usize,
pub deadband_threads: usize,
pub blackout_samples: usize,
pub ratify_samples: usize,
pub opening_policy: OpeningPolicy,
pub regression_tolerance: f64,
pub resurvey_distance: usize,
pub resurvey_persistence: Duration,
pub resize_timeout: Duration,
pub max_buffer_drift_items: u64,
pub max_buffer_drift_fraction: f64,
pub cap_history: Duration,
pub cap_persistence: Duration,
pub trace_capacity: usize,
pub min_consumer_threads: usize,
pub min_producer_slots: usize,
}
impl Default for BrokerConfig {
fn default() -> Self {
Self {
sample_interval: Duration::from_millis(100),
steady_probe_interval: None,
warmup: Duration::from_millis(400),
smoothing_windows: 3,
deadband_threads: 1,
blackout_samples: 4,
ratify_samples: 10,
opening_policy: OpeningPolicy::Fixed,
regression_tolerance: 0.05,
resurvey_distance: 1,
resurvey_persistence: Duration::from_millis(800),
resize_timeout: Duration::from_secs(2),
max_buffer_drift_items: 1 << 20,
max_buffer_drift_fraction: 0.05,
cap_history: Duration::from_millis(800),
cap_persistence: Duration::from_millis(300),
trace_capacity: 256,
min_consumer_threads: 1,
min_producer_slots: 1,
}
}
}
pub struct ThreadBroker<C: Consumer, P: Producer> {
pub(crate) consumer: Arc<C>,
pub(crate) producer: Arc<P>,
pub(crate) budget: usize,
pub(crate) config: BrokerConfig,
pub(crate) initial_producer_slots: Option<usize>,
pub(crate) steady_state_policy: SteadyStatePolicy,
}
pub struct ThreadBrokerBuilder<C: Consumer, P: Producer> {
consumer: Arc<C>,
producer: Arc<P>,
budget: Option<usize>,
config: BrokerConfig,
initial_producer_slots: Option<usize>,
steady_state_policy: SteadyStatePolicy,
}
impl<C: Consumer, P: Producer> ThreadBroker<C, P> {
pub fn builder(consumer: C, producer: P) -> ThreadBrokerBuilder<C, P> {
Self::builder_with(consumer, producer, BrokerConfig::default())
}
pub fn builder_with(
consumer: C,
producer: P,
config: BrokerConfig,
) -> ThreadBrokerBuilder<C, P> {
ThreadBrokerBuilder {
consumer: Arc::new(consumer),
producer: Arc::new(producer),
budget: None,
config,
initial_producer_slots: None,
steady_state_policy: SteadyStatePolicy::default(),
}
}
}
impl<C: Consumer, P: Producer> ThreadBrokerBuilder<C, P> {
pub fn budget(mut self, threads: usize) -> Self {
self.budget = Some(threads);
self
}
pub fn initial_producer_slots(mut self, slots: usize) -> Self {
self.initial_producer_slots = Some(slots);
self
}
pub fn steady_state_policy(mut self, policy: SteadyStatePolicy) -> Self {
self.steady_state_policy = policy;
self
}
pub fn sample_interval(mut self, interval: Duration) -> Self {
self.config.sample_interval = interval;
self
}
pub fn steady_probe_interval(mut self, interval: Duration) -> Self {
self.config.steady_probe_interval = Some(interval);
self
}
pub fn warmup(mut self, warmup: Duration) -> Self {
self.config.warmup = warmup;
self
}
pub fn smoothing_windows(mut self, windows: usize) -> Self {
self.config.smoothing_windows = windows;
self
}
pub fn deadband_threads(mut self, threads: usize) -> Self {
self.config.deadband_threads = threads;
self
}
pub fn blackout_samples(mut self, samples: usize) -> Self {
self.config.blackout_samples = samples;
self
}
pub fn ratify_samples(mut self, samples: usize) -> Self {
self.config.ratify_samples = samples;
self
}
pub fn opening_policy(mut self, policy: OpeningPolicy) -> Self {
self.config.opening_policy = policy;
self
}
pub fn regression_tolerance(mut self, fraction: f64) -> Self {
self.config.regression_tolerance = fraction;
self
}
pub fn resurvey_distance(mut self, threads: usize) -> Self {
self.config.resurvey_distance = threads;
self
}
pub fn resurvey_persistence(mut self, duration: Duration) -> Self {
self.config.resurvey_persistence = duration;
self
}
pub fn resize_timeout(mut self, timeout: Duration) -> Self {
self.config.resize_timeout = timeout;
self
}
pub fn min_consumer_threads(mut self, threads: usize) -> Self {
self.config.min_consumer_threads = threads.max(1);
self
}
pub fn min_producer_slots(mut self, slots: usize) -> Self {
self.config.min_producer_slots = slots.max(1);
self
}
pub fn build(self) -> Result<ThreadBroker<C, P>, BrokerConfigError> {
let budget = self.budget.ok_or(BrokerConfigError("budget is required"))?;
let c = &self.config;
if budget < 2 {
return Err(BrokerConfigError(
"budget must be at least 2 to divide between two sides",
));
}
if c.sample_interval.is_zero() {
return Err(BrokerConfigError("sample_interval must be non-zero"));
}
if c.steady_probe_interval
.is_some_and(|interval| interval.is_zero())
{
return Err(BrokerConfigError(
"steady_probe_interval must be non-zero when configured",
));
}
if c.warmup < c.sample_interval * 2 {
return Err(BrokerConfigError(
"warmup must span at least two sample intervals, or the first \
decision is made on the startup transient",
));
}
if !(0.0..1.0).contains(&c.regression_tolerance) {
return Err(BrokerConfigError(
"regression_tolerance must be in 0.0..1.0",
));
}
if c.smoothing_windows == 0 {
return Err(BrokerConfigError("smoothing_windows must be non-zero"));
}
if c.blackout_samples < c.smoothing_windows {
return Err(BrokerConfigError(
"blackout_samples must be at least smoothing_windows, or the \
next decision still averages in pre-move windows",
));
}
if c.ratify_samples == 0 {
return Err(BrokerConfigError("ratify_samples must be non-zero"));
}
if let OpeningPolicy::Bracket(bracket) = c.opening_policy {
if bracket.max_points == 0 {
return Err(BrokerConfigError(
"opening bracket max_points must be non-zero",
));
}
if bracket.horizon.is_zero() {
return Err(BrokerConfigError(
"opening bracket horizon must be non-zero",
));
}
if bracket.total_budget.is_zero() || bracket.total_budget < bracket.horizon {
return Err(BrokerConfigError(
"opening bracket total_budget must be at least its horizon",
));
}
}
if c.resurvey_persistence.is_zero() {
return Err(BrokerConfigError("resurvey_persistence must be non-zero"));
}
if c.resize_timeout.is_zero() {
return Err(BrokerConfigError("resize_timeout must be non-zero"));
}
if !(0.0..=1.0).contains(&c.max_buffer_drift_fraction) {
return Err(BrokerConfigError(
"max_buffer_drift_fraction must be in 0.0..=1.0",
));
}
if c.cap_history.is_zero()
|| c.cap_persistence.is_zero()
|| c.cap_persistence > c.cap_history
{
return Err(BrokerConfigError(
"cap durations must be non-zero and cap_persistence must not exceed cap_history",
));
}
if c.trace_capacity == 0 {
return Err(BrokerConfigError("trace_capacity must be non-zero"));
}
if c.resurvey_distance < c.deadband_threads {
return Err(BrokerConfigError(
"resurvey_distance must be at least deadband_threads",
));
}
if c.min_consumer_threads == 0 || c.min_producer_slots == 0 {
return Err(BrokerConfigError(
"min_consumer_threads and min_producer_slots must be non-zero",
));
}
if c.min_consumer_threads + c.min_producer_slots > budget {
return Err(BrokerConfigError(
"the two floors do not fit inside the budget",
));
}
Ok(ThreadBroker {
consumer: self.consumer,
producer: self.producer,
budget,
config: self.config,
initial_producer_slots: self.initial_producer_slots,
steady_state_policy: self.steady_state_policy,
})
}
}
pub(crate) struct Stop {
value: AtomicBool,
mutex: std::sync::Mutex<()>,
wake: std::sync::Condvar,
}
impl Stop {
pub(crate) fn new() -> Arc<Self> {
Arc::new(Self {
value: AtomicBool::new(false),
mutex: std::sync::Mutex::new(()),
wake: std::sync::Condvar::new(),
})
}
pub(crate) fn is_set(&self) -> bool {
self.value.load(Ordering::Acquire)
}
pub(crate) fn set(&self) {
self.value.store(true, Ordering::Release);
self.wake.notify_all();
}
pub(crate) fn wait_timeout(&self, timeout: Duration) -> bool {
if self.is_set() {
return true;
}
let Ok(guard) = self.mutex.lock() else {
return self.is_set();
};
if self.is_set() {
return true;
}
let _wait = self.wake.wait_timeout(guard, timeout);
self.is_set()
}
}
pub(crate) fn idle_fraction(busy_delta: u64, threads: usize, window: Duration) -> f64 {
let capacity = (threads as u128).saturating_mul(window.as_nanos());
if capacity == 0 {
return 0.0;
}
let used = (busy_delta as f64 / capacity as f64).clamp(0.0, 1.0);
1.0 - used
}
pub(crate) fn now() -> Instant {
Instant::now()
}
#[derive(Debug, Default)]
pub struct BusyMeter {
nanos: std::sync::atomic::AtomicU64,
items: std::sync::atomic::AtomicU64,
}
impl BusyMeter {
pub fn new() -> Self {
Self::default()
}
pub fn work(&self) -> Work {
Work {
busy_nanos: self.nanos.load(Ordering::Relaxed),
items: self.items.load(Ordering::Relaxed),
}
}
pub fn nanos(&self) -> u64 {
self.nanos.load(Ordering::Relaxed)
}
pub fn timer(&self) -> BatchTimer<'_> {
BatchTimer {
meter: self,
single_writer_progress: None,
single_writer_progress_total: 0,
mark: Instant::now(),
items_since_publish: 0,
time_items_since_publish: 0,
progress_every: DEFAULT_FLUSH_EVERY,
time_every: DEFAULT_FLUSH_EVERY,
}
}
}
pub const DEFAULT_FLUSH_EVERY: u64 = 256;
pub struct BatchTimer<'a> {
meter: &'a BusyMeter,
single_writer_progress: Option<&'a std::sync::atomic::AtomicU64>,
single_writer_progress_total: u64,
mark: Instant,
items_since_publish: u64,
time_items_since_publish: u64,
progress_every: u64,
time_every: u64,
}
impl<'a> BatchTimer<'a> {
pub fn single_writer_progress_counter(
mut self,
counter: &'a std::sync::atomic::AtomicU64,
) -> Self {
self.single_writer_progress_total = counter.load(Ordering::Relaxed);
self.single_writer_progress = Some(counter);
self
}
pub fn flush_every(mut self, n: u64) -> Self {
let n = n.max(1);
self.progress_every = n;
self.time_every = n;
self
}
pub fn progress_every(mut self, n: u64) -> Self {
self.progress_every = n.max(1);
self
}
pub fn time_every(mut self, n: u64) -> Self {
self.time_every = n.max(1);
self
}
#[inline]
pub fn tick(&mut self) {
self.items_since_publish += 1;
self.time_items_since_publish += 1;
if self.items_since_publish >= self.progress_every {
self.publish_items();
}
if self.time_items_since_publish >= self.time_every {
self.publish_time();
}
}
#[inline]
pub fn completed_item(&mut self) -> CompletedItem<'_, 'a> {
CompletedItem { timer: self }
}
#[inline]
fn publish_items(&mut self) {
if let Some(counter) = self.single_writer_progress {
self.single_writer_progress_total = self
.single_writer_progress_total
.saturating_add(self.items_since_publish);
counter.store(self.single_writer_progress_total, Ordering::Relaxed);
} else {
self.meter
.items
.fetch_add(self.items_since_publish, Ordering::Relaxed);
}
self.items_since_publish = 0;
}
#[inline]
fn publish_time(&mut self) {
let now = Instant::now();
self.meter
.nanos
.fetch_add((now - self.mark).as_nanos() as u64, Ordering::Relaxed);
self.mark = now;
self.time_items_since_publish = 0;
}
}
pub struct CompletedItem<'timer, 'meter> {
timer: &'timer mut BatchTimer<'meter>,
}
impl Drop for CompletedItem<'_, '_> {
#[inline]
fn drop(&mut self) {
self.timer.tick();
}
}
impl Drop for BatchTimer<'_> {
fn drop(&mut self) {
self.publish_items();
self.publish_time();
}
}