use crate::serve::history::{RUN_LOG_TRUNCATED_SEQ, RunHistory, RunLogLine};
use dashmap::DashMap;
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::OnceLock;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::time::Duration;
use tokio::sync::{broadcast, mpsc};
pub const RING_CAPACITY: usize = 10_000;
pub const BROADCAST_CAPACITY: usize = 1024;
pub const LOG_DRAIN: Duration = Duration::from_secs(60);
const PERSIST_CHANNEL_CAPACITY: usize = 16_384;
const PERSIST_BATCH: usize = 256;
enum PersistMsg {
Line {
run_id: String,
seq: u64,
ts: String,
level: String,
line: String,
},
End {
run_id: String,
},
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct LogLine {
pub seq: u64,
pub line: String,
}
#[derive(Clone, Debug)]
pub enum LogMsg {
Line(LogLine),
End,
}
struct RunBuffer {
ring: Mutex<VecDeque<LogLine>>,
seq: AtomicU64,
tx: broadcast::Sender<LogMsg>,
ended: AtomicBool,
}
impl RunBuffer {
fn new() -> Self {
let (tx, _rx) = broadcast::channel(BROADCAST_CAPACITY);
Self {
ring: Mutex::new(VecDeque::with_capacity(64)),
seq: AtomicU64::new(0),
tx,
ended: AtomicBool::new(false),
}
}
fn push(&self, line: String) -> u64 {
let seq = self.seq.fetch_add(1, Ordering::Relaxed);
let entry = LogLine { seq, line };
{
let mut ring = self.ring.lock().expect("log ring poisoned");
if ring.len() == RING_CAPACITY {
ring.pop_front();
}
ring.push_back(entry.clone());
}
let _ = self.tx.send(LogMsg::Line(entry));
seq
}
fn snapshot(&self) -> Vec<LogLine> {
self.ring
.lock()
.expect("log ring poisoned")
.iter()
.cloned()
.collect()
}
fn finish(&self) {
self.ended.store(true, Ordering::SeqCst);
let _ = self.tx.send(LogMsg::End);
}
}
#[derive(Clone, Default)]
pub struct LogHub {
inner: Arc<DashMap<String, Arc<RunBuffer>>>,
persist: Arc<OnceLock<mpsc::Sender<PersistMsg>>>,
}
impl LogHub {
pub fn new() -> Self {
Self::default()
}
fn buffer(&self, run_id: &str) -> Arc<RunBuffer> {
if let Some(b) = self.inner.get(run_id) {
return Arc::clone(b.value());
}
Arc::clone(
self.inner
.entry(run_id.to_string())
.or_insert_with(|| Arc::new(RunBuffer::new()))
.value(),
)
}
pub fn enable_persistence(&self, history: Arc<dyn RunHistory>, max_lines_per_run: usize) {
let (tx, rx) = mpsc::channel(PERSIST_CHANNEL_CAPACITY);
if self.persist.set(tx).is_err() {
return; }
tokio::spawn(persist_writer(rx, history, max_lines_per_run.max(1)));
}
pub fn capture(&self, run_id: &str, level: &str, ts: String, line: String) {
let seq = self.buffer(run_id).push(line.clone());
if let Some(tx) = self.persist.get()
&& tx
.try_send(PersistMsg::Line {
run_id: run_id.to_string(),
seq,
ts,
level: level.to_string(),
line,
})
.is_err()
{
metrics::counter!("faucet_serve_run_logs_dropped_total", "reason" => "queue_full")
.increment(1);
}
}
pub fn append(&self, run_id: &str, line: String) {
self.buffer(run_id).push(line);
}
pub fn reader(
&self,
run_id: &str,
) -> Option<(Vec<LogLine>, broadcast::Receiver<LogMsg>, bool)> {
let buf = Arc::clone(self.inner.get(run_id)?.value());
let rx = buf.tx.subscribe();
let snapshot = buf.snapshot();
let ended = buf.ended.load(Ordering::SeqCst);
Some((snapshot, rx, ended))
}
pub fn finish(&self, run_id: &str) {
if let Some(buf) = self.inner.get(run_id) {
buf.finish();
}
if let Some(tx) = self.persist.get() {
let _ = tx.try_send(PersistMsg::End {
run_id: run_id.to_string(),
});
}
}
pub fn drop_run(&self, run_id: &str) {
self.inner.remove(run_id);
}
}
#[derive(Default)]
struct RunPersistState {
pending: Vec<RunLogLine>,
persisted: u64,
truncated: bool,
}
async fn persist_writer(
mut rx: mpsc::Receiver<PersistMsg>,
history: Arc<dyn RunHistory>,
max_lines_per_run: usize,
) {
let cap = max_lines_per_run as u64;
let mut runs: HashMap<String, RunPersistState> = HashMap::new();
async fn flush(history: &Arc<dyn RunHistory>, run_id: &str, st: &mut RunPersistState) {
if st.pending.is_empty() {
return;
}
let batch = std::mem::take(&mut st.pending);
let n = batch.len() as u64;
if let Err(e) = history.record_run_logs(run_id, &batch).await {
tracing::warn!(run_id, error = %e, "persisting run logs failed");
metrics::counter!("faucet_serve_run_logs_dropped_total", "reason" => "backend_error")
.increment(n);
} else {
metrics::counter!("faucet_serve_run_log_lines_total").increment(n);
}
}
while let Some(msg) = rx.recv().await {
match msg {
PersistMsg::Line {
run_id,
seq,
ts,
level,
line,
} => {
let st = runs.entry(run_id.clone()).or_default();
if st.persisted >= cap {
if !st.truncated {
st.truncated = true;
metrics::counter!(
"faucet_serve_run_logs_dropped_total", "reason" => "per_run_cap"
)
.increment(1);
}
continue;
}
st.persisted += 1;
st.pending.push(RunLogLine {
seq,
ts,
level,
line,
});
if st.pending.len() >= PERSIST_BATCH {
flush(&history, &run_id, st).await;
}
}
PersistMsg::End { run_id } => {
if let Some(mut st) = runs.remove(&run_id) {
flush(&history, &run_id, &mut st).await;
if st.truncated {
let marker = [RunLogLine {
seq: RUN_LOG_TRUNCATED_SEQ,
ts: String::new(),
level: "WARN".to_string(),
line: "log truncated: per-run cap reached".to_string(),
}];
if let Err(e) = history.record_run_logs(&run_id, &marker).await {
tracing::warn!(run_id, error = %e, "persisting run-log truncation marker failed");
}
}
}
}
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LogEvent {
Log(String),
Truncated(u64),
End,
}
pub fn log_events(
snapshot: Vec<LogLine>,
mut rx: broadcast::Receiver<LogMsg>,
ended: bool,
) -> impl futures::Stream<Item = LogEvent> {
use tokio::sync::broadcast::error::RecvError;
async_stream::stream! {
let mut last_seq: Option<u64> = None;
for entry in snapshot {
last_seq = Some(entry.seq);
yield LogEvent::Log(entry.line);
}
if ended {
yield LogEvent::End;
return;
}
loop {
match rx.recv().await {
Ok(LogMsg::Line(entry)) => {
if last_seq.is_none_or(|s| entry.seq > s) {
last_seq = Some(entry.seq);
yield LogEvent::Log(entry.line);
}
}
Ok(LogMsg::End) => {
yield LogEvent::End;
break;
}
Err(RecvError::Lagged(n)) => {
yield LogEvent::Truncated(n);
}
Err(RecvError::Closed) => {
yield LogEvent::End;
break;
}
}
}
}
}
use tracing::field::{Field, Visit};
use tracing::span::{Attributes, Id};
use tracing::{Event, Subscriber};
use tracing_subscriber::layer::{Context, Layer};
use tracing_subscriber::registry::LookupSpan;
#[derive(Clone)]
struct RunIdExt(String);
#[derive(Default)]
struct RunIdVisitor(Option<String>);
impl Visit for RunIdVisitor {
fn record_str(&mut self, field: &Field, value: &str) {
if field.name() == "serve_run_id" {
self.0 = Some(value.to_string());
}
}
fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
if field.name() == "serve_run_id" && self.0.is_none() {
self.0 = Some(format!("{value:?}"));
}
}
}
#[derive(Default)]
struct EventLineVisitor {
message: String,
fields: String,
}
impl Visit for EventLineVisitor {
fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
use std::fmt::Write;
if field.name() == "message" {
let _ = write!(self.message, "{value:?}");
} else {
let _ = write!(self.fields, " {}={:?}", field.name(), value);
}
}
}
impl EventLineVisitor {
fn finish(self) -> String {
if self.fields.is_empty() {
self.message
} else {
format!("{}{}", self.message, self.fields)
}
}
}
pub struct RunLogLayer {
hub: LogHub,
}
impl RunLogLayer {
pub fn new(hub: LogHub) -> Self {
Self { hub }
}
}
impl<S> Layer<S> for RunLogLayer
where
S: Subscriber + for<'a> LookupSpan<'a>,
{
fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
let mut visitor = RunIdVisitor::default();
attrs.record(&mut visitor);
if let Some(run_id) = visitor.0
&& let Some(span) = ctx.span(id)
{
span.extensions_mut().insert(RunIdExt(run_id));
}
}
fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
let Some(run_id) = ctx.event_scope(event).and_then(|scope| {
scope
.from_root()
.find_map(|span| span.extensions().get::<RunIdExt>().map(|ext| ext.0.clone()))
}) else {
return;
};
let mut visitor = EventLineVisitor::default();
event.record(&mut visitor);
let meta = event.metadata();
let line = format!("{} {}: {}", meta.level(), meta.target(), visitor.finish());
let line = crate::secrets::registry::redact(&line).into_owned();
let ts = chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, true);
self.hub.capture(&run_id, meta.level().as_str(), ts, line);
}
}
#[cfg(test)]
mod tests {
use super::*;
use futures::StreamExt;
#[test]
fn ring_caps_and_orders_by_seq() {
let hub = LogHub::new();
for i in 0..(RING_CAPACITY + 5) {
hub.append("r", format!("line-{i}"));
}
let (snapshot, _rx, _ended) = hub.reader("r").unwrap();
assert_eq!(snapshot.len(), RING_CAPACITY, "ring must cap at capacity");
assert_eq!(snapshot.first().unwrap().line, "line-5");
assert!(
snapshot.windows(2).all(|w| w[0].seq < w[1].seq),
"sequence numbers must be strictly increasing"
);
}
#[test]
fn reader_none_for_unknown_run() {
let hub = LogHub::new();
assert!(hub.reader("nope").is_none());
}
#[test]
fn finish_sets_ended_flag() {
let hub = LogHub::new();
hub.append("r", "x".into());
hub.finish("r");
let (snapshot, _rx, ended) = hub.reader("r").unwrap();
assert!(ended, "reader must observe ended after finish");
assert_eq!(snapshot.len(), 1);
}
#[test]
fn drop_run_frees_buffer() {
let hub = LogHub::new();
hub.append("r", "x".into());
assert!(hub.reader("r").is_some());
hub.drop_run("r");
assert!(hub.reader("r").is_none());
}
#[tokio::test]
async fn ended_buffer_streams_snapshot_then_end() {
let hub = LogHub::new();
hub.append("r", "a".into());
hub.append("r", "b".into());
hub.finish("r");
let (snapshot, rx, ended) = hub.reader("r").unwrap();
let events: Vec<LogEvent> = log_events(snapshot, rx, ended).collect().await;
assert_eq!(
events,
vec![
LogEvent::Log("a".into()),
LogEvent::Log("b".into()),
LogEvent::End
]
);
}
#[tokio::test]
async fn snapshot_then_live_dedups_by_seq() {
let (tx, rx) = broadcast::channel(8);
let _ = tx.send(LogMsg::Line(LogLine {
seq: 0,
line: "a".into(),
}));
let _ = tx.send(LogMsg::Line(LogLine {
seq: 1,
line: "b".into(),
}));
let _ = tx.send(LogMsg::End);
let snapshot = vec![LogLine {
seq: 0,
line: "a".into(),
}];
let events: Vec<LogEvent> = log_events(snapshot, rx, false).collect().await;
assert_eq!(
events,
vec![
LogEvent::Log("a".into()),
LogEvent::Log("b".into()),
LogEvent::End
]
);
}
#[tokio::test]
async fn truncated_emitted_on_broadcast_lag() {
let (tx, rx) = broadcast::channel(2);
for i in 0..5u64 {
let _ = tx.send(LogMsg::Line(LogLine {
seq: i,
line: format!("l{i}"),
}));
}
let _ = tx.send(LogMsg::End);
let events: Vec<LogEvent> = log_events(vec![], rx, false).collect().await;
assert!(
events.iter().any(|e| matches!(e, LogEvent::Truncated(_))),
"a lagging reader must get a Truncated event: {events:?}"
);
assert_eq!(events.last(), Some(&LogEvent::End));
}
#[test]
fn layer_captures_events_in_run_span_only() {
use tracing_subscriber::layer::SubscriberExt;
let hub = LogHub::new();
let subscriber = tracing_subscriber::registry().with(RunLogLayer::new(hub.clone()));
tracing::subscriber::with_default(subscriber, || {
tracing::info!("orphan event");
let span = tracing::info_span!("faucet.serve.run", serve_run_id = "run-xyz");
let _g = span.enter();
tracing::info!("hello from the run");
});
let (snapshot, _rx, _ended) = hub.reader("run-xyz").expect("buffer for the run exists");
assert!(
snapshot
.iter()
.any(|l| l.line.contains("hello from the run")),
"in-span event must be captured: {snapshot:?}"
);
assert!(
snapshot.iter().all(|l| !l.line.contains("orphan")),
"events outside the run span must not be captured"
);
assert!(hub.reader("nonexistent").is_none());
}
}