bpf/trace.c - master - rustwizard/pgtrace

bpf/trace.c
161 строка4 KB
Rust Wizard
feat: top-N queries in summaries; trace openat/sendto/recvfrom syscalls
07 авг 2026, 14:18
07 авг 2026, 14:1858f6f7e
// pgtrace eBPF programs: trace enter/exit of selected syscalls for postgres
// processes and emit events for calls longer than the configured threshold.
//
// Attached to syscalls:sys_enter_*/sys_exit_* tracepoints, which are
// arch-independent (no CO-RE required).
#include <linux/bpf.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_tracing.h>

char LICENSE[] SEC("license") = "GPL";

#define COMM_LEN 16
#define MAX_ENTRIES 16384

// Traced syscalls. Keep in sync with syscallNames in internal/bpf/syscall.go.
enum syscall_id {
	SYS_ID_READ = 0,
	SYS_ID_PREAD64,
	SYS_ID_WRITE,
	SYS_ID_PWRITE64,
	SYS_ID_FSYNC,
	SYS_ID_FDATASYNC,
	SYS_ID_NANOSLEEP,
	SYS_ID_FUTEX,
	SYS_ID_OPENAT,
	SYS_ID_SENDTO,
	SYS_ID_RECVFROM,
	SYS_ID_MAX,
};

struct config {
	__u64 threshold_ns;
	char comm[COMM_LEN];
};

struct start_event {
	__u32 syscall_id;
	__u64 start_ts;
};

struct event {
	__u32 pid;
	__u32 syscall_id;
	__u64 duration_ns;
	__u64 ts;
};

struct {
	__uint(type, BPF_MAP_TYPE_ARRAY);
	__uint(max_entries, 1);
	__type(key, __u32);
	__type(value, struct config);
} config_map SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_HASH);
	__uint(max_entries, MAX_ENTRIES);
	__type(key, __u32); // pid
	__type(value, struct start_event);
} start_events SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_RINGBUF);
	__uint(max_entries, 256 * 1024);
} events SEC(".maps");

// Force struct event into BTF so bpf2go can generate the Go type.
const struct event *_event __attribute__((unused));

static __always_inline const struct config *get_config(void)
{
	__u32 key = 0;
	return bpf_map_lookup_elem(&config_map, &key);
}

static __always_inline int comm_matches(const char *comm, const char *target)
{
#pragma unroll
	for (int i = 0; i < COMM_LEN; i++) {
		if (comm[i] != target[i])
			return 0;
		if (comm[i] == '\0')
			return 1;
	}
	return 1;
}

static __always_inline int on_enter(__u32 syscall_id)
{
	const struct config *cfg = get_config();
	if (!cfg)
		return 0;

	char comm[COMM_LEN];
	if (bpf_get_current_comm(comm, sizeof(comm)) != 0)
		return 0;
	if (!comm_matches(comm, cfg->comm))
		return 0;

	struct start_event start = {
		.syscall_id = syscall_id,
		.start_ts = bpf_ktime_get_ns(),
	};
	__u32 pid = (__u32)bpf_get_current_pid_tgid();
	bpf_map_update_elem(&start_events, &pid, &start, BPF_ANY);
	return 0;
}

static __always_inline int on_exit(void)
{
	const struct config *cfg = get_config();
	if (!cfg)
		return 0;

	__u32 pid = (__u32)bpf_get_current_pid_tgid();
	struct start_event *start = bpf_map_lookup_elem(&start_events, &pid);
	if (!start)
		return 0;

	__u64 now = bpf_ktime_get_ns();
	__u64 duration = now - start->start_ts;
	__u32 syscall_id = start->syscall_id;
	bpf_map_delete_elem(&start_events, &pid);

	if (duration <= cfg->threshold_ns)
		return 0;

	struct event *ev = bpf_ringbuf_reserve(&events, sizeof(*ev), 0);
	if (!ev)
		return 0;
	ev->pid = pid;
	ev->syscall_id = syscall_id;
	ev->duration_ns = duration;
	ev->ts = now;
	bpf_ringbuf_submit(ev, 0);
	return 0;
}

#define TRACE_SYSCALL(name, id)                                            \
	SEC("tp/syscalls/sys_enter_" #name)                                \
	int trace_enter_##name(void *ctx)                                  \
	{                                                                  \
		return on_enter(id);                                       \
	}                                                                  \
	SEC("tp/syscalls/sys_exit_" #name)                                 \
	int trace_exit_##name(void *ctx)                                   \
	{                                                                  \
		return on_exit();                                          \
	}

TRACE_SYSCALL(read, SYS_ID_READ)
TRACE_SYSCALL(pread64, SYS_ID_PREAD64)
TRACE_SYSCALL(write, SYS_ID_WRITE)
TRACE_SYSCALL(pwrite64, SYS_ID_PWRITE64)
TRACE_SYSCALL(fsync, SYS_ID_FSYNC)
TRACE_SYSCALL(fdatasync, SYS_ID_FDATASYNC)
TRACE_SYSCALL(nanosleep, SYS_ID_NANOSLEEP)
TRACE_SYSCALL(futex, SYS_ID_FUTEX)
TRACE_SYSCALL(openat, SYS_ID_OPENAT)
TRACE_SYSCALL(sendto, SYS_ID_SENDTO)
TRACE_SYSCALL(recvfrom, SYS_ID_RECVFROM)