Skip to main content

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::OnceLock;
18use std::sync::atomic::{AtomicU64, Ordering};
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 DROPPED: AtomicU64 = AtomicU64::new(0);
39
40#[must_use]
41pub fn dropped_log_entries() -> u64 {
42    DROPPED.load(Ordering::Relaxed)
43}
44
45pub fn enqueue_background(entry: LogEntry) {
46    let Some(sender) = BACKGROUND_SENDER.get() else {
47        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        DROPPED.fetch_add(1, Ordering::Relaxed);
53        return;
54    }
55    if is_error {
56        // Why: a full queue already forces the flush this would request, and
57        // tracing here would re-enter this layer.
58        let _flush_requested = sender.try_send(LogCommand::FlushNow).is_ok();
59    }
60}
61
62struct LogChannel {
63    sender: mpsc::Sender<LogCommand>,
64}
65
66impl LogChannel {
67    fn new(capacity: usize) -> (Self, mpsc::Receiver<LogCommand>) {
68        let (sender, receiver) = mpsc::channel(capacity);
69        let channel = Self { sender };
70        (channel, receiver)
71    }
72
73    fn send(&self, command: LogCommand) {
74        if let Err(mpsc::error::TrySendError::Full(_)) = self.sender.try_send(command) {
75            DROPPED.fetch_add(1, Ordering::Relaxed);
76        }
77    }
78}
79
80pub struct DatabaseLayer {
81    channel: LogChannel,
82}
83
84impl std::fmt::Debug for DatabaseLayer {
85    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
86        f.debug_struct("DatabaseLayer")
87            .field("dropped", &dropped_log_entries())
88            .finish_non_exhaustive()
89    }
90}
91
92impl DatabaseLayer {
93    pub fn new(db_pool: &DbPool) -> (Self, LogWriterHandle) {
94        let (channel, receiver) = LogChannel::new(CHANNEL_CAPACITY);
95
96        BACKGROUND_SENDER.get_or_init(|| channel.sender.clone());
97
98        let writer = LogWriterHandle::spawn(
99            LoggingRepository::new(db_pool),
100            channel.sender.clone(),
101            receiver,
102        );
103
104        (Self { channel }, writer)
105    }
106}
107
108impl DatabaseLayer {
109    fn send_entry(&self, entry: LogEntry) {
110        let is_error = entry.level == LogLevel::Error;
111        self.channel.send(LogCommand::Entry(Box::new(entry)));
112        if is_error {
113            self.channel.send(LogCommand::FlushNow);
114        }
115    }
116}
117
118impl<S> Layer<S> for DatabaseLayer
119where
120    S: Subscriber + for<'a> LookupSpan<'a>,
121{
122    fn on_new_span(
123        &self,
124        attrs: &tracing::span::Attributes<'_>,
125        id: &tracing::span::Id,
126        ctx: Context<'_, S>,
127    ) {
128        record_span_fields(attrs, id, &ctx);
129    }
130
131    fn on_record(
132        &self,
133        id: &tracing::span::Id,
134        values: &tracing::span::Record<'_>,
135        ctx: Context<'_, S>,
136    ) {
137        update_span_fields(id, values, &ctx);
138    }
139
140    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
141        if let Some(entry) = build_log_entry(event, &ctx) {
142            self.send_entry(entry);
143        }
144    }
145}