use std::fmt::Write as _;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::mpsc::{sync_channel, Receiver, SyncSender, TrySendError};
use std::sync::{Mutex, OnceLock};
use tracing::field::{Field, Visit};
use tracing::{Event, Level, Subscriber};
use tracing_subscriber::layer::{Context, Layer};
use tracing_subscriber::registry::LookupSpan;
pub const CAPACITY: usize = 256;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LogLine {
pub level: Level,
pub target: &'static str,
pub message: String,
pub fields: Vec<(&'static str, String)>,
pub span_fields: Vec<(&'static str, String)>,
}
impl LogLine {
pub fn field(&self, name: &str) -> Option<&str> {
self.fields
.iter()
.find(|(k, _)| *k == name)
.map(|(_, v)| v.as_str())
}
pub fn get(&self, name: &str) -> Option<&str> {
self.field(name).or_else(|| {
self.span_fields
.iter()
.rev()
.find(|(k, _)| *k == name)
.map(|(_, v)| v.as_str())
})
}
pub fn event(&self) -> Option<&str> {
self.field("event")
}
}
struct Tap {
tx: SyncSender<LogLine>,
}
static TAP: OnceLock<Tap> = OnceLock::new();
static RECEIVER: OnceLock<Mutex<Option<Receiver<LogLine>>>> = OnceLock::new();
static DROPPED: AtomicU64 = AtomicU64::new(0);
pub struct TapLayer;
impl TapLayer {
pub fn new() -> Self {
TAP.get_or_init(|| {
let (tx, rx) = sync_channel(CAPACITY);
let _ = RECEIVER.set(Mutex::new(Some(rx)));
Tap { tx }
});
TapLayer
}
}
impl Default for TapLayer {
fn default() -> Self {
Self::new()
}
}
pub fn take_receiver() -> Option<Receiver<LogLine>> {
RECEIVER.get()?.lock().ok()?.take()
}
pub fn dropped() -> u64 {
DROPPED.load(Ordering::Relaxed)
}
pub fn push(line: LogLine) {
if let Some(tap) = TAP.get() {
offer(&tap.tx, line);
}
}
fn offer(tx: &SyncSender<LogLine>, line: LogLine) {
match tx.try_send(line) {
Ok(()) => {}
Err(TrySendError::Full(_)) | Err(TrySendError::Disconnected(_)) => {
DROPPED.fetch_add(1, Ordering::Relaxed);
}
}
}
#[derive(Default)]
struct SpanFields(Vec<(&'static str, String)>);
impl SpanFields {
fn put(&mut self, name: &'static str, value: String) {
if let Some(slot) = self.0.iter_mut().find(|(k, _)| *k == name) {
slot.1 = value;
} else {
self.0.push((name, value));
}
}
}
struct Collect<'a> {
message: &'a mut String,
fields: &'a mut Vec<(&'static str, String)>,
}
impl Visit for Collect<'_> {
fn record_str(&mut self, field: &Field, value: &str) {
if field.name() == "message" {
self.message.push_str(value);
} else {
self.fields.push((field.name(), value.to_owned()));
}
}
fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
if field.name() == "message" {
let _ = write!(self.message, "{value:?}");
} else {
self.fields.push((field.name(), format!("{value:?}")));
}
}
}
impl<S> Layer<S> for TapLayer
where
S: Subscriber + for<'a> LookupSpan<'a>,
{
fn on_new_span(
&self,
attrs: &tracing::span::Attributes<'_>,
id: &tracing::span::Id,
ctx: Context<'_, S>,
) {
let Some(span) = ctx.span(id) else { return };
let mut fields = SpanFields::default();
let mut message = String::new();
attrs.record(&mut Collect {
message: &mut message,
fields: &mut fields.0,
});
span.extensions_mut().insert(fields);
}
fn on_record(
&self,
id: &tracing::span::Id,
values: &tracing::span::Record<'_>,
ctx: Context<'_, S>,
) {
let Some(span) = ctx.span(id) else { return };
let mut fresh = Vec::new();
let mut message = String::new();
values.record(&mut Collect {
message: &mut message,
fields: &mut fresh,
});
let mut ext = span.extensions_mut();
let fields = ext.get_mut::<SpanFields>();
match fields {
Some(fields) => fresh.into_iter().for_each(|(k, v)| fields.put(k, v)),
None => ext.insert(SpanFields(fresh)),
}
}
fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
let Some(tap) = TAP.get() else { return };
if *event.metadata().level() > Level::INFO {
return;
}
let mut message = String::new();
let mut fields = Vec::new();
event.record(&mut Collect {
message: &mut message,
fields: &mut fields,
});
let mut span_fields: Vec<(&'static str, String)> = Vec::new();
if let Some(scope) = ctx.event_scope(event) {
for span in scope.from_root() {
if let Some(sf) = span.extensions().get::<SpanFields>() {
for (k, v) in &sf.0 {
if let Some(slot) = span_fields.iter_mut().find(|(n, _)| n == k) {
slot.1 = v.clone();
} else {
span_fields.push((k, v.clone()));
}
}
}
}
}
offer(
&tap.tx,
LogLine {
level: *event.metadata().level(),
target: event.metadata().target(),
message,
fields,
span_fields,
},
);
}
}
#[cfg(test)]
mod tests {
use super::*;
use tracing_subscriber::layer::SubscriberExt as _;
static ONE_AT_A_TIME: std::sync::Mutex<()> = std::sync::Mutex::new(());
fn with_tap<R>(f: impl FnOnce() -> R) -> R {
let _serial = ONE_AT_A_TIME.lock().unwrap_or_else(|e| e.into_inner());
let subscriber = tracing_subscriber::registry().with(TapLayer::new());
tracing::subscriber::with_default(subscriber, f)
}
#[test]
fn an_event_arrives_with_its_fields_and_the_span_fields_around_it() {
let rx = with_tap(|| {
let rx = take_receiver();
if let Some(rx) = &rx {
for _ in rx.try_iter() {}
}
let span = tracing::info_span!(
"ssh.connection",
repo = "ns/r",
principal = tracing::field::Empty
);
span.record("principal", "alice");
let _e = span.enter();
tracing::info!(target: "tap-test", event = "ssh.auth.accepted", fingerprint = "SHA256:x", "publickey accepted");
tracing::debug!(target: "tap-test", "never forwarded");
rx.map(|rx| rx.try_iter().collect::<Vec<_>>())
});
let Some(lines) = rx else { return };
let rx = lines.into_iter();
let mut rx = rx;
let line = rx
.by_ref()
.find(|l| l.target == "tap-test")
.expect("the info line was tapped");
assert_eq!(line.level, Level::INFO);
assert_eq!(line.message, "publickey accepted");
assert_eq!(line.event(), Some("ssh.auth.accepted"));
assert_eq!(line.field("fingerprint"), Some("SHA256:x"));
assert_eq!(line.get("repo"), Some("ns/r"));
assert_eq!(
line.get("principal"),
Some("alice"),
"a recorded-later span field is seen"
);
assert!(
rx.all(|l| l.target != "tap-test" || l.message != "never forwarded"),
"DEBUG is refused before the channel"
);
}
#[test]
fn a_full_buffer_drops_and_counts_rather_than_blocking() {
let before = dropped();
with_tap(|| {
for i in 0..(CAPACITY * 2) {
push(LogLine {
level: Level::INFO,
target: "tap-flood",
message: format!("{i}"),
fields: vec![],
span_fields: vec![],
});
}
});
assert!(
dropped() >= before + CAPACITY as u64,
"the surplus was counted, not waited for"
);
}
}