profile_bee/

ebpf.rs

use anyhow::anyhow;
use aya::maps::{Array, HashMap, MapData, RingBuf, StackTraceMap};
use aya::programs::{
    perf_event::{PerfEventScope, PerfTypeId, SamplePolicy},
    KProbe, PerfEvent,
};
use aya::programs::{RawTracePoint, TracePoint, UProbe};
use aya::{include_bytes_aligned, util::online_cpus};
use aya::{Btf, Ebpf, EbpfLoader};

use aya::Pod;
use profile_bee_common::{ProcInfo, ProcInfoKey, UnwindEntry};

// Create a newtype wrapper around StackInfo
#[repr(transparent)]
#[derive(Debug, Clone, Copy)]
pub struct StackInfoPod(pub profile_bee_common::StackInfo);
unsafe impl Pod for StackInfoPod {}

#[repr(transparent)]
#[derive(Debug, Clone, Copy)]
pub struct FramePointersPod(pub profile_bee_common::FramePointers);
unsafe impl Pod for FramePointersPod {}

#[repr(transparent)]
#[derive(Debug, Clone, Copy)]
pub struct UnwindEntryPod(pub UnwindEntry);
unsafe impl Pod for UnwindEntryPod {}

#[repr(transparent)]
#[derive(Debug, Clone, Copy)]
pub struct ProcInfoKeyPod(pub ProcInfoKey);
unsafe impl Pod for ProcInfoKeyPod {}

#[repr(transparent)]
#[derive(Clone, Copy)]
pub struct ProcInfoPod(pub ProcInfo);
unsafe impl Pod for ProcInfoPod {}

/// Wrapper for eBPF stuff
#[derive(Debug)]
pub struct EbpfProfiler {
    pub bpf: Ebpf,
    /// eBPF stacktrace map
    pub stack_traces: StackTraceMap<MapData>,
    /// Count stacktraces
    pub counts: HashMap<MapData, StackInfoPod, u64>,
    /// Custom storage of StackTrace IPs
    pub stacked_pointers: HashMap<MapData, StackInfoPod, FramePointersPod>,
}
/// Legacy configuration for uprobe attachment (single target).
/// Kept for backward compatibility; prefer `ResolvedProbe` for new code.
#[derive(Debug, Clone)]
pub struct UProbeConfig {
    /// Function name (with optional +offset)
    pub function: String,
    /// Path to binary/library (absolute or library name)
    pub path: String,
    /// Whether this is a return probe
    pub is_retprobe: bool,
    /// Optional PID to attach to (None = all processes)
    pub pid: Option<i32>,
}

/// Configuration for the new smart uprobe system.
/// Holds one or more resolved probe targets for multi-attach.
#[derive(Debug, Clone)]
pub struct SmartUProbeConfig {
    /// Resolved probe targets to attach to.
    pub probes: Vec<crate::probe_resolver::ResolvedProbe>,
    /// Optional PID to attach to (None = all processes).
    pub pid: Option<i32>,
}

pub struct ProfilerConfig {
    pub skip_idle: bool,
    pub stream_mode: u8,
    pub frequency: u64,
    pub kprobe: Option<String>,
    /// Legacy single-target uprobe config (backward compat).
    pub uprobe: Option<UProbeConfig>,
    /// New smart uprobe config with multi-attach support.
    pub smart_uprobe: Option<SmartUProbeConfig>,
    pub tracepoint: Option<String>,
    /// Raw tracepoint name for syscall raw_tp (e.g., "sys_enter").
    /// Uses collect_trace_raw_syscall which reads pt_regs from args[0].
    pub raw_tracepoint: Option<String>,
    /// Raw tracepoint name for generic raw_tp with task pt_regs (e.g., "sched_switch").
    /// Uses bpf_get_current_task_btf() + bpf_task_pt_regs() for full FP/DWARF unwinding.
    /// Requires kernel >= 5.15; falls back to raw_tracepoint_generic on older kernels.
    pub raw_tracepoint_task_regs: Option<String>,
    /// Raw tracepoint name for generic raw_tp (e.g., "sched_switch").
    /// Uses collect_trace_raw_tp_generic which relies on bpf_get_stackid only.
    pub raw_tracepoint_generic: Option<String>,
    /// Target syscall number for raw tracepoint filtering (-1 = all).
    pub target_syscall_nr: i64,
    pub pid: Option<u32>,
    pub cpu: Option<u32>,
    pub self_profile: bool,
    pub dwarf: bool,
}

