use super::*;
use crate::fs::{CountingFs, Fs, FsOpenOptions, MemFs};
use crate::path::Path;
const DAY: Duration = Duration::from_secs(86_400);
struct Medium {
fs: CountingFs,
files: core::cell::Cell<u32>,
}
impl Medium {
fn new() -> Self {
let fs = CountingFs::new(Arc::new(MemFs::new()));
fs.create_dir_all(Path::new("/db")).expect("mkdir");
Self {
fs,
files: core::cell::Cell::new(0),
}
}
fn ledger(&self) -> Arc<WriteIoStats> {
Arc::clone(self.fs.stats())
}
fn write(&self, origin: WriteOrigin, len: usize) {
let n = self.files.get();
self.files.set(n + 1);
let path = alloc::format!("/db/{n}");
let mut file = self
.fs
.open(
Path::new(&path),
&FsOpenOptions::new()
.write(true)
.create_new(true)
.write_origin(origin),
)
.expect("open");
file.write_all(&alloc::vec![0; len]).expect("write");
}
}
fn budget(medium: &Medium, allowance: Allowance) -> Arc<WriteBudget> {
Arc::new(WriteBudget::new_at(
medium.ledger(),
allowance,
Duration::ZERO,
))
}
#[test]
fn a_fresh_budget_grants_its_burst_and_no_more() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(1_000, DAY));
let first = budget
.reserve_at(700, Duration::ZERO)
.expect("within the burst");
assert!(budget.reserve_at(301, Duration::ZERO).is_none());
let second = budget
.reserve_at(300, Duration::ZERO)
.expect("the rest of the burst");
assert_eq!(budget.report_at(Duration::ZERO).reserved, 1_000);
drop((first, second));
assert_eq!(budget.report_at(Duration::ZERO).reserved, 0);
}
#[test]
fn the_refill_follows_wall_time_up_to_the_burst() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(1_000, DAY).with_burst(500));
assert_eq!(budget.report_at(Duration::ZERO).credit, 500);
medium.write(WriteOrigin::FlushTable, 500);
assert_eq!(budget.report_at(Duration::ZERO).credit, 0);
assert_eq!(budget.report_at(DAY / 4).credit, 250);
assert_eq!(budget.report_at(DAY * 10).credit, 500);
}
#[test]
fn polling_often_refills_as_much_as_polling_once() {
let medium = Medium::new();
let budget = budget(
&medium,
Allowance::new(3, Duration::from_secs(10)).with_burst(3),
);
medium.write(WriteOrigin::FlushTable, 3);
for ms in (0..=10_000).step_by(7) {
let _ = budget.report_at(Duration::from_millis(ms));
}
assert_eq!(budget.report_at(Duration::from_secs(10)).credit, 3);
}
#[test]
fn every_written_byte_is_charged_once_split_by_who_wrote_it() {
let medium = Medium::new();
medium.write(WriteOrigin::FlushTable, 4_000);
let budget = budget(&medium, Allowance::new(10_000, DAY));
medium.write(WriteOrigin::FlushTable, 100);
medium.write(WriteOrigin::Manifest, 20);
medium.write(WriteOrigin::CompactionTable, 300);
medium.write(WriteOrigin::CompactionBlob, 5);
let report = budget.report_at(Duration::ZERO);
assert_eq!(report.credit, 10_000 - 425);
assert_eq!(
(report.discretionary_bytes, report.mandatory_bytes),
(305, 120)
);
assert_eq!(budget.report_at(Duration::ZERO).credit, 10_000 - 425);
}
#[test]
fn a_landed_output_moves_from_the_hold_to_the_charge() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(1_000, DAY));
let hold = budget
.reserve_at(600, Duration::ZERO)
.expect("within the burst");
medium.write(WriteOrigin::CompactionTable, 250);
hold.landed(250);
let report = budget.report_at(Duration::ZERO);
assert_eq!((report.credit, report.reserved), (750, 350));
let _rest = budget
.reserve_at(400, Duration::ZERO)
.expect("what the landed part freed");
assert!(budget.reserve_at(1, Duration::ZERO).is_none());
}
#[test]
fn a_refusal_is_counted_only_when_no_override_answers() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(100, DAY));
let asked = core::cell::Cell::new(0);
let fits = budget
.reserve_or_override_at(100, Duration::ZERO, || {
asked.set(asked.get() + 1);
None
})
.expect("within the burst");
assert_eq!(asked.get(), 0, "credit covered it: nothing asked");
let over = budget
.reserve_or_override_at(50, Duration::ZERO, || Some(Override::Forced))
.expect("forced past the budget");
let report = budget.report_at(Duration::ZERO);
assert_eq!(report.refused, 0);
assert_eq!(report.forced_overrides.reservations, 1);
assert!(
budget
.reserve_or_override_at(1, Duration::ZERO, || None)
.is_none()
);
assert_eq!(budget.report_at(Duration::ZERO).refused, 1);
drop((fits, over));
}
#[test]
fn credit_freed_while_asking_grants_ordinary_work() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(100, DAY));
let first = budget.reserve_at(100, Duration::ZERO).expect("the burst");
let hold = budget
.reserve_or_override_at(100, Duration::ZERO, || {
drop(first);
None
})
.expect("freed while asked");
let report = budget.report_at(Duration::ZERO);
assert_eq!(report.refused, 0);
assert_eq!(report.forced_overrides.reservations, 0);
assert_eq!(report.backpressure_overrides.reservations, 0);
drop(hold);
}
#[test]
fn an_output_past_its_hold_is_charged_without_going_below_zero() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(1_000, DAY));
let hold = budget
.reserve_at(100, Duration::ZERO)
.expect("within the burst");
medium.write(WriteOrigin::CompactionTable, 400);
hold.landed(400);
assert_eq!(hold.held(), 0);
let report = budget.report_at(Duration::ZERO);
assert_eq!((report.credit, report.reserved), (600, 0));
}
#[test]
fn an_overrun_refuses_work_until_the_refill_repays_it() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(1_000, DAY));
medium.write(WriteOrigin::FlushTable, 1_500);
assert_eq!(budget.report_at(Duration::ZERO).credit, -500);
let early = DAY / 10;
assert!(budget.reserve_at(1, early).is_none());
let report = budget.report_at(early);
assert_eq!(report.credit, -400);
assert_eq!(report.refused, 1);
assert_eq!(report.waiting_since, Some(early));
assert!(budget.reserve_at(1, DAY / 2).is_none());
assert_eq!(budget.report_at(DAY / 2).waiting_since, Some(early));
let hold = budget
.reserve_at(100, DAY / 2 + DAY / 10)
.expect("refilled past the overrun");
assert_eq!(budget.report_at(DAY).waiting_since, None);
drop(hold);
}
#[test]
fn an_override_is_granted_charged_and_counted_by_its_reason() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(100, DAY));
let hold = budget.reserve_overriding_at(500, Override::Backpressure, Duration::ZERO);
medium.write(WriteOrigin::CompactionTable, 500);
hold.landed(500);
let forced = budget.reserve_overriding_at(7, Override::Forced, Duration::ZERO);
let report = budget.report_at(Duration::ZERO);
assert_eq!(report.credit, -400);
assert_eq!(
(report.backpressure_overrides, report.forced_overrides),
(
Overrides {
reservations: 1,
bytes: 500
},
Overrides {
reservations: 1,
bytes: 7
}
)
);
assert!(budget.reserve_at(1, Duration::ZERO).is_none());
drop(forced);
}
#[test]
fn a_resumed_budget_refills_only_the_time_since_its_checkpoint() {
let medium = Medium::new();
let before = budget(&medium, Allowance::new(1_000, DAY));
medium.write(WriteOrigin::FlushTable, 900);
let checkpoint = before.checkpoint_at(Duration::ZERO);
assert_eq!(checkpoint.credit, 100);
let restarted = Medium::new();
let resumed = Arc::new(WriteBudget::resume_at(
restarted.ledger(),
Allowance::new(1_000, DAY),
checkpoint,
50,
DAY / 10,
));
assert_eq!(resumed.report_at(DAY / 10).credit, 100 - 50 + 100);
let same_time = WriteBudget::resume_at(
restarted.ledger(),
Allowance::new(1_000, DAY),
checkpoint,
0,
Duration::ZERO,
);
assert_eq!(same_time.report_at(Duration::ZERO).credit, 100);
}
#[test]
fn unaccounted_writes_are_charged_after_the_restart_refill() {
let resumed = WriteBudget::resume_at(
Medium::new().ledger(),
Allowance::new(1_000, DAY),
BudgetCheckpoint {
credit: 1_000,
as_of: Duration::ZERO,
},
900,
DAY,
);
assert_eq!(resumed.report_at(DAY).credit, 100);
}
#[test]
fn an_overrun_survives_a_restart() {
let medium = Medium::new();
let before = budget(&medium, Allowance::new(1_000, DAY));
medium.write(WriteOrigin::FlushTable, 3_000);
let checkpoint = before.checkpoint_at(Duration::ZERO);
assert_eq!(checkpoint.credit, -2_000);
let resumed = Arc::new(WriteBudget::resume_at(
Medium::new().ledger(),
Allowance::new(1_000, DAY),
checkpoint,
0,
DAY,
));
assert!(resumed.reserve_at(1, DAY).is_none());
assert_eq!(resumed.report_at(DAY).credit, -1_000);
}
#[test]
fn a_clock_stepping_back_refills_nothing() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(1_000, DAY));
medium.write(WriteOrigin::FlushTable, 1_000);
assert_eq!(budget.report_at(Duration::ZERO).credit, 0);
assert_eq!(budget.report_at(DAY / 2).credit, 500);
assert_eq!(budget.report_at(DAY / 4).credit, 500);
assert_eq!(budget.report_at(DAY / 2).credit, 500);
assert_eq!(budget.report_at(DAY / 2 + DAY / 10).credit, 600);
}
#[test]
fn a_write_is_charged_when_the_budget_reads_it() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(1_000, DAY));
medium.write(WriteOrigin::FlushTable, 1_000);
assert_eq!(budget.report_at(DAY / 2).credit, 0);
}
#[test]
fn a_new_allowance_applies_from_its_change() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(1_000, DAY));
medium.write(WriteOrigin::FlushTable, 1_000);
assert_eq!(budget.report_at(Duration::ZERO).credit, 0);
budget.set_allowance_at(Allowance::new(100, DAY).with_burst(300), DAY / 2);
assert_eq!(budget.report_at(DAY / 2).credit, 300);
medium.write(WriteOrigin::FlushTable, 300);
assert_eq!(budget.report_at(DAY / 2).credit, 0);
assert_eq!(budget.report_at(DAY).credit, 50);
}
#[test]
fn only_mandatory_writes_past_the_allowance_are_unsustainable() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(1_000, DAY));
medium.write(WriteOrigin::CompactionTable, 5_000);
assert!(!budget.report_at(DAY).unsustainable);
medium.write(WriteOrigin::FlushTable, 1_200);
assert!(budget.report_at(DAY).unsustainable);
assert!(!budget.report_at(DAY * 3).unsustainable);
}
#[test]
fn the_last_period_slides_over_two() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(1_000, DAY).with_burst(100));
medium.write(WriteOrigin::FlushTable, 800);
let _ = budget.report_at(DAY / 2);
medium.write(WriteOrigin::FlushTable, 600);
let end_of_day = DAY.checked_sub(Duration::from_secs(1)).expect("a day");
assert!(budget.report_at(end_of_day).unsustainable);
assert!(budget.report_at(DAY + DAY / 4).unsustainable);
assert!(!budget.report_at(DAY + DAY * 3 / 4).unsustainable);
}
#[test]
fn a_cut_allowance_is_unsustainable_once_past_its_burst() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(1_000, DAY));
budget.set_allowance_at(Allowance::new(240, DAY).with_burst(10), DAY);
for hour in 1..=3u32 {
medium.write(WriteOrigin::FlushTable, 30);
let _ = budget.report_at(DAY + DAY / 24 * hour);
}
assert!(budget.report_at(DAY + DAY / 8).unsustainable);
}
#[test]
fn a_burst_within_its_refill_time_is_sustainable() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(240, DAY).with_burst(100));
medium.write(WriteOrigin::FlushTable, 100);
assert!(!budget.report_at(DAY / 24).unsustainable);
assert!(!budget.report_at(DAY * 10 / 24).unsustainable);
medium.write(WriteOrigin::FlushTable, 50);
assert!(budget.report_at(DAY * 10 / 24).unsustainable);
}
#[test]
fn the_report_forecasts_the_waiting_work() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(1_000, DAY));
let held = budget.reserve_at(200, Duration::ZERO).expect("the burst");
medium.write(WriteOrigin::FlushTable, 800);
assert!(budget.reserve_at(300, Duration::ZERO).is_none());
assert!(budget.reserve_at(100, Duration::ZERO).is_none());
let report = budget.report_at(Duration::ZERO);
assert_eq!(report.waiting_bytes, 300, "the largest that waits");
assert_eq!(report.credit_in, Some(DAY * 3 / 10));
assert_eq!(budget.report_at(DAY * 3 / 20).credit_in, Some(DAY * 3 / 20));
let granted = budget.reserve_at(300, DAY * 3 / 10).expect("covered");
let report = budget.report_at(DAY * 3 / 10);
assert_eq!(
(report.waiting_bytes, report.credit_in),
(0, Some(Duration::ZERO))
);
drop((held, granted));
}
#[test]
fn a_refused_zero_byte_request_is_forecast_from_the_overrun() {
let medium = Medium::new();
let overrun = budget(&medium, Allowance::new(1_000, DAY));
medium.write(WriteOrigin::FlushTable, 2_000);
assert!(overrun.reserve_at(0, Duration::ZERO).is_none());
let report = overrun.report_at(Duration::ZERO);
assert_eq!(report.waiting_since, Some(Duration::ZERO));
assert_eq!(report.credit_in, Some(DAY));
let dry = budget(&medium, Allowance::new(0, DAY).with_burst(500));
medium.write(WriteOrigin::FlushTable, 600);
assert!(dry.reserve_at(0, Duration::ZERO).is_none());
assert_eq!(dry.report_at(Duration::ZERO).credit_in, None);
}
#[test]
fn a_clock_stepped_back_forecasts_the_time_to_catch_up() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(1_000, DAY));
assert!(!budget.report_at(DAY).unsustainable);
medium.write(WriteOrigin::FlushTable, 1_000);
assert!(budget.reserve_at(1_000, DAY).is_none());
assert_eq!(budget.report_at(Duration::ZERO).credit_in, Some(DAY * 2));
let just_before = Duration::from_secs(2 * 86_400 - 1);
assert!(budget.reserve_at(1_000, just_before).is_none());
assert!(budget.reserve_at(1_000, DAY * 2).is_some());
}
#[test]
fn work_past_the_burst_is_forecast_never() {
let medium = Medium::new();
let narrow = budget(&medium, Allowance::new(1_000, DAY).with_burst(500));
assert!(narrow.reserve_at(501, Duration::ZERO).is_none());
assert_eq!(narrow.report_at(Duration::ZERO).credit_in, None);
let dry = budget(&medium, Allowance::new(0, DAY).with_burst(500));
medium.write(WriteOrigin::FlushTable, 500);
assert!(dry.reserve_at(1, Duration::ZERO).is_none());
assert_eq!(dry.report_at(Duration::ZERO).credit_in, None);
}
#[test]
fn a_space_override_is_counted_apart() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(10, DAY));
let hold = budget
.reserve_or_override_at(100, Duration::ZERO, || Some(Override::Space))
.expect("overridden");
let report = budget.report_at(Duration::ZERO);
assert_eq!(
report.space_overrides,
Overrides {
reservations: 1,
bytes: 100
}
);
assert_eq!(report.forced_overrides, Overrides::default());
assert_eq!(report.refused, 0);
drop(hold);
}
#[test]
fn without_a_refill_mandatory_writes_add_up_against_the_burst() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(0, DAY).with_burst(1_000));
medium.write(WriteOrigin::FlushTable, 600);
assert!(!budget.report_at(DAY / 2).unsustainable);
medium.write(WriteOrigin::FlushTable, 600);
assert!(budget.report_at(DAY * 3 / 2).unsustainable);
assert!(
budget.report_at(DAY * 10).unsustainable,
"nothing repays it"
);
}
#[test]
fn forcing_is_set_and_cleared_by_the_caller() {
let medium = Medium::new();
let budget = budget(&medium, Allowance::new(1_000, DAY));
assert!(!budget.forced());
budget.set_forced(true);
assert!(budget.forced());
budget.set_forced(false);
assert!(!budget.forced());
}