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        // Why: best-effort flush nudge; entry send above already accounts for drops
54        sender.try_send(LogCommand::FlushNow).ok();
55    }
56}
57
58enum LogCommand {
59    Entry(Box<LogEntry>),
60    FlushNow,
61}
62
63/// Bounded sender to the database writer task. On a full channel the entry is
64/// dropped and [`LogChannel::dropped`] is incremented; the send never blocks,
65/// so logging stays off the hot path even under burst.
66struct LogChannel {
67    sender: mpsc::Sender<LogCommand>,
68    dropped: Arc<AtomicU64>,
69}
70
71impl LogChannel {
72    fn new(capacity: usize) -> (Self, mpsc::Receiver<LogCommand>) {
73        let (sender, receiver) = mpsc::channel(capacity);
74        let channel = Self {
75            sender,
76            dropped: Arc::new(AtomicU64::new(0)),
77        };
78        (channel, receiver)
79    }
80
81    fn send(&self, command: LogCommand) {
82        if let Err(mpsc::error::TrySendError::Full(_)) = self.sender.try_send(command) {
83            self.dropped.fetch_add(1, Ordering::Relaxed);
84        }
85    }
86
87    fn dropped(&self) -> u64 {
88        self.dropped.load(Ordering::Relaxed)
89    }
90}
91
92pub struct DatabaseLayer {
93    channel: LogChannel,
94}
95
96impl std::fmt::Debug for DatabaseLayer {
97    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
98        f.debug_struct("DatabaseLayer")
99            .field("dropped", &self.channel.dropped())
100            .finish_non_exhaustive()
101    }
102}
103
104impl DatabaseLayer {
105    pub fn new(db_pool: DbPool) -> Self {
106        let (channel, receiver) = LogChannel::new(CHANNEL_CAPACITY);
107
108        BACKGROUND_SENDER.get_or_init(|| channel.sender.clone());
109
110        tokio::spawn(Self::batch_writer(db_pool, receiver));
111
112        Self { channel }
113    }
114
115    async fn batch_writer(db_pool: DbPool, mut receiver: mpsc::Receiver<LogCommand>) {
116        let mut buffer = Vec::with_capacity(BUFFER_FLUSH_SIZE);
117        let mut interval = tokio::time::interval(Duration::from_secs(BUFFER_FLUSH_INTERVAL_SECS));
118        let mut failed_total: u64 = 0;
119
120        loop {
121            tokio::select! {
122                Some(command) = receiver.recv() => {
123                    match command {
124                        LogCommand::Entry(entry) => {
125                            buffer.push(*entry);
126                            if buffer.len() >= BUFFER_FLUSH_SIZE {
127                                Self::flush(&db_pool, &mut buffer, &mut failed_total).await;
128                            }
129                        }
130                        LogCommand::FlushNow => {
131                            if !buffer.is_empty() {
132                                Self::flush(&db_pool, &mut buffer, &mut failed_total).await;
133                            }
134                        }
135                    }
136                }
137                _ = interval.tick() => {
138                    if !buffer.is_empty() {
139                        Self::flush(&db_pool, &mut buffer, &mut failed_total).await;
140                    }
141                }
142            }
143        }
144    }
145
146    async fn flush(db_pool: &DbPool, buffer: &mut Vec<LogEntry>, failed_total: &mut u64) {
147        if let Err(e) = Self::batch_insert(db_pool, buffer).await {
148            let lost = u64::try_from(buffer.len()).unwrap_or(u64::MAX);
149            *failed_total = failed_total.saturating_add(lost);
150            writeln!(
151                std::io::stderr(),
152                "DATABASE LOG FLUSH FAILED ({lost} entries lost this flush, {failed_total} total lost since start): {e}"
153            )
154            .ok();
155        }
156        buffer.clear();
157    }
158
159    async fn batch_insert(
160        db_pool: &DbPool,
161        entries: &[LogEntry],
162    ) -> Result<(), crate::models::LoggingError> {
163        let pool = db_pool.write_pool_arc()?;
164
165        // One commit per flush, fsync off: the audit log is best-effort, so a
166        // few buffered rows lost on an unclean shutdown is an acceptable trade.
167        let mut tx = pool.begin().await?;
168        sqlx::query!("SET LOCAL synchronous_commit = off")
169            .execute(&mut *tx)
170            .await?;
171
172        for entry in entries {
173            let metadata_json: Option<String> = entry
174                .metadata
175                .as_ref()
176                .map(serde_json::to_string)
177                .transpose()?;
178
179            let entry_id = entry.id.as_str();
180            let level_str = entry.level.to_string();
181            let user_id = entry.user_id.as_str();
182            let session_id = entry.session_id.as_str();
183            let task_id = entry.task_id.as_ref().map(TaskId::as_str);
184            let trace_id = entry.trace_id.as_str();
185            let context_id = entry.context_id.as_ref().map(ContextId::as_str);
186            let client_id = entry.client_id.as_ref().map(ClientId::as_str);
187
188            sqlx::query!(
189                r"
190                INSERT INTO logs (id, timestamp, level, module, message, metadata, user_id, session_id, task_id, trace_id, context_id, client_id)
191                VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
192                ",
193                entry_id,
194                entry.timestamp,
195                level_str,
196                entry.module,
197                entry.message,
198                metadata_json,
199                user_id,
200                session_id,
201                task_id,
202                trace_id,
203                context_id,
204                client_id
205            )
206            .execute(&mut *tx)
207            .await?;
208        }
209
210        tx.commit().await?;
211        Ok(())
212    }
213}
214
215impl DatabaseLayer {
216    fn send_entry(&self, entry: LogEntry) {
217        let is_error = entry.level == LogLevel::Error;
218        self.channel.send(LogCommand::Entry(Box::new(entry)));
219        if is_error {
220            self.channel.send(LogCommand::FlushNow);
221        }
222    }
223}
224
225impl<S> Layer<S> for DatabaseLayer
226where
227    S: Subscriber + for<'a> LookupSpan<'a>,
228{
229    fn on_new_span(
230        &self,
231        attrs: &tracing::span::Attributes<'_>,
232        id: &tracing::span::Id,
233        ctx: Context<'_, S>,
234    ) {
235        record_span_fields(attrs, id, &ctx);
236    }
237
238    fn on_record(
239        &self,
240        id: &tracing::span::Id,
241        values: &tracing::span::Record<'_>,
242        ctx: Context<'_, S>,
243    ) {
244        update_span_fields(id, values, &ctx);
245    }
246
247    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
248        if let Some(entry) = build_log_entry(event, &ctx) {
249            self.send_entry(entry);
250        }
251    }
252}
systemprompt_logging/layer/mod.rs

systemprompt_logging/layer/
mod.rs
