add CString support and optimise string allocations

dtrace-parser/src/lib.rs

@@ -156,6 +156,7 @@ impl Integer {
 pub enum DataType {
     Integer(Integer),
     Pointer(Integer),
+    CString,
     String,
 }
 
@@ -246,7 +247,7 @@ impl DataType {
         match self {
             DataType::Integer(int) => int.to_c_type(),
             DataType::Pointer(int) => format!("{}*", int.to_c_type()),
-            DataType::String => String::from("char*"),
+            DataType::CString | DataType::String => String::from("char*"),
         }
     }
 
@@ -255,7 +256,7 @@ impl DataType {
         match self {
             DataType::Integer(int) => int.to_rust_ffi_type(),
             DataType::Pointer(int) => format!("*const {}", int.to_rust_ffi_type()),
-            DataType::String => format!("*const {RUST_TYPE_PREFIX}char"),
+            DataType::CString | DataType::String => format!("*const {RUST_TYPE_PREFIX}char"),
         }
     }
 
@@ -265,6 +266,7 @@ impl DataType {
             DataType::Integer(int) => int.to_rust_type(),
             DataType::Pointer(int) => format!("*const {}", int.to_rust_type()),
             DataType::String => String::from("&str"),
+            DataType::CString => String::from("&::core::ffi::CStr"),
         }
     }
 }

probe-test-attr/src/main.rs

@@ -82,6 +82,9 @@ mod test {
 
     /// Constant pointers to integer types are also supported
     fn work_with_pointer(_buffer: *const u8, _: u64) {}
+
+    /// Avoid expensive string allocation with cstring support
+    fn cstring(_: &CStr, _: CString) {}
 }
 
 fn main() {
@@ -106,5 +109,6 @@ fn main() {
         test::arg_as_tuple!(|| (arg.x, &arg.buffer[..]));
         test::not_json_serializable!(|| Whoops::NoBueno(0));
         test::work_with_pointer!(|| (buffer.as_ptr(), buffer.len() as u64));
+        test::cstring!(|| (c"hello world", c"and when owned".to_owned()));
     }
 }

usdt-attr-macro/src/lib.rs

@@ -303,6 +303,8 @@ fn is_simple_type(ident: &syn::Ident) -> bool {
             | "i64"
             | "String"
             | "str"
+            | "CString"
+            | "CStr"
             | "usize"
             | "isize"
     )
