attributes: permit setting parent span via #[instrument(parent = …)]

Fixes: tokio-rs#2021
Refs: tokio-rs#879

tracing-attributes/src/attr.rs

@@ -11,6 +11,7 @@ pub(crate) struct InstrumentArgs {
     level: Option<Level>,
     pub(crate) name: Option<LitStr>,
     target: Option<LitStr>,
+    pub(crate) parent: Option<Expr>,
     pub(crate) skips: HashSet<Ident>,
     pub(crate) fields: Option<Fields>,
     pub(crate) err_mode: Option<FormatMode>,
@@ -122,6 +123,12 @@ impl Parse for InstrumentArgs {
                 }
                 let target = input.parse::<StrArg<kw::target>>()?.value;
                 args.target = Some(target);
+            } else if lookahead.peek(kw::parent) {
+                if args.target.is_some() {
+                    return Err(input.error("expected only a single `parent` argument"));
+                }
+                let parent = input.parse::<ExprArg<kw::parent>>()?;
+                args.parent = Some(parent.value);
             } else if lookahead.peek(kw::level) {
                 if args.level.is_some() {
                     return Err(input.error("expected only a single `level` argument"));
@@ -180,6 +187,23 @@ impl<T: Parse> Parse for StrArg<T> {
     }
 }
 
+struct ExprArg<T> {
+    value: Expr,
+    _p: std::marker::PhantomData<T>,
+}
+
+impl<T: Parse> Parse for ExprArg<T> {
+    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
+        let _ = input.parse::<T>()?;
+        let _ = input.parse::<Token![=]>()?;
+        let value = input.parse()?;
+        Ok(Self {
+            value,
+            _p: std::marker::PhantomData,
+        })
+    }
+}
+
 struct Skips(HashSet<Ident>);
 
 impl Parse for Skips {
@@ -360,6 +384,7 @@ mod kw {
     syn::custom_keyword!(skip);
     syn::custom_keyword!(level);
     syn::custom_keyword!(target);
+    syn::custom_keyword!(parent);
     syn::custom_keyword!(name);
     syn::custom_keyword!(err);
     syn::custom_keyword!(ret);

tracing-attributes/src/expand.rs

@@ -135,6 +135,8 @@ fn gen_block<B: ToTokens>(
 
         let target = args.target();
 
+        let parent = args.parent.iter();
+
         // filter out skipped fields
         let quoted_fields: Vec<_> = param_names
             .iter()
@@ -182,6 +184,7 @@ fn gen_block<B: ToTokens>(
 
         quote!(tracing::span!(
             target: #target,
+            #(parent: #parent,)*
             #level,
             #span_name,
             #(#quoted_fields,)*

tracing-attributes/src/lib.rs

@@ -153,6 +153,14 @@ mod expand;
 ///     // ...
 /// }
 /// ```
+/// Overriding the generated span's parent:
+/// ```
+/// # use tracing_attributes::instrument;
+/// #[instrument(parent = None)]
+/// pub fn my_function() {
+///     // ...
+/// }
+/// ```
 ///
 /// To skip recording an argument, pass the argument's name to the `skip`:
 ///

tracing-attributes/tests/parents.rs

@@ -0,0 +1,99 @@
+use tracing::{collect::with_default, Level, Span};
+use tracing_attributes::instrument;
+use tracing_mock::*;
+
+#[instrument]
+fn with_default_parent() {}
+
+#[instrument(parent = parent_span)]
+fn with_explicit_parent(parent_span: &Span) {}
+
+#[test]
+fn default_parent_test() {
+    let contextual_parent = span::mock().named("contextual_parent");
+    let child = span::mock().named("with_default_parent");
+
+    let (collector, handle) = collector::mock()
+        .new_span(
+            contextual_parent
+                .clone()
+                .with_contextual_parent(None)
+                .with_explicit_parent(None),
+        )
+        .new_span(
+            child
+                .clone()
+                .with_contextual_parent(Some("contextual_parent"))
+                .with_explicit_parent(None),
+        )
+        .enter(child.clone())
+        .exit(child.clone())
+        .enter(contextual_parent.clone())
+        .new_span(
+            child
+                .clone()
+                .with_contextual_parent(Some("contextual_parent"))
+                .with_explicit_parent(None),
+        )
+        .enter(child.clone())
+        .exit(child.clone())
+        .exit(contextual_parent.clone())
+        .done()
+        .run_with_handle();
+
+    with_default(collector, || {
+        let contextual_parent = tracing::span!(Level::TRACE, "contextual_parent");
+
+        with_default_parent();
+
+        contextual_parent.in_scope(|| {
+            with_default_parent();
+        });
+    });
+
+    handle.assert_finished();
+}
+
+#[test]
+fn explicit_parent_test() {
+    let contextual_parent = span::mock().named("contextual_parent");
+    let explicit_parent = span::mock().named("explicit_parent");
+    let child = span::mock().named("with_explicit_parent");
+
+    let (collector, handle) = collector::mock()
+        .new_span(
+            contextual_parent
+                .clone()
+                .with_contextual_parent(None)
+                .with_explicit_parent(None),
+        )
+        .new_span(
+            explicit_parent
+                .clone()
+                .with_contextual_parent(None)
+                .with_explicit_parent(None),
+        )
+        .enter(contextual_parent.clone())
+        .new_span(
+            child
+                .clone()
+                .with_contextual_parent(Some("contextual_parent"))
+                .with_explicit_parent(Some("explicit_parent")),
+        )
+        .enter(child.clone())
+        .exit(child.clone())
+        .exit(contextual_parent.clone())
+        .done()
+        .run_with_handle();
+
+    with_default(collector, || {
+        let contextual_parent = tracing::span!(Level::INFO, "contextual_parent");
+        let explicit_parent = tracing::span!(Level::INFO, "explicit_parent");
+
+        contextual_parent.in_scope(|| {
+            with_explicit_parent(&explicit_parent);
+        });
+    });
+
+    handle.assert_finished();
+}