#!/usr/bin/python3
#
# drsnoopv2  Trace direct reclaim and print details including issuing PID.
#       For Linux, uses BCC, eBPF.
#
# This uses in-kernel eBPF maps to cache process details (PID and comm) by
# direct reclaim begin, as well as a starting timestamp for calculating
# latency.
#

from __future__ import print_function
from bcc import ArgString, BPF
import argparse
from datetime import datetime, timedelta
import os
import math
import sys

# symbols
kallsyms = "/proc/kallsyms"

# arguments
examples = """examples:
    ./drsnoopv2           # trace all direct reclaim
    ./drsnoopv2 -T        # include timestamps
    ./drsnoopv2 -U        # include UID
    ./drsnoopv2 -O        # include page order
    ./drsnoopv2 -P 181    # only trace PID 181
    ./drsnoopv2 -t 123    # only trace TID 123
    ./drsnoopv2 -u 1000   # only trace UID 1000
    ./drsnoopv2 -d 10     # trace for 10 seconds only
    ./drsnoopv2 -n main   # only print process names containing "main"
    ./drsnoopv2 -S        # print summary instead of event stream
    ./drsnoopv2 -v        # print current vmstat
"""
parser = argparse.ArgumentParser(
    description="Trace direct reclaim",
    formatter_class=argparse.RawDescriptionHelpFormatter,
    epilog=examples)
parser.add_argument("-T", "--timestamp", action="store_true",
                    help="include timestamp on output")
parser.add_argument("-U", "--print-uid", action="store_true",
                    help="print UID column")
parser.add_argument("-O", "--print-order", action="store_true",
                    help="print ORDER column")
parser.add_argument("-p", "--pid",
                    help="trace this PID only")
parser.add_argument("-t", "--tid",
                    help="trace this TID only")
parser.add_argument("-u", "--uid",
                    help="trace this UID only")
parser.add_argument("-d", "--duration",
                    help="total duration of trace in seconds")
parser.add_argument("-n", "--name",
                    type=ArgString,
                    help="only print process names containing this name")
parser.add_argument("-v", "--verbose", action="store_true",
                    help="show system memory state")
parser.add_argument("--ebpf", action="store_true",
                    help=argparse.SUPPRESS)
parser.add_argument("-S", "--summary", action="store_true",
                    help="print summary instead of event stream")
args = parser.parse_args()
debug = 0
if args.duration:
    args.duration = timedelta(seconds=int(args.duration))


# vm_stat
vm_stat_addr = ''
with open(kallsyms) as syms:
    for line in syms:
        (addr, size, name) = line.rstrip().split(" ", 2)
        name = name.split("\t")[0]
        if name == "vm_stat":
            vm_stat_addr = "0x" + addr
            break
        if name == "vm_zone_stat":
            vm_stat_addr = "0x" + addr
            break
    if vm_stat_addr == '':
        print("ERROR: no vm_stat or vm_zone_stat in /proc/kallsyms. Exiting.")
        print("HINT: the kernel should be built with CONFIG_KALLSYMS_ALL.")
        exit()

NR_FREE_PAGES = 0

PAGE_SIZE = os.sysconf("SC_PAGE_SIZE")
PAGE_SHIFT = int(math.log(PAGE_SIZE) / math.log(2))

def K(x):
    return x << (PAGE_SHIFT - 10)

