systemprompt_logging/layer/
mod.rs1mod 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 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}