AkurAI Build
Menu

Public repository / 262e2146713f325f0b45d3cb01d407461d774a9a

akurai-log / src/lib.rs

Latest change dfa0605e64c2086eca2341e817eaecfa115486c8 - AKURAI-LOG-2 consolidate log shipping contract by Ólafur Búi Ólafsson

//! `akurai-log` — a [`tracing`] layer that ships structured log events to
//! AkurAI-Monitor's `POST /api/ingest` endpoint, in addition to normal stdout.
//!
//! Design goals:
//! - **Never stall the app.** Events are handed to a bounded channel with
//!   `try_send`; if the channel is full (Monitor slow/down) events are dropped,
//!   never blocking the caller.
//! - **Runtime-agnostic.** A dedicated OS thread owns a blocking HTTP client and
//!   batches events, so consumers need no particular async runtime.
//! - **No feedback loop.** Events from the HTTP stack (`reqwest`/`hyper`) are
//!   ignored so shipping logs can't generate more logs to ship.
//!
//! Configuration (environment):
//! - `MONITOR_INGEST_URL`   base URL of Monitor (e.g. `https://monitor.olibuijr.com`)
//! - `MONITOR_INGEST_TOKEN` bearer token for `/api/ingest`
//! - `MONITOR_HOST_ID`      optional host id to attribute logs to
//! - `MONITOR_APP_ID`       optional application id to attribute logs to
//! - `MONITOR_LOG_LEVEL`    filter for what reaches Monitor (default `debug`)
//! - `RUST_LOG`             standard filter for stdout (unchanged)
//!
//! Usage:
//! ```no_run
//! akurai_log::init("akurai-dns");
//! tracing::info!(user = "olibuijr", "started");
//! ```

use std::sync::mpsc::{sync_channel, SyncSender};
use std::sync::OnceLock;
use std::time::{Duration, Instant};

use tracing::field::{Field, Visit};
use tracing_subscriber::layer::{Context, Layer};
use tracing_subscriber::prelude::*;
use tracing_subscriber::{fmt, EnvFilter};

/// One structured log record queued for shipment.
#[derive(Clone)]
struct Record {
    source: String,
    level: &'static str,
    target: String,
    message: String,
    fields: serde_json::Map<String, serde_json::Value>,
    pid: i64,
    ts: i64,
}

/// Sender is populated at `init()` from runtime env — kept as a `OnceLock`
/// because the initializer needs runtime input (the configured channel).
static SENDER: OnceLock<SyncSender<Record>> = OnceLock::new();

const CHANNEL_CAPACITY: usize = 8192;
const BATCH_MAX: usize = 128;
const MAX_LINE_LENGTH: usize = 8192;
const FLUSH_INTERVAL: Duration = Duration::from_secs(2);

/// Ingestion targets that must be ignored to avoid a log→ship→log loop.
fn is_noise(target: &str) -> bool {
    target.starts_with("reqwest")
        || target.starts_with("hyper")
        || target.starts_with("rustls")
        || target.starts_with("h2")
        || target.starts_with("want")
        || target.starts_with("mio")
}

fn level_str(level: &tracing::Level) -> &'static str {
    match *level {
        tracing::Level::ERROR => "error",
        tracing::Level::WARN => "warn",
        tracing::Level::INFO => "info",
        tracing::Level::DEBUG => "debug",
        tracing::Level::TRACE => "trace",
    }
}

/// Collects the event's `message` plus any structured key/value fields.
#[derive(Default)]
struct FieldVisitor {
    message: String,
    fields: serde_json::Map<String, serde_json::Value>,
}

impl Visit for FieldVisitor {
    fn record_str(&mut self, field: &Field, value: &str) {
        if field.name() == "message" {
            self.message = value.to_string();
        } else {
            self.fields
                .insert(field.name().to_string(), serde_json::Value::from(value));
        }
    }

    fn record_i64(&mut self, field: &Field, value: i64) {
        self.fields
            .insert(field.name().to_string(), serde_json::Value::from(value));
    }

    fn record_u64(&mut self, field: &Field, value: u64) {
        self.fields
            .insert(field.name().to_string(), serde_json::Value::from(value));
    }

    fn record_bool(&mut self, field: &Field, value: bool) {
        self.fields
            .insert(field.name().to_string(), serde_json::Value::from(value));
    }

    fn record_f64(&mut self, field: &Field, value: f64) {
        self.fields
            .insert(field.name().to_string(), serde_json::Value::from(value));
    }

    fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
        let rendered = format!("{value:?}");
        if field.name() == "message" {
            self.message = rendered;
        } else {
            self.fields
                .insert(field.name().to_string(), serde_json::Value::from(rendered));
        }
    }
}

/// The tracing layer that forwards events into the shipment channel.
struct MonitorLayer {
    source: String,
    pid: i64,
}

impl<S: tracing::Subscriber> Layer<S> for MonitorLayer {
    fn on_event(&self, event: &tracing::Event<'_>, _ctx: Context<'_, S>) {
        let target = event.metadata().target();
        if is_noise(target) {
            return;
        }
        let Some(sender) = SENDER.get() else {
            return;
        };
        let mut visitor = FieldVisitor::default();
        event.record(&mut visitor);
        let record = Record {
            source: self.source.clone(),
            level: level_str(event.metadata().level()),
            target: target.to_string(),
            message: visitor.message.chars().take(MAX_LINE_LENGTH).collect(),
            fields: visitor.fields,
            pid: self.pid,
            ts: now_secs(),
        };
        // Drop on backpressure — shipping logs must never block the app.
        let _ = sender.try_send(record);
    }
}