# load BPF program
bpf_text = """
#include <uapi/linux/ptrace.h>
#include <linux/sched.h>
#include <linux/mmzone.h>

struct val_t {
    u64 id;
    u64 ts; // start time
    u32 order;
    char name[TASK_COMM_LEN];
    u64 vm_stat[NR_VM_ZONE_STAT_ITEMS];
};

struct data_t {
    u64 id;
    u32 uid;
    u64 nr_reclaimed;
    u64 delta;
    u64 ts;    // end time
    u32 order;
    char name[TASK_COMM_LEN];
    u64 vm_stat[NR_VM_ZONE_STAT_ITEMS];
};

struct info_t {
    u64 count;
    u64 total_lat;
    u64 total_pages;
    u64 total_free;
    char name[TASK_COMM_LEN];
};

BPF_HASH(start, u64, struct val_t);
BPF_HASH(counts, u32, struct info_t);
BPF_HISTOGRAM(dist);
BPF_HISTOGRAM(order_dist);
BPF_PERF_OUTPUT(events);

TRACEPOINT_PROBE(vmscan, mm_vmscan_direct_reclaim_begin) {
    struct val_t val = {};
    u64 id = bpf_get_current_pid_tgid();
    u32 pid = id >> 32; // PID is higher part
    u32 tid = id;       // Cast and get the lower part
    u32 uid = bpf_get_current_uid_gid();
    u64 ts;

    PID_TID_FILTER
    UID_FILTER
    if (bpf_get_current_comm(&val.name, sizeof(val.name)) == 0) {
        val.id = id;
        val.ts = bpf_ktime_get_ns();
        bpf_probe_read_kernel(&val.order, sizeof(val.order), &args->order);
        bpf_probe_read_kernel(&val.vm_stat, sizeof(val.vm_stat), (const void *)%s);
        start.update(&id, &val);
    }
    return 0;
}

TRACEPOINT_PROBE(vmscan, mm_vmscan_direct_reclaim_end) {
    u64 id = bpf_get_current_pid_tgid();
    u32 pid = id >> 32;
    struct val_t *valp;
    struct data_t data = {};
    u64 ts = bpf_ktime_get_ns();

    valp = start.lookup(&id);
    if (valp == NULL) {
        // missed entry
        return 0;
    }

    u64 delta = ts - valp->ts;

    SUMMARY_FILTER

    data.delta = delta;
    data.ts = ts / 1000;
    data.id = valp->id;
    data.uid = bpf_get_current_uid_gid();
    data.order = valp->order;
    bpf_probe_read_kernel(&data.name, sizeof(data.name), valp->name);
    bpf_probe_read_kernel(&data.vm_stat, sizeof(data.vm_stat), valp->vm_stat);
    data.nr_reclaimed = args->nr_reclaimed;

    events.perf_submit(args, &data, sizeof(data));
    start.delete(&id);

    return 0;
}
""" % vm_stat_addr

if args.tid:  # TID trumps PID
    bpf_text = bpf_text.replace('PID_TID_FILTER',
                                'if (tid != %s) { return 0; }' % args.tid)
elif args.pid:
    bpf_text = bpf_text.replace('PID_TID_FILTER',
                                'if (pid != %s) { return 0; }' % args.pid)
else:
    bpf_text = bpf_text.replace('PID_TID_FILTER', '')
if args.uid:
    bpf_text = bpf_text.replace('UID_FILTER',
                                'if (uid != %s) { return 0; }' % args.uid)
else:
    bpf_text = bpf_text.replace('UID_FILTER', '')

if args.summary:
    bpf_text = bpf_text.replace('SUMMARY_FILTER', """
    struct info_t *infop, info = {};
    infop = counts.lookup(&pid);
    if (infop == NULL) {
        info.count = 1;
        info.total_lat = delta;
        info.total_pages = args->nr_reclaimed;
        bpf_probe_read_kernel(&info.name, sizeof(info.name), valp->name);
        %s
        counts.update(&pid, &info);
    } else {
        infop->count++;
        infop->total_lat += delta;
        infop->total_pages += args->nr_reclaimed;
        %s
    }
    dist.increment(bpf_log2l(delta / 1000));
    order_dist.increment(valp->order);

    start.delete(&id);
    return 0;
    """ % ("info.total_free = valp->vm_stat[NR_FREE_PAGES];" if args.verbose else "info.total_free = 0;",
           "infop->total_free += valp->vm_stat[NR_FREE_PAGES];" if args.verbose else ""))
else:
    bpf_text = bpf_text.replace('SUMMARY_FILTER', '')

if debug or args.ebpf:
    print(bpf_text)
    if args.ebpf:
        exit()

# initialize BPF
b = BPF(text=bpf_text)

initial_ts = 0

# Track event count and header printed flag
event_count = 0
header_printed = False

