feat(heaptrack): allow enabling/disabling via IPC

not-matthias · not-matthias · commit 47d5bb086986 · 2025-11-26T14:14:31.000+01:00
diff --git a/crates/heaptrack/src/bpf.rs b/crates/heaptrack/src/bpf.rs
@@ -55,6 +55,38 @@ impl HeaptrackBpf {
         Ok(())
     }
 
+    /// Enable event tracking
+    pub fn enable_tracking(&mut self) -> Result<()> {
+        let key = 0u32;
+        let value = true as u8;
+        self.skel
+            .maps
+            .tracking_enabled
+            .update(
+                &key.to_le_bytes(),
+                &value.to_le_bytes(),
+                libbpf_rs::MapFlags::ANY,
+            )
+            .context("Failed to enable tracking")?;
+        Ok(())
+    }
+
+    /// Disable event tracking
+    pub fn disable_tracking(&mut self) -> Result<()> {
+        let key = 0u32;
+        let value = false as u8;
+        self.skel
+            .maps
+            .tracking_enabled
+            .update(
+                &key.to_le_bytes(),
+                &value.to_le_bytes(),
+                libbpf_rs::MapFlags::ANY,
+            )
+            .context("Failed to disable tracking")?;
+        Ok(())
+    }
+
     pub fn attach_malloc(&mut self, libc_path: &Path) -> Result<()> {
         let malloc_opts = UprobeOpts {
             func_name: Some("malloc".to_string()),
diff --git a/crates/heaptrack/src/bpf/heaptrack.bpf.c b/crates/heaptrack/src/bpf/heaptrack.bpf.c
@@ -29,6 +29,26 @@ struct {
     __uint(max_entries, 256 * 1024);
 } events SEC(".maps");
 
+/* Map to control whether tracking is enabled (0 = disabled, 1 = enabled) */
+struct {
+    __uint(type, BPF_MAP_TYPE_ARRAY);
+    __uint(max_entries, 1);
+    __type(key, __u32);
+    __type(value, __u8);
+} tracking_enabled SEC(".maps");
+
+/* Helper to check if tracking is currently enabled */
+static __always_inline int is_enabled(void) {
+    __u32 key = 0;
+    __u8 *enabled = bpf_map_lookup_elem(&tracking_enabled, &key);
+
+    /* Default to enabled if map not initialized */
+    if (!enabled) {
+        return 1;
+    }
+
+    return *enabled;
+}
 
 /* Helper to check if a PID or any of its ancestors should be tracked */
 static __always_inline int is_tracked(__u32 pid) {
@@ -76,14 +96,16 @@ int tracepoint_sched_fork(struct trace_event_raw_sched_process_fork *ctx) {
 
 SEC("tracepoint/syscalls/sys_enter_execve")
 int tracepoint_sys_execve(struct trace_event_raw_sys_enter *ctx) {
-    __u32 pid = bpf_get_current_pid_tgid() >> 32;
+    __u64 tid = bpf_get_current_pid_tgid();
+    __u32 pid = tid >> 32;
 
     /* Check if this process or any parent is being tracked */
-    if (is_tracked(pid)) {
+    if (is_tracked(pid) && is_enabled()) {
         struct event *e = bpf_ringbuf_reserve(&events, sizeof(*e), 0);
         if (e) {
             e->timestamp = bpf_ktime_get_ns();
             e->pid = pid;
+            e->tid = tid & 0xFFFFFFFF;
             e->event_type = EVENT_TYPE_EXECVE;
             e->addr = 0;
             e->size = 0;
@@ -123,7 +145,7 @@ int uretprobe_malloc(struct pt_regs *ctx) {
     __u32 pid = tid >> 32;
 
     /* Check if this PID is being tracked */
-    if (is_tracked(pid)) {
+    if (is_tracked(pid) && is_enabled()) {
         __u64 *size_ptr = bpf_map_lookup_elem(&malloc_size, &tid);
         if (size_ptr) {
             __u64 addr = PT_REGS_RC(ctx);
@@ -135,6 +157,7 @@ int uretprobe_malloc(struct pt_regs *ctx) {
                 if (e) {
                     e->timestamp = bpf_ktime_get_ns();
                     e->pid = pid;
+                    e->tid = tid & 0xFFFFFFFF;
                     e->event_type = EVENT_TYPE_MALLOC;
                     e->addr = addr;
                     e->size = size;
@@ -151,10 +174,11 @@ int uretprobe_malloc(struct pt_regs *ctx) {
 
 SEC("uprobe")
 int uprobe_free(struct pt_regs *ctx) {
-    __u32 pid = bpf_get_current_pid_tgid() >> 32;
+    __u64 tid = bpf_get_current_pid_tgid();
+    __u32 pid = tid >> 32;
 
     /* Check if this PID is being tracked */
-    if (is_tracked(pid)) {
+    if (is_tracked(pid) && is_enabled()) {
         __u64 addr = PT_REGS_PARM1(ctx);
 
         /* Only track non-NULL frees */
@@ -163,6 +187,7 @@ int uprobe_free(struct pt_regs *ctx) {
             if (e) {
                 e->timestamp = bpf_ktime_get_ns();
                 e->pid = pid;
+                e->tid = tid & 0xFFFFFFFF;
                 e->event_type = EVENT_TYPE_FREE;
                 e->addr = addr;
                 e->size = 0;  /* size unknown for free */
@@ -205,7 +230,7 @@ int uretprobe_calloc(struct pt_regs *ctx) {
     __u32 pid = tid >> 32;
 
     /* Check if this PID is being tracked */
-    if (is_tracked(pid)) {
+    if (is_tracked(pid) && is_enabled()) {
         __u64 *size_ptr = bpf_map_lookup_elem(&calloc_size, &tid);
         if (size_ptr) {
             __u64 addr = PT_REGS_RC(ctx);
@@ -217,6 +242,7 @@ int uretprobe_calloc(struct pt_regs *ctx) {
                 if (e) {
                     e->timestamp = bpf_ktime_get_ns();
                     e->pid = pid;
+                    e->tid = tid & 0xFFFFFFFF;
                     e->event_type = EVENT_TYPE_CALLOC;
                     e->addr = addr;
                     e->size = size;
@@ -260,7 +286,7 @@ int uretprobe_realloc(struct pt_regs *ctx) {
     __u32 pid = tid >> 32;
 
     /* Check if this PID is being tracked */
-    if (is_tracked(pid)) {
+    if (is_tracked(pid) && is_enabled()) {
         __u64 *size_ptr = bpf_map_lookup_elem(&realloc_size, &tid);
         if (size_ptr) {
             __u64 addr = PT_REGS_RC(ctx);
@@ -272,6 +298,7 @@ int uretprobe_realloc(struct pt_regs *ctx) {
                 if (e) {
                     e->timestamp = bpf_ktime_get_ns();
                     e->pid = pid;
+                    e->tid = tid & 0xFFFFFFFF;
                     e->event_type = EVENT_TYPE_REALLOC;
                     e->addr = addr;
                     e->size = size;
@@ -315,7 +342,7 @@ int uretprobe_aligned_alloc(struct pt_regs *ctx) {
     __u32 pid = tid >> 32;
 
     /* Check if this PID is being tracked */
-    if (is_tracked(pid)) {
+    if (is_tracked(pid) && is_enabled()) {
         __u64 *size_ptr = bpf_map_lookup_elem(&aligned_alloc_size, &tid);
         if (size_ptr) {
             __u64 addr = PT_REGS_RC(ctx);
@@ -327,6 +354,7 @@ int uretprobe_aligned_alloc(struct pt_regs *ctx) {
                 if (e) {
                     e->timestamp = bpf_ktime_get_ns();
                     e->pid = pid;
+                    e->tid = tid & 0xFFFFFFFF;
                     e->event_type = EVENT_TYPE_ALIGNED_ALLOC;
                     e->addr = addr;
                     e->size = size;
diff --git a/crates/heaptrack/src/ipc.rs b/crates/heaptrack/src/ipc.rs
@@ -0,0 +1,115 @@
+use anyhow::{Context, Result};
+use ipc_channel::ipc::{self, IpcOneShotServer, IpcSender};
+use log::debug;
+use serde::{Deserialize, Serialize};
+use std::sync::{Arc, Mutex};
+
+pub type HeaptrackIpcServer = IpcOneShotServer<IpcSender<IpcMessage>>;
+
+use crate::Tracker;
+
+/// Commands sent from the runner to control heaptrack
+#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
+pub enum IpcCommand {
+    Enable,
+    Disable,
+    Ping,
+}
+
+/// Response sent back to runner
+#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
+pub enum IpcResponse {
+    Ack,
+    Err,
+}
+
+/// Message combining command and response channel
+#[derive(Debug, Serialize, Deserialize)]
+pub struct IpcMessage {
+    pub command: IpcCommand,
+    pub response_channel: IpcSender<IpcResponse>,
+}
+
+/// Client for sending commands from runner to heaptrack
+pub struct HeaptrackIpcClient {
+    sender: IpcSender<IpcMessage>,
+}
+
+impl HeaptrackIpcClient {
+    pub fn new(sender: IpcSender<IpcMessage>) -> Self {
+        Self { sender }
+    }
+
+    /// Create from the connection accepted by the runner
+    pub fn from_accepted(sender: IpcSender<IpcMessage>) -> Self {
+        Self { sender }
+    }
+
+    fn send_command(&self, cmd: IpcCommand) -> Result<IpcResponse> {
+        let (response_tx, response_rx) = ipc::channel::<IpcResponse>()?;
+
+        let msg = IpcMessage {
+            command: cmd,
+            response_channel: response_tx,
+        };
+
+        self.sender.send(msg).context("Failed to send command")?;
+        let response = response_rx.recv().context("Failed to receive response")?;
+
+        Ok(response)
+    }
+
+    pub fn enable(&self) -> Result<()> {
+        let response = self.send_command(IpcCommand::Enable)?;
+        match response {
+            IpcResponse::Ack => Ok(()),
+            IpcResponse::Err => anyhow::bail!("Failed to enable tracking"),
+        }
+    }
+
+    pub fn disable(&self) -> Result<()> {
+        let response = self.send_command(IpcCommand::Disable)?;
+        match response {
+            IpcResponse::Ack => Ok(()),
+            IpcResponse::Err => anyhow::bail!("Failed to disable tracking"),
+        }
+    }
+
+    pub fn ping(&self) -> Result<()> {
+        let response = self.send_command(IpcCommand::Ping)?;
+        match response {
+            IpcResponse::Ack => Ok(()),
+            IpcResponse::Err => anyhow::bail!("Failed to ping heaptrack"),
+        }
+    }
+}
+
+/// Handle incoming IPC messages in heaptrack
+pub fn handle_ipc_message(msg: IpcMessage, tracker: &Arc<Mutex<Tracker>>) {
+    let response = match msg.command {
+        IpcCommand::Enable => match tracker.lock() {
+            Ok(mut t) => match t.enable() {
+                Ok(_) => {
+                    debug!("Tracking enabled");
+                    IpcResponse::Ack
+                }
+                Err(_) => IpcResponse::Err,
+            },
+            Err(_) => IpcResponse::Err,
+        },
+        IpcCommand::Disable => match tracker.lock() {
+            Ok(mut t) => match t.disable() {
+                Ok(_) => {
+                    debug!("Tracking disabled");
+                    IpcResponse::Ack
+                }
+                Err(_) => IpcResponse::Err,
+            },
+            Err(_) => IpcResponse::Err,
+        },
+        IpcCommand::Ping => IpcResponse::Ack,
+    };
+
+    // Send response back
+    let _ = msg.response_channel.send(response);
+}
diff --git a/crates/heaptrack/src/lib.rs b/crates/heaptrack/src/lib.rs
@@ -1,9 +1,14 @@
 mod bpf;
 mod events;
+mod ipc;
 mod libc;
 mod poller;
 mod tracker;
 
 pub use events::{Event, EventType};
+pub use ipc::{
+    HeaptrackIpcClient, HeaptrackIpcServer, IpcCommand as HeaptrackIpcCommand,
+    IpcMessage as HeaptrackIpcMessage, IpcResponse as HeaptrackIpcResponse, handle_ipc_message,
+};
 pub use poller::{EventHandler, RingBufferPoller};
 pub use tracker::Tracker;
diff --git a/crates/heaptrack/src/main.rs b/crates/heaptrack/src/main.rs
@@ -28,6 +28,10 @@ enum Commands {
         /// Output folder for the allocations data
         #[arg(short, long, default_value = ".")]
         output: PathBuf,
+
+        /// Optional IPC server name for receiving control commands
+        #[arg(long)]
+        ipc_server: Option<String>,
     },
 }
 
diff --git a/crates/heaptrack/src/tracker.rs b/crates/heaptrack/src/tracker.rs
@@ -80,4 +80,14 @@ impl Tracker {
 
         Ok(())
     }
+
+    /// Enable event tracking in the BPF program
+    pub fn enable(&mut self) -> anyhow::Result<()> {
+        self.heaptrack.enable_tracking()
+    }
+
+    /// Disable event tracking in the BPF program
+    pub fn disable(&mut self) -> anyhow::Result<()> {
+        self.heaptrack.disable_tracking()
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -80,4 +80,14 @@ impl Tracker {`
`80`	`80`
`81`	`81`	`Ok(())`
`82`	`82`	`}`
	`83`	`+`
	`84`	`+ /// Enable event tracking in the BPF program`
	`85`	`+ pub fn enable(&mut self) -> anyhow::Result<()> {`
	`86`	`+ self.heaptrack.enable_tracking()`
	`87`	`+ }`
	`88`	`+`
	`89`	`+ /// Disable event tracking in the BPF program`
	`90`	`+ pub fn disable(&mut self) -> anyhow::Result<()> {`
	`91`	`+ self.heaptrack.disable_tracking()`
	`92`	`+ }`
`83`	`93`	`}`