add retained flag to notes, allowing us to inline and monomorphise again

Author: conradludgate

tokio/src/util/usdt/macos.rs

@@ -83,61 +83,31 @@ pub(super) fn task_start(task_id: u64, spawned: u8, size: usize, original_size:
 #[inline(always)]
 pub(super) fn task_poll_start(task_id: u64) {
     extern "C" {
-        #[link_name = "__dtrace_isenabled$tokio$task__poll__start$v1"]
-        fn task_poll_start_enabled() -> i32;
-
         #[link_name = "__dtrace_probe$tokio$task__poll__start$v1$75696e7436345f74"]
         fn __task_poll_start(task_id: core::ffi::c_ulonglong);
     }
 
-    #[inline(never)]
-    fn probe_inner(task_id: u64) {
-        unsafe { __task_poll_start(task_id) }
-    }
-
-    if unsafe { task_poll_start_enabled() } != 0 {
-        probe_inner(task_id);
-    }
+    unsafe { __task_poll_start(task_id) }
 }
 
 #[inline(always)]
 pub(super) fn task_poll_end(task_id: u64) {
     extern "C" {
-        #[link_name = "__dtrace_isenabled$tokio$task__poll__end$v1"]
-        fn task_poll_end_enabled() -> i32;
-
         #[link_name = "__dtrace_probe$tokio$task__poll__end$v1$75696e7436345f74"]
         fn __task_poll_end(task_id: core::ffi::c_ulonglong);
     }
 
-    #[inline(never)]
-    fn probe_inner(task_id: u64) {
-        unsafe { __task_poll_end(task_id) }
-    }
-
-    if unsafe { task_poll_end_enabled() } != 0 {
-        probe_inner(task_id);
-    }
+    unsafe { __task_poll_end(task_id) }
 }
 
 #[inline(always)]
 pub(crate) fn task_terminate(task_id: u64, reason: u8) {
     extern "C" {
-        #[link_name = "__dtrace_isenabled$tokio$task__terminate$v1"]
-        fn task_terminate_enabled() -> i32;
-
         #[link_name = "__dtrace_probe$tokio$task__terminate$v1$75696e7436345f74$75696e74385f74"]
         fn __task_terminate(task_id: core::ffi::c_ulonglong, reason: core::ffi::c_uchar);
     }
 
-    #[inline(never)]
-    fn probe_inner(task_id: u64, reason: u8) {
-        unsafe { __task_terminate(task_id, reason) }
-    }
-
-    if unsafe { task_terminate_enabled() } != 0 {
-        probe_inner(task_id, reason);
-    }
+    unsafe { __task_terminate(task_id, reason) }
 }
 
 #[inline(always)]

tokio/src/util/usdt/stapsdt.rs

@@ -24,14 +24,6 @@
 //! into a seperate function and marking it as `#[cold]` or `#[inline(never)]`.
 //! This ensures the hot path (probe is disabled) doesn't need to skip over a large
 //! number of instructions, and the branch predictor can make better predictions.
-//!
-//! ## Avoid monomorphisation
-//!
-//! We also use semaphores to ensure that probe callsites are **NOT** monomorphised.
-//! When a probe is monomorphised in many callsites, it has been observed to
-//! produce weird results after linking where the probe addresses are not correct.
-//!
-//! Be very mindful of the above when you add a probe.
 
 #[cfg(target_arch = "x86_64")]
 macro_rules! call_probe {
@@ -42,9 +34,16 @@ macro_rules! call_probe {
         $($args:tt)*
     ) => {
         // <https://sourceware.org/systemtap/wiki/UserSpaceProbeImplementation>
+        // <https://github.com/cdkey/systemtap/blob/cbb34b7244ba60cb0904d61dc9167290855106aa/includes/sys/sdt.h#L240-L254>
+        //
+        // The `R` here means "retained". It prevents the linker from GC the note.
+        // This additionally means that the function containing the `nop` is also not GCd.
+        //
+        // This is not great, but without this we get notes that are not GCd but the functions
+        // they reference are GCd. This causes notes to have incorrect addresses.
         std::arch::asm!(
             "990:   nop
-                    .pushsection .note.stapsdt, \"\", \"note\"
+                    .pushsection .note.stapsdt, \"R\", \"note\"
                     .balign 4
                     .4byte 992f-991f, 994f-993f, 3
             991:
@@ -76,9 +75,16 @@ macro_rules! call_probe {
         $($args:tt)*
     ) => {
         // <https://sourceware.org/systemtap/wiki/UserSpaceProbeImplementation>
+        // <https://github.com/cdkey/systemtap/blob/cbb34b7244ba60cb0904d61dc9167290855106aa/includes/sys/sdt.h#L240-L254>
+        //
+        // The `R` here means "retained". It prevents the linker from GC the note.
+        // This additionally means that the function containing the `nop` is also not GCd.
+        //
+        // This is not great, but without this we get notes that are not GCd but the functions
+        // they reference are GCd. This causes notes to have incorrect addresses.
         std::arch::asm!(
             "990:   nop
-                    .pushsection .note.stapsdt, \"\", \"note\"
+                    .pushsection .note.stapsdt, \"R\", \"note\"
                     .balign 4
                     .4byte 992f-991f, 994f-993f, 3
             991:
@@ -171,73 +177,43 @@ pub(super) fn task_start(task_id: u64, kind: u8, size: usize, original_size: usi
     }
 }
 
-declare_semaphore!(__usdt_sema_tokio_task__poll__start);
-
 #[inline(always)]
 pub(super) fn task_poll_start(task_id: u64) {
-    // `inline(never)` since poll_start probes are monomorphised otherwise
-    #[inline(never)]
-    fn probe_inner(task_id: u64) {
-        unsafe {
-            call_probe!(
-                "task-poll-start",
-                x86_64: "8@{task_id:r}",
-                aarch64: "8@{task_id}",
-                semaphore = sym __usdt_sema_tokio_task__poll__start,
-                task_id = in(reg) task_id,
-            );
-        }
-    }
-
-    if unsafe { __usdt_sema_tokio_task__poll__start } != 0 {
-        probe_inner(task_id);
+    unsafe {
+        call_probe!(
+            "task-poll-start",
+            x86_64: "8@{task_id:r}",
+            aarch64: "8@{task_id}",
+            semaphore = const 0,
+            task_id = in(reg) task_id,
+        );
     }
 }
 
-declare_semaphore!(__usdt_sema_tokio_task__poll__end);
-
 #[inline(always)]
 pub(super) fn task_poll_end(task_id: u64) {
-    // `inline(never)` since poll_end probes are monomorphised otherwise
-    #[inline(never)]
-    fn probe_inner(task_id: u64) {
-        unsafe {
-            call_probe!(
-                "task-poll-end",
-                x86_64: "8@{task_id:r}",
-                aarch64: "8@{task_id}",
-                semaphore = sym __usdt_sema_tokio_task__poll__end,
-                task_id = in(reg) task_id,
-            );
-        }
-    }
-
-    if unsafe { __usdt_sema_tokio_task__poll__end } != 0 {
-        probe_inner(task_id);
+    unsafe {
+        call_probe!(
+            "task-poll-end",
+            x86_64: "8@{task_id:r}",
+            aarch64: "8@{task_id}",
+            semaphore = const 0,
+            task_id = in(reg) task_id,
+        );
     }
 }
 
-declare_semaphore!(__usdt_sema_tokio_task__terminate);
-
 #[inline(always)]
 pub(super) fn task_terminate(task_id: u64, reason: u8) {
-    // `inline(never)` since terminate probes are monomorphised otherwise
-    #[inline(never)]
-    fn probe_inner(task_id: u64, reason: u8) {
-        unsafe {
-            call_probe!(
-                "task-terminate",
-                x86_64: "8@{task_id:r} 1@{reason:l}",
-                aarch64: "8@{task_id} 1@{reason:w}",
-                semaphore = sym __usdt_sema_tokio_task__terminate,
-                task_id = in(reg) task_id,
-                reason = in(reg) reason as u32,
-            );
-        }
-    }
-
-    if unsafe { __usdt_sema_tokio_task__terminate } != 0 {
-        probe_inner(task_id, reason);
+    unsafe {
+        call_probe!(
+            "task-terminate",
+            x86_64: "8@{task_id:r} 1@{reason:l}",
+            aarch64: "8@{task_id} 1@{reason:w}",
+            semaphore = const 0,
+            task_id = in(reg) task_id,
+            reason = in(reg) reason as u32,
+        );
     }
 }