buzz/crates/buzz-datastore-tracing/src/lib.rs

//! Policy-enforcing instrumentation for logical datastore operations.

use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::{format_ident, quote};
use syn::parse::{Parse, ParseStream};
use syn::{
    parenthesized, parse_macro_input, Error, Ident, ItemFn, LitStr, Result, ReturnType, Token, Type,
};

struct DatastoreArgs {
    name: LitStr,
    system: LitStr,
    fields: Option<TokenStream2>,
}

impl Parse for DatastoreArgs {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        let mut name = None;
        let mut system = None;
        let mut fields = None;

        while !input.is_empty() {
            let key: Ident = input.parse()?;
            match key.to_string().as_str() {
                "name" => {
                    if name.is_some() {
                        return Err(Error::new(key.span(), "duplicate `name` argument"));
                    }
                    input.parse::<Token![=]>()?;
                    name = Some(input.parse()?);
                }
                "system" => {
                    if system.is_some() {
                        return Err(Error::new(key.span(), "duplicate `system` argument"));
                    }
                    input.parse::<Token![=]>()?;
                    system = Some(input.parse()?);
                }
                "fields" => {
                    if fields.is_some() {
                        return Err(Error::new(key.span(), "duplicate `fields` argument"));
                    }
                    let content;
                    parenthesized!(content in input);
                    fields = Some(content.parse()?);
                }
                _ => {
                    return Err(Error::new(
                        key.span(),
                        "expected `name`, `system`, or `fields`",
                    ))
                }
            }
            if !input.is_empty() {
                input.parse::<Token![,]>()?;
            }
        }

        Ok(Self {
            name: name.ok_or_else(|| Error::new(input.span(), "missing `name`"))?,
            system: system.ok_or_else(|| Error::new(input.span(), "missing `system`"))?,
            fields,
        })
    }
}

/// Instruments an async logical datastore operation according to Buzz policy.
///
/// PostgreSQL spans always omit function arguments, use the `buzz_datastore`
/// target, and expose only canonical semantic fields plus explicitly supplied
/// safe fields. An `Err` sets `otel.status_code` without inspecting the error.
/// The literal `name` also labels a logical-operation duration histogram. Slow
/// completions are sampled and logged with only that name, outcome, and elapsed
/// time; arguments, error values, and return values are never formatted.
#[proc_macro_attribute]
pub fn datastore_span(args: TokenStream, item: TokenStream) -> TokenStream {
    let args = parse_macro_input!(args as DatastoreArgs);
    let mut function = parse_macro_input!(item as ItemFn);

    if function.sig.asyncness.is_none() {
        return Error::new_spanned(
            function.sig.fn_token,
            "`datastore_span` requires an async function",
        )
        .into_compile_error()
        .into();
    }
    if args.system.value() != "postgresql" {
        return Error::new_spanned(
            args.system,
            "unsupported datastore system; only `postgresql` is currently supported",
        )
        .into_compile_error()
        .into();
    }

    let name = args.name;
    let extra_fields = args.fields.map(|tokens| quote!(, #tokens));
    function.attrs.push(syn::parse_quote!(
        #[::tracing::instrument(
            target = "buzz_datastore",
            name = #name,
            skip_all,
            fields(
                otel.kind = "client",
                db.system.name = "postgresql",
                otel.status_code = ::tracing::field::Empty
                #extra_fields
            )
        )]
    ));

    let return_type: Type = match &function.sig.output {
        ReturnType::Type(_, ty) => *ty.clone(),
        ReturnType::Default => syn::parse_quote!(()),
    };
    let returns_result = match &return_type {
        Type::Path(path) => path
            .path
            .segments
            .last()
            .is_some_and(|segment| segment.ident == "Result"),
        _ => false,
    };
    let original_body = function.block;
    let result = format_ident!("__buzz_datastore_result_7f3a9c");
    let record_error = returns_result.then(|| {
        quote! {
            if #result.is_err() {
                ::tracing::Span::current().record("otel.status_code", "ERROR");
            }
        }
    });
    let outcome = if returns_result {
        quote! {
            if #result.is_err() { "error" } else { "success" }
        }
    } else {
        quote!("success")
    };
    function.block = Box::new(syn::parse_quote!({
        let __buzz_datastore_started_7f3a9c = ::std::time::Instant::now();
        let #result: #return_type = (async #original_body).await;
        #record_error
        let __buzz_datastore_outcome_7f3a9c = #outcome;
        let __buzz_datastore_elapsed_7f3a9c = __buzz_datastore_started_7f3a9c.elapsed();
        ::metrics::histogram!(
            "buzz_db_operation_duration_seconds",
            "operation" => #name,
            "outcome" => __buzz_datastore_outcome_7f3a9c,
        )
        .record(__buzz_datastore_elapsed_7f3a9c.as_secs_f64());
        if __buzz_datastore_elapsed_7f3a9c >= ::std::time::Duration::from_millis(500) {
            static __BUZZ_DATASTORE_SLOW_SAMPLE_7F3A9C:
                ::std::sync::atomic::AtomicU64 = ::std::sync::atomic::AtomicU64::new(0);
            if __BUZZ_DATASTORE_SLOW_SAMPLE_7F3A9C.fetch_add(
                1,
                ::std::sync::atomic::Ordering::Relaxed,
            ) % 100 == 0 {
                let __buzz_datastore_elapsed_ms_7f3a9c =
                    __buzz_datastore_elapsed_7f3a9c
                        .as_millis()
                        .min(::std::primitive::u64::MAX as u128) as u64;
                ::tracing::warn!(
                    target: "buzz_datastore",
                    parent: None,
                    operation = #name,
                    outcome = __buzz_datastore_outcome_7f3a9c,
                    elapsed_ms = __buzz_datastore_elapsed_ms_7f3a9c,
                    "slow datastore operation"
                );
            }
        }
        #result
    }));

    quote!(#function).into()
}

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

    #[test]
    fn parses_safe_tracing_fields() {
        let args: DatastoreArgs = syn::parse_str(
            r#"name = "audit", system = "postgresql", fields(action = %entry.action, from_seq = from_seq, limit = limit)"#,
        )
        .expect("valid arguments");
        assert!(args.fields.is_some());
    }

    #[test]
    fn rejects_duplicate_arguments() {
        for args in [
            r#"name = "one", name = "two", system = "postgresql""#,
            r#"name = "one", system = "postgresql", system = "postgresql""#,
            r#"name = "one", system = "postgresql", fields(a = 1), fields(b = 2)"#,
        ] {
            let error = match syn::parse_str::<DatastoreArgs>(args) {
                Ok(_) => panic!("duplicate accepted"),
                Err(error) => error,
            };
            assert!(error.to_string().starts_with("duplicate `"));
        }
    }
}