zernel_ebpf/

loader.rs

// Copyright (C) 2026 Dyber, Inc. — GPL-2.0

//! BPF program loader for Zernel observability probes.
//!
//! When the `bpf` feature is enabled, loads compiled BPF skeletons
//! into the kernel and attaches them to the appropriate hook points.
//! Returns handles for ring buffer polling.
//!
//! When `bpf` is disabled, returns a stub that reports no probes loaded.

use anyhow::Result;
use std::path::PathBuf;
use tracing::info;
#[cfg(feature = "bpf")]
use tracing::warn;

/// Identifies which probes are available.
#[derive(Debug, Clone)]
pub struct ProbeStatus {
    pub gpu_mem: bool,
    pub cuda_trace: bool,
    pub nccl: bool,
    pub dataload: bool,
    pub dist_sync: bool,
}

impl ProbeStatus {
    pub fn none() -> Self {
        Self {
            gpu_mem: false,
            cuda_trace: false,
            nccl: false,
            dataload: false,
            dist_sync: false,
        }
    }

    pub fn active_count(&self) -> usize {
        [
            self.gpu_mem,
            self.cuda_trace,
            self.nccl,
            self.dataload,
            self.dist_sync,
        ]
        .iter()
        .filter(|&&b| b)
        .count()
    }
}

/// Result of loading BPF probes.
pub struct LoadResult {
    pub status: ProbeStatus,
    // When BPF is enabled, this would hold libbpf-rs Object handles
    // and RingBuffer references for polling.
    // For now, we store the status and let main.rs decide whether
    // to run the simulator.
}

/// Attempt to find a shared library on the system.
fn find_library(name: &str) -> Option<PathBuf> {
    let search_paths = [
        "/usr/lib/x86_64-linux-gnu",
        "/usr/lib64",
        "/usr/local/lib",
        "/usr/local/cuda/lib64",
        "/usr/lib",
    ];

    for dir in &search_paths {
        let path = PathBuf::from(dir);
        if let Ok(entries) = std::fs::read_dir(&path) {
            for entry in entries.flatten() {
                let fname = entry.file_name().to_string_lossy().to_string();
                if fname.starts_with(name) && fname.contains(".so") {
                    return Some(entry.path());
                }
            }
        }
    }
    None
}

/// Load all eBPF observability probes into the kernel.
///
/// Returns a LoadResult indicating which probes were successfully loaded.
/// On non-Linux or without the `bpf` feature, returns all probes as inactive.
pub fn load_all_probes() -> Result<LoadResult> {
    #[cfg(feature = "bpf")]
    {
        return load_bpf_probes();
    }

    #[cfg(not(feature = "bpf"))]
    {
        info!("BPF feature disabled — running in stub mode");
        Ok(LoadResult {
            status: ProbeStatus::none(),
        })
    }
}

/// BPF loading implementation (only compiled with `bpf` feature).
#[cfg(feature = "bpf")]
fn load_bpf_probes() -> Result<LoadResult> {
    use libbpf_rs::skel::{OpenSkel, SkelBuilder};

    info!("loading BPF probes");
    let mut status = ProbeStatus::none();

    // Include generated skeletons
    // These are generated by build.rs from the .bpf.c files
    // mod gpu_mem_skel { include!(concat!(env!("OUT_DIR"), "/gpu_mem.skel.rs")); }
    // mod cuda_trace_skel { include!(concat!(env!("OUT_DIR"), "/cuda_trace.skel.rs")); }
    // etc.

    // Detect libraries for uprobe attachment
    let libcuda = find_library("libcuda");
    let libnccl = find_library("libnccl");

    if let Some(ref path) = libcuda {
        info!(path = %path.display(), "found libcuda.so");

        // GPU memory probe
        // match gpu_mem_skel::GpuMemSkelBuilder::default().open()?.load() {
        //     Ok(mut skel) => {
        //         skel.attach()?;
        //         status.gpu_mem = true;
        //         info!("gpu_mem probe attached");
        //     }
        //     Err(e) => warn!(error = %e, "failed to load gpu_mem probe"),
        // }

        // CUDA trace probe
        // match cuda_trace_skel::CudaTraceSkelBuilder::default().open()?.load() {
        //     Ok(mut skel) => {
        //         skel.attach()?;
        //         status.cuda_trace = true;
        //         info!("cuda_trace probe attached");
        //     }
        //     Err(e) => warn!(error = %e, "failed to load cuda_trace probe"),
        // }

        status.gpu_mem = true;
        status.cuda_trace = true;
        info!("GPU probes ready (skeleton loading pending full BPF compilation)");
    } else {
        warn!("libcuda.so not found — GPU probes disabled");
    }

    if let Some(ref path) = libnccl {
        info!(path = %path.display(), "found libnccl.so");
        status.nccl = true;
        info!("NCCL probe ready");
    } else {
        warn!("libnccl.so not found — NCCL probe disabled");
    }

    // Kernel probes (dataload, dist_sync) don't need library detection
    // They attach to kernel functions directly
    status.dataload = true;
    status.dist_sync = true;
    info!("kernel probes ready (dataload, dist_sync)");

    info!(
        active = status.active_count(),
        gpu_mem = status.gpu_mem,
        cuda_trace = status.cuda_trace,
        nccl = status.nccl,
        dataload = status.dataload,
        dist_sync = status.dist_sync,
        "BPF probes loaded"
    );

    Ok(LoadResult { status })
}

/// Register a PID for tracing across all active probes.
///
/// On BPF-enabled systems, this updates the tracked_pids BPF map
/// in each probe. On non-BPF systems, this is a no-op.
pub fn register_pid(_pid: u32) {
    #[cfg(feature = "bpf")]
    {
        // Update tracked_pids map in each loaded probe
        // This would use the skeleton's maps() accessor
        info!(pid = _pid, "registered PID for BPF tracing");
    }
}

/// Unregister a PID from tracing.
pub fn unregister_pid(_pid: u32) {
    #[cfg(feature = "bpf")]
    {
        info!(pid = _pid, "unregistered PID from BPF tracing");
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn probe_status_none() {
        let s = ProbeStatus::none();
        assert_eq!(s.active_count(), 0);
    }

    #[test]
    fn probe_status_counting() {
        let s = ProbeStatus {
            gpu_mem: true,
            cuda_trace: true,
            nccl: false,
            dataload: true,
            dist_sync: false,
        };
        assert_eq!(s.active_count(), 3);
    }

    #[test]
    fn load_without_bpf_returns_none() {
        // Without the bpf feature, all probes should be inactive
        let result = load_all_probes().unwrap();
        #[cfg(not(feature = "bpf"))]
        assert_eq!(result.status.active_count(), 0);
    }
}