From a528226022e7dea01e6335ab489f381097958f4d Mon Sep 17 00:00:00 2001
From: Alice Ryhl <aliceryhl@google.com>
Date: Fri, 31 May 2024 16:52:37 +0200
Subject: [PATCH] rust: add tracepoint support

Make it possible to have Rust code call into tracepoints defined by C
code. It is still required that the tracepoint is declared in a C
header, and that this header is included in the input to bindgen.

Signed-off-by: Alice Ryhl <aliceryhl@google.com>
---
 rust/bindings/bindings_helper.h |  1 +
 rust/kernel/lib.rs              |  1 +
 rust/kernel/tracepoint.rs       | 80 +++++++++++++++++++++++++++++++++
 3 files changed, 82 insertions(+)
 create mode 100644 rust/kernel/tracepoint.rs

diff --git a/rust/bindings/bindings_helper.h b/rust/bindings/bindings_helper.h
index ddb5644d4fd900..d442f9ccfc2c43 100644
--- a/rust/bindings/bindings_helper.h
+++ b/rust/bindings/bindings_helper.h
@@ -15,6 +15,7 @@
 #include <linux/refcount.h>
 #include <linux/sched.h>
 #include <linux/slab.h>
+#include <linux/tracepoint.h>
 #include <linux/wait.h>
 #include <linux/workqueue.h>
 
diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index 208d4c4f9cdc8e..adf8152d039195 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -48,6 +48,7 @@ pub mod str;
 pub mod sync;
 pub mod task;
 pub mod time;
+pub mod tracepoint;
 pub mod types;
 pub mod workqueue;
 
diff --git a/rust/kernel/tracepoint.rs b/rust/kernel/tracepoint.rs
new file mode 100644
index 00000000000000..32743c1f24233f
--- /dev/null
+++ b/rust/kernel/tracepoint.rs
@@ -0,0 +1,80 @@
+// SPDX-License-Identifier: GPL-2.0
+
+// Copyright (C) 2024 Google LLC.
+
+//! Logic for tracepoints.
+
+/// Declare the Rust entry point for a tracepoint.
+#[macro_export]
+macro_rules! declare_trace {
+    ($(#[$attr:meta])* $pub:vis fn $name:ident($($argname:ident : $argtyp:ty),* $(,)?);) => {
+        $crate::macros::paste! {
+            $( #[$attr] )*
+            #[inline(always)]
+            $pub fn [< trace_ $name >]($($argname : $argtyp),*) {
+                #[cfg(CONFIG_TRACEPOINTS)]
+                {
+                    use $crate::bindings::*;
+
+                    let should_trace = unsafe {
+                        $crate::static_key::static_key_false!(
+                            [< __tracepoint_ $name >],
+                            $crate::bindings::tracepoint,
+                            key
+                        )
+                    };
+
+                    if should_trace {
+                        // TODO: cpu_online(raw_smp_processor_id())
+                        let cond = true;
+                        $crate::tracepoint::do_trace!($name($($argname : $argtyp),*), cond);
+                    }
+                }
+            }
+        }
+    }
+}
+
+#[doc(hidden)]
+#[macro_export]
+macro_rules! do_trace {
+    ($name:ident($($argname:ident : $argtyp:ty),* $(,)?), $cond:expr) => {{
+        if !$cond {
+            return;
+        }
+
+        // SAFETY: This call is balanced with the call below.
+        unsafe { $crate::bindings::preempt_disable_notrace() };
+
+        #[cfg(CONFIG_HAVE_STATIC_CALL)]
+        unsafe {
+            let it_func_ptr: *mut $crate::bindings::tracepoint_func =
+                $crate::bindings::rcu_dereference_raw(
+                    ::core::ptr::addr_of!(
+                        $crate::macros::concat_idents!(__tracepoint_, $name).funcs
+                    )
+                );
+
+            if !it_func_ptr.is_null() {
+                let __data = (*it_func_ptr).data;
+                $crate::macros::paste! {
+                    $crate::static_call::static_call! {
+                        [< tp_func_ $name >] (__data, $($argname),*)
+                    };
+                }
+            }
+        }
+
+        #[cfg(not(CONFIG_HAVE_STATIC_CALL))]
+        unsafe {
+            $crate::macros::concat_idents!(__traceiter_, $name)(
+                ::core::ptr::null_mut(),
+                $($argname),*
+            );
+        }
+
+        unsafe { $crate::bindings::preempt_enable_notrace() };
+    }}
+}
+
+pub use {declare_trace, do_trace};