#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
pub enum FlushPolicy {
#[default]
Manual,
EveryNRecords(u64),
EveryNBytes(u64),
}
impl FlushPolicy {
pub fn is_auto(&self) -> bool {
!matches!(self, FlushPolicy::Manual)
}
}
#[derive(Debug, Clone, Default)]
pub struct StreamingConfig {
pub policy: FlushPolicy,
}
impl StreamingConfig {
pub fn new() -> Self {
Self::default()
}
pub fn every_n_records(n: u64) -> Self {
Self {
policy: FlushPolicy::EveryNRecords(n),
}
}
pub fn every_n_bytes(n: u64) -> Self {
Self {
policy: FlushPolicy::EveryNBytes(n),
}
}
}
#[derive(Debug, Default)]
pub(super) struct FlushState {
pub records_since_flush: u64,
pub bytes_since_flush: u64,
pub total_records: u64,
pub total_bytes: u64,
pub flush_count: u64,
}
impl FlushState {
pub fn record_write(&mut self, records: u64, bytes: u64) {
self.records_since_flush += records;
self.bytes_since_flush += bytes;
self.total_records += records;
self.total_bytes += bytes;
}
pub fn should_flush(&self, policy: &FlushPolicy) -> bool {
match policy {
FlushPolicy::Manual => false,
FlushPolicy::EveryNRecords(n) => self.records_since_flush >= *n,
FlushPolicy::EveryNBytes(n) => self.bytes_since_flush >= *n,
}
}
pub fn on_flush(&mut self) {
self.records_since_flush = 0;
self.bytes_since_flush = 0;
self.flush_count += 1;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_flush_policy_default() {
assert_eq!(FlushPolicy::default(), FlushPolicy::Manual);
}
#[test]
fn test_flush_policy_is_auto() {
assert!(!FlushPolicy::Manual.is_auto());
assert!(FlushPolicy::EveryNRecords(100).is_auto());
assert!(FlushPolicy::EveryNBytes(1024).is_auto());
}
#[test]
fn test_flush_state_should_flush() {
let mut state = FlushState::default();
state.record_write(1000, 10000);
assert!(!state.should_flush(&FlushPolicy::Manual));
assert!(!state.should_flush(&FlushPolicy::EveryNRecords(1001)));
assert!(state.should_flush(&FlushPolicy::EveryNRecords(1000)));
assert!(state.should_flush(&FlushPolicy::EveryNRecords(500)));
assert!(!state.should_flush(&FlushPolicy::EveryNBytes(10001)));
assert!(state.should_flush(&FlushPolicy::EveryNBytes(10000)));
assert!(state.should_flush(&FlushPolicy::EveryNBytes(5000)));
}
#[test]
fn test_flush_state_reset() {
let mut state = FlushState::default();
state.record_write(100, 1000);
assert_eq!(state.records_since_flush, 100);
assert_eq!(state.bytes_since_flush, 1000);
state.on_flush();
assert_eq!(state.records_since_flush, 0);
assert_eq!(state.bytes_since_flush, 0);
assert_eq!(state.total_records, 100); assert_eq!(state.total_bytes, 1000);
assert_eq!(state.flush_count, 1);
}
#[test]
fn test_streaming_config_constructors() {
let config = StreamingConfig::every_n_records(500);
assert_eq!(config.policy, FlushPolicy::EveryNRecords(500));
let config = StreamingConfig::every_n_bytes(1024);
assert_eq!(config.policy, FlushPolicy::EveryNBytes(1024));
}
}