#!/usr/bin/python3

from asyncore import write
from bcc import BPF
from time import sleep
import argparse

parser = argparse.ArgumentParser(
	description="Measure time interval between receiving a message and handling it in a client thread in glib",
	formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--no-method-calls", action="store_true",
	help="disable detecting D-Bus method calls")
parser.add_argument("--no-method-returns-and-errors", action="store_true",
	help="disable detecting D-Bus method returns and errors")
parser.add_argument("--no-signals", action="store_true",
	help="disable detecting D-Bus signals")
parser.add_argument("--raw", action="store_true",
	help="put data into .csv file")
parser.add_argument("-o", "--output", default=None)
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)

b = BPF(src_file="dbus-glib-receiving.c")

# Setup uprobes

# ****** General

# this uprobe function detects an incoming, constructed message and stores a reference to it
# it also stores time of detection of the message
b.attach_uprobe(name="gio-2.0", sym="on_worker_message_received", fn_name="handle_on_worker_received")

# this uprobe function detects destroying of a message and removes any existing references
b.attach_uprobe(name="gio-2.0", sym="g_dbus_message_finalize", fn_name="handle_message_delete")

# ****** Receiving method calls
if not args.no_method_calls:
	# this uprobe function detects when a GDBusMethodInvocation is created, and stores a reference to it
	b.attach_uretprobe(name="gio-2.0", sym="_g_dbus_method_invocation_new", fn_name="handle_ret_method_invocation_new")

	# this uprobe function detects when a GDBusMethodInvocation is handled in a client thread
	# it removes any existing references to the invocation and the message associated with the invocation
	# it also records time interval between detection of the message and this point
	b.attach_uprobe(name="gio-2.0", sym="call_in_idle_cb", fn_name="handle_call_in_idle_cb")

	# this uprobe function detects non-handled method invocations for removing references
	b.attach_uprobe(name="gio-2.0", sym="g_dbus_method_invocation_finalize", fn_name="handle_finalize_invocation")

# ***** Receiving method returns and errors

if not args.no_method_returns_and_errors:
	# this uprobe function detects when a task is associated with a message, and stores a reference to it
	b.attach_uprobe(name="gio-2.0", sym="g_task_return_pointer", fn_name="handle_task_return_pointer")

	# this uprobe function detects when a task is handled in a client thread
	# it removes any existing references to the task and the message associated with the task
	# it also records time interval between detection of the message and this point
	b.attach_uprobe(name="gio-2.0", sym="g_task_return_now", fn_name="handle_task_return_now")

	# this uprobe function detects non-handled tasks for removing references
	b.attach_uprobe(name="gio-2.0", sym="g_task_finalize", fn_name="handle_finalize_task")

# ***** Receiving signals
if not args.no_method_returns_and_errors:
	# this uprobe function detects when a signal is about to be delivered in a client thread
	# it removes any existing references to the message
	# it also records time interval between detection of the message and this point
	# thus, it records only timings for the very first subscriber that receives the signal
	b.attach_uprobe(name="gio-2.0", sym="emit_signal_instance_in_idle_cb", fn_name="handle_emit_signal")

# ***** Create .csv with data
# process event put data into .csv file while in --raw mode.
def new_timestamp_to_file(cpu, data, size):
	# this func is triggered when new timestamp is put into timestamp_out
	# buffer. Data is then put into .csv file
	event = b["timestamp_out"].event(data)
	global f
	f.write(str(event.timestamp) + '\n')

def new_timestamp_to_std(cpu, data, size):
	# this func is triggered when new timestamp is put into timestamp_out
	# buffer. Data is displayed on std output.
	event = b["timestamp_out"].event(data)
	print("Measured Latency: " + str(event.timestamp) + "\n")

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

if args.raw :
	# loop with callback to print_event
	if args.output != None :
		print("File mode active\n")
		b["timestamp_out"].open_perf_buffer(new_timestamp_to_file)
		with open(args.output, "w") as f:
			f.write("Latency\n")
			while (1):
				b.perf_buffer_poll()
	else:
		print("STD mode active\n")
		b["timestamp_out"].open_perf_buffer(new_timestamp_to_std)

	while (1):
		b.perf_buffer_poll()

while (1):
	sleep(interval)
	loop = loop + 1
	print("\n%d:" % loop)
	b["timings"].print_log2_hist("us")

	# Kinda debug part, but important - it shows any leftovers in hashmaps.
	# This way, if it shows up, we can tell how bad it is.
	msgs = len(b["msg_on_worker_received"].items())
	invocations = len(b["invocation_to_message"].items())
	tasks = len(b["task_to_message"].items())

	if msgs + invocations + tasks > 0:
		print(" containers stats:")
		if msgs > 0:        print('     msg_on_worker_received:', msgs)
		if invocations > 0: print('     invocation_to_message :', invocations)
		if tasks > 0:       print('     task_to_message       :', tasks)

	if b["count_averages"][1].value > 0:
		print('Average:', (b["count_averages"][0].value // b["count_averages"][1].value) // 1000)

	if count == loop:
		print('Finished')
		break
