bpf/offcpu.c - master - rustwizard/pgtrace

bpf/offcpu.c
146 строк4 KB
Rust Wizard
fix(offcpu): capture stack at block time via kprobe:schedule
08 авг 2026, 10:55
08 авг 2026, 10:55d84b1f0
// pgtrace off-CPU module: trace time postgres backends spend off-CPU
// (blocked on locks, I/O, network).
//
// kprobe:schedule fires when the current task is about to block, so its stack
// shows the reason (futex_wait, io_schedule, do_nanosleep, ...). We record
// {ts, stack} for the task. When sched:sched_switch sees the task coming back,
// we emit an event with the off-CPU duration and the stack captured at block
// time. Wait-type classification happens in user space (Go) on the symbols.
#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
#define MAX_STACK_DEPTH 8

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

struct off_cpu_start {
	__u64 start_ts;
	// Kernel stack captured when the task blocked, leaf first, zero-padded.
	__u64 stack[MAX_STACK_DEPTH];
};

struct offcpu_event {
	__u32 pid;
	__u32 wait_type; // filled by user space, kept for BTF shape
	__u64 duration_ns;
	__u64 ts;
	__u64 stack[MAX_STACK_DEPTH];
};

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

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

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

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

static __always_inline const struct offcpu_config *get_offcpu_config(void)
{
	__u32 key = 0;
	return bpf_map_lookup_elem(&offcpu_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;
}

// Fires when the current task calls schedule(): it is about to yield the CPU.
// The current task's stack is still intact here, so we capture the kernel
// stack that shows WHY it is blocking. The comm filter keeps this cheap for
// non-postgres processes.
SEC("kprobe/schedule")
int BPF_KPROBE(trace_schedule)
{
	const struct offcpu_config *cfg = get_offcpu_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;

	__u32 pid = (__u32)bpf_get_current_pid_tgid();

	// Only record the first block in a chain of nested schedule() calls, so
	// start_ts is not pushed forward by inner schedules during one sleep.
	struct off_cpu_start *existing = bpf_map_lookup_elem(&off_cpu_start, &pid);
	if (existing)
		return 0;

	struct off_cpu_start start = {
		.start_ts = bpf_ktime_get_ns(),
	};
	bpf_get_stack(ctx, start.stack, sizeof(start.stack), 0);
	bpf_map_update_elem(&off_cpu_start, &pid, &start, BPF_ANY);
	return 0;
}

// Fires when the scheduler switches tasks. If the incoming task was recorded
// as blocked by kprobe:schedule, emit the off-CPU event with the block stack.
SEC("tp/sched/sched_switch")
int trace_sched_switch(void *ctx)
{
	const struct offcpu_config *cfg = get_offcpu_config();
	if (!cfg)
		return 0;

	// Tracepoint layout: 8-byte trace_entry, prev_comm[16], prev_pid(4),
	// prev_prio(4), prev_state(8), next_comm[16], next_pid(4), next_prio(4).
	__u32 next_pid = 0;
	bpf_probe_read_kernel(&next_pid, sizeof(next_pid), (void *)ctx + 56);

	struct off_cpu_start *start = bpf_map_lookup_elem(&off_cpu_start, &next_pid);
	if (!start)
		return 0;
	bpf_map_delete_elem(&off_cpu_start, &next_pid);

	__u64 now = bpf_ktime_get_ns();
	__u64 duration = now - start->start_ts;
	if (duration <= cfg->threshold_ns)
		return 0;

	struct offcpu_event *ev = bpf_ringbuf_reserve(&offcpu_events, sizeof(*ev), 0);
	if (!ev)
		return 0;
	ev->pid = next_pid;
	ev->wait_type = 0;
	ev->duration_ns = duration;
	ev->ts = now;
	__builtin_memcpy(ev->stack, start->stack, sizeof(ev->stack));

	bpf_ringbuf_submit(ev, 0);
	return 0;
}