#!/usr/bin/python

from bcc import BPF
from time import sleep
from ctypes import c_ushort, c_int, c_ulonglong
from sys import argv


print("Hit Ctrl-C to end.")
count = 100
interval = 5
loop = 0
# load BPF program
bpf_text="""
#include <hostcompat/linux/sched.h>
#include <hostcompat/uapi/linux/ptrace.h>

typedef struct pdata_T {
	u32 pid;
	char buf[TASK_COMM_LEN];
	u64 drun;
	u64 dwait;
	u64 grun;
	u64 gwait;
	u64 usage;
}pdata_t;

BPF_HASH(timestamp_b, u64);
BPF_HASH(timestamp_e, u64);
BPF_HASH(stats, u32, pdata_t, sizeof(pdata_t));
BPF_HASH(run_start, u32);
BPF_HASH(wait_start, u32);

BPF_HISTOGRAM(dist);

int do_entry_dbus(struct pt_regs *ctx) {
	u64 time = 0;
	pdata_t pdata;
	pdata_t new_pdata = {0, ' ', 0, 0, 0, 0, 0};
	u64 ts, *tsp, delta;
	u32 pid;

	ts = bpf_ktime_get_ns() / 1000;
	if(timestamp_b.lookup(&time) == 0){
		timestamp_b.update(&time, &ts);
	}
	timestamp_e.update(&time, &ts);

	pid = bpf_get_current_pid_tgid();
	tsp = run_start.lookup(&pid);
	run_start.delete(&pid);
	wait_start.update(&pid, &ts);
	if(tsp != 0) {
		delta = ts -*tsp;
	}
	else {
		return 0;
	}
	new_pdata.pid = pid;
	pdata = *(stats.lookup_or_init(&pid, &new_pdata));
	pdata.drun += delta;
	pdata.usage += delta;
	bpf_get_current_comm(&(pdata.buf), sizeof(pdata.buf));
	stats.update(&pid, &pdata);
	return 0;
}

int do_return_dbus(struct pt_regs *ctx) {
	u64 time = 0;
	pdata_t pdata;
	pdata_t new_pdata = {0, ' ', 0, 0, 0, 0, 0};
	u64 ts, *tsp, delta;
	u32 pid;

	ts = bpf_ktime_get_ns() / 1000;
	if(timestamp_b.lookup(&time) == 0){
		timestamp_b.update(&time, &ts);
	}
	timestamp_e.update(&time, &ts);

	pid = bpf_get_current_pid_tgid();
	tsp = wait_start.lookup(&pid);
	wait_start.delete(&pid);
	run_start.update(&pid, &ts);
	if(tsp != 0) {
		delta = ts - *tsp;
	}
	else {
		return 0;
	}
	new_pdata.pid = pid;
	pdata = *(stats.lookup_or_init(&pid, &new_pdata));
	pdata.dwait += delta;
	bpf_get_current_comm(&(pdata.buf), sizeof(pdata.buf));
	stats.update(&pid, &pdata);
	return 0;
}

int do_entry_glib(struct pt_regs *ctx) {
	u64 time = 0;
	pdata_t pdata;
	pdata_t new_pdata = {0, ' ', 0, 0, 0, 0, 0};
	u64 ts, *tsp, delta;
	u32 pid;

	ts = bpf_ktime_get_ns() / 1000;
	if(timestamp_b.lookup(&time) == 0){
		timestamp_b.update(&time, &ts);
	}
	timestamp_e.update(&time, &ts);

	pid = bpf_get_current_pid_tgid();
	tsp = run_start.lookup(&pid);
	run_start.delete(&pid);
	wait_start.update(&pid, &ts);
	if(tsp != 0) {
		delta = ts -*tsp;
	}
	else {
		return 0;
	}
	new_pdata.pid = pid;
	pdata = *(stats.lookup_or_init(&pid, &new_pdata));
	pdata.grun += delta;
	pdata.usage += delta;
	bpf_get_current_comm(&(pdata.buf), sizeof(pdata.buf));
	stats.update(&pid, &pdata);
	return 0;
}

int do_return_glib(struct pt_regs *ctx) {
	u64 time = 0;
	pdata_t pdata;
	pdata_t new_pdata = {0, ' ', 0, 0, 0, 0, 0};
	u64 ts, *tsp, delta;
	u32 pid;

	ts = bpf_ktime_get_ns() / 1000;
	if(timestamp_b.lookup(&time) == 0){
		timestamp_b.update(&time, &ts);
	}
	timestamp_e.update(&time, &ts);

	pid = bpf_get_current_pid_tgid();
	tsp = wait_start.lookup(&pid);
	wait_start.delete(&pid);
	run_start.update(&pid, &ts);
	if(tsp != 0) {
		delta = ts - *tsp;
	}
	else {
		return 0;
	}
	new_pdata.pid = pid;
	pdata = *(stats.lookup_or_init(&pid, &new_pdata));
	pdata.gwait += delta;
	bpf_get_current_comm(&(pdata.buf), sizeof(pdata.buf));
	stats.update(&pid, &pdata);
	return 0;
}
"""

b = BPF(text=bpf_text)

b.attach_uprobe(name="glib-2.0", sym="g_poll", fn_name="do_entry_glib")
b.attach_uprobe(name="dbus-1", sym="_dbus_poll", fn_name="do_entry_dbus")

b.attach_uretprobe(name="dbus-1", sym="_dbus_poll", fn_name="do_return_dbus")
b.attach_uretprobe(name="glib-2.0", sym="g_poll", fn_name="do_return_glib")

while (1):
	if count > 0:

		loop += 1
		if loop > count:
			exit()
	sleep(interval)
	print ("%d : \n" % loop)
	b["dist"].print_log2_hist("usecs")
	b["dist"].clear()
	stats_t = b["stats"]
	timestamp = b["timestamp_e"]
	timeframe = 0
	for v in timestamp.values():
		timeframe = v.value
	timestamp = b["timestamp_b"]
	for v in timestamp.values():
		timeframe -= v.value
	#print("%d\n" % stats_t[0].pid)
	print ("%10s %20s %10s %10s" % ("PID", "NAME", "DBUS", "GLIB"))
	for v in sorted(stats_t.values(), key=lambda stats_t: stats_t.usage, reverse=True):
		print("%10d %20s %10.6f %10.6f" % (v.pid, v.buf.encode('string-escape'), float(v.drun) / float(timeframe), float(v.grun) / float(timeframe)))
