lb_lat.py

#!/usr/bin/python
# @lint-avoid-python-3-compatibility-imports

from __future__ import print_function
from bcc import BPF
from time import sleep, strftime
import argparse

examples = """examples:
"""

parser = argparse.ArgumentParser(
    description="Summarize load balance time per task as a histogram.",
    formatter_class=argparse.RawDescriptionHelpFormatter,
    epilog=examples)
parser.add_argument("-T", "--timestamp", action="store_true",
    help="include timestamp on output")
parser.add_argument("-m", "--milliseconds", action="store_true",
    help="millisecond histogram")
parser.add_argument("-P", "--pids", action="store_true",
    help="print a histogram per process ID")
parser.add_argument("-L", "--tids", action="store_true",
    help="print a histogram per thread ID")
parser.add_argument("-p", "--pid",
    help="trace this PID only")
parser.add_argument("-t", "--tid",
    help="trace this TID only")
parser.add_argument("interval", nargs="?", default=99999999,
    help="output interval, in seconds")
parser.add_argument("count", nargs="?", default=99999999,
    help="number of outputs")
parser.add_argument("--ebpf", action="store_true",
    help=argparse.SUPPRESS)
args = parser.parse_args()
countdown = int(args.count)
debug = 0

bpf_text = """
#include <uapi/linux/ptrace.h>
#include <linux/sched/topology.h>
#include <linux/sched.h>

struct pid_key {
    u64 id;
    u64 slot;
};

BPF_HASH(start, u32, u64, MAX_PID);
STORAGE

int trace_load_balance_entry(struct pt_regs *ctx, int this_cpu, struct rq *this_rq,
                             struct sched_domain *sd, enum cpu_idle_type idle,
                             int *continue_balancing)
{
    u64 pid_tgid = bpf_get_current_pid_tgid();
    u32 tgid = pid_tgid >> 32, pid = pid_tgid;
    u64 ts = bpf_ktime_get_ns();

    if (FILTER)
        return 0;
    start.update(&pid, &ts);
    return 0;
}

int trace_load_balance_return(struct pt_regs *ctx)
{
    u64 pid_tgid = bpf_get_current_pid_tgid();
    u32 tgid = pid_tgid >> 32, pid = pid_tgid;
    u64 ts = bpf_ktime_get_ns();
    u64 *tsp, delta;

    if (FILTER)
        return 0;
    tsp = start.lookup(&pid);
    if (!tsp)
        return 0;
    if (ts < *tsp)
        return 0;
    delta = ts - *tsp;
    FACTOR
    STORE
    return 0;
}

"""

if args.pid:
    bpf_text = bpf_text.replace('FILTER', 'tgid != %s' % args.pid)
elif args.tid:
    bpf_text = bpf_text.replace('FILTER', 'pid != %s' % args.tid)
else:
    bpf_text = bpf_text.replace('FILTER', '0')
if args.milliseconds:
    bpf_text = bpf_text.replace('FACTOR', 'delta /= 1000000;')
    label = "msecs"
else:
    bpf_text = bpf_text.replace('FACTOR', 'delta /= 1000;')
    label = "usecs"
if args.pids or args.tids:
    section = "pid"
    pid = "tgid"
    if args.tids:
        pid = "pid"
        section = "tid"
    bpf_text = bpf_text.replace('STORAGE',
        'BPF_HISTOGRAM(dist, struct pid_key, MAX_PID);')
    bpf_text = bpf_text.replace('STORE',
        'struct pid_key key = {.id = ' + pid + ', .slot = bpf_log2l(delta)}; ' +
        'dist.increment(key);')
else:
    section = ""
    bpf_text = bpf_text.replace('STORAGE', 'BPF_HISTOGRAM(dist);')
    bpf_text = bpf_text.replace('STORE',
        'dist.increment(bpf_log2l(delta));')
if debug or args.ebpf:
    print(bpf_text)
    if args.ebpf:
        exit()

max_pid = int(open("/proc/sys/kernel/pid_max").read())

b = BPF(text=bpf_text, cflags=["-DMAX_PID=%d" % max_pid])
b.attach_kprobe(event="load_balance", fn_name="trace_load_balance_entry")
b.attach_kretprobe(event="load_balance", fn_name="trace_load_balance_return")

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

exiting = 0 if args.interval else 1
dist = b.get_table("dist")
while (1):
    try:
        sleep(int(args.interval))
    except KeyboardInterrupt:
        exiting = 1

    print()
    if args.timestamp:
        print("%-8s\n" % strftime("%H:%M:%S"), end="")

    def pid_to_comm(pid):
        try:
            comm = open("/proc/%d/comm" % pid, "r").read()
            return "%d %s" % (pid, comm)
        except IOError:
            return str(pid)

    dist.print_log2_hist(label, section, section_print_fn=pid_to_comm)
    dist.clear()

    countdown -= 1
    if exiting or countdown == 0:
        exit()