From 5cebb656842152581a19dadefed0c4c54a6c569a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dawid=20Ci=C4=99=C5=BCarkiewicz?= <dpc@dpc.pw>
Date: Sat, 25 May 2024 11:08:34 -0700
Subject: [PATCH] tracing-attributes: support const values for `target` and
 `name` (#2941)

Fixes #2960

Co-authored-by: Matthijs Brobbel <m1brobbel@gmail.com>
Co-authored-by: Hayden Stainsby <hds@caffeineconcepts.com>
Co-authored-by: Eliza Weisman <eliza@buoyant.io>
---
 tracing-attributes/src/attr.rs         | 32 ++++++++++--
 tracing-attributes/tests/instrument.rs | 72 ++++++++++++++++++++++++++
 2 files changed, 101 insertions(+), 3 deletions(-)

diff --git a/tracing-attributes/src/attr.rs b/tracing-attributes/src/attr.rs
index 8b4526164b..12525ec9e6 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>,
@@ -87,6 +87,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"));
                 }
@@ -211,8 +213,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 d01df0c313..402be8572a 100644
--- a/tracing-attributes/tests/instrument.rs
+++ b/tracing-attributes/tests/instrument.rs
@@ -252,3 +252,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 (subscriber, handle) = subscriber::mock()
+        .new_span(span_name.clone())
+        .enter(span_name.clone())
+        .exit(span_name.clone())
+        .drop_span(span_name)
+        .only()
+        .run_with_handle();
+
+    with_default(subscriber, || {
+        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 (subscriber, handle) = subscriber::mock()
+        .new_span(span_target.clone())
+        .enter(span_target.clone())
+        .exit(span_target.clone())
+        .drop_span(span_target)
+        .only()
+        .run_with_handle();
+
+    with_default(subscriber, || {
+        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 (subscriber, handle) = subscriber::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(subscriber, || {
+        name_target();
+    });
+
+    handle.assert_finished();
+}