use commonware_utils::time::NANOS_PER_SEC;
use std::time::Duration;
const BYTE_SCALE: u128 = 1 << 32;
#[derive(Clone, Copy, Debug, Default)]
pub(super) struct Score {
averages: Option<(u128, u128)>,
}
impl Score {
pub(super) fn record(&mut self, elapsed: Duration, bytes: usize) {
let bytes = (bytes as u128) * BYTE_SCALE;
let elapsed_ns = elapsed.as_nanos().max(1);
self.averages = Some(match self.averages {
Some((mean_bytes, mean_ns)) => ((mean_bytes + bytes) / 2, (mean_ns + elapsed_ns) / 2),
None => (bytes, elapsed_ns),
});
}
pub(super) fn weight(&self) -> u64 {
let Some((bytes, elapsed_ns)) = self.averages else {
return 1;
};
let rate = bytes * NANOS_PER_SEC / (elapsed_ns * BYTE_SCALE);
u64::try_from(rate).unwrap_or(u64::MAX).saturating_add(1)
}
}
#[cfg(test)]
mod tests {
use super::*;
const FAILURE_DURATION: Duration = Duration::from_secs(2);
#[test]
fn throughput_accounts_for_bytes_and_duration() {
for (bytes, elapsed_ms, expected) in [
(1, 1, 1_001),
(1_024, 1, 1_024_001),
(1_024, 2, 512_001),
(10_240, 20, 512_001),
(10 * 1024 * 1024, 10, 1_048_576_001),
(10 * 1024 * 1024, 5, 2_097_152_001),
] {
let mut score = Score::default();
assert_eq!(score.weight(), 1);
score.record(Duration::from_millis(elapsed_ms), bytes);
assert_eq!(score.weight(), expected);
}
}
#[test]
fn throughput_averages_bytes_and_time_before_dividing() {
let mut score = Score::default();
score.record(Duration::from_millis(1), 1_000);
score.record(Duration::from_millis(999), 1);
assert_eq!(score.weight(), 1_002);
}
#[test]
fn failures_reduce_score_and_service_recovers() {
let mut score = Score::default();
score.record(Duration::from_millis(100), 1_000_000);
assert_eq!(score.weight(), 10_000_001);
score.record(FAILURE_DURATION, 0);
assert_eq!(score.weight(), 476_191);
score.record(Duration::from_millis(100), 1_000_000);
assert_eq!(score.weight(), 1_304_348);
}
#[test]
fn empty_responses_cannot_restore_old_byte_credit() {
let mut score = Score::default();
score.record(Duration::from_millis(100), 1_000_000);
for _ in 0..10 {
let previous = score.weight();
score.record(FAILURE_DURATION, 0);
assert!(score.weight() < previous);
}
for elapsed in [
Duration::from_millis(100),
Duration::from_nanos(1),
Duration::ZERO,
] {
let previous = score.weight();
score.record(elapsed, 0);
assert!(score.weight() <= previous);
assert!(score.weight() < 1_000);
}
}
#[test]
fn small_responses_recover_after_repeated_failures() {
let mut score = Score::default();
score.record(Duration::from_nanos(1), usize::MAX);
let mut previous = score.weight();
for _ in 0..128 {
score.record(FAILURE_DURATION, 0);
assert!(score.weight() <= previous);
previous = score.weight();
}
assert_eq!(previous, 1);
for _ in 0..16 {
score.record(Duration::from_millis(1), 1);
assert!(score.weight() >= previous);
previous = score.weight();
}
assert!(previous > 900);
assert!(previous < 1_001);
}
#[test]
fn failure_before_first_success_is_retained() {
let mut score = Score::default();
score.record(FAILURE_DURATION, 0);
assert_eq!(score.weight(), 1);
score.record(Duration::from_millis(100), 1_000_000);
assert_eq!(score.weight(), 476_191);
}
#[test]
fn zero_and_maximum_inputs_remain_bounded() {
for bytes in [0, 1, usize::MAX] {
for elapsed in [Duration::ZERO, Duration::from_nanos(1), Duration::MAX] {
let mut score = Score::default();
let rate = (bytes as u128) * NANOS_PER_SEC / elapsed.as_nanos().max(1);
let expected = u64::try_from(rate + 1).unwrap_or(u64::MAX);
for _ in 0..2 {
score.record(elapsed, bytes);
assert_eq!(score.weight(), expected);
}
score.record(Duration::MAX, 0);
assert!(score.weight() <= expected);
score.record(Duration::ZERO, usize::MAX);
assert!(score.weight() >= 1);
}
}
}
}