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

systemprompt_logging/layer/
mod.rs
