bpf script for block read into shared_buffer

blk_read_pgsql.py

#!/usr/bin/python

from bcc import BPF, USDT
import argparse
import re
import ctypes as ct
import subprocess

examples = """examples:
    blk_read_pgsql -p 188 # trace specific PostgreSQL processes
"""
parser = argparse.ArgumentParser(
    description="",
    formatter_class=argparse.RawDescriptionHelpFormatter,
    epilog=examples)
parser.add_argument("-p", "--pid", type=int, nargs='*',
    dest="pids", metavar="PID", help="the pid(s) to trace")
args = parser.parse_args()

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

struct data_t {
    u64 timestamp;
    int pid;
    int oid;
    int blkno;
    int requested_size;
    int act_size;
};

BPF_PERF_OUTPUT(events);

int query_blk_read(struct pt_regs *ctx) {
    struct data_t data = {};
    data.timestamp = bpf_ktime_get_ns();
    bpf_usdt_readarg(5, ctx, &data.oid);
    bpf_usdt_readarg(2, ctx, &data.blkno);
    bpf_usdt_readarg(8, ctx, &data.requested_size);
    bpf_usdt_readarg(7, ctx, &data.act_size);
    data.pid = bpf_get_current_pid_tgid();
    events.perf_submit(ctx, &data, sizeof(data));
    return 0;
}
"""

usdts = map(lambda pid: USDT(pid=pid), args.pids)
for usdt in usdts:
    usdt.enable_probe("smgr__md__read__done", "query_blk_read")

bpf = BPF(text=program, usdt_contexts=usdts)

class Data(ct.Structure):
    _fields_ = [
        ("timestamp", ct.c_ulonglong),
        ("pid", ct.c_int),
        ("oid", ct.c_int),
        ("blkno", ct.c_int),
        ("requested_size", ct.c_int),
        ("act_size", ct.c_int),
    ]

start = BPF.monotonic_time()

def print_event(cpu, data, size):
    event = ct.cast(data, ct.POINTER(Data)).contents
    print("%-14.6f %-6d %-6d %-6d %-6d %-6d" % (
        float(event.timestamp - start) / 1000000000,
        event.pid,
        event.oid,
        event.blkno,
        event.requested_size,
        event.act_size))

print("%-14s %-6s %-6s %-6s %-6s %-6s" % ("TIME(s)", "PID", "OID", "BLKNO", "REQ", "ACT"))

bpf["events"].open_perf_buffer(print_event, page_cnt=64)
while True:
    bpf.perf_buffer_poll()