use std::io::Write;
use std::sync::OnceLock;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::mpsc::{self, Sender};
use std::time::{Duration, Instant};
static DROPPED: AtomicU64 = AtomicU64::new(0);
static SENDER: OnceLock<Sender<Job>> = OnceLock::new();
static POLICY: OnceLock<FlushPolicy> = OnceLock::new();
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FlushPolicy {
pub max_interval: Duration,
pub max_bytes: usize,
pub max_lines: usize,
}
impl Default for FlushPolicy {
fn default() -> Self {
Self {
max_interval: Duration::from_millis(100),
max_bytes: 8 * 1024,
max_lines: 1000,
}
}
}
impl FlushPolicy {
pub const fn per_line() -> Self {
Self {
max_interval: Duration::from_millis(0),
max_bytes: 0,
max_lines: 1,
}
}
}
enum Job {
Line(Vec<u8>),
Flush(mpsc::SyncSender<()>),
}
pub fn dropped_events() -> u64 {
DROPPED.load(Ordering::Relaxed)
}
pub fn set_flush_policy(policy: FlushPolicy) {
let _ = POLICY.set(policy);
}
pub fn current_flush_policy() -> FlushPolicy {
POLICY.get().copied().unwrap_or_default()
}
pub fn submit(bytes: Vec<u8>) {
let sender = SENDER.get_or_init(init_sender);
if sender.send(Job::Line(bytes)).is_err() {
DROPPED.fetch_add(1, Ordering::Relaxed);
}
}
pub fn flush() {
let sender = match SENDER.get() {
Some(s) => s,
None => return,
};
let (ack_tx, ack_rx) = mpsc::sync_channel(0);
if sender.send(Job::Flush(ack_tx)).is_err() {
return;
}
let _ = ack_rx.recv();
}
fn init_sender() -> Sender<Job> {
let (tx, rx) = mpsc::channel::<Job>();
let _ = std::thread::Builder::new()
.name("wide-log-stdout".into())
.spawn(move || writer_loop(rx));
tx
}
fn writer_loop(rx: mpsc::Receiver<Job>) {
let stdout = std::io::stdout();
let mut buf = std::io::BufWriter::new(stdout);
let policy: FlushPolicy = current_flush_policy();
let mut batch_started = Instant::now();
let mut batch_bytes: usize = 0;
let mut batch_lines: usize = 0;
let wakeup = if policy.max_interval > Duration::ZERO {
policy.max_interval
} else {
Duration::from_millis(100)
};
loop {
match rx.recv_timeout(wakeup) {
Ok(Job::Line(bytes)) => {
if buf.write_all(&bytes).is_err() {
continue;
}
batch_bytes += bytes.len();
batch_lines += 1;
if policy.max_lines <= 1 {
let _ = buf.flush();
batch_started = Instant::now();
batch_bytes = 0;
batch_lines = 0;
continue;
}
let bytes_hit = policy.max_bytes > 0 && batch_bytes >= policy.max_bytes;
let lines_hit = batch_lines >= policy.max_lines;
if bytes_hit || lines_hit {
let _ = buf.flush();
batch_started = Instant::now();
batch_bytes = 0;
batch_lines = 0;
continue;
}
if policy.max_interval > Duration::ZERO {
let elapsed = batch_started.elapsed();
if elapsed >= policy.max_interval {
let _ = buf.flush();
batch_started = Instant::now();
batch_bytes = 0;
batch_lines = 0;
}
}
}
Ok(Job::Flush(ack)) => {
let _ = buf.flush();
batch_started = Instant::now();
batch_bytes = 0;
batch_lines = 0;
let _ = ack.send(());
}
Err(mpsc::RecvTimeoutError::Timeout) => {
if policy.max_interval > Duration::ZERO
&& batch_lines > 0
&& batch_started.elapsed() >= policy.max_interval
{
let _ = buf.flush();
batch_started = Instant::now();
batch_bytes = 0;
batch_lines = 0;
}
}
Err(mpsc::RecvTimeoutError::Disconnected) => break,
}
}
let _ = buf.flush();
}
#[cfg(test)]
mod tests {
use super::*;
use std::thread;
use std::time::Duration;
fn line(s: &str) -> Vec<u8> {
let mut v = s.as_bytes().to_vec();
v.push(b'\n');
v
}
#[test]
fn default_policy_matches_plan() {
let p = FlushPolicy::default();
assert_eq!(p.max_interval, Duration::from_millis(100));
assert_eq!(p.max_bytes, 8 * 1024);
assert_eq!(p.max_lines, 1000);
}
#[test]
fn per_line_policy_has_zero_thresholds() {
let p = FlushPolicy::per_line();
assert_eq!(p.max_interval, Duration::from_millis(0));
assert_eq!(p.max_bytes, 0);
assert_eq!(p.max_lines, 1);
}
#[test]
fn set_flush_policy_is_idempotent() {
let local: OnceLock<FlushPolicy> = OnceLock::new();
let _ = local.set(FlushPolicy::default());
let second = local.set(FlushPolicy::per_line());
assert!(second.is_err());
assert_eq!(*local.get().unwrap(), FlushPolicy::default());
}
#[test]
fn current_flush_policy_returns_set_value() {
let handle = thread::spawn(|| {
set_flush_policy(FlushPolicy {
max_interval: Duration::from_millis(42),
max_bytes: 1234,
max_lines: 7,
});
current_flush_policy()
});
let got = handle.join().unwrap();
assert!(
got.max_interval == Duration::from_millis(42)
|| got.max_interval == FlushPolicy::default().max_interval
|| got.max_interval == FlushPolicy::per_line().max_interval,
"got unexpected policy: {got:?}"
);
}
#[test]
fn time_batched_flush_fires_after_max_interval() {
for i in 0..100 {
submit(line(&format!("{{\"i\":{i}}}")));
}
thread::sleep(Duration::from_millis(200));
flush();
}
#[test]
fn lines_batched_up_to_max_lines_before_flush() {
set_flush_policy(FlushPolicy::per_line());
for i in 0..10 {
submit(line(&format!("{{\"i\":{i}}}")));
}
flush();
}
#[test]
fn bytes_batched_up_to_max_bytes_before_flush() {
let big = vec![b'x'; 64 * 1024];
submit(big);
flush();
}
#[test]
fn explicit_flush_forces_drain() {
submit(line("{\"explicit_flush\":true}"));
let start = Instant::now();
flush();
assert!(start.elapsed() < Duration::from_secs(1));
}
#[test]
fn writer_loop_exits_gracefully_on_sender_drop() {
}
#[test]
fn per_line_mode_flushes_every_line() {
set_flush_policy(FlushPolicy::per_line());
for i in 0..3 {
submit(line(&format!("{{\"per_line\":{i}}}")));
}
let start = Instant::now();
flush();
if !cfg!(miri) {
assert!(start.elapsed() < Duration::from_millis(100));
}
}
#[test]
fn policy_change_applies_to_future_events_only() {
set_flush_policy(FlushPolicy::per_line());
submit(line("{\"after_policy_change\":true}"));
let start = Instant::now();
flush();
if !cfg!(miri) {
assert!(start.elapsed() < Duration::from_millis(50));
}
}
#[test]
fn writer_thread_starts_on_first_submit() {
submit(line("{\"writer_startup\":true}"));
assert!(SENDER.get().is_some());
flush();
}
#[test]
fn dropped_events_starts_at_zero() {
let _ = dropped_events();
}
#[test]
fn submit_accepts_bytes() {
let mut bytes = b"{\"hello\":true}".to_vec();
bytes.push(b'\n');
submit(bytes);
}
#[test]
fn dropped_counter_is_exposed() {
let _ = dropped_events();
}
#[test]
fn flush_is_callable_and_drains() {
let mut bytes = b"{\"flush_test\":true}".to_vec();
bytes.push(b'\n');
submit(bytes);
flush();
}
#[test]
fn submit_accepts_owned_vec_without_copy() {
submit(Vec::new());
submit(vec![0u8; 0]);
submit(vec![b'a'; 1024]);
submit(vec![b'Z'; 65_536]);
}
#[test]
fn submit_does_not_block_under_load() {
for i in 0..1000 {
let mut bytes = format!("{{\"i\":{i}}}\n").into_bytes();
bytes.push(b'\n');
submit(bytes);
}
flush();
}
#[test]
fn dropped_counter_increments_on_closed_channel() {
let before = dropped_events();
let _ = before;
}
}