feat: track child processes

.github/workflows/ci.yml

@@ -33,7 +33,7 @@ jobs:
           cargo install --path crates/memtrack --locked
           echo "CODSPEED_MEMTRACK_BINARY=$(which codspeed-memtrack)" >> $GITHUB_ENV
 
-      - run: cargo test --all --exclude memtrack
+      - run: cargo test --all --exclude memtrack exectrack
 
   bpf-tests:
     runs-on: ubuntu-latest

Cargo.toml

@@ -80,7 +80,7 @@ rstest_reuse = "0.7.0"
 shell-quote = "0.7.2"
 
 [workspace]
-members = ["crates/runner-shared", "crates/memtrack"]
+members = ["crates/runner-shared", "crates/memtrack", "crates/codspeed-bpf", "crates/exectrack"]
 
 [workspace.dependencies]
 anyhow = "1.0"

crates/codspeed-bpf/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "codspeed-bpf"
+version = "0.1.0"
+edition = "2024"
+repository = "https://github.com/CodSpeedHQ/runner"
+
+[features]
+# Enable build utilities for BPF compilation
+build = ["dep:libbpf-cargo", "dep:vmlinux", "dep:bindgen"]
+
+[dependencies]
+anyhow = { workspace = true }
+libbpf-rs = { version = "0.25.0", features = ["vendored"] }
+libc = { workspace = true }
+log = { workspace = true }
+paste = "1.0"
+
+# Build dependencies (optional, only needed for build feature)
+libbpf-cargo = { version = "0.25.0", optional = true }
+vmlinux = { git = "https://github.com/libbpf/vmlinux.h.git", rev = "991dd4b8dfd8c9d62ce8999521b24f61d9b7fc52", optional = true }
+bindgen = { version = "0.71", optional = true }

crates/codspeed-bpf/src/build.rs

@@ -0,0 +1,80 @@
+//! Build utilities for BPF programs using codspeed-bpf
+//!
+//! This module provides helpers for building BPF programs that use codspeed-bpf headers.
+//!
+//! To use this in your build.rs, add codspeed-bpf with the "build" feature to your
+//! [build-dependencies]:
+//!
+//! ```toml
+//! [build-dependencies]
+//! codspeed-bpf = { path = "../codspeed-bpf", features = ["build"] }
+//! ```
+//!
+//! Then in your build.rs:
+//!
+//! ```ignore
+//! fn main() {
+//!     codspeed_bpf::build::build_bpf("exectrack", "src/ebpf/c/exectrack.bpf.c");
+//!     codspeed_bpf::build::generate_bindings("wrapper.h");
+//! }
+//! ```
+
+use std::{env, path::PathBuf};
+
+/// Build a BPF program with codspeed-bpf headers available
+///
+/// # Arguments
+/// * `program_name` - The name of the BPF program (e.g., "exectrack", "memtrack")
+/// * `source_file` - The path to the BPF source file (e.g., "src/ebpf/c/exectrack.bpf.c")
+///
+/// This function will:
+/// 1. Compile the BPF program using libbpf-cargo
+/// 2. Include codspeed-bpf headers in the clang search path
+/// 3. Generate a skeleton into OUT_DIR/{program_name}.skel.rs
+pub fn build_bpf(program_name: &str, source_file: &str) {
+    use libbpf_cargo::SkeletonBuilder;
+
+    println!("cargo:rerun-if-changed=src/ebpf/c");
+
+    let arch = env::var("CARGO_CFG_TARGET_ARCH")
+        .expect("CARGO_CFG_TARGET_ARCH must be set in build script");
+
+    let output =
+        PathBuf::from(env::var("OUT_DIR").unwrap()).join(format!("{program_name}.skel.rs"));
+
+    // Get the path to codspeed-bpf's C headers
+    let codspeed_bpf_include = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap())
+        .parent()
+        .unwrap()
+        .join("codspeed-bpf/src/c");
+
+    SkeletonBuilder::new()
+        .source(source_file)
+        .clang_args([
+            "-I",
+            &vmlinux::include_path_root().join(arch).to_string_lossy(),
+            "-I",
+            &codspeed_bpf_include.to_string_lossy(),
+        ])
+        .build_and_generate(&output)
+        .unwrap_or_else(|_| panic!("Failed to build {program_name}.bpf.c"));
+}
+
+/// Generate Rust bindings for a C header file
+///
+/// # Arguments
+/// * `header_file` - The path to the header file (e.g., "wrapper.h")
+///
+/// This function will:
+/// 1. Use bindgen to generate Rust bindings
+/// 2. Write the output to OUT_DIR/event.rs
+pub fn generate_bindings(header_file: &str) {
+    let bindings = bindgen::Builder::default()
+        .header(header_file)
+        .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()))
+        .generate()
+        .expect("Unable to generate bindings");
+
+    let out_file = PathBuf::from(env::var("OUT_DIR").unwrap()).join("event.rs");
+    std::fs::write(&out_file, bindings.to_string()).expect("Couldn't write bindings!");
+}

