usdt_macro/
lib.rs

1//! Prototype proc-macro crate for parsing a DTrace provider definition into Rust code.
2
3// Copyright 2021 Oxide Computer Company
4//
5// Licensed under the Apache License, Version 2.0 (the "License");
6// you may not use this file except in compliance with the License.
7// You may obtain a copy of the License at
8//
9//     http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing, software
12// distributed under the License is distributed on an "AS IS" BASIS,
13// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14// See the License for the specific language governing permissions and
15// limitations under the License.
16
17use std::iter::FromIterator;
18use std::{fs, path::Path};
19
20use quote::quote;
21use syn::{parse_macro_input, Lit};
22
23use usdt_impl::compile_provider_source;
24
25/// Generate DTrace probe macros from a provider definition file.
26///
27/// This macro parses a DTrace provider.d file, given as a single literal string path. It then
28/// generates a Rust macro for each of the DTrace probe definitions.
29///
30/// For example, let's say the file `"test.d"` has the following contents:
31///
32/// ```ignore
33/// provider test {
34///     probe start(uint8_t);
35///     probe stop(char *, uint8_t);
36/// };
37/// ```
38///
39/// In a Rust library or application, write:
40///
41/// ```ignore
42/// dtrace_provider!("test.d");
43/// ```
44///
45/// The macro looks for the file relative to the root of the package, so `"test.d"`
46/// in this case would be in the same directory as `"Cargo.toml"`.
47///
48/// By default probe macros are named `{provider}_{probe}!`. Arguments are passed
49/// via a closure that returns a tuple. Note that the provided closure is only
50/// evaluated when the probe is enabled. One can then add points of instrumentation
51/// by invoking the macro:
52///
53/// ```ignore
54/// fn do_stuff(count: u8, name: String) {
55///     // doing stuff
56///     test_stop!(|| (name, count));
57/// }
58/// ```
59///
60/// The probe macro names can be customized by adding `, format =
61/// my_prefix_{provider}_{probe}` to the macro invocation where `{provider}` and
62/// `{probe}` are optional and will be substituted with the actual provider and
63/// probe names:
64///
65/// ```ignore
66/// dtrace_provider!("test.d", format = "dtrace_{provider}_{probe}");
67/// ```
68///
69/// Note
70/// ----
71/// The only supported types are integers of specific bit-width (e.g., `uint16_t`),
72/// pointers to integers, and `char *`.
73#[proc_macro]
74pub fn dtrace_provider(item: proc_macro::TokenStream) -> proc_macro::TokenStream {
75    let mut tokens = item.into_iter().collect::<Vec<proc_macro::TokenTree>>();
76
77    let comma_index = tokens
78        .iter()
79        .enumerate()
80        .find_map(|(i, token)| match token {
81            proc_macro::TokenTree::Punct(p) if p.as_char() == ',' => Some(i),
82            _ => None,
83        });
84
85    // Split off the tokens after the comma if there is one.
86    let rest = if let Some(index) = comma_index {
87        let mut rest = tokens.split_off(index);
88        let _ = rest.remove(0);
89        rest
90    } else {
91        Vec::new()
92    };
93
94    // Parse the config from the remaining tokens.
95    let config: usdt_impl::CompileProvidersConfig = serde_tokenstream::from_tokenstream(
96        &proc_macro2::TokenStream::from(proc_macro::TokenStream::from_iter(rest)),
97    )
98    .unwrap();
99
100    let first_item = proc_macro::TokenStream::from_iter(tokens);
101    let tok = parse_macro_input!(first_item as Lit);
102    let filename = match tok {
103        Lit::Str(f) => f.value(),
104        _ => panic!("DTrace provider must be a single literal string filename"),
105    };
106    let source = if filename.ends_with(".d") {
107        let dir = std::env::var("CARGO_MANIFEST_DIR").map_or_else(
108            |_| std::env::current_dir().unwrap(),
109            |s| Path::new(&s).to_path_buf(),
110        );
111
112        let path = dir.join(&filename);
113        fs::read_to_string(path).unwrap_or_else(|_| {
114            panic!(
115                "Could not read D source file \"{}\" in {:?}",
116                &filename, dir,
117            )
118        })
119    } else {
120        filename.clone()
121    };
122    match compile_provider_source(&source, &config) {
123        Ok(provider) => provider.into(),
124        Err(e) => {
125            let message = format!(
126                "Error building provider definition in \"{}\"\n\n{}",
127                filename, e
128            );
129            let out = quote! {
130                compile_error!(#message);
131            };
132            out.into()
133        }
134    }
135}
usdt_macro/lib.rs

usdt_macro/
lib.rs
