use std::time::Duration;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct WorkBudget {
pub max_input_bytes: usize,
pub max_output_bytes: usize,
pub max_items: usize,
pub max_cpu_time: Duration,
}
impl WorkBudget {
pub fn new(
max_input_bytes: usize,
max_output_bytes: usize,
max_items: usize,
max_cpu_time: Duration,
) -> Self {
Self {
max_input_bytes,
max_output_bytes,
max_items,
max_cpu_time,
}
}
pub fn unlimited() -> Self {
Self {
max_input_bytes: usize::MAX,
max_output_bytes: usize::MAX,
max_items: usize::MAX,
max_cpu_time: Duration::from_secs(3600),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Step<Progress> {
Yield(Progress),
Complete,
}
pub trait BoundedJob: Send {
type Progress: Send;
fn step(&mut self, budget: WorkBudget) -> anyhow::Result<Step<Self::Progress>>;
fn estimated_next_bytes(&self) -> usize;
}
#[cfg(test)]
mod test {
use super::*;
struct CountingJob {
total: usize,
processed: usize,
bytes_per_item: usize,
}
impl CountingJob {
fn new(total: usize) -> Self {
Self {
total,
processed: 0,
bytes_per_item: 10,
}
}
}
impl BoundedJob for CountingJob {
type Progress = usize;
fn step(&mut self, budget: WorkBudget) -> anyhow::Result<Step<Self::Progress>> {
if self.processed >= self.total {
return Ok(Step::Complete);
}
let take = budget.max_items.min(self.total - self.processed);
self.processed += take;
if self.processed >= self.total {
Ok(Step::Complete)
} else {
Ok(Step::Yield(self.processed))
}
}
fn estimated_next_bytes(&self) -> usize {
if self.processed >= self.total {
0
} else {
(self.total - self.processed) * self.bytes_per_item
}
}
}
fn budget_max_items(items: usize) -> WorkBudget {
WorkBudget::new(usize::MAX, usize::MAX, items, Duration::from_secs(3600))
}
#[test]
fn test_work_budget_bounds_step() {
let mut job = CountingJob::new(100);
let budget = budget_max_items(30);
let seen = [
job.step(budget).unwrap(),
job.step(budget).unwrap(),
job.step(budget).unwrap(),
job.step(budget).unwrap(),
];
assert_eq!(seen[0], Step::Yield(30));
assert_eq!(seen[1], Step::Yield(60));
assert_eq!(seen[2], Step::Yield(90));
assert_eq!(seen[3], Step::Complete);
assert_eq!(job.processed, 100);
}
#[test]
fn test_work_budget_never_oversteps() {
let budget = budget_max_items(7);
let mut job = CountingJob::new(23);
assert_eq!(job.step(budget).unwrap(), Step::Yield(7));
assert_eq!(job.step(budget).unwrap(), Step::Yield(14));
assert_eq!(job.step(budget).unwrap(), Step::Yield(21));
assert_eq!(job.step(budget).unwrap(), Step::Complete);
}
#[test]
fn test_step_yield_complete_semantics() {
let mut job = CountingJob::new(10);
let budget = budget_max_items(5);
assert_eq!(job.step(budget).unwrap(), Step::Yield(5));
assert_eq!(job.step(budget).unwrap(), Step::Complete);
let processed_before = job.processed;
assert_eq!(job.step(budget).unwrap(), Step::Complete);
assert_eq!(job.processed, processed_before);
}
#[test]
fn test_complete_only_when_all_done() {
let mut job = CountingJob::new(5);
let budget = budget_max_items(2);
assert_eq!(job.step(budget).unwrap(), Step::Yield(2));
assert_eq!(job.step(budget).unwrap(), Step::Yield(4));
assert_eq!(job.step(budget).unwrap(), Step::Complete);
}
#[test]
fn test_estimated_next_bytes() {
let mut job = CountingJob::new(5);
assert!(
job.estimated_next_bytes() > 0,
"non-complete job must estimate > 0"
);
let budget = WorkBudget::unlimited();
loop {
if job.step(budget).unwrap() == Step::Complete {
break;
}
}
assert_eq!(job.estimated_next_bytes(), 0, "complete job estimates 0");
}
#[test]
fn test_estimate_reflects_remaining() {
let mut job = CountingJob::new(10);
let first = job.estimated_next_bytes();
job.step(budget_max_items(4)).unwrap();
let second = job.estimated_next_bytes();
assert!(second < first, "estimate shrinks as work advances");
}
}