expose an eager macro that skips the branch. make the lazy macro branch cold

Author: conradludgate

probe-test-attr/src/main.rs

@@ -96,6 +96,7 @@ fn main() {
     let buffer = [2; 4];
     loop {
         test::start_work!(|| arg.x);
+        test::start_work_eager!(|| arg.x);
         std::thread::sleep(std::time::Duration::from_secs(1));
         arg.x = arg.x.wrapping_add(1);
         test::stop_work!(|| { (format!("the probe has fired {}", arg.x), &arg) });

usdt-impl/src/common.rs

@@ -205,30 +205,47 @@ pub(crate) fn build_probe_macro(
     config: &crate::CompileProvidersConfig,
     probe_name: &str,
     types: &[DataType],
-    impl_block: TokenStream,
+    eager: TokenStream,
+    lazy: TokenStream,
 ) -> TokenStream {
     let module = config.module_ident();
-    let macro_name = config.probe_ident(probe_name);
+    let lazy_macro_name = config.probe_ident(probe_name);
+    let eager_macro_name = config.probe_ident_eager(probe_name);
+
     let no_args_match = if types.is_empty() {
-        quote! { () => { crate::#module::#macro_name!(|| ()) }; }
+        quote! { () => { crate::#module::#eager_macro_name!(|| ()) }; }
     } else {
         quote! {}
     };
     quote! {
         #[allow(unused_macros)]
-        macro_rules! #macro_name {
+        macro_rules! #lazy_macro_name {
+            #no_args_match
+            ($tree:tt) => {
+                compile_error!("USDT probe macros should be invoked with a closure returning the arguments");
+            };
+            ($args_lambda:expr) => {
+                {
+                    #lazy
+                }
+            };
+        }
+        #[allow(unused_imports)]
+        pub(crate) use #lazy_macro_name;
+        #[allow(unused_macros)]
+        macro_rules! #eager_macro_name {
             #no_args_match
             ($tree:tt) => {
                 compile_error!("USDT probe macros should be invoked with a closure returning the arguments");
             };
             ($args_lambda:expr) => {
                 {
-                    #impl_block
+                    #eager
                 }
             };
         }
         #[allow(unused_imports)]
-        pub(crate) use #macro_name;
+        pub(crate) use #eager_macro_name;
     }
 }
 

usdt-impl/src/lib.rs

@@ -100,6 +100,11 @@ impl CompileProvidersConfig {
         quote::format_ident!("{}", self.format_probe(probe_name))
     }
 
+    /// Return the formatted name of the probe as an identifier.
+    pub fn probe_ident_eager(&self, probe_name: &str) -> proc_macro2::Ident {
+        quote::format_ident!("{}_eager", self.format_probe(probe_name))
+    }
+
     /// Return the formatted module name as an identifier.
     pub fn module_ident(&self) -> proc_macro2::Ident {
         let name = self.module.as_ref().unwrap_or_else(|| {

usdt-impl/src/linker.rs

@@ -180,7 +180,7 @@ fn compile_probe(
     #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
     compile_error!("USDT only supports x86_64 and AArch64 architectures");
 
-    let impl_block = quote! {
+    let extern_block = quote! {
         unsafe extern "C" {
             #[allow(unused)]
             #[link_name = #stability]
@@ -198,26 +198,43 @@ fn compile_probe(
             #[link_name = #probe]
             fn #extern_probe_fn(#(#ffi_param_list,)*);
         }
+    };
+
+    let call = quote! {
         unsafe {
-            if #is_enabled_fn() != 0 {
-                #arg_lambda
-                #type_check_fn
-                #unpacked_args
-                ::std::arch::asm!(
-                    ".reference {typedefs}",
-                    #call_instruction,
-                    ".reference {stability}",
-                    typedefs = sym #typedef_fn,
-                    extern_probe_fn = sym #extern_probe_fn,
-                    stability = sym #stability_fn,
-                    #in_regs
-                    options(nomem, nostack, preserves_flags)
-                );
-            }
+            ::std::arch::asm!(
+                ".reference {typedefs}",
+                #call_instruction,
+                ".reference {stability}",
+                typedefs = sym #typedef_fn,
+                extern_probe_fn = sym #extern_probe_fn,
+                stability = sym #stability_fn,
+                #in_regs
+                options(nomem, nostack, preserves_flags)
+            );
+        }
+    };
+
+    let eager = quote! {
+        #extern_block
+        #arg_lambda
+        #type_check_fn
+        #unpacked_args
+        #call
+    };
+
+    let lazy = quote! {
+        #extern_block
+        if unsafe { #is_enabled_fn() } != 0 {
+            ::usdt::cold();
+            #arg_lambda
+            #type_check_fn
+            #unpacked_args
+            #call
         }
     };
 
-    common::build_probe_macro(config, probe_name, types, impl_block)
+    common::build_probe_macro(config, probe_name, types, eager, lazy)
 }
 
 #[derive(Debug, Default, Clone)]

usdt-impl/src/no-linker.rs

@@ -89,7 +89,8 @@ fn compile_probe(
         &probe.types,
     );
 
-    let impl_block = quote! {
+    // I don't know whether it's safe to remove the is_enabled check here.
+    let eager = quote! {
         {
             let mut is_enabled: u64;
             unsafe {
@@ -116,7 +117,37 @@ fn compile_probe(
             }
         }
     };
-    common::build_probe_macro(config, &probe.name, &probe.types, impl_block)
+
+    let lazy = quote! {
+        {
+            let mut is_enabled: u64;
+            unsafe {
+                ::std::arch::asm!(
+                    "990:   clr rax",
+                    #is_enabled_rec,
+                    out("rax") is_enabled,
+                    options(nomem, nostack)
+                );
+            }
+
+            if is_enabled != 0 {
+                ::usdt::cold();
+                #arg_lambda
+                #type_check_fn
+                #unpacked_args
+                unsafe {
+                    ::std::arch::asm!(
+                        "990:   nop",
+                        #probe_rec,
+                        #in_regs
+                        options(nomem, nostack, preserves_flags)
+                    );
+                }
+            }
+        }
+    };
+
+    common::build_probe_macro(config, &probe.name, &probe.types, eager, lazy)
 }
 
 fn extract_probe_records_from_section() -> Result<Section, crate::Error> {

usdt-impl/src/stapsdt.rs

@@ -202,7 +202,8 @@ fn compile_probe(
     );
 
     let sema_name = format_ident!("__usdt_sema_{}_{}", provider.name, probe.name);
-    let impl_block = quote! {
+
+    let extern_block = quote! {
         unsafe extern "C" {
             // Note: C libraries use a struct containing an unsigned short
             // for the semaphore counter. Using just a u16 here directly
@@ -211,28 +212,46 @@ fn compile_probe(
             // knowledge.
             static #sema_name: u16;
         }
+    };
+
+    let call = quote! {
+        #[allow(named_asm_labels)]
+        unsafe {
+            ::std::arch::asm!(
+                "990:   nop",
+                #probe_rec,
+                #in_regs
+                options(nomem, nostack, preserves_flags)
+            );
+        }
+    };
+
+    let eager = quote! {
+        #extern_block
+        #arg_lambda
+        #type_check_fn
+        #unpacked_args
+        #call
+    };
+
+    let lazy =quote! {
+        #extern_block
 
         let is_enabled: u16;
         unsafe {
             is_enabled = (&raw const #sema_name).read_volatile();
         }
 
         if is_enabled != 0 {
+            ::usdt::cold();
             #arg_lambda
             #type_check_fn
             #unpacked_args
-            #[allow(named_asm_labels)]
-            unsafe {
-                ::std::arch::asm!(
-                    "990:   nop",
-                    #probe_rec,
-                    #in_regs
-                    options(nomem, nostack, preserves_flags)
-                );
-            }
+            #call
         }
     };
-    common::build_probe_macro(config, &probe.name, &probe.types, impl_block)
+
+    common::build_probe_macro(config, &probe.name, &probe.types, eager, lazy)
 }
 
 pub fn register_probes() -> Result<(), crate::Error> {

usdt/src/lib.rs

@@ -299,6 +299,10 @@ pub use usdt_impl::{to_json, SerializeString};
 pub use usdt_impl::{Error, UniqueId};
 pub use usdt_macro::dtrace_provider;
 
+#[doc(hidden)]
+#[cold]
+pub fn cold() {}
+
 /// A simple struct used to build DTrace probes into Rust code in a build.rs script.
 #[derive(Debug)]
 pub struct Builder {