fn now_secs() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs() as i64)
        .unwrap_or(0)
}

/// Initialize logging: stdout (via `RUST_LOG`) plus shipment to AkurAI-Monitor.
///
/// Safe to call once at startup. If `MONITOR_INGEST_URL`/`MONITOR_INGEST_TOKEN`
/// are unset, only stdout logging is installed and nothing is shipped.
pub fn init(app_name: &str) {
    let stdout_filter = EnvFilter::try_from_default_env()
        .unwrap_or_else(|_| EnvFilter::new(format!("{}=info", app_name.replace('-', "_"))));

    let url = std::env::var("MONITOR_INGEST_URL")
        .ok()
        .filter(|s| !s.is_empty());
    let token = std::env::var("MONITOR_INGEST_TOKEN")
        .ok()
        .filter(|s| !s.is_empty());

    let monitor_layer = match (url, token) {
        (Some(url), Some(token)) => {
            let host_id = std::env::var("MONITOR_HOST_ID")
                .ok()
                .and_then(|s| s.parse::<i64>().ok());
            let app_id = std::env::var("MONITOR_APP_ID")
                .ok()
                .and_then(|s| s.parse::<i64>().ok());
            let ship_level =
                std::env::var("MONITOR_LOG_LEVEL").unwrap_or_else(|_| "debug".to_string());

            let (tx, rx) = sync_channel::<Record>(CHANNEL_CAPACITY);
            let _ = SENDER.set(tx);
            spawn_shipper(rx, url, token, host_id, app_id);

            // Ship every event at/above the configured level (noise targets are
            // excluded in the layer itself), regardless of the emitting module.
            let filter = EnvFilter::new(&ship_level);
            Some(
                MonitorLayer {
                    source: app_name.to_string(),
                    pid: std::process::id() as i64,
                }
                .with_filter(filter),
            )
        }
        _ => None,
    };

    tracing_subscriber::registry()
        .with(
            fmt::layer()
                .with_writer(std::io::stdout)
                .with_filter(stdout_filter),
        )
        .with(monitor_layer)
        .init();
}

/// Background thread: batch records and POST them to `/api/ingest`.
fn spawn_shipper(
    rx: std::sync::mpsc::Receiver<Record>,
    base_url: String,
    token: String,
    host_id: Option<i64>,
    app_id: Option<i64>,
) {
    let endpoint = format!("{}/api/ingest", base_url.trim_end_matches('/'));
    std::thread::Builder::new()
        .name("akurai-log-shipper".into())
        .spawn(move || {
            let client = match reqwest::blocking::Client::builder()
                .timeout(Duration::from_secs(5))
                .build()
            {
                Ok(c) => c,
                Err(_) => return,
            };
            let mut batch: Vec<Record> = Vec::with_capacity(BATCH_MAX);
            let mut last_flush = Instant::now();
            loop {
                match rx.recv_timeout(FLUSH_INTERVAL) {
                    Ok(record) => {
                        batch.push(record);
                        // Drain anything else already queued without blocking.
                        while batch.len() < BATCH_MAX {
                            match rx.try_recv() {
                                Ok(r) => batch.push(r),
                                Err(_) => break,
                            }
                        }
                    }
                    Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {}
                    Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => {
                        flush(&client, &endpoint, &token, host_id, app_id, &mut batch);
                        return;
                    }
                }
                if !batch.is_empty()
                    && (batch.len() >= BATCH_MAX || last_flush.elapsed() >= FLUSH_INTERVAL)
                {
                    flush(&client, &endpoint, &token, host_id, app_id, &mut batch);
                    last_flush = Instant::now();
                }
            }
        })
        .ok();
}

fn flush(
    client: &reqwest::blocking::Client,
    endpoint: &str,
    token: &str,
    host_id: Option<i64>,
    app_id: Option<i64>,
    batch: &mut Vec<Record>,
) {
    if batch.is_empty() {
        return;
    }
    let logs: Vec<serde_json::Value> = batch
        .iter()
        .map(|r| {
            serde_json::json!({
                "source": r.source,
                "line": r.message,
                "level": r.level,
                "target": r.target,
                "pid": r.pid,
                "origin": "server",
                "fields": if r.fields.is_empty() { serde_json::Value::Null }
                          else { serde_json::Value::Object(r.fields.clone()) },
                "host_id": host_id,
                "application_id": app_id,
                "ts": r.ts,
            })
        })
        .collect();
    let body = serde_json::json!({ "logs": logs });
    // Best-effort; a failed ship is dropped, never retried into unbounded memory.
    let _ = client.post(endpoint).bearer_auth(token).json(&body).send();
    batch.clear();
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn noise_targets_are_filtered() {
        assert!(is_noise("reqwest::connect"));
        assert!(is_noise("hyper::proto::h1"));
        assert!(is_noise("rustls::client"));
        assert!(!is_noise("akurai_dns::server"));
        assert!(!is_noise("notes_api::auth"));
    }

    #[test]
    fn levels_map_to_lowercase_strings() {
        assert_eq!(level_str(&tracing::Level::ERROR), "error");
        assert_eq!(level_str(&tracing::Level::WARN), "warn");
        assert_eq!(level_str(&tracing::Level::INFO), "info");
        assert_eq!(level_str(&tracing::Level::DEBUG), "debug");
        assert_eq!(level_str(&tracing::Level::TRACE), "trace");
    }

    #[test]
    fn visitor_default_is_empty_and_json_safe() {
        let v = FieldVisitor::default();
        assert!(v.message.is_empty());
        assert!(v.fields.is_empty());
        let value = serde_json::to_value(&v.fields).expect("serializable");
        assert!(value.is_object());
    }
}