The average time elapsed and P99 elapsed time for a function in the analysis process

func_time.py

from __future__ import print_function
import os
from bcc import BPF
from time import sleep
import argparse
import ctypes

# Argument parsing
parser = argparse.ArgumentParser(description="Measure function duration for a specific PID")
parser.add_argument("pid", help="The PID of the process")
parser.add_argument("function_name", help="The function to trace")
parser.add_argument("interval", type=int, help="The interval in seconds for statistics")
args = parser.parse_args()

# Define BPF program
bpf_text = """
#include <uapi/linux/ptrace.h>

struct key_t {
    u64 pid;
    char comm[16];
};

struct data_t {
    u64 pid;
    u64 duration;
};

BPF_HASH(start, struct key_t);
BPF_PERF_OUTPUT(times);

int trace_start(struct pt_regs *ctx) {
    struct key_t key = {};
    u64 ts;

    key.pid = bpf_get_current_pid_tgid();
    bpf_get_current_comm(&(key.comm), sizeof(key.comm));

    ts = bpf_ktime_get_ns();
    start.update(&key, &ts);

    return 0;
}

int trace_end(struct pt_regs *ctx) {
    struct key_t key = {};
    struct data_t data = {};
    u64 *tsp, delta;

    key.pid = bpf_get_current_pid_tgid();
    bpf_get_current_comm(&(key.comm), sizeof(key.comm));

    tsp = start.lookup(&key);
    if (tsp != 0) {
        delta = bpf_ktime_get_ns() - *tsp;
        data.pid = key.pid;
        data.duration = delta;
        times.perf_submit(ctx, &data, sizeof(data));
        start.delete(&key);
    }

    return 0;
}
"""

# Load BPF program
b = BPF(text=bpf_text)

pid = int(args.pid)
mangled_name = args.function_name
interval = args.interval

binary_path = os.readlink(f"/proc/{pid}/exe")

b.attach_uprobe(name=binary_path, sym=mangled_name, fn_name="trace_start")
b.attach_uretprobe(name=binary_path, sym=mangled_name, fn_name="trace_end")

times = []

# Define output data structure
class Data(ctypes.Structure):
    _fields_ = [
        ("pid", ctypes.c_ulonglong),
        ("duration", ctypes.c_ulonglong)
    ]

# Define callback
def print_event(cpu, data, size):
    event = ctypes.cast(data, ctypes.POINTER(Data)).contents
    times.append(event.duration)

# Attach callback to BPF_PERF_OUTPUT
b["times"].open_perf_buffer(print_event)

print("Tracing... Hit Ctrl-C to end.")

try:
    while True:
        while b.perf_buffer_poll():
            pass
        sleep(interval)

        if times:
            average = sum(times) / len(times) / 1e6  # convert ns to ms

            times.sort()
            p99 = times[int(len(times) * 0.99)] / 1e6  # convert ns to ms
            print("Avg time: {} ms, P99: {} ms".format(average, p99))

        times = []

except KeyboardInterrupt:
    pass