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

systemprompt_logging/layer/
mod.rs
