Skip to content

Commit

Permalink
tracing-attributes: support const values for target and name
Browse files Browse the repository at this point in the history
  • Loading branch information
mbrobbel authored and dpc committed Apr 23, 2024
1 parent 908cc43 commit 2ddad6f
Show file tree
Hide file tree
Showing 2 changed files with 101 additions and 3 deletions.
32 changes: 29 additions & 3 deletions tracing-attributes/src/attr.rs
Expand Up @@ -17,8 +17,8 @@ pub(crate) struct EventArgs {
#[derive(Clone, Default, Debug)]
pub(crate) struct InstrumentArgs {
level: Option<Level>,
pub(crate) name: Option<LitStr>,
target: Option<LitStr>,
pub(crate) name: Option<LitStrOrIdent>,
target: Option<LitStrOrIdent>,
pub(crate) parent: Option<Expr>,
pub(crate) follows_from: Option<Expr>,
pub(crate) skips: HashSet<Ident>,
Expand Down Expand Up @@ -86,6 +86,8 @@ impl Parse for InstrumentArgs {
// XXX: apparently we support names as either named args with an
// sign, _or_ as unnamed string literals. That's weird, but
// changing it is apparently breaking.
// This also means that when using idents for name, it must be via
// a named arg, i.e. `#[instrument(name = SOME_IDENT)]`.
if args.name.is_some() {
return Err(input.error("expected only a single `name` argument"));
}
Expand Down Expand Up @@ -198,8 +200,32 @@ impl Parse for EventArgs {
}
}

#[derive(Debug, Clone)]
pub(super) enum LitStrOrIdent {
LitStr(LitStr),
Ident(Ident),
}

impl ToTokens for LitStrOrIdent {
fn to_tokens(&self, tokens: &mut TokenStream) {
match self {
LitStrOrIdent::LitStr(target) => target.to_tokens(tokens),
LitStrOrIdent::Ident(ident) => ident.to_tokens(tokens),
}
}
}

impl Parse for LitStrOrIdent {
fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
input
.parse::<LitStr>()
.map(LitStrOrIdent::LitStr)
.or_else(|_| input.parse::<Ident>().map(LitStrOrIdent::Ident))
}
}

struct StrArg<T> {
value: LitStr,
value: LitStrOrIdent,
_p: std::marker::PhantomData<T>,
}

Expand Down
72 changes: 72 additions & 0 deletions tracing-attributes/tests/instrument.rs
Expand Up @@ -238,3 +238,75 @@ fn impl_trait_return_type() {

handle.assert_finished();
}

#[test]
fn name_ident() {
const MY_NAME: &str = "my_name";
#[instrument(name = MY_NAME)]
fn name() {}

let span_name = expect::span().named(MY_NAME);

let (collector, handle) = collector::mock()
.new_span(span_name.clone())
.enter(span_name.clone())
.exit(span_name.clone())
.drop_span(span_name)
.only()
.run_with_handle();

with_default(collector, || {
name();
});

handle.assert_finished();
}

#[test]
fn target_ident() {
const MY_TARGET: &str = "my_target";

#[instrument(target = MY_TARGET)]
fn target() {}

let span_target = expect::span().named("target").with_target(MY_TARGET);

let (collector, handle) = collector::mock()
.new_span(span_target.clone())
.enter(span_target.clone())
.exit(span_target.clone())
.drop_span(span_target)
.only()
.run_with_handle();

with_default(collector, || {
target();
});

handle.assert_finished();
}

#[test]
fn target_name_ident() {
const MY_NAME: &str = "my_name";
const MY_TARGET: &str = "my_target";

#[instrument(target = MY_TARGET, name = MY_NAME)]
fn name_target() {}

let span_name_target = expect::span().named(MY_NAME).with_target(MY_TARGET);

let (collector, handle) = collector::mock()
.new_span(span_name_target.clone())
.enter(span_name_target.clone())
.exit(span_name_target.clone())
.drop_span(span_name_target)
.only()
.run_with_handle();

with_default(collector, || {
name_target();
});

handle.assert_finished();
}

0 comments on commit 2ddad6f

Please sign in to comment.