khadhroony-solana-project/crates/ksp-logging-lib/src/domain.rs

// file: crates/ksp-logging-lib/src/domain.rs
// version: 1

std::thread_local! {
    static CURRENT_DOMAIN: std::cell::RefCell<std::option::Option<std::string::String>> = const { std::cell::RefCell::new(std::option::Option::None) };
}

#[derive(Clone, Debug, Default, Eq, PartialEq)]
struct SpanDomain {
    value: std::option::Option<std::string::String>,
}

#[derive(Default)]
struct DomainVisitor {
    value: std::option::Option<std::string::String>,
}

impl tracing::field::Visit for DomainVisitor {
    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
        if field.name() == "domain" {
            self.value = std::option::Option::Some(value.to_string());
        }
    }

    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
        if field.name() == "domain" {
            self.value = std::option::Option::Some(format!("{value:?}"));
        }
    }
}

pub(crate) struct DomainContextLayer;

impl DomainContextLayer {
    pub(crate) const fn new() -> Self {
        return Self;
    }
}

impl<S> tracing_subscriber::Layer<S> for DomainContextLayer
where
    S: tracing::Subscriber + for<'lookup> tracing_subscriber::registry::LookupSpan<'lookup>,
{
    fn on_new_span(&self, attrs: &tracing::span::Attributes<'_>, id: &tracing::span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
        let mut visitor = DomainVisitor::default();
        attrs.record(&mut visitor);
        let effective_domain = match visitor.value {
            std::option::Option::Some(domain) => std::option::Option::Some(domain),
            std::option::Option::None => span_parent_domain(id, &ctx),
        };
        if let std::option::Option::Some(span) = ctx.span(id) {
            span.extensions_mut().insert(SpanDomain { value: effective_domain.clone() });
        }
        set_current_domain(effective_domain.as_deref());
    }

    fn on_record(&self, id: &tracing::span::Id, values: &tracing::span::Record<'_>, ctx: tracing_subscriber::layer::Context<'_, S>) {
        let mut visitor = DomainVisitor::default();
        values.record(&mut visitor);
        if let std::option::Option::Some(domain) = visitor.value {
            if let std::option::Option::Some(span) = ctx.span(id) {
                let mut extensions = span.extensions_mut();
                if let std::option::Option::Some(stored) = extensions.get_mut::<SpanDomain>() {
                    stored.value = std::option::Option::Some(domain.clone());
                } else {
                    extensions.insert(SpanDomain { value: std::option::Option::Some(domain.clone()) });
                }
            }
            set_current_domain(std::option::Option::Some(domain.as_str()));
            return;
        }
        let domain = span_domain(id, &ctx);
        set_current_domain(domain.as_deref());
    }

    fn on_event(&self, event: &tracing::Event<'_>, ctx: tracing_subscriber::layer::Context<'_, S>) {
        let mut visitor = DomainVisitor::default();
        event.record(&mut visitor);
        let effective_domain = match visitor.value {
            std::option::Option::Some(domain) => std::option::Option::Some(domain),
            std::option::Option::None => event_parent_domain(event, &ctx),
        };
        set_current_domain(effective_domain.as_deref());
    }

    fn on_enter(&self, id: &tracing::span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
        let domain = span_domain(id, &ctx);
        set_current_domain(domain.as_deref());
    }

    fn on_exit(&self, id: &tracing::span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
        let domain = span_domain(id, &ctx);
        set_current_domain(domain.as_deref());
    }

    fn on_close(&self, id: tracing::span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
        let domain = span_domain(&id, &ctx);
        set_current_domain(domain.as_deref());
    }
}

pub(crate) fn current_domain_matches(selectors: &[std::string::String]) -> bool {
    if let [selector] = selectors
        && selector == "*"
    {
        return true;
    }
    return CURRENT_DOMAIN.with(|current| -> bool {
        let borrow_result = current.try_borrow();
        let current = match borrow_result {
            std::result::Result::Ok(current) => current,
            std::result::Result::Err(_) => return false,
        };
        let domain = match current.as_ref() {
            std::option::Option::Some(domain) => domain,
            std::option::Option::None => return false,
        };
        return selectors.iter().any(|selector| -> bool {
            return domain.starts_with(selector.as_str());
        });
    });
}

fn set_current_domain(domain: std::option::Option<&str>) {
    CURRENT_DOMAIN.with(|current| {
        let borrow_result = current.try_borrow_mut();
        if let std::result::Result::Ok(mut current) = borrow_result {
            *current = domain.map(std::borrow::ToOwned::to_owned);
        }
    });
}

fn span_parent_domain<S>(id: &tracing::span::Id, ctx: &tracing_subscriber::layer::Context<'_, S>) -> std::option::Option<std::string::String>
where
    S: tracing::Subscriber + for<'lookup> tracing_subscriber::registry::LookupSpan<'lookup>,
{
    let span = match ctx.span(id) {
        std::option::Option::Some(span) => span,
        std::option::Option::None => return std::option::Option::None,
    };
    let parent = match span.parent() {
        std::option::Option::Some(parent) => parent,
        std::option::Option::None => return std::option::Option::None,
    };
    let extensions = parent.extensions();
    return extensions.get::<SpanDomain>().and_then(|domain| -> std::option::Option<std::string::String> {
        return domain.value.clone();
    });
}

fn span_domain<S>(id: &tracing::span::Id, ctx: &tracing_subscriber::layer::Context<'_, S>) -> std::option::Option<std::string::String>
where
    S: tracing::Subscriber + for<'lookup> tracing_subscriber::registry::LookupSpan<'lookup>,
{
    let span = match ctx.span(id) {
        std::option::Option::Some(span) => span,
        std::option::Option::None => return std::option::Option::None,
    };
    let extensions = span.extensions();
    return extensions.get::<SpanDomain>().and_then(|domain| -> std::option::Option<std::string::String> {
        return domain.value.clone();
    });
}

fn event_parent_domain<S>(event: &tracing::Event<'_>, ctx: &tracing_subscriber::layer::Context<'_, S>) -> std::option::Option<std::string::String>
where
    S: tracing::Subscriber + for<'lookup> tracing_subscriber::registry::LookupSpan<'lookup>,
{
    let parent = match ctx.event_span(event) {
        std::option::Option::Some(parent) => parent,
        std::option::Option::None => return std::option::Option::None,
    };
    let extensions = parent.extensions();
    return extensions.get::<SpanDomain>().and_then(|domain| -> std::option::Option<std::string::String> {
        return domain.value.clone();
    });
}

#[cfg(test)]
#[path = "../unit_tests/domain.rs"]
mod tests;