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. The task belongs to the [`LogWriterHandle`]
6//! returned alongside the layer; awaiting [`LogWriterHandle::shutdown`] flushes
7//! whatever is still buffered or queued and stops the task.
8//! [`ProxyDatabaseLayer`] is the proxy-side variant.
9//!
10//! Copyright (c) systemprompt.io — Business Source License 1.1.
11//! See <https://systemprompt.io> for licensing details.
12
13mod proxy;
14mod visitor;
15mod writer;
16
17use std::sync::atomic::{AtomicU64, Ordering};
18use std::sync::{Arc, OnceLock};
19
20use tokio::sync::mpsc;
21use tracing::{Event, Subscriber};
22use tracing_subscriber::Layer;
23use tracing_subscriber::layer::Context;
24use tracing_subscriber::registry::LookupSpan;
25
26pub use proxy::ProxyDatabaseLayer;
27use proxy::{build_log_entry, record_span_fields, update_span_fields};
28pub use writer::{LogWriterHandle, LogWriterShutdownError};
29
30use self::writer::LogCommand;
31use crate::models::{LogEntry, LogLevel};
32use crate::repository::LoggingRepository;
33use systemprompt_database::DbPool;
34
35const CHANNEL_CAPACITY: usize = 8192;
36
37static BACKGROUND_SENDER: OnceLock<mpsc::Sender<LogCommand>> = OnceLock::new();
38static BACKGROUND_DROPPED: AtomicU64 = AtomicU64::new(0);
39
40pub fn enqueue_background(entry: LogEntry) {
41    let Some(sender) = BACKGROUND_SENDER.get() else {
42        BACKGROUND_DROPPED.fetch_add(1, Ordering::Relaxed);
43        return;
44    };
45    let is_error = entry.level == LogLevel::Error;
46    if sender.try_send(LogCommand::Entry(Box::new(entry))).is_err() {
47        BACKGROUND_DROPPED.fetch_add(1, Ordering::Relaxed);
48        return;
49    }
50    if is_error {
51        // Why: a full queue already forces the flush this would request, and
52        // tracing here would re-enter this layer.
53        let _flush_requested = sender.try_send(LogCommand::FlushNow).is_ok();
54    }
55}
56
57struct LogChannel {
58    sender: mpsc::Sender<LogCommand>,
59    dropped: Arc<AtomicU64>,
60}
61
62impl LogChannel {
63    fn new(capacity: usize) -> (Self, mpsc::Receiver<LogCommand>) {
64        let (sender, receiver) = mpsc::channel(capacity);
65        let channel = Self {
66            sender,
67            dropped: Arc::new(AtomicU64::new(0)),
68        };
69        (channel, receiver)
70    }
71
72    fn send(&self, command: LogCommand) {
73        if let Err(mpsc::error::TrySendError::Full(_)) = self.sender.try_send(command) {
74            self.dropped.fetch_add(1, Ordering::Relaxed);
75        }
76    }
77
78    fn dropped(&self) -> u64 {
79        self.dropped.load(Ordering::Relaxed)
80    }
81}
82
83pub struct DatabaseLayer {
84    channel: LogChannel,
85}
86
87impl std::fmt::Debug for DatabaseLayer {
88    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
89        f.debug_struct("DatabaseLayer")
90            .field("dropped", &self.channel.dropped())
91            .finish_non_exhaustive()
92    }
93}
94
95impl DatabaseLayer {
96    pub fn new(db_pool: &DbPool) -> (Self, LogWriterHandle) {
97        let (channel, receiver) = LogChannel::new(CHANNEL_CAPACITY);
98
99        BACKGROUND_SENDER.get_or_init(|| channel.sender.clone());
100
101        let writer = LogWriterHandle::spawn(
102            LoggingRepository::new(db_pool),
103            channel.sender.clone(),
104            receiver,
105        );
106
107        (Self { channel }, writer)
108    }
109}
110
111impl DatabaseLayer {
112    fn send_entry(&self, entry: LogEntry) {
113        let is_error = entry.level == LogLevel::Error;
114        self.channel.send(LogCommand::Entry(Box::new(entry)));
115        if is_error {
116            self.channel.send(LogCommand::FlushNow);
117        }
118    }
119}
120
121impl<S> Layer<S> for DatabaseLayer
122where
123    S: Subscriber + for<'a> LookupSpan<'a>,
124{
125    fn on_new_span(
126        &self,
127        attrs: &tracing::span::Attributes<'_>,
128        id: &tracing::span::Id,
129        ctx: Context<'_, S>,
130    ) {
131        record_span_fields(attrs, id, &ctx);
132    }
133
134    fn on_record(
135        &self,
136        id: &tracing::span::Id,
137        values: &tracing::span::Record<'_>,
138        ctx: Context<'_, S>,
139    ) {
140        update_span_fields(id, values, &ctx);
141    }
142
143    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
144        if let Some(entry) = build_log_entry(event, &ctx) {
145            self.send_entry(entry);
146        }
147    }
148}
systemprompt_logging/layer/mod.rs

systemprompt_logging/layer/
mod.rs
