use std::time::{Duration, Instant};
use tracing::debug;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FlushReason {
Time,
Count,
Size,
Manual,
Shutdown,
}
#[derive(Debug, Clone)]
pub struct BatchConfig {
pub flush_ms: u64,
pub flush_count: usize,
pub flush_bytes: usize,
}
impl Default for BatchConfig {
fn default() -> Self {
Self {
flush_ms: 100,
flush_count: 1000,
flush_bytes: 1024 * 1024, }
}
}
#[derive(Debug)]
pub struct FlushBatch<T> {
pub items: Vec<T>,
pub total_bytes: usize,
pub reason: FlushReason,
}
#[derive(Debug)]
pub struct Batch<T> {
pub items: Vec<T>,
pub total_bytes: usize,
pub created_at: Instant,
}
impl<T> Batch<T> {
pub fn new() -> Self {
Self {
items: Vec::new(),
total_bytes: 0,
created_at: Instant::now(),
}
}
pub fn is_empty(&self) -> bool {
self.items.is_empty()
}
pub fn len(&self) -> usize {
self.items.len()
}
pub fn age(&self) -> Duration {
self.created_at.elapsed()
}
pub fn push(&mut self, item: T, size_bytes: usize) {
self.items.push(item);
self.total_bytes += size_bytes;
}
pub fn take(&mut self) -> Vec<T> {
self.total_bytes = 0;
self.created_at = Instant::now();
std::mem::take(&mut self.items)
}
}
impl<T> Default for Batch<T> {
fn default() -> Self {
Self::new()
}
}
pub struct HybridBatcher<T> {
config: BatchConfig,
batch: Batch<T>,
}
impl<T> HybridBatcher<T> {
pub fn new(config: BatchConfig) -> Self {
Self {
config,
batch: Batch::new(),
}
}
pub fn push(&mut self, item: T, size_bytes: usize) -> Option<FlushReason> {
self.batch.push(item, size_bytes);
if self.batch.len() >= self.config.flush_count {
Some(FlushReason::Count)
} else if self.batch.total_bytes >= self.config.flush_bytes {
Some(FlushReason::Size)
} else {
None
}
}
#[must_use]
pub fn should_flush_time(&self) -> bool {
!self.batch.is_empty()
&& self.batch.age() >= Duration::from_millis(self.config.flush_ms)
}
pub fn take_batch(&mut self) -> Vec<T> {
let count = self.batch.len();
let bytes = self.batch.total_bytes;
let items = self.batch.take();
debug!(count, bytes, "Batch taken for flush");
items
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.batch.is_empty()
}
#[must_use]
pub fn stats(&self) -> (usize, usize, Duration) {
(self.batch.len(), self.batch.total_bytes, self.batch.age())
}
}
impl<T: SizedItem> HybridBatcher<T> {
pub fn add(&mut self, item: T) -> Option<FlushReason> {
let size = item.size_bytes();
self.push(item, size)
}
pub fn add_batch(&mut self, items: Vec<T>) {
for item in items {
self.add(item);
}
}
pub fn take_if_ready(&mut self) -> Option<FlushBatch<T>> {
let reason = if self.batch.len() >= self.config.flush_count {
Some(FlushReason::Count)
} else if self.batch.total_bytes >= self.config.flush_bytes {
Some(FlushReason::Size)
} else if self.should_flush_time() {
Some(FlushReason::Time)
} else {
None
};
reason.map(|r| {
let total_bytes = self.batch.total_bytes;
FlushBatch {
items: self.batch.take(),
total_bytes,
reason: r,
}
})
}
pub fn force_flush(&mut self) -> Option<FlushBatch<T>> {
self.force_flush_with_reason(FlushReason::Manual)
}
pub fn force_flush_with_reason(&mut self, reason: FlushReason) -> Option<FlushBatch<T>> {
if self.batch.is_empty() {
return None;
}
let total_bytes = self.batch.total_bytes;
Some(FlushBatch {
items: self.batch.take(),
total_bytes,
reason,
})
}
}
impl<T: BatchableItem> HybridBatcher<T> {
#[must_use]
pub fn contains(&self, id: &str) -> bool {
self.batch.items.iter().any(|item| item.id() == id)
}
}
pub trait SizedItem {
#[must_use]
fn size_bytes(&self) -> usize;
}
pub trait BatchableItem: SizedItem {
fn id(&self) -> &str;
}
#[cfg(test)]
mod tests {
use super::*;
use std::thread::sleep;
#[derive(Debug, Clone)]
#[allow(dead_code)]
struct TestItem {
id: String,
size: usize,
}
impl SizedItem for TestItem {
fn size_bytes(&self) -> usize {
self.size
}
}
fn item(id: &str, size: usize) -> TestItem {
TestItem { id: id.to_string(), size }
}
#[test]
fn test_batch_empty_initially() {
let batcher: HybridBatcher<TestItem> = HybridBatcher::new(BatchConfig::default());
assert!(batcher.is_empty());
let (count, bytes, _) = batcher.stats();
assert_eq!(count, 0);
assert_eq!(bytes, 0);
}
#[test]
fn test_batch_tracks_items_and_bytes() {
let mut batcher = HybridBatcher::new(BatchConfig::default());
batcher.add(item("a", 100));
batcher.add(item("b", 200));
batcher.add(item("c", 150));
let (count, bytes, _) = batcher.stats();
assert_eq!(count, 3);
assert_eq!(bytes, 450);
assert!(!batcher.is_empty());
}
#[test]
fn test_flush_on_count_threshold() {
let config = BatchConfig {
flush_count: 3,
flush_bytes: 1_000_000,
flush_ms: 10_000,
};
let mut batcher = HybridBatcher::new(config);
assert!(batcher.add(item("a", 100)).is_none());
assert!(batcher.add(item("b", 100)).is_none());
let reason = batcher.add(item("c", 100));
assert_eq!(reason, Some(FlushReason::Count));
}
#[test]
fn test_flush_on_size_threshold() {
let config = BatchConfig {
flush_count: 1000,
flush_bytes: 500,
flush_ms: 10_000,
};
let mut batcher = HybridBatcher::new(config);
assert!(batcher.add(item("a", 200)).is_none());
assert!(batcher.add(item("b", 200)).is_none());
let reason = batcher.add(item("c", 200));
assert_eq!(reason, Some(FlushReason::Size));
}
#[test]
fn test_flush_on_time_threshold() {
let config = BatchConfig {
flush_count: 1000,
flush_bytes: 1_000_000,
flush_ms: 10, };
let mut batcher = HybridBatcher::new(config);
batcher.add(item("a", 100));
assert!(!batcher.should_flush_time());
sleep(Duration::from_millis(15));
assert!(batcher.should_flush_time());
}
#[test]
fn test_take_if_ready_returns_batch() {
let config = BatchConfig {
flush_count: 2,
flush_bytes: 1_000_000,
flush_ms: 10_000,
};
let mut batcher = HybridBatcher::new(config);
batcher.add(item("a", 100));
assert!(batcher.take_if_ready().is_none());
batcher.add(item("b", 200));
let batch = batcher.take_if_ready().unwrap();
assert_eq!(batch.items.len(), 2);
assert_eq!(batch.total_bytes, 300);
assert_eq!(batch.reason, FlushReason::Count);
assert!(batcher.is_empty());
}
#[test]
fn test_force_flush() {
let mut batcher = HybridBatcher::new(BatchConfig::default());
batcher.add(item("a", 100));
batcher.add(item("b", 200));
let batch = batcher.force_flush().unwrap();
assert_eq!(batch.items.len(), 2);
assert_eq!(batch.total_bytes, 300);
assert_eq!(batch.reason, FlushReason::Manual);
assert!(batcher.is_empty());
assert!(batcher.force_flush().is_none());
batcher.add(item("c", 100));
let batch = batcher.force_flush_with_reason(FlushReason::Shutdown).unwrap();
assert_eq!(batch.reason, FlushReason::Shutdown);
}
#[test]
fn test_take_resets_batch() {
let mut batcher = HybridBatcher::new(BatchConfig::default());
batcher.add(item("a", 100));
batcher.add(item("b", 200));
let items = batcher.take_batch();
assert_eq!(items.len(), 2);
let (count, bytes, age) = batcher.stats();
assert_eq!(count, 0);
assert_eq!(bytes, 0);
assert!(age < Duration::from_millis(10)); }
#[test]
fn test_add_batch() {
let mut batcher = HybridBatcher::new(BatchConfig::default());
let items = vec![item("a", 100), item("b", 200), item("c", 300)];
batcher.add_batch(items);
let (count, bytes, _) = batcher.stats();
assert_eq!(count, 3);
assert_eq!(bytes, 600);
}
#[test]
fn test_count_beats_size_on_simultaneous_threshold() {
let config = BatchConfig {
flush_count: 2,
flush_bytes: 200,
flush_ms: 10_000,
};
let mut batcher = HybridBatcher::new(config);
batcher.add(item("a", 100));
let reason = batcher.add(item("b", 100));
assert_eq!(reason, Some(FlushReason::Count));
}
}