use std::time::Duration;
#[cfg(test)]
mod tests;
pub const DEFAULT_SMOOTHING: f64 = 0.3;
const STAGGER_SCALE: f64 = 0.5;
pub const DEFAULT_MAX_CHUNK_SAMPLES: u32 = 65_536;
#[must_use]
pub fn marginal_rate(delta_samples: u32, elapsed: Duration) -> Option<f64> {
if delta_samples == 0 {
return None;
}
let secs = elapsed.as_secs_f64();
if secs <= 0.0 {
return None;
}
Some(f64::from(delta_samples) / secs)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ChunkPolicy {
pub target: Duration,
pub min_samples: u32,
pub max_samples: u32,
pub calibration_samples: u32,
}
impl ChunkPolicy {
pub const EXPORT: Self = Self {
target: Duration::from_secs(22),
min_samples: 32,
max_samples: DEFAULT_MAX_CHUNK_SAMPLES,
calibration_samples: 8,
};
pub const INTERACTIVE: Self = Self {
target: Duration::from_millis(1500),
min_samples: 1,
max_samples: DEFAULT_MAX_CHUNK_SAMPLES,
calibration_samples: 1,
};
#[must_use]
pub const fn fixed(samples: u32) -> Self {
let samples = if samples == 0 { 1 } else { samples };
Self {
target: Duration::ZERO,
min_samples: samples,
max_samples: samples,
calibration_samples: samples,
}
}
#[must_use]
pub fn samples_for_rate(&self, rate: f64) -> u32 {
self.clamp_samples(rate * self.target.as_secs_f64())
}
#[must_use]
pub(crate) fn scaled_samples(&self, samples: u32, scale: f64) -> u32 {
self.clamp_samples(f64::from(samples) * scale)
}
fn clamp_samples(&self, wanted: f64) -> u32 {
let (min, max) = self.bounds();
if !wanted.is_finite() || wanted < f64::from(min) {
return min;
}
if wanted >= f64::from(max) {
return max;
}
(wanted.round() as u32).clamp(min, max)
}
fn bounds(&self) -> (u32, u32) {
let max = self.max_samples.max(1);
(self.min_samples.clamp(1, max), max)
}
#[must_use]
pub fn first_chunk_samples(&self) -> u32 {
self.calibration_samples.clamp(1, self.bounds().1)
}
#[must_use]
pub fn tail_aware_samples(&self, rate: f64, remaining: u32, sum_rates: f64) -> u32 {
let by_target = self.samples_for_rate(rate);
if remaining == 0 || !sum_rates.is_finite() || sum_rates <= 0.0 {
return by_target;
}
let share = (f64::from(remaining) * rate / sum_rates).ceil();
if !share.is_finite() {
return by_target;
}
let (min, max) = self.bounds();
let share_samples = if share <= 0.0 {
min
} else if share >= f64::from(max) {
max
} else {
(share as u32).clamp(min, max)
};
by_target.min(share_samples).max(min)
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct RateModel {
guess: f64,
estimate: Option<f64>,
smoothing: f64,
observations: u32,
first_chunk_scale: Option<f64>,
}
impl RateModel {
#[must_use]
pub const fn new(initial_guess: f64) -> Self {
Self {
guess: initial_guess,
estimate: None,
smoothing: DEFAULT_SMOOTHING,
observations: 0,
first_chunk_scale: None,
}
}
#[must_use]
pub fn calibrated(rate: f64) -> Self {
let mut model = Self::new(rate);
if valid_rate(rate) {
model.estimate = Some(rate);
}
model
}
#[must_use]
pub fn staggered(rate: f64) -> Self {
let mut model = Self::calibrated(rate);
if model.is_calibrated() {
model.first_chunk_scale = Some(STAGGER_SCALE);
}
model
}
pub const fn take_first_chunk_scale(&mut self) -> Option<f64> {
self.first_chunk_scale.take()
}
#[must_use]
pub const fn with_smoothing(mut self, weight_of_new: f64) -> Self {
self.smoothing = if weight_of_new.is_finite() {
weight_of_new.clamp(0.01, 1.0)
} else {
DEFAULT_SMOOTHING
};
self
}
#[must_use]
pub const fn is_calibrated(&self) -> bool {
self.estimate.is_some()
}
#[must_use]
pub fn rate(&self) -> f64 {
self.estimate.unwrap_or(self.guess)
}
#[must_use]
pub const fn estimate(&self) -> Option<f64> {
self.estimate
}
#[must_use]
pub const fn observations(&self) -> u32 {
self.observations
}
pub fn observe(&mut self, done: u32, elapsed: Duration, reported: Option<f64>) -> Option<f64> {
let measured = reported
.filter(|&rate| valid_rate(rate))
.or_else(|| marginal_rate(done, elapsed))
.filter(|&rate| valid_rate(rate))?;
self.estimate = Some(
self.estimate
.map_or(measured, |old| self.smoothing.mul_add(measured - old, old)),
);
self.observations += 1;
Some(measured)
}
#[must_use]
pub fn chunk_samples(&self, policy: &ChunkPolicy) -> u32 {
self.estimate.map_or_else(
|| policy.first_chunk_samples(),
|rate| policy.samples_for_rate(rate),
)
}
}
fn valid_rate(rate: f64) -> bool {
rate.is_finite() && rate > 0.0
}