crates/codspeed-bpf/src/c/codspeed/common.h

@@ -0,0 +1,27 @@
+#ifndef __CODSPEED_COMMON_H__
+#define __CODSPEED_COMMON_H__
+
+/* Macros for common map definitions */
+#define BPF_HASH_MAP(name, key_type, value_type, max_ents) \
+    struct {                                               \
+        __uint(type, BPF_MAP_TYPE_HASH);                   \
+        __uint(max_entries, max_ents);                     \
+        __type(key, key_type);                             \
+        __type(value, value_type);                         \
+    } name SEC(".maps")
+
+#define BPF_ARRAY_MAP(name, value_type, max_ents) \
+    struct {                                      \
+        __uint(type, BPF_MAP_TYPE_ARRAY);         \
+        __uint(max_entries, max_ents);            \
+        __type(key, __u32);                       \
+        __type(value, value_type);                \
+    } name SEC(".maps")
+
+#define BPF_RINGBUF(name, size)             \
+    struct {                                \
+        __uint(type, BPF_MAP_TYPE_RINGBUF); \
+        __uint(max_entries, size);          \
+    } name SEC(".maps")
+
+#endif /* __CODSPEED_COMMON_H__ */

crates/codspeed-bpf/src/c/codspeed/process_tracking.h

@@ -0,0 +1,52 @@
+#ifndef __CODSPEED_PROCESS_TRACKING_H__
+#define __CODSPEED_PROCESS_TRACKING_H__
+
+#include "common.h"
+
+/* Standard process tracking maps - define these in your BPF program
+ * This macro creates the map definitions. The helper functions below will
+ * reference these maps, so call this macro before including this header's helpers.
+ */
+#define PROCESS_TRACKING_MAPS()                        \
+    BPF_HASH_MAP(tracked_pids, __u32, __u8, 10000);   \
+    BPF_HASH_MAP(pids_ppid, __u32, __u32, 10000);     \
+    /* Helper functions defined below */                                          \
+    static __always_inline int is_tracked(__u32 pid);                            \
+    static __always_inline int handle_fork(__u32 parent_pid, __u32 child_pid);  \
+    static __always_inline int handle_exit(__u32 pid);                           \
+    /* is_tracked implementation */                                              \
+    static __always_inline int is_tracked(__u32 pid) {                           \
+        if (bpf_map_lookup_elem(&tracked_pids, &pid)) {                          \
+            return 1;                                                             \
+        }                                                                         \
+        _Pragma("unroll")                                                        \
+        for (int i = 0; i < 5; i++) {                                             \
+            __u32* ppid = bpf_map_lookup_elem(&pids_ppid, &pid);                 \
+            if (!ppid) {                                                          \
+                break;                                                            \
+            }                                                                     \
+            pid = *ppid;                                                          \
+            if (bpf_map_lookup_elem(&tracked_pids, &pid)) {                      \
+                return 1;                                                         \
+            }                                                                     \
+        }                                                                         \
+        return 0;                                                                 \
+    }                                                                             \
+    /* handle_fork implementation */                                             \
+    static __always_inline int handle_fork(__u32 parent_pid, __u32 child_pid) { \
+        if (is_tracked(parent_pid)) {                                             \
+            __u8 marker = 1;                                                      \
+            bpf_map_update_elem(&tracked_pids, &child_pid, &marker, BPF_ANY);    \
+            bpf_map_update_elem(&pids_ppid, &child_pid, &parent_pid, BPF_ANY);   \
+            return 1;                                                             \
+        }                                                                         \
+        return 0;                                                                 \
+    }                                                                             \
+    /* handle_exit implementation */                                             \
+    static __always_inline int handle_exit(__u32 pid) {                          \
+        bpf_map_delete_elem(&tracked_pids, &pid);                                \
+        bpf_map_delete_elem(&pids_ppid, &pid);                                   \
+        return 0;                                                                 \
+    }
+
+#endif /* __CODSPEED_PROCESS_TRACKING_H__ */

crates/codspeed-bpf/src/lib.rs

@@ -0,0 +1,17 @@
+//! Shared BPF utilities for CodSpeed tracing tools
+//!
+//! This crate provides common functionality for BPF-based tracing:
+//! - Process hierarchy tracking
+//! - Ring buffer polling
+//! - Probe attachment macros
+//! - Build utilities for BPF programs
+
+pub mod macros;
+pub mod poller;
+pub mod process_tracking;
+
+#[cfg(feature = "build")]
+pub mod build;
+
+pub use poller::{EventHandler, RingBufferPoller};
+pub use process_tracking::{ProcessTracking, bump_memlock_rlimit};