1#!/usr/bin/python 2import argparse 3from time import sleep, strftime 4from sys import argv 5import ctypes as ct 6from bcc import BPF, USDT 7import inspect 8import os 9 10# Parse command line arguments 11parser = argparse.ArgumentParser(description="Trace the latency distribution of an operation using usdt probes.", 12 formatter_class=argparse.RawDescriptionHelpFormatter) 13parser.add_argument("-p", "--pid", type=int, help="The id of the process to trace.") 14parser.add_argument("-i", "--interval", type=int, help="The interval in seconds on which to report the latency distribution.") 15parser.add_argument("-f", "--filterstr", type=str, default="", help="The prefix filter for the operation input. If specified, only operations for which the input string starts with the filterstr are traced.") 16parser.add_argument("-v", "--verbose", dest="verbose", action="store_true", help="If true, will output generated bpf program and verbose logging information.") 17parser.add_argument("-s", "--sdt", dest="sdt", action="store_true", help="If true, will use the probes, created by systemtap's dtrace.") 18 19parser.set_defaults(verbose=False) 20args = parser.parse_args() 21this_pid = int(args.pid) 22this_interval = int(args.interval) 23this_filter = str(args.filterstr) 24 25if this_interval < 1: 26 print("Invalid value for interval, using 1.") 27 this_interval = 1 28 29debugLevel=0 30if args.verbose: 31 debugLevel=4 32 33# BPF program 34bpf_text_shared = "%s/bpf_text_shared.c" % os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) 35bpf_text = open(bpf_text_shared, 'r').read() 36bpf_text += """ 37 38/** 39 * @brief The key to use for the latency histogram. 40 */ 41struct dist_key_t 42{ 43 char input[64]; ///< The input string of the request. 44 u64 slot; ///< The histogram slot. 45}; 46 47/** 48 * @brief Contains the histogram for the operation latencies. 49 */ 50BPF_HISTOGRAM(dist, struct dist_key_t); 51 52/** 53 * @brief Reads the operation response arguments, calculates the latency, and stores it in the histogram. 54 * @param ctx The BPF context. 55 */ 56int trace_operation_end(struct pt_regs* ctx) 57{ 58 u64 operation_id; 59 bpf_usdt_readarg(1, ctx, &operation_id); 60 61 struct start_data_t* start_data = start_hash.lookup(&operation_id); 62 if (0 == start_data) { 63 return 0; 64 } 65 66 u64 duration = bpf_ktime_get_ns() - start_data->start; 67 struct dist_key_t dist_key = {}; 68 __builtin_memcpy(&dist_key.input, start_data->input, sizeof(dist_key.input)); 69 dist_key.slot = bpf_log2l(duration / 1000); 70 start_hash.delete(&operation_id); 71 72 dist.atomic_increment(dist_key); 73 return 0; 74} 75""" 76 77bpf_text = bpf_text.replace("FILTER_STRING", this_filter) 78if this_filter: 79 bpf_text = bpf_text.replace("FILTER_STATEMENT", "if (!filter(start_data.input)) { return 0; }") 80else: 81 bpf_text = bpf_text.replace("FILTER_STATEMENT", "") 82 83# Create USDT context 84print("lat_dist.py - Attaching probes to pid: %d; filter: %s" % (this_pid, this_filter)) 85usdt_ctx = USDT(pid=this_pid) 86 87if args.sdt: 88 usdt_ctx.enable_probe(probe="usdt_sample_lib1_sdt:operation_start_sdt", fn_name="trace_operation_start") 89 usdt_ctx.enable_probe(probe="usdt_sample_lib1_sdt:operation_end_sdt", fn_name="trace_operation_end") 90else: 91 usdt_ctx.enable_probe(probe="usdt_sample_lib1:operation_start", fn_name="trace_operation_start") 92 usdt_ctx.enable_probe(probe="usdt_sample_lib1:operation_end", fn_name="trace_operation_end") 93 94# Create BPF context, load BPF program 95bpf_ctx = BPF(text=bpf_text, usdt_contexts=[usdt_ctx], debug=debugLevel) 96 97print("Tracing... Hit Ctrl-C to end.") 98 99start = 0 100dist = bpf_ctx.get_table("dist") 101while (1): 102 try: 103 sleep(this_interval) 104 except KeyboardInterrupt: 105 exit() 106 107 print("[%s]" % strftime("%H:%M:%S")) 108 dist.print_log2_hist("latency (us)") 109