#
# iosize.py   Block I/O size histogram.
#               For Linux, uses BCC, eBPF. See .c file.
#
# USAGE: iosize
#
# Based on bisize of bcc/tools.
#
# Ctrl-C will print the partially gathered histogram then exit.

from bcc import BPF
from time import sleep

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

struct proc_key_t {
    char name[DISK_NAME_LEN];
    u64 slot;
};

BPF_HISTOGRAM(bio_dist, struct proc_key_t);
BPF_HISTOGRAM(rq_dist, struct proc_key_t);

#if LINUX_VERSION_CODE < KERNEL_VERSION(3, 11, 0)
int trace_submit_bio(struct pt_regs *ctx, int rw, struct bio *bio)
#else
int trace_submit_bio(struct pt_regs *ctx, struct bio *bio)
#endif
{
#if LINUX_VERSION_CODE < KERNEL_VERSION(3, 11, 0)
    struct proc_key_t key = {.slot = bpf_log2l(bio->bi_size / 512)};
#else
    struct proc_key_t key = {.slot = bpf_log2l(bio->bi_iter.bi_size / 512)};
#endif
#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 14, 0)
    struct gendisk *disk = bio->bi_disk;
#else
    struct gendisk *disk = bio->bi_bdev->bd_disk;
#endif

    if (disk) {
       bpf_probe_read(&key.name, sizeof(key.name), disk->disk_name);
       bio_dist.increment(key);
    }
    return 0;
}
int trace_blk_mq_start_request(struct pt_regs *ctx, struct request *rq)
{
    struct proc_key_t key = {.slot = bpf_log2l(rq->__data_len / 512)};
    struct gendisk *disk = rq->rq_disk;

    if (disk) {
       bpf_probe_read(&key.name, sizeof(key.name), disk->disk_name);
       rq_dist.increment(key);
    }
    return 0;
}
"""

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

b.attach_kprobe(event="submit_bio", fn_name="trace_submit_bio")
b.attach_kprobe(event="blk_mq_start_request", fn_name="trace_blk_mq_start_request")


print("Tracing block I/O... Hit Ctrl-C to end.")

# trace until Ctrl-C
bio_dist = b.get_table("bio_dist")
rq_dist = b.get_table("rq_dist")

try:
    sleep(99999999)
except KeyboardInterrupt:
    bio_dist.print_log2_hist("bio sects", "Disk Name",
            section_print_fn=bytes.decode)
    rq_dist.print_log2_hist("rq sects", "Disk Name",
            section_print_fn=bytes.decode)