use core::cell::Cell;
use std::io::{self, BufWriter, Write};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::mpsc::{self, Receiver, SyncSender, TryRecvError, TrySendError};
use std::sync::{Arc, Mutex, PoisonError, RwLock, RwLockReadGuard, TryLockError};
use std::thread::JoinHandle;
use super::{Sink, lock_recover, trim_buffer};
use crate::error::{Error, Result};
use crate::format::Format;
use crate::record::Record;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[non_exhaustive]
pub enum Overflow {
#[default]
Block,
DropNewest,
}
pub type WorkerErrorHandler = Arc<dyn Fn(&io::Error) + Send + Sync>;
enum Message {
Record(String),
Flush(SyncSender<io::Result<()>>),
FlushRequest,
Shutdown,
}
thread_local! {
static WORKER_OF: Cell<usize> = const { Cell::new(0) };
}
pub(crate) fn on_background_worker() -> bool {
WORKER_OF.try_with(Cell::get).unwrap_or(0) != 0
}
struct Shared {
queued_bytes: AtomicUsize,
dropped: AtomicU64,
write_errors: AtomicU64,
deferred: Mutex<Option<io::Error>>,
on_error: Option<WorkerErrorHandler>,
}
struct Channel {
sender: Option<SyncSender<Message>>,
worker: Option<JoinHandle<Box<dyn Write + Send>>>,
}
#[cfg_attr(all(feature = "json"), doc = "```")]
#[cfg_attr(not(all(feature = "json")), doc = "```ignore")]
pub struct NonBlocking<F: Format> {
format: F,
overflow: Overflow,
channel: RwLock<Channel>,
fallback: Mutex<Option<Box<dyn Write + Send>>>,
pool: Mutex<Receiver<String>>,
shared: Arc<Shared>,
max_queued_bytes: usize,
}
impl<F: Format> NonBlocking<F> {
pub fn new<W: Write + Send + 'static>(writer: W, format: F) -> Result<Self> {
Self::builder(writer, format).build()
}
pub fn builder<W: Write + Send + 'static>(writer: W, format: F) -> NonBlockingBuilder<F> {
NonBlockingBuilder {
writer: Box::new(writer),
format,
capacity: 8192,
max_queued_bytes: DEFAULT_MAX_QUEUED_BYTES,
overflow: Overflow::Block,
buffer_capacity: 64 * 1024,
thread_name: String::from("log-io-writer"),
on_error: None,
}
}
#[must_use]
pub fn dropped(&self) -> u64 {
self.shared.dropped.load(Ordering::Relaxed)
}
#[must_use]
pub fn write_errors(&self) -> u64 {
self.shared.write_errors.load(Ordering::Relaxed)
}
#[must_use]
pub fn queued_bytes(&self) -> usize {
self.shared.queued_bytes.load(Ordering::Relaxed)
}
fn fits(&self, len: usize) -> bool {
let queued = self.shared.queued_bytes.load(Ordering::Relaxed);
queued == 0 || queued.saturating_add(len) <= self.max_queued_bytes
}
fn unqueue(&self, len: usize) {
let _previous = self.shared.queued_bytes.fetch_sub(len, Ordering::Relaxed);
}
fn take_buffer(&self) -> String {
self.pool
.try_lock()
.ok()
.and_then(|rx| rx.try_recv().ok())
.unwrap_or_default()
}
fn on_own_worker(&self) -> bool {
let me = Arc::as_ptr(&self.shared).addr();
WORKER_OF.try_with(Cell::get).is_ok_and(|w| w == me)
}
fn count_drop(&self) {
let _previous = self.shared.dropped.fetch_add(1, Ordering::Relaxed);
}
fn channel(&self, may_wait: bool) -> Option<RwLockReadGuard<'_, Channel>> {
match self.channel.try_read() {
Ok(channel) => Some(channel),
Err(TryLockError::Poisoned(p)) => Some(p.into_inner()),
Err(TryLockError::WouldBlock) if may_wait => {
Some(self.channel.read().unwrap_or_else(PoisonError::into_inner))
}
Err(TryLockError::WouldBlock) => None,
}
}
fn queue_flush(&self) -> Result<()> {
let Some(channel) = self.channel(false) else {
return Ok(());
};
if let Some(sender) = &channel.sender {
match sender.try_send(Message::FlushRequest) {
Ok(()) | Err(TrySendError::Full(_)) => {}
Err(TrySendError::Disconnected(_)) => return Err(Error::Closed),
}
drop(channel);
return self.take_deferred();
}
drop(channel);
if let Some(writer) = lock_recover(&self.fallback).as_mut() {
writer.flush()?;
}
self.take_deferred()
}
fn take_deferred(&self) -> Result<()> {
match lock_recover(&self.shared.deferred).take() {
Some(e) => Err(Error::Io(e)),
None => Ok(()),
}
}
}
impl<F: Format> Sink for NonBlocking<F> {
fn write_record(&self, record: &Record<'_>) -> Result<()> {
let mut text = self.take_buffer();
text.clear();
self.format
.write_record_generic(record, &mut text)
.map_err(Error::Format)?;
let Some(channel) = self.channel(!self.on_own_worker()) else {
self.count_drop();
return Ok(());
};
if let Some(sender) = &channel.sender {
let len = text.len();
let may_block = matches!(self.overflow, Overflow::Block) && !self.on_own_worker();
while !self.fits(len) {
if !may_block {
self.count_drop();
return Ok(());
}
let (ack_tx, ack_rx) = mpsc::sync_channel(1);
sender
.send(Message::Flush(ack_tx))
.map_err(|_| Error::Closed)?;
let _flushed = ack_rx.recv().map_err(|_| Error::Closed)?;
}
let _previous = self.shared.queued_bytes.fetch_add(len, Ordering::Relaxed);
return match sender.try_send(Message::Record(text)) {
Ok(()) => Ok(()),
Err(TrySendError::Full(message)) => {
if may_block {
sender.send(message).map_err(|_| {
self.unqueue(len);
Error::Closed
})
} else {
self.unqueue(len);
self.count_drop();
Ok(())
}
}
Err(TrySendError::Disconnected(_)) => {
self.unqueue(len);
Err(Error::Closed)
}
};
}
drop(channel);
match lock_recover(&self.fallback).as_mut() {
Some(writer) => writer.write_all(text.as_bytes()).map_err(Error::from),
None => Err(Error::Closed),
}
}
fn flush(&self) -> Result<()> {
if self.on_own_worker() {
return Err(Error::Closed);
}
let channel = self.channel.read().unwrap_or_else(PoisonError::into_inner);
if let Some(sender) = &channel.sender {
let (ack_tx, ack_rx) = mpsc::sync_channel(1);
sender
.send(Message::Flush(ack_tx))
.map_err(|_| Error::Closed)?;
drop(channel);
let flushed = ack_rx.recv().map_err(|_| Error::Closed)?;
self.take_deferred()?;
return flushed.map_err(Error::from);
}
drop(channel);
if let Some(writer) = lock_recover(&self.fallback).as_mut() {
writer.flush()?;
}
self.take_deferred()
}
fn request_flush(&self) -> Result<()> {
if self.on_own_worker() {
return Ok(());
}
match self.overflow {
Overflow::DropNewest => self.queue_flush(),
_ => self.flush(),
}
}
fn shutdown(&self) -> Result<()> {
if self.on_own_worker() {
return Err(Error::Closed);
}
let mut channel = self.channel.write().unwrap_or_else(PoisonError::into_inner);
let Some(sender) = channel.sender.take() else {
drop(channel);
return self.flush();
};
let sent = sender.send(Message::Shutdown);
drop(sender);
let joined = channel.worker.take().map(JoinHandle::join);
match joined {
Some(Ok(writer)) => *lock_recover(&self.fallback) = Some(writer),
Some(Err(_)) => return Err(Error::Closed),
None => {}
}
drop(channel);
sent.map_err(|_| Error::Closed)?;
self.take_deferred()
}
}
impl<F: Format> Drop for NonBlocking<F> {
fn drop(&mut self) {
let _shutdown_in_drop = self.shutdown();
}
}
impl<F: Format + core::fmt::Debug> core::fmt::Debug for NonBlocking<F> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("NonBlocking")
.field("format", &self.format)
.field("overflow", &self.overflow)
.field("dropped", &self.dropped())
.field("write_errors", &self.write_errors())
.finish_non_exhaustive()
}
}
#[must_use = "call `build()` to start the background writer"]
pub struct NonBlockingBuilder<F> {
writer: Box<dyn Write + Send>,
format: F,
capacity: usize,
max_queued_bytes: usize,
overflow: Overflow,
buffer_capacity: usize,
thread_name: String,
on_error: Option<WorkerErrorHandler>,
}
const MAX_CAPACITY: usize = 1 << 20;
const DEFAULT_MAX_QUEUED_BYTES: usize = 16 * 1024 * 1024;
const MAX_POOLED: usize = 256;
impl<F: Format> NonBlockingBuilder<F> {
pub fn capacity(mut self, records: usize) -> Self {
self.capacity = records.clamp(1, MAX_CAPACITY);
self
}
pub fn max_queued_bytes(mut self, bytes: usize) -> Self {
self.max_queued_bytes = bytes.max(1);
self
}
pub fn overflow(mut self, overflow: Overflow) -> Self {
self.overflow = overflow;
self
}
pub fn buffer_capacity(mut self, bytes: usize) -> Self {
self.buffer_capacity = bytes;
self
}
pub fn thread_name(mut self, name: impl Into<String>) -> Self {
self.thread_name = name.into();
self
}
pub fn on_error(mut self, handler: WorkerErrorHandler) -> Self {
self.on_error = Some(handler);
self
}
pub fn build(self) -> Result<NonBlocking<F>> {
if self.thread_name.contains('\0') {
return Err(Error::Configuration(
"NonBlocking thread name must not contain a NUL byte",
));
}
let (tx, rx) = mpsc::sync_channel::<Message>(self.capacity);
let (pool_tx, pool_rx) = mpsc::sync_channel::<String>(self.capacity.min(MAX_POOLED));
let shared = Arc::new(Shared {
queued_bytes: AtomicUsize::new(0),
dropped: AtomicU64::new(0),
write_errors: AtomicU64::new(0),
deferred: Mutex::new(None),
on_error: self.on_error,
});
let worker_shared = Arc::clone(&shared);
let identity = Arc::as_ptr(&shared).addr();
let writer = BufWriter::with_capacity(self.buffer_capacity, self.writer);
let worker = std::thread::Builder::new()
.name(self.thread_name)
.spawn(move || {
WORKER_OF.with(|w| w.set(identity));
run_worker(writer, &rx, &pool_tx, &worker_shared)
})?;
Ok(NonBlocking {
format: self.format,
overflow: self.overflow,
channel: RwLock::new(Channel {
sender: Some(tx),
worker: Some(worker),
}),
fallback: Mutex::new(None),
pool: Mutex::new(pool_rx),
shared,
max_queued_bytes: self.max_queued_bytes,
})
}
}
impl<F: core::fmt::Debug> core::fmt::Debug for NonBlockingBuilder<F> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("NonBlockingBuilder")
.field("format", &self.format)
.field("capacity", &self.capacity)
.field("max_queued_bytes", &self.max_queued_bytes)
.field("overflow", &self.overflow)
.field("thread_name", &self.thread_name)
.finish_non_exhaustive()
}
}
fn report(shared: &Shared, error: io::Error) {
let _previous = shared.write_errors.fetch_add(1, Ordering::Relaxed);
if let Some(handler) = &shared.on_error {
handler(&error);
}
let mut deferred = lock_recover(&shared.deferred);
if deferred.is_none() {
*deferred = Some(error);
}
}
fn run_worker(
mut writer: BufWriter<Box<dyn Write + Send>>,
rx: &Receiver<Message>,
pool: &SyncSender<String>,
shared: &Shared,
) -> Box<dyn Write + Send> {
let mut dirty = false;
loop {
let message = match rx.try_recv() {
Ok(m) => m,
Err(TryRecvError::Empty) => {
if dirty {
if let Err(e) = writer.flush() {
report(shared, e);
}
}
match rx.recv() {
Ok(m) => m,
Err(_) => break,
}
}
Err(TryRecvError::Disconnected) => break,
};
match message {
Message::Record(mut text) => {
if let Err(e) = writer.write_all(text.as_bytes()) {
report(shared, e);
}
let _previous = shared.queued_bytes.fetch_sub(text.len(), Ordering::Relaxed);
dirty = true;
trim_buffer(&mut text);
let _recycled = pool.try_send(text);
}
Message::Flush(ack) => {
dirty = false;
let _requester_may_be_gone = ack.send(writer.flush());
}
Message::FlushRequest => {
dirty = false;
if let Err(e) = writer.flush() {
report(shared, e);
}
}
Message::Shutdown => break,
}
}
if let Err(e) = writer.flush() {
report(shared, e);
}
match writer.into_inner() {
Ok(inner) => inner,
Err(e) => {
let (error, writer) = e.into_parts();
report(shared, error);
Box::new(writer)
}
}
}
#[cfg(all(test, feature = "json"))]
mod tests {
use super::*;
use crate::format::JsonFormat;
use crate::{Level, Metadata};
use std::sync::Barrier;
#[derive(Clone, Default)]
struct Capture(Arc<Mutex<Vec<u8>>>);
impl Write for Capture {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.0.lock().unwrap().extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
impl Capture {
fn lines(&self) -> usize {
String::from_utf8(self.0.lock().unwrap().clone())
.unwrap()
.lines()
.count()
}
}
fn record(msg: &str) -> Record<'_> {
Record::new(Metadata::new(Level::Info, "t"), msg, &[])
}
#[test]
fn test_non_blocking_flush_waits_for_queued_records() {
let out = Capture::default();
let sink = NonBlocking::new(out.clone(), JsonFormat::new()).unwrap();
for _ in 0..1_000 {
sink.write_record(&record("x")).unwrap();
}
sink.flush().unwrap();
assert_eq!(out.lines(), 1_000);
}
#[test]
fn test_non_blocking_shutdown_drains_then_writes_synchronously() {
let out = Capture::default();
let sink = NonBlocking::new(out.clone(), JsonFormat::new()).unwrap();
for _ in 0..100 {
sink.write_record(&record("queued")).unwrap();
}
sink.shutdown().unwrap();
assert_eq!(out.lines(), 100);
sink.write_record(&record("late")).unwrap();
sink.flush().unwrap();
sink.shutdown().unwrap();
assert_eq!(out.lines(), 101);
}
#[test]
fn test_non_blocking_drop_newest_counts_drops() {
struct Gate(Arc<Barrier>, Capture, bool);
impl Write for Gate {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
if !self.2 {
self.2 = true;
let _leader = self.0.wait();
}
self.1.write(buf)
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
let gate = Arc::new(Barrier::new(2));
let out = Capture::default();
let sink = NonBlocking::builder(Gate(gate.clone(), out.clone(), false), JsonFormat::new())
.capacity(2)
.overflow(Overflow::DropNewest)
.buffer_capacity(0)
.build()
.unwrap();
for _ in 0..50 {
sink.write_record(&record("burst")).unwrap();
}
let _leader = gate.wait();
sink.flush().unwrap();
let written = out.lines() as u64;
assert!(sink.dropped() > 0, "the queue must have overflowed");
assert_eq!(
written + sink.dropped(),
50,
"every record is written or counted"
);
}
#[test]
fn test_non_blocking_write_errors_are_counted_and_deferred_to_flush() {
struct Failing;
impl Write for Failing {
fn write(&mut self, _: &[u8]) -> io::Result<usize> {
Err(io::Error::other("disk full"))
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
let seen = Arc::new(AtomicU64::new(0));
let seen2 = Arc::clone(&seen);
let sink = NonBlocking::builder(Failing, JsonFormat::new())
.buffer_capacity(0)
.thread_name("log-io-test-writer")
.on_error(Arc::new(move |_| {
let _n = seen2.fetch_add(1, Ordering::Relaxed);
}))
.build()
.unwrap();
sink.write_record(&record("lost")).unwrap();
sink.write_record(&record("lost too")).unwrap();
let err = sink.flush().unwrap_err();
assert!(matches!(err, Error::Io(_)), "{err}");
assert_eq!(sink.write_errors(), 2);
assert_eq!(seen.load(Ordering::Relaxed), 2);
assert!(sink.flush().is_ok(), "the deferred error is reported once");
assert!(format!("{sink:?}").contains("write_errors: 2"));
}
#[test]
fn test_non_blocking_many_producers_lose_nothing_with_block() {
let out = Capture::default();
let sink = Arc::new(
NonBlocking::builder(out.clone(), JsonFormat::new())
.capacity(16)
.build()
.unwrap(),
);
std::thread::scope(|s| {
for _ in 0..8 {
let sink = Arc::clone(&sink);
let _handle = s.spawn(move || {
for _ in 0..500 {
sink.write_record(&record("p")).unwrap();
}
});
}
});
sink.shutdown().unwrap();
assert_eq!(out.lines(), 8 * 500);
assert_eq!(sink.dropped(), 0);
}
#[test]
fn test_non_blocking_worker_panic_reports_closed() {
struct Panics;
impl Write for Panics {
fn write(&mut self, _: &[u8]) -> io::Result<usize> {
panic!("writer exploded");
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
let sink = NonBlocking::builder(Panics, JsonFormat::new())
.buffer_capacity(0)
.build()
.unwrap();
let _queued = sink.write_record(&record("boom"));
assert!(matches!(sink.shutdown(), Err(Error::Closed)));
assert!(matches!(
sink.write_record(&record("after")),
Err(Error::Closed)
));
let builder = NonBlocking::builder(Vec::<u8>::new(), JsonFormat::new()).capacity(0);
assert!(format!("{builder:?}").contains("capacity: 1"));
}
#[test]
fn test_non_blocking_builder_rejects_bad_options_without_panicking() {
let err = NonBlocking::builder(Vec::<u8>::new(), JsonFormat::new())
.thread_name("a\0b")
.build()
.unwrap_err();
assert!(matches!(err, Error::Configuration(_)), "{err}");
let builder =
NonBlocking::builder(Vec::<u8>::new(), JsonFormat::new()).capacity(usize::MAX / 2);
assert!(format!("{builder:?}").contains("capacity: 1048576"));
}
struct Stall(Option<mpsc::Receiver<()>>, Capture);
impl Write for Stall {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
if let Some(gate) = self.0.take() {
let _released = gate.recv();
}
self.1.write(buf)
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
#[test]
fn test_non_blocking_queue_is_bounded_in_bytes() {
let (release, gate) = mpsc::channel();
let out = Capture::default();
let sink = NonBlocking::builder(Stall(Some(gate), out.clone()), JsonFormat::new())
.capacity(1024)
.overflow(Overflow::DropNewest)
.buffer_capacity(0)
.build()
.unwrap();
let big = "x".repeat(256 * 1024);
for _ in 0..256 {
sink.write_record(&record(&big)).unwrap();
}
let queued = sink.queued_bytes();
let limit = DEFAULT_MAX_QUEUED_BYTES + big.len() + 1024;
release.send(()).unwrap();
sink.flush().unwrap();
assert!(queued <= limit, "{queued} bytes queued");
assert!(sink.dropped() > 0, "the byte budget must have been hit");
assert_eq!(out.lines() as u64 + sink.dropped(), 256);
assert_eq!(sink.queued_bytes(), 0);
}
#[test]
fn test_non_blocking_block_waits_for_the_byte_budget_and_loses_nothing() {
let out = Capture::default();
let sink = NonBlocking::builder(out.clone(), JsonFormat::new())
.max_queued_bytes(4 * 1024)
.build()
.unwrap();
let text = "y".repeat(1024);
for _ in 0..200 {
sink.write_record(&record(&text)).unwrap();
}
sink.flush().unwrap();
assert_eq!(out.lines(), 200);
assert_eq!(sink.dropped(), 0);
assert_eq!(sink.queued_bytes(), 0);
}
#[test]
fn test_non_blocking_drop_newest_flush_level_record_does_not_wait() {
let (release, gate) = mpsc::channel();
let out = Capture::default();
let sink = NonBlocking::builder(Stall(Some(gate), out.clone()), JsonFormat::new())
.capacity(4)
.overflow(Overflow::DropNewest)
.buffer_capacity(0)
.build()
.unwrap();
let logger = crate::Logger::builder()
.level(Level::Trace)
.no_timestamps()
.with_sink(sink)
.build();
for _ in 0..20 {
logger.log(Level::Info, "queued", &[]);
}
let (done, returned) = mpsc::channel();
let error_logger = logger.clone();
let handle = std::thread::spawn(move || {
error_logger.log(Level::Error, "error", &[]);
let _sent = done.send(());
});
let waited = returned.recv_timeout(std::time::Duration::from_secs(10));
release.send(()).unwrap();
handle.join().unwrap();
assert!(
waited.is_ok(),
"a flush-level record must not wait for the writer under DropNewest"
);
logger.flush().unwrap();
assert!(out.lines() >= 1);
}
#[test]
fn test_non_blocking_worker_logging_into_its_own_sink_never_deadlocks() {
struct Echo {
logger: Arc<std::sync::OnceLock<crate::Logger>>,
echoes: usize,
flush_from_worker: Arc<Mutex<Option<bool>>>,
}
impl Write for Echo {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
if let Some(logger) = self.logger.get() {
if self.echoes < 50 {
self.echoes += 1;
logger.log(Level::Error, "echo from the writer", &[]);
let closed = matches!(logger.flush(), Err(Error::Closed));
*self.flush_from_worker.lock().unwrap() = Some(closed);
}
}
Ok(buf.len())
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
let cell = Arc::new(std::sync::OnceLock::new());
let flush_from_worker = Arc::new(Mutex::new(None));
let sink = NonBlocking::builder(
Echo {
logger: Arc::clone(&cell),
echoes: 0,
flush_from_worker: Arc::clone(&flush_from_worker),
},
JsonFormat::new(),
)
.capacity(1)
.buffer_capacity(0)
.build()
.unwrap();
let logger = crate::Logger::builder()
.level(Level::Trace)
.no_timestamps()
.with_sink(sink)
.build();
let _set = cell.set(logger.clone());
let (done, finished) = mpsc::channel();
let _detached = std::thread::spawn(move || {
for _ in 0..200 {
logger.log(Level::Info, "from a producer", &[]);
}
let shut = logger.shutdown();
let _sent = done.send(shut.is_ok());
});
let outcome = finished.recv_timeout(std::time::Duration::from_secs(30));
assert_eq!(outcome, Ok(true), "the worker waited for itself");
assert_eq!(
*flush_from_worker.lock().unwrap(),
Some(true),
"a flush from the worker thread reports `Closed` instead of waiting"
);
drop(cell);
}
#[test]
fn test_non_blocking_drop_without_logger_drains_the_queue() {
let out = Capture::default();
let sink = NonBlocking::new(out.clone(), JsonFormat::new()).unwrap();
for _ in 0..100 {
sink.write_record(&record("x")).unwrap();
}
drop(sink);
assert_eq!(out.lines(), 100);
}
}