/// Creates an aya Ebpf object
pub fn load_ebpf(config: &ProfilerConfig) -> Result<Ebpf, anyhow::Error> {
    // The eBPF object file is selected by build.rs: it uses a freshly-built
    // binary from `cargo xtask build-ebpf` if available, otherwise the prebuilt
    // binary shipped in ebpf-bin/. Either way it ends up in OUT_DIR.
    let data = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/profile-bee.bpf.o"));

    let skip_idle = if config.skip_idle { 1u8 } else { 0u8 };
    let dwarf_enabled = if config.dwarf { 1u8 } else { 0u8 };
    let target_syscall_nr: i64 = config.target_syscall_nr;

    let bpf = EbpfLoader::new()
        .set_global("SKIP_IDLE", &skip_idle, true)
        .set_global("NOTIFY_TYPE", &config.stream_mode, true)
        .set_global("DWARF_ENABLED", &dwarf_enabled, true)
        .set_global("TARGET_SYSCALL_NR", &target_syscall_nr, true)
        .btf(Btf::from_sys_fs().ok().as_ref())
        .load(data)
        .map_err(|e| {
            println!("{:?}", e);
            e
        })?;

    // // this might be useful for debugging, but definitely disable bpf logging for performance purposes
    // aya_log::EbpfLogger::init(&mut bpf)?;

    Ok(bpf)
}

