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

systemprompt_logging/layer/
mod.rs
