systemprompt_logging/layer/
writer.rs

1//! The background task that batch-inserts buffered log entries, and the
2//! [`LogWriterHandle`] its owner awaits to flush and stop it.
3//!
4//! Copyright (c) systemprompt.io — Business Source License 1.1.
5//! See <https://systemprompt.io> for licensing details.
6
7use std::io::Write;
8use std::time::Duration;
9
10use systemprompt_traits::OwnedTask;
11use tokio::sync::mpsc;
12
13use crate::models::LogEntry;
14use crate::repository::LoggingRepository;
15
16const BUFFER_FLUSH_SIZE: usize = 100;
17const BUFFER_FLUSH_INTERVAL_SECS: u64 = 10;
18
19pub(super) enum LogCommand {
20    Entry(Box<LogEntry>),
21    FlushNow,
22    Shutdown,
23}
24
25#[derive(Debug, thiserror::Error)]
26pub enum LogWriterShutdownError {
27    #[error("the log writer had already stopped before the shutdown flush was requested")]
28    WriterGone,
29
30    #[error("the log writer task failed: {0}")]
31    Join(#[from] tokio::task::JoinError),
32}
33
34#[derive(Debug)]
35pub struct LogWriterHandle {
36    sender: mpsc::Sender<LogCommand>,
37    task: OwnedTask<()>,
38}
39
40impl LogWriterHandle {
41    pub(super) fn spawn(
42        repository: LoggingRepository,
43        sender: mpsc::Sender<LogCommand>,
44        receiver: mpsc::Receiver<LogCommand>,
45    ) -> Self {
46        Self {
47            sender,
48            task: OwnedTask::spawn("log_batch_writer", batch_writer(repository, receiver)),
49        }
50    }
51
52    pub async fn shutdown(self) -> Result<(), LogWriterShutdownError> {
53        let requested = self.sender.send(LogCommand::Shutdown).await;
54        self.task.join().await?;
55        requested.map_err(|_closed| LogWriterShutdownError::WriterGone)
56    }
57}
58
59async fn batch_writer(repository: LoggingRepository, mut receiver: mpsc::Receiver<LogCommand>) {
60    let mut buffer = Vec::with_capacity(BUFFER_FLUSH_SIZE);
61    let mut interval = tokio::time::interval(Duration::from_secs(BUFFER_FLUSH_INTERVAL_SECS));
62    let mut failed_total: u64 = 0;
63
64    loop {
65        tokio::select! {
66            command = receiver.recv() => match command {
67                Some(LogCommand::Entry(entry)) => {
68                    buffer.push(*entry);
69                    if buffer.len() >= BUFFER_FLUSH_SIZE {
70                        flush(&repository, &mut buffer, &mut failed_total).await;
71                    }
72                }
73                Some(LogCommand::FlushNow) => {
74                    if !buffer.is_empty() {
75                        flush(&repository, &mut buffer, &mut failed_total).await;
76                    }
77                }
78                Some(LogCommand::Shutdown) | None => {
79                    drain_queued(&mut receiver, &mut buffer);
80                    if !buffer.is_empty() {
81                        flush(&repository, &mut buffer, &mut failed_total).await;
82                    }
83                    return;
84                }
85            },
86            _ = interval.tick() => {
87                if !buffer.is_empty() {
88                    flush(&repository, &mut buffer, &mut failed_total).await;
89                }
90            }
91        }
92    }
93}
94
95fn drain_queued(receiver: &mut mpsc::Receiver<LogCommand>, buffer: &mut Vec<LogEntry>) {
96    while let Ok(command) = receiver.try_recv() {
97        if let LogCommand::Entry(entry) = command {
98            buffer.push(*entry);
99        }
100    }
101}
102
103async fn flush(repository: &LoggingRepository, buffer: &mut Vec<LogEntry>, failed_total: &mut u64) {
104    if let Err(e) = repository.insert_batch(buffer).await {
105        let lost = u64::try_from(buffer.len()).unwrap_or(u64::MAX);
106        *failed_total = failed_total.saturating_add(lost);
107        // Why: stderr is the last-resort sink once the database layer has
108        // failed; tracing here would feed back into the failing layer.
109        let _stderr_written = writeln!(
110            std::io::stderr(),
111            "DATABASE LOG FLUSH FAILED ({lost} entries lost this flush, {failed_total} total lost since start): {e}"
112        )
113        .is_ok();
114    }
115    buffer.clear();
116}
systemprompt_logging/layer/writer.rs

systemprompt_logging/layer/
writer.rs
