attributes: permit #[instrument(follows_from = …)]

The optional `follows_from` argument allows users to specify any
number of `Span::follows_from` relationships. The value to this
argument is an expression of type:
  `impl IntoIterator<Item = impl Into<Option<Id>>>`

Fixes: tokio-rs#879

tracing-attributes/src/attr.rs

@@ -12,6 +12,7 @@ pub(crate) struct InstrumentArgs {
     pub(crate) name: Option<LitStr>,
     target: Option<LitStr>,
     pub(crate) parent: Option<Expr>,
+    pub(crate) follows_from: Option<Expr>,
     pub(crate) skips: HashSet<Ident>,
     pub(crate) fields: Option<Fields>,
     pub(crate) err_mode: Option<FormatMode>,
@@ -129,6 +130,12 @@ impl Parse for InstrumentArgs {
                 }
                 let parent = input.parse::<ExprArg<kw::parent>>()?;
                 args.parent = Some(parent.value);
+            } else if lookahead.peek(kw::follows_from) {
+                if args.target.is_some() {
+                    return Err(input.error("expected only a single `follows_from` argument"));
+                }
+                let follows_from = input.parse::<ExprArg<kw::follows_from>>()?;
+                args.follows_from = Some(follows_from.value);
             } else if lookahead.peek(kw::level) {
                 if args.level.is_some() {
                     return Err(input.error("expected only a single `level` argument"));
@@ -385,6 +392,7 @@ mod kw {
     syn::custom_keyword!(level);
     syn::custom_keyword!(target);
     syn::custom_keyword!(parent);
+    syn::custom_keyword!(follows_from);
     syn::custom_keyword!(name);
     syn::custom_keyword!(err);
     syn::custom_keyword!(ret);

tracing-attributes/src/expand.rs

@@ -88,6 +88,13 @@ fn gen_block<B: ToTokens>(
 
     let level = args.level();
 
+    let follows_from = args.follows_from.iter();
+    let follows_from = quote! {
+        #(for cause in #follows_from {
+            __tracing_attr_span.follows_from(cause);
+        })*
+    };
+
     // generate this inside a closure, so we can return early on errors.
     let span = (|| {
         // Pull out the arguments-to-be-skipped first, so we can filter results
@@ -261,6 +268,7 @@ fn gen_block<B: ToTokens>(
             let __tracing_attr_span = #span;
             let __tracing_instrument_future = #mk_fut;
             if !__tracing_attr_span.is_disabled() {
+                #follows_from
                 tracing::Instrument::instrument(
                     __tracing_instrument_future,
                     __tracing_attr_span
@@ -287,6 +295,7 @@ fn gen_block<B: ToTokens>(
         let __tracing_attr_guard;
         if tracing::level_enabled!(#level) {
             __tracing_attr_span = #span;
+            #follows_from
             __tracing_attr_guard = __tracing_attr_span.enter();
         }
     );

tracing-attributes/src/lib.rs

@@ -176,6 +176,14 @@ mod expand;
 ///     fn my_method(&self) {}
 /// }
 /// ```
+/// Specifying `follows_from` relationships:
+/// ```
+/// # use tracing_attributes::instrument;
+/// #[instrument(follows_from = causes)]
+/// pub fn my_function(causes: &[tracing::Id]) {
+///     // ...
+/// }
+/// ```
 ///
 /// To skip recording an argument, pass the argument's name to the `skip`:
 ///

tracing-attributes/tests/follows_from.rs

@@ -0,0 +1,73 @@
+use tracing::{collect::with_default, Id, Level};
+use tracing_attributes::instrument;
+use tracing_mock::*;
+
+#[instrument(follows_from = causes, skip(causes))]
+fn with_follows_from_sync(causes: impl IntoIterator<Item = impl Into<Option<Id>>>) {}
+
+#[instrument(follows_from = causes, skip(causes))]
+async fn with_follows_from_async(causes: impl IntoIterator<Item = impl Into<Option<Id>>>) {}
+
+#[test]
+fn follows_from_sync_test() {
+    let cause_a = span::mock().named("cause_a");
+    let cause_b = span::mock().named("cause_b");
+    let cause_c = span::mock().named("cause_c");
+    let consequence = span::mock().named("with_follows_from_sync");
+
+    let (collector, handle) = collector::mock()
+        .new_span(cause_a.clone())
+        .new_span(cause_b.clone())
+        .new_span(cause_c.clone())
+        .new_span(consequence.clone())
+        .follows_from(consequence.clone(), cause_a)
+        .follows_from(consequence.clone(), cause_b)
+        .follows_from(consequence.clone(), cause_c)
+        .enter(consequence.clone())
+        .exit(consequence)
+        .done()
+        .run_with_handle();
+
+    with_default(collector, || {
+        let cause_a = tracing::span!(Level::TRACE, "cause_a");
+        let cause_b = tracing::span!(Level::TRACE, "cause_b");
+        let cause_c = tracing::span!(Level::TRACE, "cause_c");
+
+        with_follows_from_sync(&[cause_a, cause_b, cause_c])
+    });
+
+    handle.assert_finished();
+}
+
+#[test]
+fn follows_from_async_test() {
+    let cause_a = span::mock().named("cause_a");
+    let cause_b = span::mock().named("cause_b");
+    let cause_c = span::mock().named("cause_c");
+    let consequence = span::mock().named("with_follows_from_async");
+
+    let (collector, handle) = collector::mock()
+        .new_span(cause_a.clone())
+        .new_span(cause_b.clone())
+        .new_span(cause_c.clone())
+        .new_span(consequence.clone())
+        .follows_from(consequence.clone(), cause_a)
+        .follows_from(consequence.clone(), cause_b)
+        .follows_from(consequence.clone(), cause_c)
+        .enter(consequence.clone())
+        .exit(consequence)
+        .done()
+        .run_with_handle();
+
+    with_default(collector, || {
+        block_on_future(async {
+            let cause_a = tracing::span!(Level::TRACE, "cause_a");
+            let cause_b = tracing::span!(Level::TRACE, "cause_b");
+            let cause_c = tracing::span!(Level::TRACE, "cause_c");
+
+            with_follows_from_async(&[cause_a, cause_b, cause_c]).await
+        })
+    });
+
+    handle.assert_finished();
+}