From 1286fdf54908e4fb0cb2266735c0ecf5da43e8aa Mon Sep 17 00:00:00 2001
From: Matthijs Brobbel <m1brobbel@gmail.com>
Date: Tue, 26 Mar 2024 13:11:54 +0100
Subject: [PATCH] tracing-attributes: support const values for `target` and
 `name`

Fixes: #2460
---
 tracing-attributes/src/attr.rs             | 32 +++++++++-
 tracing-attributes/tests/instrument.rs     | 72 ++++++++++++++++++++++
 tracing-subscriber/src/registry/sharded.rs |  1 +
 3 files changed, 102 insertions(+), 3 deletions(-)

diff --git a/tracing-attributes/src/attr.rs b/tracing-attributes/src/attr.rs
index 902ca1450b..4684326d62 100644
--- a/tracing-attributes/src/attr.rs
+++ b/tracing-attributes/src/attr.rs
@@ -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>,
@@ -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"));
                 }
@@ -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>,
 }
 
diff --git a/tracing-attributes/tests/instrument.rs b/tracing-attributes/tests/instrument.rs
index 0376be4069..7bc8e962f9 100644
--- a/tracing-attributes/tests/instrument.rs
+++ b/tracing-attributes/tests/instrument.rs
@@ -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();
+}
diff --git a/tracing-subscriber/src/registry/sharded.rs b/tracing-subscriber/src/registry/sharded.rs
index a6c06391fa..c915109fe5 100644
--- a/tracing-subscriber/src/registry/sharded.rs
+++ b/tracing-subscriber/src/registry/sharded.rs
@@ -592,6 +592,7 @@ mod tests {
         closed: Vec<(&'static str, Weak<()>)>,
     }
 
+    #[allow(dead_code)]
     struct SetRemoved(Arc<()>);
 
     impl<C> Subscribe<C> for CloseSubscriber