pub fn setup_ebpf_profiler(config: &ProfilerConfig) -> Result<EbpfProfiler, anyhow::Error> {
    let mut bpf = load_ebpf(config)?;

    if let Some(kprobe) = &config.kprobe {
        let program: &mut KProbe = bpf.program_mut("kprobe_profile").unwrap().try_into()?;
        program.load()?;
        program.attach(kprobe, 0)?;
    } else if let Some(smart) = &config.smart_uprobe {
        // New smart uprobe path: attach to all resolved probe targets
        if smart.probes.is_empty() {
            return Err(anyhow!("no probe targets resolved — nothing to attach"));
        }

        // Separate probes by type (uprobe vs uretprobe)
        let has_uprobe = smart.probes.iter().any(|p| !p.is_ret);
        let has_uretprobe = smart.probes.iter().any(|p| p.is_ret);

        if has_uprobe {
            let program: &mut UProbe = bpf.program_mut("uprobe_profile").unwrap().try_into()?;
            program.load()?;

            for probe in smart.probes.iter().filter(|p| !p.is_ret) {
                let display_name = probe.demangled.as_deref().unwrap_or(&probe.symbol_name);
                eprintln!(
                    "  attaching uprobe: {}:{} (0x{:x})",
                    probe.library_path.display(),
                    display_name,
                    probe.address,
                );
                program.attach(
                    Some(probe.symbol_name.as_str()),
                    probe.offset,
                    probe.library_path.to_str().ok_or_else(|| {
                        anyhow!("non-UTF8 library path: {}", probe.library_path.display())
                    })?,
                    smart.pid,
                )?;
            }
        }

        if has_uretprobe {
            let program: &mut UProbe = bpf.program_mut("uretprobe_profile").unwrap().try_into()?;
            program.load()?;

            for probe in smart.probes.iter().filter(|p| p.is_ret) {
                let display_name = probe.demangled.as_deref().unwrap_or(&probe.symbol_name);
                eprintln!(
                    "  attaching uretprobe: {}:{} (0x{:x})",
                    probe.library_path.display(),
                    display_name,
                    probe.address,
                );
                program.attach(
                    Some(probe.symbol_name.as_str()),
                    probe.offset,
                    probe.library_path.to_str().ok_or_else(|| {
                        anyhow!("non-UTF8 library path: {}", probe.library_path.display())
                    })?,
                    smart.pid,
                )?;
            }
        }
    } else if let Some(uprobe_config) = &config.uprobe {
        // Choose the right program based on is_retprobe flag
        let program_name = if uprobe_config.is_retprobe {
            "uretprobe_profile"
        } else {
            "uprobe_profile"
        };

        let program: &mut UProbe = bpf.program_mut(program_name).unwrap().try_into()?;
        program.load()?;

        // Parse function name and offset (format: "function_name" or "function_name+offset")
        let (fn_name, offset) = if let Some(plus_pos) = uprobe_config.function.find('+') {
            let (name, offset_str) = uprobe_config.function.split_at(plus_pos);
            let offset = offset_str[1..]
                .parse::<u64>()
                .map_err(|e| anyhow::anyhow!("Invalid offset in uprobe function: {}", e))?;
            (Some(name), offset)
        } else {
            (Some(uprobe_config.function.as_str()), 0)
        };

        program.attach(
            fn_name,
            offset,
            uprobe_config.path.as_str(),
            uprobe_config.pid,
        )?;
    } else if let Some(raw_tp) = &config.raw_tracepoint {
        // Pick the correct syscall raw_tp program based on enter vs exit
        let prog_name = if raw_tp == "sys_exit" {
            "raw_tp_sys_exit"
        } else {
            "raw_tp_sys_enter"
        };
        let program: &mut RawTracePoint = bpf.program_mut(prog_name).unwrap().try_into()?;
        program.load()?;
        program.attach(raw_tp)?;
    } else if let Some(raw_tp) = &config.raw_tracepoint_task_regs {
        let program: &mut RawTracePoint =
            bpf.program_mut("raw_tp_with_regs").unwrap().try_into()?;
        program.load()?;
        program.attach(raw_tp)?;
    } else if let Some(raw_tp) = &config.raw_tracepoint_generic {
        let program: &mut RawTracePoint = bpf.program_mut("raw_tp_generic").unwrap().try_into()?;
        program.load()?;
        program.attach(raw_tp)?;
    } else if let Some(tracepoint) = &config.tracepoint {
        let program: &mut TracePoint = bpf.program_mut("tracepoint_profile").unwrap().try_into()?;
        program.load()?;

        let mut split = tracepoint.split(':');
        let category = split.next().expect("category");
        let name = split.next().expect("name");

        program.attach(category, name)?;
    } else {
        let program: &mut PerfEvent = bpf.program_mut("profile_cpu").unwrap().try_into()?;

        program.load()?;

        // https://elixir.bootlin.com/linux/v4.2/source/include/uapi/linux/perf_event.h#L103
        const PERF_COUNT_SW_CPU_CLOCK: u64 = 0;

        // could change this to Hardware if your system supports
        // `lscpu | grep -i pmu`
        let perf_type = PerfTypeId::Software;

        if config.self_profile {
            program.attach(
                perf_type,
                PERF_COUNT_SW_CPU_CLOCK,
                PerfEventScope::CallingProcessAnyCpu,
                SamplePolicy::Frequency(config.frequency),
                true,
            )?;
        } else if config.pid.is_some() || config.cpu.is_some() {
            // When filtering by PID or CPU, attach to all/specific CPUs
            // and let eBPF filter by tgid. This allows profiling child processes.
            let cpus = if let Some(cpu) = config.cpu {
                vec![cpu]
            } else {
                online_cpus().map_err(|(_, error)| error)?
            };

            let nprocs = cpus.len();
            if let Some(pid) = config.pid {
                eprintln!(
                    "Profiling PID {} and child processes across {} CPUs",
                    pid, nprocs
                );
            } else if let Some(cpu) = config.cpu {
                eprintln!("Profiling CPU {}", cpu);
            }

            for cpu in cpus {
                program.attach(
                    perf_type.clone(),
                    PERF_COUNT_SW_CPU_CLOCK,
                    PerfEventScope::AllProcessesOneCpu { cpu },
                    SamplePolicy::Frequency(config.frequency),
                    true,
                )?;
            }
        } else {
            let cpus = online_cpus().map_err(|(_, error)| error)?;
            let nprocs = cpus.len();
            eprintln!("CPUs: {}", nprocs);

            for cpu in cpus {
                program.attach(
                    perf_type.clone(),
                    PERF_COUNT_SW_CPU_CLOCK,
                    PerfEventScope::AllProcessesOneCpu { cpu },
                    SamplePolicy::Frequency(config.frequency),
                    true,
                )?;
            }
        }
    }

    let stack_traces = StackTraceMap::try_from(
        bpf.take_map("stack_traces")
            .ok_or(anyhow!("stack_traces not found"))?,
    )?;

    let counts = bpf
        .take_map("counts")
        .ok_or(anyhow!("counts not found"))?
        .try_into()?;

    let stacked_pointers = bpf
        .take_map("stacked_pointers")
        .ok_or(anyhow!("stacked_pointers not found"))?
        .try_into()?;

    Ok(EbpfProfiler {
        bpf,
        stack_traces,
        counts,
        stacked_pointers,
    })
}