@@ -363,6 +365,8 @@ fn data_type_from_path(path: &syn::Path, pointer: bool) -> DataType {
         }))
     } else if path.is_ident("String") || path.is_ident("str") {
         DataType::Native(DType::String)
+    } else if path.is_ident("CString") || path.is_ident("CStr") {
+        DataType::Native(DType::CString)
     } else if path.is_ident("isize") {
         DataType::Native(variant(Integer {
             sign: Sign::Signed,
@@ -462,6 +466,10 @@ mod tests {
     #[case("String", DType::String)]
     #[case("&&str", DType::String)]
     #[case("&String", DType::String)]
+    #[case("&CStr", DType::CString)]
+    #[case("CString", DType::CString)]
+    #[case("&&CStr", DType::CString)]
+    #[case("&CString", DType::CString)]
     fn test_parse_probe_argument_native(#[case] name: &str, #[case] ty: dtrace_parser::DataType) {
         let arg = syn::parse_str(name).unwrap();
         let out = parse_probe_argument(&arg, 0, 0).unwrap();

usdt-impl/src/common.rs

@@ -43,33 +43,36 @@ pub fn construct_type_check(
                 match ty {
                     syn::Type::Reference(reference) => {
                         if let Some(elem) = shared_slice_elem_type(reference) {
-                            quote! { _: impl AsRef<[#elem]> }
+                            quote! { _: &impl AsRef<[#elem]> }
                         } else {
                             let elem = &*reference.elem;
-                            quote! { _: impl ::std::borrow::Borrow<#elem> }
+                            quote! { _: &impl ::std::borrow::Borrow<#elem> }
                         }
                     }
                     syn::Type::Slice(slice) => {
                         let elem = &*slice.elem;
-                        quote! { _: impl AsRef<[#elem]> }
+                        quote! { _: &impl AsRef<[#elem]> }
                     }
                     syn::Type::Array(array) => {
                         let elem = &*array.elem;
-                        quote! { _: impl AsRef<[#elem]> }
+                        quote! { _: &impl AsRef<[#elem]> }
                     }
                     syn::Type::Path(_) => {
-                        quote! { _: impl ::std::borrow::Borrow<#ty> }
+                        quote! { _: &impl ::std::borrow::Borrow<#ty> }
                     }
                     _ => {
                         // Any other type must be specified exactly as given in the probe parameter
-                        quote! { _: #ty }
+                        quote! { _: &#ty }
                     }
                 }
             }
-            DataType::Native(dtrace_parser::DataType::String) => quote! { _: impl AsRef<str> },
+            DataType::Native(dtrace_parser::DataType::String) => quote! { _: &impl AsRef<str> },
+            DataType::Native(dtrace_parser::DataType::CString) => {
+                quote! { _: &impl AsRef<::core::ffi::CStr> }
+            }
             _ => {
                 let arg = typ.to_rust_type();
-                quote! { _: impl ::std::borrow::Borrow<#arg> }
+                quote! { _: &impl ::std::borrow::Borrow<#arg> }
             }
         })
         .collect::<Vec<_>>();
@@ -79,7 +82,7 @@ pub fn construct_type_check(
     let type_check_args = (0..types.len())
         .map(|i| {
             let index = syn::Index::from(i);
-            quote! { args.#index }
+            quote! { &args.#index }
         })
         .collect::<Vec<_>>();
 
@@ -103,7 +106,7 @@ fn shared_slice_elem_type(reference: &syn::TypeReference) -> Option<&syn::Type>
 
 // Return code to destructure a probe arguments into identifiers, and to pass those to ASM
 // registers.
-pub fn construct_probe_args(types: &[DataType]) -> (TokenStream, TokenStream) {
+pub fn construct_probe_args(types: &[DataType]) -> (TokenStream, TokenStream, TokenStream) {
     // x86_64 passes the first 6 arguments in registers, with the rest on the stack.
     // We limit this to 6 arguments in all cases for now, as handling those stack
     // arguments would be challenging with the current `asm!` macro implementation.
@@ -141,12 +144,9 @@ pub fn construct_probe_args(types: &[DataType]) -> (TokenStream, TokenStream) {
         })
         .unzip();
     let arg_lambda = call_argument_closure(types);
-    let unpacked_args = quote! {
-        #arg_lambda
-        #(#unpacked_args)*
-    };
+    let unpacked_args = quote! { #(#unpacked_args)* };
     let in_regs = quote! { #(#in_regs,)* };
-    (unpacked_args, in_regs)
+    (arg_lambda, unpacked_args, in_regs)
 }
 
 /// Call the argument closure, assigning its output to `args`.
@@ -171,20 +171,19 @@ fn asm_type_convert(typ: &DataType, input: TokenStream) -> (TokenStream, TokenSt
             // data in a result-like JSON blob, mapping the `Result`'s variants to the keys "ok"
             // and "err".
             quote! {
-                [
-                    match ::usdt::to_json(&#input) {
-                        Ok(json) => format!("{{\"ok\":{}}}", json),
-                        Err(e) => format!("{{\"err\":\"{}\"}}", e.to_string()),
-                    }.as_bytes(),
-                    &[0_u8]
-                ].concat()
+                match ::usdt::to_json(&#input) {
+                    Ok(json) => format!("{{\"ok\":{}}}\0", json).into_bytes(),
+                    Err(e) => format!("{{\"err\":\"{}\"}}\0", e.to_string()).into_bytes(),
+                }
             },
             quote! { .as_ptr() as usize },
         ),
+        DataType::Native(dtrace_parser::DataType::CString) => (
+            quote! { #input.as_ref() as &::core::ffi::CStr },
+            quote! { .as_ptr() as usize },
+        ),
         DataType::Native(dtrace_parser::DataType::String) => (
-            quote! {
-                [(#input.as_ref() as &str).as_bytes(), &[0_u8]].concat()
-            },
+            quote! { ::usdt::SerializeString::serialize(#input) },
             quote! { .as_ptr() as usize },
         ),
         DataType::Native(_) => {
@@ -269,11 +268,11 @@ mod tests {
             #[allow(unused_imports)]
             #[allow(non_snake_case)]
             fn __usdt_private_provider_probe_type_check(
-                _: impl ::std::borrow::Borrow<u8>,
-                _: impl ::std::borrow::Borrow<i64>
+                _: &impl ::std::borrow::Borrow<u8>,
+                _: &impl ::std::borrow::Borrow<i64>
             ) { }
             let _ = || {
-                __usdt_private_provider_probe_type_check(args.0, args.1);
+                __usdt_private_provider_probe_type_check(&args.0, &args.1);
             };
         };
         let block = construct_type_check(provider, probe, &[], types);
@@ -289,9 +288,9 @@ mod tests {
         let expected = quote! {
             #[allow(unused_imports)]
             #[allow(non_snake_case)]
-            fn __usdt_private_provider_probe_type_check(_: impl AsRef<str>) { }
+            fn __usdt_private_provider_probe_type_check(_: &impl AsRef<str>) { }
             let _ = || {
-                __usdt_private_provider_probe_type_check(args.0);
+                __usdt_private_provider_probe_type_check(&args.0);
             };
         };
         let block = construct_type_check(provider, probe, &use_statements, types);
@@ -307,9 +306,9 @@ mod tests {
         let expected = quote! {
             #[allow(unused_imports)]
             #[allow(non_snake_case)]
-            fn __usdt_private_provider_probe_type_check(_: impl AsRef<[u8]>) { }
+            fn __usdt_private_provider_probe_type_check(_: &impl AsRef<[u8]>) { }
             let _ = || {
-                __usdt_private_provider_probe_type_check(args.0);
+                __usdt_private_provider_probe_type_check(&args.0);
             };
         };
         let block = construct_type_check(provider, probe, &use_statements, types);
@@ -326,9 +325,9 @@ mod tests {
             #[allow(unused_imports)]
             use my_module::MyType;
             #[allow(non_snake_case)]
-            fn __usdt_private_provider_probe_type_check(_: impl ::std::borrow::Borrow<MyType>) { }
+            fn __usdt_private_provider_probe_type_check(_: &impl ::std::borrow::Borrow<MyType>) { }
             let _ = || {
-                __usdt_private_provider_probe_type_check(args.0);
+                __usdt_private_provider_probe_type_check(&args.0);
             };
         };
         let block = construct_type_check(provider, probe, &use_statements, types);
@@ -348,11 +347,17 @@ mod tests {
         let registers = ["rdi", "rsi"];
         #[cfg(target_arch = "aarch64")]
         let registers = ["x0", "x1"];
-        let (args, regs) = construct_probe_args(types);
+
+        let (lambda, args, regs) = construct_probe_args(types);
+
         let expected = quote! {
             let args = ($args_lambda)();
+        };
+        assert_eq!(lambda.to_string(), expected.to_string());
+
+        let expected = quote! {
             let arg_0 = (*<_ as ::std::borrow::Borrow<*const u8>>::borrow(&args.0) as usize);
-            let arg_1 = [(args.1.as_ref() as &str).as_bytes(), &[0_u8]].concat();
+            let arg_1 = ::usdt::SerializeString::serialize(args.1);
         };
         assert_eq!(args.to_string(), expected.to_string());
 
@@ -394,7 +399,17 @@ mod tests {
         );
         assert_eq!(
             out.to_string(),
-            quote! { [(foo.as_ref() as &str).as_bytes(), &[0_u8]].concat() }.to_string()
+            quote! { ::usdt::SerializeString::serialize(foo) }.to_string()
+        );
+        assert_eq!(post.to_string(), quote! { .as_ptr() as usize }.to_string());
+
+        let (out, post) = asm_type_convert(
+            &DataType::Native(dtrace_parser::DataType::CString),
+            TokenStream::from_str("foo").unwrap(),
+        );
+        assert_eq!(
+            out.to_string(),
+            quote! { foo.as_ref() as &::core::ffi::CStr }.to_string()
         );
         assert_eq!(post.to_string(), quote! { .as_ptr() as usize }.to_string());
     }

usdt-impl/src/lib.rs

@@ -421,6 +421,27 @@ impl Clone for UniqueId {
     }
 }
 
+#[doc(hidden)]
+pub trait SerializeString {
+    fn serialize(self) -> Vec<u8>;
+}
+
+#[doc(hidden)]
+impl<S: AsRef<str> + ?Sized> SerializeString for &S {
+    #[inline]
+    fn serialize(self) -> Vec<u8> {
+        [self.as_ref().as_bytes(), b"0"].concat()
+    }
+}
+
+#[doc(hidden)]
+impl SerializeString for String {
+    fn serialize(mut self) -> Vec<u8> {
+        self.push('\0');
+        self.into_bytes()
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;

usdt-impl/src/linker.rs

@@ -169,7 +169,7 @@ fn compile_probe(
         let ty = typ.to_rust_ffi_type();
         syn::parse2::<syn::FnArg>(quote! { _: #ty }).unwrap()
     });
-    let (unpacked_args, in_regs) = common::construct_probe_args(types);
+    let (arg_lambda, unpacked_args, in_regs) = common::construct_probe_args(types);
     let type_check_fn =
         common::construct_type_check(&provider.name, probe_name, &provider.use_statements, types);
 
@@ -200,8 +200,9 @@ fn compile_probe(
         }
         unsafe {
             if #is_enabled_fn() != 0 {
-                #unpacked_args
+                #arg_lambda
                 #type_check_fn
+                #unpacked_args
                 ::std::arch::asm!(
                     ".reference {typedefs}",
                     #call_instruction,

usdt-impl/src/no-linker.rs

@@ -79,7 +79,7 @@ fn compile_probe(
     probe: &Probe,
     config: &crate::CompileProvidersConfig,
 ) -> TokenStream {
-    let (unpacked_args, in_regs) = common::construct_probe_args(&probe.types);
+    let (arg_lambda, unpacked_args, in_regs) = common::construct_probe_args(&probe.types);
     let is_enabled_rec = emit_probe_record(&provider.name, &probe.name, None);
     let probe_rec = emit_probe_record(&provider.name, &probe.name, Some(&probe.types));
     let type_check_fn = common::construct_type_check(
@@ -102,8 +102,9 @@ fn compile_probe(
             }
 
             if is_enabled != 0 {
-                #unpacked_args
+                #arg_lambda
                 #type_check_fn
+                #unpacked_args
                 unsafe {
                     ::std::arch::asm!(
                         "990:   nop",

usdt-impl/src/stapsdt.rs

@@ -192,7 +192,7 @@ fn compile_probe(
     probe: &Probe,
     config: &crate::CompileProvidersConfig,
 ) -> TokenStream {
-    let (unpacked_args, in_regs) = common::construct_probe_args(&probe.types);
+    let (arg_lambda, unpacked_args, in_regs) = common::construct_probe_args(&probe.types);
     let probe_rec = emit_probe_record(&provider.name, &probe.name, Some(&probe.types));
     let type_check_fn = common::construct_type_check(
         &provider.name,
@@ -218,8 +218,9 @@ fn compile_probe(
         }
 
         if is_enabled != 0 {
-            #unpacked_args
+            #arg_lambda
             #type_check_fn
+            #unpacked_args
             #[allow(named_asm_labels)]
             unsafe {
                 ::std::arch::asm!(

usdt/src/lib.rs

@@ -295,7 +295,7 @@ use std::{env, fs};
 
 pub use usdt_attr_macro::provider;
 #[doc(hidden)]
-pub use usdt_impl::to_json;
+pub use usdt_impl::{to_json, SerializeString};
 pub use usdt_impl::{Error, UniqueId};
 pub use usdt_macro::dtrace_provider;