Merge pull request #29 from palantir/macros

Add a macro to span functions and methods

Cargo.toml

@@ -2,5 +2,6 @@
 members = [
     "http-zipkin",
     "zipkin",
+    "zipkin-macros",
     "zipkin-types",
 ]

zipkin-macros/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "zipkin-macros"
+version = "0.1.0"
+authors = ["Steven Fackler <sfackler@palantir.com>"]
+edition = "2018"
+license = "Apache-2.0"
+description = "Macros for use with `zipkin`"
+repository = "https://github.com/palantir/rust-zipkin"
+readme = "../README.md"
+
+[lib]
+proc-macro = true
+
+[dependencies]
+quote = "1.0"
+proc-macro2 = "1.0"
+syn = { version = "1.0", features = ["full"] }

zipkin-macros/src/lib.rs

@@ -0,0 +1,123 @@
+//  Copyright 2020 Palantir Technologies, Inc.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//! Macros for use with `zipkin`.
+//!
+//! You should not depend on this crate directly.
+extern crate proc_macro;
+
+use proc_macro::TokenStream;
+use proc_macro2::Span;
+use quote::{quote, ToTokens};
+use syn::{
+    parse_macro_input, AttributeArgs, Error, ImplItemMethod, Lit, LitStr, Meta, NestedMeta, Stmt,
+};
+
+struct Options {
+    name: LitStr,
+}
+
+/// Wraps the execution of a function or method in a span.
+///
+/// Both normal and `async` methods and functions are supported. The name of the span is specified as an argument
+/// to the macro attribute.
+///
+/// Requires the `macros` Cargo feature.
+///
+/// # Examples
+///
+/// ```ignore
+/// #[zipkin::spanned(name = "shave yaks")]
+/// fn shave_some_yaks(yaks: &mut [Yak]) {
+///     // ...
+/// }
+///
+/// #[zipkin::spanned(name = "asynchronously shave yaks")]
+/// async fn shave_some_other_yaks(yaks: &mut [Yak]) {
+///     // ...
+/// }
+///
+/// struct Yak;
+///
+/// impl Yak {
+///     #[zipkin::spanned(name = "shave a yak")]
+///     fn shave(&mut self) {
+///         // ...
+///     }
+///
+///     #[zipkin::spanned(name = "asynchronously shave a yak")]
+///     async fn shave_nonblocking(&mut self) {
+///          // ...
+///     }
+/// }
+/// ```
+#[proc_macro_attribute]
+pub fn spanned(args: TokenStream, item: TokenStream) -> TokenStream {
+    let args = parse_macro_input!(args as AttributeArgs);
+    let func = parse_macro_input!(item as ImplItemMethod);
+
+    spanned_impl(args, func).unwrap_or_else(|e| e.to_compile_error().into())
+}
+
+fn spanned_impl(args: AttributeArgs, mut func: ImplItemMethod) -> Result<TokenStream, Error> {
+    let options = parse_options(args)?;
+    let name = &options.name;
+
+    if func.sig.asyncness.is_some() {
+        let stmts = &func.block.stmts;
+        let stmt = quote! {
+            zipkin::next_span()
+                .with_name(#name)
+                .detach()
+                .bind(async move { #(#stmts)* })
+                .await
+        };
+        func.block.stmts = vec![Stmt::Expr(syn::parse2(stmt).unwrap())];
+    } else {
+        let stmt = quote! {
+            let __macro_impl_span = zipkin::next_span().with_name(#name);
+        };
+        func.block.stmts.insert(0, syn::parse2(stmt).unwrap());
+    };
+
+    Ok(func.into_token_stream().into())
+}
+
+fn parse_options(args: AttributeArgs) -> Result<Options, Error> {
+    let mut name = None;
+
+    for arg in args {
+        let meta = match arg {
+            NestedMeta::Meta(meta) => meta,
+            NestedMeta::Lit(lit) => {
+                return Err(Error::new_spanned(&lit, "invalid attribute syntax"))
+            }
+        };
+
+        if meta.path().is_ident("name") {
+            match meta {
+                Meta::NameValue(meta) => match meta.lit {
+                    Lit::Str(lit) => name = Some(lit),
+                    lit => return Err(Error::new_spanned(&lit, "expected a string literal")),
+                },
+                _ => return Err(Error::new_spanned(meta, "expected `name = \"...\"`")),
+            }
+        } else {
+            return Err(Error::new_spanned(meta.path(), "unknown option"));
+        }
+    }
+
+    Ok(Options {
+        name: name.ok_or_else(|| Error::new(Span::call_site(), "missing `name` option"))?,
+    })
+}

zipkin/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "zipkin"
-version = "0.4.0"
+version = "0.4.1"
 authors = ["Steven Fackler <sfackler@palantir.com>"]
 edition = "2018"
 license = "Apache-2.0"
@@ -10,15 +10,20 @@ readme = "../README.md"
 categories = ["network-programming", "web-programming"]
 keywords = ["zipkin", "tracing"]
 
+[package.metadata.docs.rs]
+all-features = true
+
 [features]
 serde = ["zipkin-types/serde"]
+macros = ["zipkin-macros"]
 
 [dependencies]
 log = "0.4"
 lazycell = "1.0"
 pin-project-lite = "0.1"
 rand = "0.7"
 
+zipkin-macros = { version = "0.1.0", optional = true, path = "../zipkin-macros" }
 zipkin-types = { version = "0.1.0", path = "../zipkin-types" }
 
 [dev-dependencies]

zipkin/src/lib.rs

@@ -27,6 +27,8 @@
 #![doc(html_root_url = "https://docs.rs/zipkin/0.4")]
 #![warn(missing_docs)]
 
+#[cfg(feature = "macros")]
+pub use zipkin_macros::*;
 #[doc(inline)]
 pub use zipkin_types::{
     annotation, endpoint, span, span_id, trace_id, Annotation, Endpoint, Kind, Span, SpanId,

zipkin/src/test/macros.rs

@@ -0,0 +1,185 @@
+//  Copyright 2020 Palantir Technologies, Inc.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+use crate as zipkin; // hack to get the macro codegen to work in the same crate
+use crate::{spanned, test};
+use futures::executor;
+
+#[test]
+fn blocking_free_function() {
+    #[spanned(name = "foobar")]
+    fn foo() {
+        zipkin::next_span().with_name("fizzbuzz");
+    }
+
+    test::init();
+
+    let span = zipkin::next_span().with_name("root");
+    foo();
+    drop(span);
+
+    let spans = test::take();
+    assert_eq!(spans.len(), 3);
+    assert_eq!(spans[0].name(), Some("fizzbuzz"));
+    assert_eq!(spans[1].name(), Some("foobar"));
+    assert_eq!(spans[2].name(), Some("root"));
+    assert_eq!(spans[0].trace_id(), spans[2].trace_id());
+    assert_eq!(spans[1].trace_id(), spans[2].trace_id());
+    assert_eq!(spans[0].parent_id(), Some(spans[1].id()));
+    assert_eq!(spans[1].parent_id(), Some(spans[2].id()));
+    assert_eq!(spans[2].parent_id(), None);
+}
+
+#[test]
+fn blocking_associated_function() {
+    struct Foo;
+
+    impl Foo {
+        #[spanned(name = "foobar")]
+        fn foo() {
+            zipkin::next_span().with_name("fizzbuzz");
+        }
+    }
+
+    test::init();
+
+    let span = zipkin::next_span().with_name("root");
+    Foo::foo();
+    drop(span);
+
+    let spans = test::take();
+    assert_eq!(spans.len(), 3);
+    assert_eq!(spans[0].name(), Some("fizzbuzz"));
+    assert_eq!(spans[1].name(), Some("foobar"));
+    assert_eq!(spans[2].name(), Some("root"));
+    assert_eq!(spans[0].trace_id(), spans[2].trace_id());
+    assert_eq!(spans[1].trace_id(), spans[2].trace_id());
+    assert_eq!(spans[0].parent_id(), Some(spans[1].id()));
+    assert_eq!(spans[1].parent_id(), Some(spans[2].id()));
+    assert_eq!(spans[2].parent_id(), None);
+}
+
+#[test]
+fn blocking_method() {
+    struct Foo;
+
+    impl Foo {
+        #[spanned(name = "foobar")]
+        fn foo(&self) {
+            zipkin::next_span().with_name("fizzbuzz");
+        }
+    }
+
+    test::init();
+
+    let span = zipkin::next_span().with_name("root");
+    Foo.foo();
+    drop(span);
+
+    let spans = test::take();
+    assert_eq!(spans.len(), 3);
+    assert_eq!(spans[0].name(), Some("fizzbuzz"));
+    assert_eq!(spans[1].name(), Some("foobar"));
+    assert_eq!(spans[2].name(), Some("root"));
+    assert_eq!(spans[0].trace_id(), spans[2].trace_id());
+    assert_eq!(spans[1].trace_id(), spans[2].trace_id());
+    assert_eq!(spans[0].parent_id(), Some(spans[1].id()));
+    assert_eq!(spans[1].parent_id(), Some(spans[2].id()));
+    assert_eq!(spans[2].parent_id(), None);
+}
+
+#[test]
+fn async_free_function() {
+    #[spanned(name = "foobar")]
+    async fn foo() {
+        zipkin::next_span().with_name("fizzbuzz");
+    }
+
+    test::init();
+
+    let future = zipkin::next_span().with_name("root").detach().bind(foo());
+    executor::block_on(future);
+
+    let spans = test::take();
+    assert_eq!(spans.len(), 3);
+    assert_eq!(spans[0].name(), Some("fizzbuzz"));
+    assert_eq!(spans[1].name(), Some("foobar"));
+    assert_eq!(spans[2].name(), Some("root"));
+    assert_eq!(spans[0].trace_id(), spans[2].trace_id());
+    assert_eq!(spans[1].trace_id(), spans[2].trace_id());
+    assert_eq!(spans[0].parent_id(), Some(spans[1].id()));
+    assert_eq!(spans[1].parent_id(), Some(spans[2].id()));
+    assert_eq!(spans[2].parent_id(), None);
+}
+
+#[test]
+fn async_associated_function() {
+    struct Foo;
+
+    impl Foo {
+        #[spanned(name = "foobar")]
+        async fn foo() {
+            zipkin::next_span().with_name("fizzbuzz");
+        }
+    }
+
+    test::init();
+
+    let future = zipkin::next_span()
+        .with_name("root")
+        .detach()
+        .bind(Foo::foo());
+    executor::block_on(future);
+
+    let spans = test::take();
+    assert_eq!(spans.len(), 3);
+    assert_eq!(spans[0].name(), Some("fizzbuzz"));
+    assert_eq!(spans[1].name(), Some("foobar"));
+    assert_eq!(spans[2].name(), Some("root"));
+    assert_eq!(spans[0].trace_id(), spans[2].trace_id());
+    assert_eq!(spans[1].trace_id(), spans[2].trace_id());
+    assert_eq!(spans[0].parent_id(), Some(spans[1].id()));
+    assert_eq!(spans[1].parent_id(), Some(spans[2].id()));
+    assert_eq!(spans[2].parent_id(), None);
+}
+
+#[test]
+fn async_method() {
+    struct Foo;
+
+    impl Foo {
+        #[spanned(name = "foobar")]
+        async fn foo(&self) {
+            zipkin::next_span().with_name("fizzbuzz");
+        }
+    }
+
+    test::init();
+
+    let future = zipkin::next_span()
+        .with_name("root")
+        .detach()
+        .bind(Foo.foo());
+    executor::block_on(future);
+
+    let spans = test::take();
+    assert_eq!(spans.len(), 3);
+    assert_eq!(spans[0].name(), Some("fizzbuzz"));
+    assert_eq!(spans[1].name(), Some("foobar"));
+    assert_eq!(spans[2].name(), Some("root"));
+    assert_eq!(spans[0].trace_id(), spans[2].trace_id());
+    assert_eq!(spans[1].trace_id(), spans[2].trace_id());
+    assert_eq!(spans[0].parent_id(), Some(spans[1].id()));
+    assert_eq!(spans[1].parent_id(), Some(spans[2].id()));
+    assert_eq!(spans[2].parent_id(), None);
+}