#!/usr/bin/python3

from bcc import BPF
from time import sleep
from ctypes import c_ushort, c_int, c_ulonglong
import argparse

parser = argparse.ArgumentParser(
	description="Measure time interval needed for transporting D-Bus messages",
	formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--process", action="store_true",
	help="change output format to top-like")
parser.add_argument("--no-libdbus", action="store_true",
	help="exclude messages sent with libdbus")
parser.add_argument("--no-gio", action="store_true",
	help="exclude messages sent with gio")
parser.add_argument("interval", nargs="?", default=2,
	help="output interval, in seconds")
parser.add_argument("count", nargs="?", default=-1,
	help="number of outputs")

args = parser.parse_args()
interval = int(args.interval)
loop = 0
count = int(args.count)

if args.process:
	mode = "process"
else:
	mode = "histogram"

b = BPF(src_file="dbus-latency.c")

# catch sending
if not args.no_libdbus:
	b.attach_uprobe(name="dbus-1", sym="_dbus_message_get_network_data", fn_name="get_msg_addr")
if not args.no_gio:
	b.attach_uprobe(name="gio-2.0", sym="on_worker_message_about_to_be_sent", fn_name="g_get_msg_addr")

# catch receiving - don't disable libdbus or libgio, we still want to know that a gio-based
# program received a message sent by a libdbus-based program
b.attach_uprobe(name="dbus-1", sym="_dbus_connection_queue_received_message_link", fn_name="get_msg_latency")
b.attach_uprobe(name="gio-2.0", sym="on_worker_message_received", fn_name="g_get_msg_latency")

print("Hit Ctrl-C to end.")

while (1):
	sleep(interval)
	loop = loop + 1
	print("\n%d:" % loop)
	latency = b["msg_latency"]
	if mode =="histogram":
		b["latency_histo"].print_log2_hist("us")
		avgcnt = 0
		avgval = 0
		for v, p in latency.items():
			avgcnt += long(p.count)
			avgval += long(p.count) * long(p.avg)
		if avgcnt > 0:
			print 'Average:', avgval // avgcnt
	else:
		print ("%10s %20s %10s %10s %10s" % ("PID", "COMM", "MIN", "AVG", "MAX"))
		for v, p in sorted(latency.items(), key=lambda latency: latency[1].avg, reverse=True):
			print("%10d %20s %10d %10d %10d" % (v.pid, v.comm.encode('string-escape'), int(p.min), int(p.avg), int(p.max)))

	print(" messages being transferred: %d" % len(b["msg_sent_addr"].items()))

	if count == loop:
		print 'Finished'
		break
