1#!/usr/bin/env 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 verbose logging information.") 17parser.set_defaults(verbose=False) 18args = parser.parse_args() 19this_pid = int(args.pid) 20this_interval = int(args.interval) 21this_filter = str(args.filterstr) 22 23if this_interval < 1: 24 print("Invalid value for interval, using 1.") 25 this_interval = 1 26 27debugLevel=0 28if args.verbose: 29 debugLevel=4 30 31# BPF program 32bpf_text_shared = "%s/bpf_text_shared.c" % os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) 33bpf_text = open(bpf_text_shared, 'r').read() 34bpf_text += """ 35 36/** 37 * @brief The key to use for the latency histogram. 38 */ 39struct dist_key_t 40{ 41 char input[64]; ///< The input string of the request. 42 u64 slot; ///< The histogram slot. 43}; 44 45/** 46 * @brief Contains the histogram for the operation latencies. 47 */ 48BPF_HISTOGRAM(dist, struct dist_key_t); 49 50/** 51 * @brief Reads the operation response arguments, calculates the latency, and stores it in the histogram. 52 * @param ctx The BPF context. 53 */ 54int trace_operation_end(struct pt_regs* ctx) 55{ 56 u64 operation_id; 57 bpf_usdt_readarg(1, ctx, &operation_id); 58 59 struct start_data_t* start_data = start_hash.lookup(&operation_id); 60 if (0 == start_data) { 61 return 0; 62 } 63 64 u64 duration = bpf_ktime_get_ns() - start_data->start; 65 struct dist_key_t dist_key = {}; 66 __builtin_memcpy(&dist_key.input, start_data->input, sizeof(dist_key.input)); 67 dist_key.slot = bpf_log2l(duration / 1000); 68 start_hash.delete(&operation_id); 69 70 dist.increment(dist_key); 71 return 0; 72} 73""" 74 75bpf_text = bpf_text.replace("FILTER_STRING", this_filter) 76if this_filter: 77 bpf_text = bpf_text.replace("FILTER", "if (!filter(start_data.input)) { return 0; }") 78else: 79 bpf_text = bpf_text.replace("FILTER", "") 80 81# Create USDT context 82print("Attaching probes to pid %d" % this_pid) 83usdt_ctx = USDT(pid=this_pid) 84usdt_ctx.enable_probe(probe="operation_start", fn_name="trace_operation_start") 85usdt_ctx.enable_probe(probe="operation_end", fn_name="trace_operation_end") 86 87# Create BPF context, load BPF program 88bpf_ctx = BPF(text=bpf_text, usdt_contexts=[usdt_ctx], debug=debugLevel) 89 90start = 0 91dist = bpf_ctx.get_table("dist") 92while (1): 93 try: 94 sleep(this_interval) 95 except KeyboardInterrupt: 96 exit() 97 98 print("[%s]" % strftime("%H:%M:%S")) 99 dist.print_log2_hist("latency (us)") 100