#!/usr/bin/python3

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

debug = 0

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

typedef struct aio_key {
    u32 id;
    u32 slot;
} aio_key_t;

BPF_HISTOGRAM(dist_nr, aio_key_t);
BPF_HISTOGRAM(dist_events, aio_key_t);
BPF_HISTOGRAM(dist_depth, aio_key_t);

BPF_ARRAY(depth, unsigned long, 1);


static int handle_submit(int v)
{
    aio_key_t key = {0};
    unsigned int dk = 0;
    unsigned long *qd;

    qd = depth.lookup(&dk);
    if (!qd)
       return -1;

    key.slot = bpf_log2l(v);
    dist_nr.increment(key);

    *qd = *qd + v;
   key.slot = *qd;
   dist_depth.increment(key);

   return 0;
}

static int handle_getevents(int v)
{
    aio_key_t key = {0};
    unsigned int dk = 0;
    unsigned long *qd;

    qd = depth.lookup(&dk);
    if (!qd)
       return -1;

    key.slot = bpf_log2l(v);
    dist_events.increment(key);

    if (*qd > v)
        *qd = *qd - v;
    key.slot = *qd;
    dist_depth.increment(key);
    return 0;
}

#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 10, 0)
TRACEPOINT_PROBE(syscalls, sys_enter_io_submit)
{
    int v = (int)args->nr;

    handle_submit(v);
    return 0;
}

TRACEPOINT_PROBE(syscalls, sys_exit_io_getevents)
{
    int v = (int)args->ret;
    handle_getevents(v);
    return 0;
}
#else
int trace_do_io_submit(struct pt_regs *ctx, unsigned long id, long nr)
{
    handle_submit(nr);
    return 0;
}

int trace_do_io_getevents(struct pt_regs *ctx)
{
    int v = PT_REGS_RC(ctx);
    handle_getevents(v);
    return 0;
}
#endif

"""

# code substitutions
if debug:
    print(bpf_text)
    exit()

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

release=str(subprocess.check_output(['uname', '-a']))
if release.find("3.10.0") > 0:
    b.attach_kprobe(event="do_io_submit", fn_name="trace_do_io_submit")
    b.attach_kretprobe(event="read_events", fn_name="trace_do_io_getevents")

# these should really use irq:irq_handler_entry/exit tracepoints:
print("Tracing request size... Hit Ctrl-C to end.")

# output
dist_nr = b.get_table("dist_nr")
dist_events = b.get_table("dist_events")
dist_depth = b.get_table("dist_depth")

while (1):
    try:
        sleep(1000000000)
    except KeyboardInterrupt:
        exiting = 1

    dist_nr.print_log2_hist("nr in io_submit", "aio")
    dist_events.print_log2_hist("events in io_getevents", "aio")
    dist_depth.print_linear_hist("io depth", "aio")
    
    dist_nr.clear()
    dist_events.clear()
    dist_depth.clear()

    if exiting:
        exit()