use alloc::{borrow::Cow, boxed::Box, sync::Arc};
use core::any::Any;
use ax_errno::{AxError, AxResult};
use axpoll::Pollable;
use kbpf_basic::raw_tracepoint::BpfRawTracePointArg;
use ktracepoint::{RawTraceEventFunc, TraceCallbackType};
use crate::{
file::{FileLike, add_file_like, get_file_like},
perf::bpf::OwnedEbpfVm,
tracepoint::{KernelExtTracePoint, find_ext_tracepoint_by_name},
};
type RawTpCallback = Box<dyn Fn(&[u64], &(dyn Any + Send + Sync)) + Send + Sync>;
pub struct RawTracepointPerfEvent {
ext_tp: KernelExtTracePoint,
callback: Arc<RawTraceEventFunc>,
}
impl Pollable for RawTracepointPerfEvent {
fn poll(&self) -> axpoll::IoEvents {
axpoll::IoEvents::empty()
}
fn register(&self, _context: &mut core::task::Context<'_>, _events: axpoll::IoEvents) {
}
}
impl FileLike for RawTracepointPerfEvent {
fn read(&self, _dst: &mut crate::file::IoDst) -> AxResult<usize> {
Err(AxError::Unsupported)
}
fn write(&self, _src: &mut crate::file::IoSrc) -> AxResult<usize> {
Err(AxError::Unsupported)
}
fn stat(&self) -> AxResult<crate::file::Kstat> {
Ok(crate::file::Kstat::default())
}
fn path(&self) -> Cow<'_, str> {
"anon_inode:[raw_tracepoint_perf_event]".into()
}
}
impl Drop for RawTracepointPerfEvent {
fn drop(&mut self) {
self.ext_tp
.lock()
.unregister(TraceCallbackType::RawEvent(self.callback.clone()));
}
}
impl RawTracepointPerfEvent {
pub fn new(ext_tp: KernelExtTracePoint, bpf_prog: Arc<dyn FileLike>) -> AxResult<Self> {
struct Ctx {
vm: OwnedEbpfVm,
}
let ctx = Box::new(Ctx {
vm: OwnedEbpfVm::new(bpf_prog)?,
});
let func: RawTpCallback = Box::new(|args: &[u64], data: &(dyn Any + Send + Sync)| {
let ctx = data
.downcast_ref::<Ctx>()
.expect("raw_tracepoint Ctx mismatch");
let arg_bytes = unsafe {
core::slice::from_raw_parts_mut(
args.as_ptr() as *mut u8,
core::mem::size_of_val(args),
)
};
if let Err(e) = ctx.vm.execute_program(arg_bytes) {
error!("raw_tracepoint BPF program failed: {e:?}");
}
});
let callback = Arc::new(RawTraceEventFunc::new(func, ctx));
ext_tp
.lock()
.register(TraceCallbackType::RawEvent(callback.clone()));
Ok(Self { ext_tp, callback })
}
}
pub fn bpf_raw_tracepoint_open(arg: BpfRawTracePointArg) -> AxResult<isize> {
let ext_tp = find_ext_tracepoint_by_name(&arg.name).ok_or(AxError::InvalidInput)?;
let prog = get_file_like(arg.prog_fd as _)?;
let event = RawTracepointPerfEvent::new(ext_tp, prog)?;
let fd = add_file_like(Arc::new(event), true)?;
Ok(fd as isize)
}