refactor stapsdt handling

Author: conradludgate

tokio/src/util/usdt.rs

@@ -29,12 +29,8 @@ cfg_rt! {
         #[path = "usdt/macos.rs"]
         pub(crate) mod usdt_impl;
 
-        #[cfg(all(target_os = "linux", target_arch = "x86_64"))]
-        #[path = "usdt/stapsdt_x86_64.rs"]
-        pub(crate) mod usdt_impl;
-
-        #[cfg(all(target_os = "linux", target_arch = "aarch64"))]
-        #[path = "usdt/stapsdt_aarch64.rs"]
+        #[cfg(all(target_os = "linux", any(target_arch = "x86_64", target_arch = "aarch64")))]
+        #[path = "usdt/stapsdt.rs"]
         pub(crate) mod usdt_impl;
 
         #[cfg(not(any(

tokio/src/util/usdt/stapsdt.rs

@@ -0,0 +1,291 @@
+//! USDT support on linux, and other platforms that use the SystemTap SDT V3 system.
+
+#[cfg(target_arch = "x86_64")]
+macro_rules! call_probe {
+    (
+        $name:literal,
+        x86_64: $x86_64:literal,
+        aarch64: $aarch64:literal,
+        $($args:tt)*
+    ) => {
+        std::arch::asm!(
+            "990:   nop
+                    .pushsection .note.stapsdt, \"\", \"note\"
+                    .balign 4
+                    .4byte 992f-991f, 994f-993f, 3
+            991:
+                    .asciz \"stapsdt\"
+            992:
+                    .balign 4
+            993:
+                    .8byte 990b
+                    .8byte _.stapsdt.base
+                    .8byte {semaphore}
+                    .asciz \"tokio\"",
+            concat!(".asciz \"", $name, "\""),
+            concat!(".asciz \"", $x86_64, "\""),
+            "994:
+                    .balign 4
+                    .popsection",
+            $($args)*
+            options(att_syntax, readonly, nostack, preserves_flags),
+        );
+    };
+}
+
+#[cfg(target_arch = "aarch64")]
+macro_rules! call_probe {
+    (
+        $name:literal,
+        x86_64: $x86_64:literal,
+        aarch64: $aarch64:literal,
+        $($args:tt)*
+    ) => {
+        std::arch::asm!(
+            "990:   nop
+                    .pushsection .note.stapsdt, \"\", \"note\"
+                    .balign 4
+                    .4byte 992f-991f, 994f-993f, 3
+            991:
+                    .asciz \"stapsdt\"
+            992:
+                    .balign 4
+            993:
+                    .8byte 990b
+                    .8byte _.stapsdt.base
+                    .8byte {semaphore}
+                    .asciz \"tokio\"",
+            concat!(".asciz \"", $name, "\""),
+            concat!(".asciz \"", $aarch64, "\""),
+            "994:
+                    .balign 4
+                    .popsection",
+            $($args)*
+            options(readonly, nostack, preserves_flags),
+        );
+    };
+}
+
+#[inline(always)]
+pub(super) fn task_details(task_id: u64, name: &str, file: &str, line: u32, col: u32) {
+    #[cold]
+    fn task_details_inner(task_id: u64, name: &str, file: &str, line: u32, col: u32) {
+        // add nul bytes
+        let name0 = [name.as_bytes(), b"\0"].concat();
+        let file0 = [file.as_bytes(), b"\0"].concat();
+
+        unsafe {
+            call_probe!(
+                "task-details",
+                x86_64: "8@{task_id:r} 8@{name} 8@{file} 4@{line:e} 4@{col:e}",
+                aarch64: "8@{task_id} 8@{name} 8@{file} 4@{line:w} 4@{col:w}",
+                semaphore = sym __usdt_sema_tokio_task__details,
+                task_id = in(reg) task_id,
+                name = in(reg) name0.as_ptr(),
+                file = in(reg) file0.as_ptr(),
+                line = in(reg) line,
+                col = in(reg) col,
+            );
+        }
+    }
+
+    if unsafe { __usdt_sema_tokio_task__details } != 0 {
+        task_details_inner(task_id, name, file, line, col);
+    }
+}
+
+#[inline(always)]
+pub(super) fn task_start(task_id: u64, kind: u8, size: usize, original_size: usize) {
+    unsafe {
+        call_probe!(
+            "task-start",
+            x86_64: "8@{task_id:r} 1@{kind:l} 8@{size} 8@{original_size}",
+            aarch64: "8@{task_id} 1@{kind:w} 8@{size} 8@{original_size}",
+            semaphore = const 0,
+            task_id = in(reg) task_id,
+            kind = in(reg) kind as u32,
+            size = in(reg) size,
+            original_size = in(reg) original_size,
+        );
+    }
+}
+
+#[inline(always)]
+pub(super) fn task_poll_start(task_id: u64) {
+    #[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);
+    }
+}
+
+#[inline(always)]
+pub(super) fn task_poll_end(task_id: u64) {
+    #[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);
+    }
+}
+
+#[inline(always)]
+pub(super) fn task_terminate(task_id: u64, reason: u8) {
+    #[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);
+    }
+}
+
+#[inline(always)]
+pub(crate) fn waker_clone(task_id: u64) {
+    unsafe {
+        call_probe!(
+            "task-waker-clone",
+            x86_64: "8@{task_id:r}",
+            aarch64: "8@{task_id}",
+            semaphore = const 0,
+            task_id = in(reg) task_id,
+        );
+    }
+}
+
+#[inline(always)]
+pub(crate) fn waker_drop(task_id: u64) {
+    unsafe {
+        call_probe!(
+            "task-waker-drop",
+            x86_64: "8@{task_id:r}",
+            aarch64: "8@{task_id}",
+            semaphore = const 0,
+            task_id = in(reg) task_id,
+        );
+    }
+}
+
+#[inline(always)]
+pub(crate) fn waker_wake(task_id: u64) {
+    unsafe {
+        call_probe!(
+            "task-waker-wake",
+            x86_64: "8@{task_id:r}",
+            aarch64: "8@{task_id}",
+            semaphore = const 0,
+            task_id = in(reg) task_id,
+        );
+    }
+}
+
+// if defined, attached probes will increment these counters.
+unsafe extern "C" {
+    static mut __usdt_sema_tokio_task__details: u16;
+    static mut __usdt_sema_tokio_task__poll__start: u16;
+    static mut __usdt_sema_tokio_task__poll__end: u16;
+    static mut __usdt_sema_tokio_task__terminate: u16;
+}
+
+std::arch::global_asm!(
+    ".ifndef {semaphore}
+            .pushsection .probes, \"aw\", \"progbits\"
+            .weak {semaphore}
+            .hidden {semaphore}
+    {semaphore}:
+            .zero 2
+            .type {semaphore}, @object
+            .size {semaphore}, 2
+            .popsection
+    .endif",
+    semaphore = sym __usdt_sema_tokio_task__details
+);
+
+std::arch::global_asm!(
+    ".ifndef {semaphore}
+            .pushsection .probes, \"aw\", \"progbits\"
+            .weak {semaphore}
+            .hidden {semaphore}
+    {semaphore}:
+            .zero 2
+            .type {semaphore}, @object
+            .size {semaphore}, 2
+            .popsection
+    .endif",
+    semaphore = sym __usdt_sema_tokio_task__poll__start
+);
+
+std::arch::global_asm!(
+    ".ifndef {semaphore}
+            .pushsection .probes, \"aw\", \"progbits\"
+            .weak {semaphore}
+            .hidden {semaphore}
+    {semaphore}:
+            .zero 2
+            .type {semaphore}, @object
+            .size {semaphore}, 2
+            .popsection
+    .endif",
+    semaphore = sym __usdt_sema_tokio_task__poll__end
+);
+
+std::arch::global_asm!(
+    ".ifndef {semaphore}
+            .pushsection .probes, \"aw\", \"progbits\"
+            .weak {semaphore}
+            .hidden {semaphore}
+    {semaphore}:
+            .zero 2
+            .type {semaphore}, @object
+            .size {semaphore}, 2
+            .popsection
+    .endif",
+    semaphore = sym __usdt_sema_tokio_task__terminate
+);
+
+std::arch::global_asm!(
+    ".ifndef _.stapsdt.base
+            .pushsection .stapsdt.base, \"aGR\", \"progbits\", .stapsdt.base, comdat
+            .weak _.stapsdt.base
+            .hidden _.stapsdt.base
+    _.stapsdt.base:
+            .space 1
+            .size _.stapsdt.base, 1
+            .popsection
+    .endif",
+);
+
+pub(crate) fn register_probes() -> std::io::Result<()> {
+    Ok(())
+}