systemprompt_logging/layer/
mod.rs

1//! `tracing` subscriber layer that persists events to the database.
2//!
3//! [`DatabaseLayer`] buffers log events off the hot path and batch-inserts them
4//! from a background task, flushing on a size threshold, a timer, or
5//! immediately on an error. [`ProxyDatabaseLayer`] is the proxy-side variant.
6
7mod proxy;
8mod visitor;
9
10use std::io::Write;
11use std::sync::atomic::{AtomicU64, Ordering};
12use std::sync::{Arc, OnceLock};
13use std::time::Duration;
14
15use tokio::sync::mpsc;
16use tracing::{Event, Subscriber};
17use tracing_subscriber::Layer;
18use tracing_subscriber::layer::Context;
19use tracing_subscriber::registry::LookupSpan;
20
21pub use proxy::ProxyDatabaseLayer;
22use proxy::{build_log_entry, record_span_fields, update_span_fields};
23
24use crate::models::{LogEntry, LogLevel};
25use systemprompt_database::DbPool;
26use systemprompt_identifiers::{ClientId, ContextId, TaskId};
27
28const BUFFER_FLUSH_SIZE: usize = 100;
29const BUFFER_FLUSH_INTERVAL_SECS: u64 = 10;
30
31/// Bounded capacity of the log channel. Beyond this depth (a sustained burst
32/// the database writer cannot drain) entries are dropped rather than queued, so
33/// a logging backlog cannot grow the heap without bound.
34const CHANNEL_CAPACITY: usize = 8192;
35
36static BACKGROUND_SENDER: OnceLock<mpsc::Sender<LogCommand>> = OnceLock::new();
37static BACKGROUND_DROPPED: AtomicU64 = AtomicU64::new(0);
38
39/// Non-blocking and off the caller's hot path: the entry is dropped (and
40/// counted) if the sink is unattached or the channel is full. Error entries
41/// also request an immediate flush.
42pub fn enqueue_background(entry: LogEntry) {
43    let Some(sender) = BACKGROUND_SENDER.get() else {
44        BACKGROUND_DROPPED.fetch_add(1, Ordering::Relaxed);
45        return;
46    };
47    let is_error = entry.level == LogLevel::Error;
48    if sender.try_send(LogCommand::Entry(Box::new(entry))).is_err() {
49        BACKGROUND_DROPPED.fetch_add(1, Ordering::Relaxed);
50        return;
51    }
52    if is_error {
53        sender.try_send(LogCommand::FlushNow).ok();
54    }
55}
56
57enum LogCommand {
58    Entry(Box<LogEntry>),
59    FlushNow,
60}
61
62/// Bounded sender to the database writer task. On a full channel the entry is
63/// dropped and [`LogChannel::dropped`] is incremented; the send never blocks,
64/// so logging stays off the hot path even under burst.
65struct LogChannel {
66    sender: mpsc::Sender<LogCommand>,
67    dropped: Arc<AtomicU64>,
68}
69
70impl LogChannel {
71    fn new(capacity: usize) -> (Self, mpsc::Receiver<LogCommand>) {
72        let (sender, receiver) = mpsc::channel(capacity);
73        let channel = Self {
74            sender,
75            dropped: Arc::new(AtomicU64::new(0)),
76        };
77        (channel, receiver)
78    }
79
80    fn send(&self, command: LogCommand) {
81        if let Err(mpsc::error::TrySendError::Full(_)) = self.sender.try_send(command) {
82            self.dropped.fetch_add(1, Ordering::Relaxed);
83        }
84    }
85
86    fn dropped(&self) -> u64 {
87        self.dropped.load(Ordering::Relaxed)
88    }
89}
90
91pub struct DatabaseLayer {
92    channel: LogChannel,
93}
94
95impl std::fmt::Debug for DatabaseLayer {
96    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
97        f.debug_struct("DatabaseLayer")
98            .field("dropped", &self.channel.dropped())
99            .finish_non_exhaustive()
100    }
101}
102
103impl DatabaseLayer {
104    pub fn new(db_pool: DbPool) -> Self {
105        let (channel, receiver) = LogChannel::new(CHANNEL_CAPACITY);
106
107        BACKGROUND_SENDER.get_or_init(|| channel.sender.clone());
108
109        tokio::spawn(Self::batch_writer(db_pool, receiver));
110
111        Self { channel }
112    }
113
114    async fn batch_writer(db_pool: DbPool, mut receiver: mpsc::Receiver<LogCommand>) {
115        let mut buffer = Vec::with_capacity(BUFFER_FLUSH_SIZE);
116        let mut interval = tokio::time::interval(Duration::from_secs(BUFFER_FLUSH_INTERVAL_SECS));
117        let mut failed_total: u64 = 0;
118
119        loop {
120            tokio::select! {
121                Some(command) = receiver.recv() => {
122                    match command {
123                        LogCommand::Entry(entry) => {
124                            buffer.push(*entry);
125                            if buffer.len() >= BUFFER_FLUSH_SIZE {
126                                Self::flush(&db_pool, &mut buffer, &mut failed_total).await;
127                            }
128                        }
129                        LogCommand::FlushNow => {
130                            if !buffer.is_empty() {
131                                Self::flush(&db_pool, &mut buffer, &mut failed_total).await;
132                            }
133                        }
134                    }
135                }
136                _ = interval.tick() => {
137                    if !buffer.is_empty() {
138                        Self::flush(&db_pool, &mut buffer, &mut failed_total).await;
139                    }
140                }
141            }
142        }
143    }
144
145    async fn flush(db_pool: &DbPool, buffer: &mut Vec<LogEntry>, failed_total: &mut u64) {
146        if let Err(e) = Self::batch_insert(db_pool, buffer).await {
147            let lost = u64::try_from(buffer.len()).unwrap_or(u64::MAX);
148            *failed_total = failed_total.saturating_add(lost);
149            writeln!(
150                std::io::stderr(),
151                "DATABASE LOG FLUSH FAILED ({lost} entries lost this flush, {failed_total} total lost since start): {e}"
152            )
153            .ok();
154        }
155        buffer.clear();
156    }
157
158    async fn batch_insert(
159        db_pool: &DbPool,
160        entries: &[LogEntry],
161    ) -> Result<(), crate::models::LoggingError> {
162        let pool = db_pool.write_pool_arc()?;
163
164        // One commit per flush, fsync off: the audit log is best-effort, so a
165        // few buffered rows lost on an unclean shutdown is an acceptable trade.
166        let mut tx = pool.begin().await?;
167        sqlx::query!("SET LOCAL synchronous_commit = off")
168            .execute(&mut *tx)
169            .await?;
170
171        for entry in entries {
172            let metadata_json: Option<String> = entry
173                .metadata
174                .as_ref()
175                .map(serde_json::to_string)
176                .transpose()?;
177
178            let entry_id = entry.id.as_str();
179            let level_str = entry.level.to_string();
180            let user_id = entry.user_id.as_str();
181            let session_id = entry.session_id.as_str();
182            let task_id = entry.task_id.as_ref().map(TaskId::as_str);
183            let trace_id = entry.trace_id.as_str();
184            let context_id = entry.context_id.as_ref().map(ContextId::as_str);
185            let client_id = entry.client_id.as_ref().map(ClientId::as_str);
186
187            sqlx::query!(
188                r"
189                INSERT INTO logs (id, timestamp, level, module, message, metadata, user_id, session_id, task_id, trace_id, context_id, client_id)
190                VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
191                ",
192                entry_id,
193                entry.timestamp,
194                level_str,
195                entry.module,
196                entry.message,
197                metadata_json,
198                user_id,
199                session_id,
200                task_id,
201                trace_id,
202                context_id,
203                client_id
204            )
205            .execute(&mut *tx)
206            .await?;
207        }
208
209        tx.commit().await?;
210        Ok(())
211    }
212}
213
214impl DatabaseLayer {
215    fn send_entry(&self, entry: LogEntry) {
216        let is_error = entry.level == LogLevel::Error;
217        self.channel.send(LogCommand::Entry(Box::new(entry)));
218        if is_error {
219            self.channel.send(LogCommand::FlushNow);
220        }
221    }
222}
223
224impl<S> Layer<S> for DatabaseLayer
225where
226    S: Subscriber + for<'a> LookupSpan<'a>,
227{
228    fn on_new_span(
229        &self,
230        attrs: &tracing::span::Attributes<'_>,
231        id: &tracing::span::Id,
232        ctx: Context<'_, S>,
233    ) {
234        record_span_fields(attrs, id, &ctx);
235    }
236
237    fn on_record(
238        &self,
239        id: &tracing::span::Id,
240        values: &tracing::span::Record<'_>,
241        ctx: Context<'_, S>,
242    ) {
243        update_span_fields(id, values, &ctx);
244    }
245
246    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
247        if let Some(entry) = build_log_entry(event, &ctx) {
248            self.send_entry(entry);
249        }
250    }
251}
systemprompt_logging/layer/mod.rs

systemprompt_logging/layer/
mod.rs
