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

systemprompt_logging/layer/
mod.rs