def print_header():
    global header_printed
    if header_printed:
        return

    # header
    if args.timestamp:
        print("%-14s" % ("TIME(s)"), end="")
    if args.print_uid:
        print("%-6s" % ("UID"), end="")
    if args.print_order:
        print("%-8s" % ("ORDER"), end="")

    print("%-14s %-6s %8s %5s" %
          ("COMM", "TID" if args.tid else "PID", "LAT(ms)", "PAGES"), end="")
    if args.verbose:
        print("%10s" % ("FREE(KB)"))
    else:
        print("")

    header_printed = True

def print_event(cpu, data, size):
    # Print header before first event
    print_header()

    event = b["events"].event(data)

    global initial_ts

    if not initial_ts:
        initial_ts = event.ts

    if args.name and bytes(args.name) not in event.name:
        return

    if args.timestamp:
        delta = event.ts - initial_ts
        print("%-14.9f" % (float(delta) / 1000000), end="")

    if args.print_uid:
        print("%-6d" % event.uid, end="")

    if args.print_order:
        print("%-5d" % event.order, end="")

    print("%-14.14s %-6s %8.2f %5d" %
          (event.name.decode('utf-8', 'replace'),
           event.id & 0xffffffff if args.tid else event.id >> 32,
           float(event.delta) / 1000000, event.nr_reclaimed), end="")
    if args.verbose:
        print("%10d" % K(event.vm_stat[NR_FREE_PAGES]))
    else:
        print("")

    sys.stdout.flush()

def print_event_wrapper(cpu, data, size):
    global event_count
    event_count += 1
    print_event(cpu, data, size)

print("Tracing direct reclaim events... Hit Ctrl-C to end.")
if not args.summary:
    # loop with callback to print_event
    b["events"].open_perf_buffer(print_event_wrapper, page_cnt=64)
    start_time = datetime.now()
    try:
        while not args.duration or datetime.now() - start_time < args.duration:
            b.perf_buffer_poll(timeout=1000)
    except KeyboardInterrupt:
        # Print message if no events were found
        if event_count == 0:
            print("No any allocstall found!")
        exit()

    # Print message if no events were found
    if event_count == 0:
        print("No any allocstall found!")
else:
    try:
        # sleep until interrupted or duration expires
        if args.duration:
            import time
            time.sleep(args.duration.total_seconds())
        else:
            try:
                b.trace_print() # This is a blocking call if we were using trace_print, but we are not.
                # Just wait forever
                while True:
                    pass
            except KeyboardInterrupt:
                pass
    except KeyboardInterrupt:
        if event_count == 0:
            print("No any allocstall found!")
        exit()

    # Sort and limit output
    output_items = []
    counts = b.get_table("counts")
    for k, v in counts.items():
        event_count += 1
        output_items.append((k, v))

    # Print message if no events were found
    if event_count == 0:
        print("No any allocstall found!")
        exit()

    print("\n%-10s %-16s %8s %14s %8s %10s" % ("PID", "COMM", "COUNT", "TOTAL_LAT(ms)", "AVG_LAT", "PAGES"))

    # If more than 50 items, sort by AVG_LAT descending
    if len(output_items) > 50:
        print("WARNING: Too many processes, showing top 50 by Average Latency")
        output_items.sort(key=lambda item: (item[1].total_lat / item[1].count), reverse=True)
        output_items = output_items[:50]
    else:
        # standard sort by count
        output_items.sort(key=lambda item: item[1].count, reverse=True)

    total_reclaimed=0
    total_free=0
    total_count=0
    for k, v in output_items:
        print("%-10d %-16s %8d %14.2f %8.2f %10d" % (
            k.value,
            v.name.decode('utf-8', 'replace'),
            v.count,
            v.total_lat / 1e6,
            (v.total_lat / 1e6) / v.count,
            v.total_pages
        ))
        total_count += v.count
        total_reclaimed += v.total_pages
        total_free += v.total_free

    print("\nDirect Reclaim Latency Histogram (usecs):")
    b["dist"].print_log2_hist("usecs")

    print("\nDirect Reclaim Order Histogram:")
    b["order_dist"].print_linear_hist("Order")

    total_reclaimed_kb = K(total_reclaimed)
    print("\nTotal Reclaimed Pages: %d KB" % total_reclaimed_kb)

    if args.verbose:
        avg_free_kb = K(total_free // total_count) if total_count > 0 else 0
        print("Avg Free Pages: %d KB" % avg_free_kb)