pub fn setup_ring_buffer(bpf: &mut Ebpf) -> Result<RingBuf<MapData>, anyhow::Error> {
    let ring_buf = RingBuf::try_from(
        bpf.take_map("RING_BUF_STACKS")
            .ok_or(anyhow!("RING_BUF_STACKS not found"))?,
    )?;

    Ok(ring_buf)
}

impl EbpfProfiler {
    /// Set the target PID for eBPF filtering (0 = profile all processes)
    pub fn set_target_pid(&mut self, pid: u32) -> Result<(), anyhow::Error> {
        let mut target_pid_map: Array<&mut MapData, u32> = Array::try_from(
            self.bpf
                .map_mut("target_pid_map")
                .ok_or(anyhow!("target_pid_map not found"))?,
        )?;
        target_pid_map.set(0, pid, 0)?;
        Ok(())
    }

    /// Load DWARF unwind tables into eBPF maps for a process
    pub fn load_dwarf_unwind_tables(
        &mut self,
        manager: &crate::dwarf_unwind::DwarfUnwindManager,
    ) -> Result<(), anyhow::Error> {
        // Load all shards
        let all_shard_ids: Vec<u8> = manager.binary_tables.keys().copied().collect();
        self.update_dwarf_tables(manager, &all_shard_ids)
    }

    /// Load a single shard's unwind entries into its eBPF Array map.
    fn load_shard(&mut self, shard_id: u8, entries: &[UnwindEntry]) -> Result<(), anyhow::Error> {
        let map_name = format!("shard_{}", shard_id);
        let mut arr: Array<&mut MapData, UnwindEntryPod> = Array::try_from(
            self.bpf
                .map_mut(&map_name)
                .ok_or_else(|| anyhow!("{} map not found", map_name))?,
        )?;
        for (idx, entry) in entries.iter().enumerate() {
            arr.set(idx as u32, UnwindEntryPod(*entry), 0)?;
        }
        Ok(())
    }

    /// Incrementally update eBPF maps with new unwind shard entries,
    /// and refresh all proc_info entries.
    pub fn update_dwarf_tables(
        &mut self,
        manager: &crate::dwarf_unwind::DwarfUnwindManager,
        new_shard_ids: &[u8],
    ) -> Result<(), anyhow::Error> {
        if !new_shard_ids.is_empty() {
            let mut total_entries = 0usize;
            for &shard_id in new_shard_ids {
                if let Some(entries) = manager.binary_tables.get(&shard_id) {
                    self.load_shard(shard_id, entries)?;
                    total_entries += entries.len();
                    tracing::info!("Loaded shard {} with {} entries", shard_id, entries.len());
                }
            }

            tracing::info!(
                "Loaded {} total unwind entries across {} shards",
                total_entries,
                new_shard_ids.len()
            );
        }

        let mut proc_info_map: HashMap<&mut MapData, ProcInfoKeyPod, ProcInfoPod> =
            HashMap::try_from(
                self.bpf
                    .map_mut("proc_info")
                    .ok_or(anyhow!("proc_info map not found"))?,
            )?;

        for (&tgid, proc_info) in &manager.proc_info {
            let key = ProcInfoKeyPod(ProcInfoKey { tgid, _pad: 0 });
            let value = ProcInfoPod(*proc_info);
            proc_info_map.insert(key, value, 0)?;

            tracing::info!(
                "Loaded process info for tgid {} ({} mappings)",
                tgid,
                proc_info.mapping_count,
            );
        }

        Ok(())
    }
}