use std::fmt::Write as _;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use rmpv::Value as MpVal;
use tokio::io::AsyncWriteExt;
use tokio::net::UnixStream;
use tokio::sync::mpsc;
use tracing::field::{Field, Visit};
use tracing::{Event, Level as TraceLevel, Subscriber};
use tracing_subscriber::layer::{Context, Layer};
use tracing_subscriber::registry::LookupSpan;
use crate::logd::{Fields, IngestFrame, Level, LogFrame};
pub use crate::logd::emitter::DEFAULT_INGEST_SOCK;
pub const CHANNEL_CAPACITY: usize = 1024;
const BATCH_MAX_FRAMES: usize = 128;
const BATCH_LINGER: Duration = Duration::from_millis(100);
const RECONNECT_PACE: crate::retry::RetryPace = crate::retry::LOCAL_SOCKET;
#[derive(Debug, Default)]
pub struct LayerStats {
enqueued: AtomicU64,
dropped: AtomicU64,
}
impl LayerStats {
fn record_enqueued(&self) {
self.enqueued.fetch_add(1, Ordering::Relaxed);
}
fn record_dropped(&self) {
self.dropped.fetch_add(1, Ordering::Relaxed);
}
pub fn enqueued(&self) -> u64 {
self.enqueued.load(Ordering::Relaxed)
}
pub fn dropped(&self) -> u64 {
self.dropped.load(Ordering::Relaxed)
}
}
fn spawn_shipper(rx: mpsc::Receiver<LogFrame>, path: std::path::PathBuf) {
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(shipper_task(rx, path));
return;
}
let _ = std::thread::Builder::new()
.name("logd-shipper".to_string())
.spawn(move || {
if let Ok(rt) = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
{
rt.block_on(shipper_task(rx, path));
}
});
}
pub struct LogdLayer<S> {
source: String,
tx: mpsc::Sender<LogFrame>,
stats: Arc<LayerStats>,
_marker: std::marker::PhantomData<fn(S)>,
}
impl<S> LogdLayer<S> {
pub fn new(source: impl Into<String>) -> Self {
Self::with_socket(source, DEFAULT_INGEST_SOCK)
}
pub fn with_socket(source: impl Into<String>, socket: impl AsRef<Path>) -> Self {
let (tx, rx) = mpsc::channel::<LogFrame>(CHANNEL_CAPACITY);
let stats = Arc::new(LayerStats::default());
let path = socket.as_ref().to_path_buf();
spawn_shipper(rx, path);
Self {
source: source.into(),
tx,
stats,
_marker: std::marker::PhantomData,
}
}
pub fn stats(&self) -> Arc<LayerStats> {
Arc::clone(&self.stats)
}
fn enqueue(&self, frame: LogFrame) {
match self.tx.try_send(frame) {
Ok(()) => self.stats.record_enqueued(),
Err(mpsc::error::TrySendError::Full(_)) => self.stats.record_dropped(),
Err(mpsc::error::TrySendError::Closed(_)) => self.stats.record_dropped(),
}
}
}
fn map_level(level: &TraceLevel) -> Level {
match *level {
TraceLevel::TRACE => Level::Trace,
TraceLevel::DEBUG => Level::Debug,
TraceLevel::INFO => Level::Info,
TraceLevel::WARN => Level::Warn,
TraceLevel::ERROR => Level::Error,
}
}
fn now_us() -> i64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_micros() as i64)
.unwrap_or(0)
}
struct FieldVisitor {
message: String,
fields: Fields,
}
impl FieldVisitor {
fn new() -> Self {
Self {
message: String::new(),
fields: Fields::new(),
}
}
fn insert(&mut self, field: &Field, value: MpVal) {
let name = field.name();
if name == "message" {
return;
}
self.fields.insert(name.to_string(), value);
}
}
impl Visit for FieldVisitor {
fn record_str(&mut self, field: &Field, value: &str) {
if field.name() == "message" {
self.message.push_str(value);
} else {
self.insert(field, MpVal::from(value));
}
}
fn record_i64(&mut self, field: &Field, value: i64) {
self.insert(field, MpVal::from(value));
}
fn record_u64(&mut self, field: &Field, value: u64) {
self.insert(field, MpVal::from(value));
}
fn record_f64(&mut self, field: &Field, value: f64) {
self.insert(field, MpVal::from(value));
}
fn record_bool(&mut self, field: &Field, value: bool) {
self.insert(field, MpVal::from(value));
}
fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
if field.name() == "message" {
let _ = write!(self.message, "{value:?}");
} else {
self.insert(field, MpVal::from(format!("{value:?}")));
}
}
}
impl<S> Layer<S> for LogdLayer<S>
where
S: Subscriber + for<'a> LookupSpan<'a>,
{
fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) {
let meta = event.metadata();
let mut visitor = FieldVisitor::new();
event.record(&mut visitor);
let mut frame = LogFrame::new(
now_us(),
self.source.clone(),
map_level(meta.level()),
visitor.message,
);
frame.target = Some(meta.target().to_string());
frame.fields = visitor.fields;
frame.redact_fields();
self.enqueue(frame);
}
}
struct SocketWriter {
path: PathBuf,
stream: Option<UnixStream>,
}
impl SocketWriter {
fn new(path: PathBuf) -> Self {
Self { path, stream: None }
}
async fn ensure_connected(&mut self) -> bool {
if self.stream.is_some() {
return true;
}
match UnixStream::connect(&self.path).await {
Ok(s) => {
self.stream = Some(s);
true
}
Err(_) => {
tokio::time::sleep(RECONNECT_PACE.wait()).await;
false
}
}
}
async fn write_all(&mut self, buf: &[u8]) -> bool {
let Some(stream) = self.stream.as_mut() else {
return false;
};
match stream.write_all(buf).await {
Ok(()) => true,
Err(_) => {
self.stream = None;
false
}
}
}
}
fn encode_batch(batch: &[LogFrame]) -> Vec<u8> {
let mut buf = Vec::new();
for frame in batch {
if let Ok(bytes) = IngestFrame::Log(frame.clone()).encode() {
buf.extend_from_slice(&bytes);
}
}
buf
}
async fn shipper_task(mut rx: mpsc::Receiver<LogFrame>, path: PathBuf) {
let mut writer = SocketWriter::new(path);
let mut batch: Vec<LogFrame> = Vec::with_capacity(BATCH_MAX_FRAMES);
loop {
let Some(first) = rx.recv().await else {
break;
};
batch.clear();
batch.push(first);
let deadline = tokio::time::Instant::now() + BATCH_LINGER;
while batch.len() < BATCH_MAX_FRAMES {
match tokio::time::timeout_at(deadline, rx.recv()).await {
Ok(Some(frame)) => batch.push(frame),
Ok(None) => break,
Err(_) => break,
}
}
if writer.ensure_connected().await {
let buf = encode_batch(&batch);
if !buf.is_empty() {
let _ = writer.write_all(&buf).await;
}
}
}
if !batch.is_empty() && writer.stream.is_some() {
let buf = encode_batch(&batch);
if !buf.is_empty() {
let _ = writer.write_all(&buf).await;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::io::AsyncReadExt;
use tokio::net::UnixListener;
use tracing::Level as TLevel;
use tracing_subscriber::prelude::*;
use crate::frame::{decode_len, HEADER_SIZE};
use crate::logd::{IngestFrame, Level, LOGD_MAX_FRAME};
#[test]
fn constructing_the_layer_outside_a_runtime_does_not_panic() {
assert!(
tokio::runtime::Handle::try_current().is_err(),
"this test is only meaningful without an ambient runtime"
);
let layer: LogdLayer<tracing_subscriber::Registry> =
LogdLayer::with_socket("sync-binary", tmp_sock());
assert_eq!(layer.stats().enqueued(), 0);
assert_eq!(layer.stats().dropped(), 0);
}
fn tmp_sock() -> PathBuf {
let mut p = std::env::temp_dir();
let uniq = format!(
"ados-logd-layer-test-{}-{}.sock",
std::process::id(),
now_us()
);
p.push(uniq);
let _ = std::fs::remove_file(&p);
p
}
async fn read_one(stream: &mut UnixStream) -> IngestFrame {
let mut header = [0u8; HEADER_SIZE];
stream.read_exact(&mut header).await.unwrap();
let len = decode_len(header, LOGD_MAX_FRAME, true).unwrap();
let mut body = vec![0u8; len];
stream.read_exact(&mut body).await.unwrap();
IngestFrame::decode(&body).unwrap()
}
#[test]
fn level_mapping_covers_every_level() {
assert_eq!(map_level(&TLevel::TRACE), Level::Trace);
assert_eq!(map_level(&TLevel::DEBUG), Level::Debug);
assert_eq!(map_level(&TLevel::INFO), Level::Info);
assert_eq!(map_level(&TLevel::WARN), Level::Warn);
assert_eq!(map_level(&TLevel::ERROR), Level::Error);
}
#[test]
fn field_visitor_separates_message_from_fields() {
let out: Arc<std::sync::Mutex<Option<(String, Fields)>>> =
Arc::new(std::sync::Mutex::new(None));
let sub = CapturingSub {
out: Arc::clone(&out),
};
tracing::subscriber::with_default(sub, || {
tracing::info!(count = 7u64, name = "rig", "hello world");
});
let (msg, fields) = out.lock().unwrap().take().expect("event captured");
assert_eq!(msg, "hello world");
assert_eq!(fields.get("count").and_then(|v| v.as_u64()), Some(7));
assert_eq!(fields.get("name").and_then(|v| v.as_str()), Some("rig"));
assert!(!fields.contains_key("message"));
}
struct CapturingSub {
out: Arc<std::sync::Mutex<Option<(String, Fields)>>>,
}
impl tracing::Subscriber for CapturingSub {
fn enabled(&self, _m: &tracing::Metadata<'_>) -> bool {
true
}
fn new_span(&self, _a: &tracing::span::Attributes<'_>) -> tracing::span::Id {
tracing::span::Id::from_u64(1)
}
fn record(&self, _s: &tracing::span::Id, _v: &tracing::span::Record<'_>) {}
fn record_follows_from(&self, _s: &tracing::span::Id, _f: &tracing::span::Id) {}
fn event(&self, event: &Event<'_>) {
let mut v = FieldVisitor::new();
event.record(&mut v);
if let Ok(mut guard) = self.out.lock() {
*guard = Some((v.message, v.fields));
}
}
fn enter(&self, _s: &tracing::span::Id) {}
fn exit(&self, _s: &tracing::span::Id) {}
}
#[tokio::test]
async fn ships_a_redacted_frame_to_a_live_socket() {
let path = tmp_sock();
let listener = UnixListener::bind(&path).unwrap();
let accept = {
let listener = listener;
tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
read_one(&mut stream).await
})
};
let layer: LogdLayer<_> = LogdLayer::with_socket("ados-test", &path);
let subscriber = tracing_subscriber::registry().with(layer);
tracing::subscriber::with_default(subscriber, || {
tracing::warn!(
api_key = "ABCDEFGHIJ1234567890",
attempt = 2u64,
"link down"
);
});
let frame = tokio::time::timeout(Duration::from_secs(5), accept)
.await
.expect("frame delivered within the deadline")
.expect("accept task ok");
match frame {
IngestFrame::Log(log) => {
assert_eq!(log.source, "ados-test");
assert_eq!(log.level, Level::Warn);
assert_eq!(log.msg, "link down");
assert_eq!(log.fields.get("attempt").and_then(|v| v.as_u64()), Some(2));
assert_eq!(
log.fields.get("api_key").and_then(|v| v.as_str()),
Some("redacted:len=20")
);
}
other => panic!("expected a log frame, got {other:?}"),
}
let _ = std::fs::remove_file(&path);
}
#[tokio::test]
async fn absent_socket_never_blocks_or_panics() {
let path = tmp_sock(); let layer: LogdLayer<_> = LogdLayer::with_socket("ados-test", &path);
let stats = layer.stats();
let subscriber = tracing_subscriber::registry().with(layer);
tracing::subscriber::with_default(subscriber, || {
for i in 0..50 {
tracing::info!(i, "tick");
}
});
assert!(stats.enqueued() >= 1, "at least one frame was enqueued");
tokio::time::sleep(Duration::from_millis(50)).await;
let _ = std::fs::remove_file(&path);
}
#[tokio::test]
async fn full_channel_drops_without_blocking_and_counts_it() {
let path = tmp_sock(); let layer: LogdLayer<tracing_subscriber::Registry> =
LogdLayer::with_socket("ados-test", &path);
let stats = layer.stats();
let total = CHANNEL_CAPACITY * 2;
for n in 0..total {
let frame = LogFrame::new(now_us(), "ados-test", Level::Debug, format!("m{n}"));
layer.enqueue(frame);
}
let enq = stats.enqueued();
let drp = stats.dropped();
assert_eq!(enq + drp, total as u64, "every frame accounted for");
assert!(drp >= 1, "overflow past capacity was dropped and counted");
let _ = std::fs::remove_file(&path);
}
#[tokio::test]
async fn warn_and_error_are_kept_until_the_channel_saturates() {
let path = tmp_sock();
let listener = UnixListener::bind(&path).unwrap();
tokio::spawn(async move {
loop {
let Ok((mut stream, _)) = listener.accept().await else {
break;
};
let mut sink = [0u8; 4096];
while stream.read(&mut sink).await.unwrap_or(0) > 0 {}
}
});
let layer: LogdLayer<tracing_subscriber::Registry> =
LogdLayer::with_socket("ados-test", &path);
let stats = layer.stats();
for lvl in [Level::Debug, Level::Info, Level::Warn, Level::Error] {
layer.enqueue(LogFrame::new(now_us(), "ados-test", lvl, "x"));
}
assert_eq!(stats.dropped(), 0);
assert_eq!(stats.enqueued(), 4);
let _ = std::fs::remove_file(&path);
}
}