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

systemprompt_logging/layer/
mod.rs
