use std::time::{Duration, Instant};
#[derive(Debug)]
pub struct Batcher<T> {
items: Vec<T>,
max_size: usize,
timeout: Duration,
last_flush: Instant,
}
impl<T> Batcher<T> {
pub fn new(max_size: usize, timeout: Duration) -> Self {
Self {
items: Vec::with_capacity(max_size),
max_size,
timeout,
last_flush: Instant::now(),
}
}
pub fn add(&mut self, item: T) {
self.items.push(item);
}
pub fn should_flush(&self) -> bool {
self.items.len() >= self.max_size || self.last_flush.elapsed() >= self.timeout
}
pub fn is_empty(&self) -> bool {
self.items.is_empty()
}
pub fn len(&self) -> usize {
self.items.len()
}
pub fn flush(&mut self) -> Vec<T> {
self.last_flush = Instant::now();
std::mem::take(&mut self.items)
}
pub fn items(&self) -> &[T] {
&self.items
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::thread;
#[test]
fn test_batcher_size() {
let mut batcher = Batcher::new(3, Duration::from_secs(10));
assert!(!batcher.should_flush());
batcher.add(1);
batcher.add(2);
assert!(!batcher.should_flush());
batcher.add(3);
assert!(batcher.should_flush());
let items = batcher.flush();
assert_eq!(items, vec![1, 2, 3]);
assert!(!batcher.should_flush());
}
#[test]
fn test_batcher_timeout() {
let mut batcher = Batcher::new(10, Duration::from_millis(50));
batcher.add(1);
assert!(!batcher.should_flush());
thread::sleep(Duration::from_millis(60));
assert!(batcher.should_flush());
let items = batcher.flush();
assert_eq!(items, vec![1]);
}
#[test]
fn test_batcher_empty() {
let batcher: Batcher<i32> = Batcher::new(10, Duration::from_secs(1));
assert!(batcher.is_empty());
assert_eq!(batcher.len(), 0);
}
}