kill_me.rs

use std::any::TypeId;
use std::cell::RefCell;
use std::fs::{File, OpenOptions};
use std::io;
use std::io::Write;
use std::path::PathBuf;
use std::rc::Rc;
use std::sync::Mutex;
use tracing::field::{Field, ValueSet, Visit};
use tracing::span::{Attributes, Record};
use tracing::subscriber::Interest;
use tracing::{error_span, span, Event, Id, Instrument, Level, Metadata, Subscriber};
use tracing_subscriber::layer::{Context, Layer, SubscriberExt};
use tracing_subscriber::registry::LookupSpan;

use std::time::{Duration, Instant};
use tracing::level_filters::LevelFilter;
use tracing_appender::non_blocking::{NonBlocking, WorkerGuard};
use tracing_subscriber::fmt::writer::{
    BoxMakeWriter, EitherWriter, MutexGuardWriter, OptionalWriter,
};
use tracing_subscriber::fmt::{format, FormatEvent, FormatFields, FormattedFields, MakeWriter};
use tracing_subscriber::{fmt, layer};
use uuid::Uuid;

thread_local! {
    static WRITER: RefCell<Writer> = RefCell::new(Writer(None));
}

struct CustomLayer<S, N, E, W> {
    base: PathBuf,
    inner: fmt::Layer<S, N, E, W>,
}

#[derive(Clone)]
struct Writer(Option<NonBlocking>);

impl<'w> MakeWriter<'w> for Writer {
    type Writer = EitherWriter<NonBlocking, io::Sink>;

    fn make_writer(&'w self) -> Self::Writer {
        self.0
            .clone()
            .map(EitherWriter::A)
            .unwrap_or(EitherWriter::B(io::sink()))
    }
}

impl<S, N, E> CustomLayer<S, N, E, RefCellWriter> {
    pub fn new<W>(base: PathBuf, layer: fmt::Layer<S, N, E, W>) -> Self {
        let inner = layer.map_writer(|_| RefCellWriter(RefCell::new(Writer(None))));
        Self { base, inner }
    }
}

#[derive(Clone)]
struct RefCellWriter(RefCell<Writer>);

impl<'w> MakeWriter<'w> for RefCellWriter {
    type Writer = EitherWriter<NonBlocking, io::Sink>;

    fn make_writer(&'w self) -> Self::Writer {
        self.0
            .borrow()
            .clone()
            .0
            .map(EitherWriter::A)
            .unwrap_or(EitherWriter::B(io::sink()))
    }
}

#[derive(Debug)]
struct CustomLayerEnabled {
    writer: NonBlocking,
    _guard: WorkerGuard,
}

impl<S, N, E> Layer<S> for CustomLayer<S, N, E, RefCellWriter>
where
    S: Subscriber + for<'a> LookupSpan<'a>,
    N: for<'writer> FormatFields<'writer> + 'static,
    E: FormatEvent<S, N> + 'static,
{
    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
        let f = || {
            let Some(__id) = attrs.fields().field("__id") else {
                return;
            };

            let mut visitor = MatchStrVisitor { found: None };
            attrs.values().record(&mut visitor);

            let Some(__id) = visitor.found else { return };
            let __id = __id.to_string() + ".log";
            let file = OpenOptions::new()
                .write(true)
                .create(true)
                .truncate(false)
                .open(self.base.join(__id))
                .unwrap();

            let (writer, _guard) = NonBlocking::new(file);

            ctx.span(id)
                .unwrap()
                .extensions_mut()
                .insert(CustomLayerEnabled { writer, _guard });
        };

        f();
        self.inner.on_new_span(attrs, id, ctx);

        struct MatchStrVisitor {
            found: Option<String>,
        }

        impl Visit for MatchStrVisitor {
            fn record_str(&mut self, field: &Field, value: &str) {
                if field.name() == "__id" {
                    self.found = Some(value.to_string());
                }
            }
            fn record_debug(&mut self, _field: &Field, _value: &dyn std::fmt::Debug) {}
        }
    }

    fn on_record(&self, id: &Id, values: &Record<'_>, ctx: Context<'_, S>) {
        self.inner.on_record(id, values, ctx)
    }

    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
        let f = || {
            let Some(span) = ctx.lookup_current() else {
                return;
            };
            let mut extensions = span.extensions_mut();
            let Some(extension) = extensions.get_mut::<CustomLayerEnabled>() else {
                return;
            };

            WRITER.set(Writer(Some(extension.writer.clone())));
        };

        f();
        WRITER.with(|f| {
            let writer = f.clone().borrow().clone();
            *self.inner.writer().0.borrow_mut() = writer.clone();
            self.inner.on_event(event, ctx);
            
            // let _ = writer.flush();
            // *lock = OptionalWriter::none();
        });

    }

    fn on_enter(&self, id: &Id, ctx: Context<'_, S>) {
        self.inner.on_enter(id, ctx)
    }

    fn on_exit(&self, id: &Id, ctx: Context<'_, S>) {
        self.inner.on_exit(id, ctx)
    }

    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
        self.inner.on_close(id, ctx)
    }

    unsafe fn downcast_raw(&self, id: TypeId) -> Option<*const ()> {
        self.inner.downcast_raw(id)
    }
}

async fn a(id: String) {
    tracing::info!("{id} | a");
    for c in 1..=5 {
        tracing::info!("{id} | a{c}");
        tokio::time::sleep(Duration::from_secs(1)).await;
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let base = std::env::current_dir()?.join("logs");

    let format = tracing_subscriber::fmt::format()
        .with_level(true) // don't include levels in formatted output
        .with_target(true) // don't include targets
        .with_thread_ids(true) // include the thread ID of the current thread
        .with_thread_names(false) // include the name of the current thread
        .with_ansi(false)
        .pretty(); // use the `Compact` formatting style.

    let inner = tracing_subscriber::fmt::layer().event_format(format.clone());
    let registry = tracing_subscriber::Registry::default().with(CustomLayer::new(base, inner));

    tracing::subscriber::set_global_default(registry)?;

    let id = Uuid::new_v4().to_string();
    let span = error_span!("bot", __id = id);
    tokio::spawn(a(id.to_string()).instrument(span));

    let id = Uuid::new_v4().to_string();
    let span = error_span!("bot", __id = id);
    tokio::spawn(a(id.to_string()).instrument(span));

    tokio::time::sleep(Duration::from_secs(5)).await;
    Ok(())
}