1#!/usr/bin/python 2import argparse 3from time import sleep 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 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("-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.") 15parser.add_argument("-v", "--verbose", dest="verbose", action="store_true", help="If true, will output generated bpf program and verbose logging information.") 16parser.add_argument("-s", "--sdt", dest="sdt", action="store_true", help="If true, will use the probes, created by systemtap's dtrace.") 17 18parser.set_defaults(verbose=False) 19args = parser.parse_args() 20this_pid = int(args.pid) 21this_filter = str(args.filterstr) 22 23debugLevel=0 24if args.verbose: 25 debugLevel=4 26 27# BPF program 28bpf_text_shared = "%s/bpf_text_shared.c" % os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) 29bpf_text = open(bpf_text_shared, 'r').read() 30bpf_text += """ 31 32/** 33 * @brief Contains the latency data w.r.t. the complete operation from request to response. 34 */ 35struct end_data_t 36{ 37 u64 operation_id; ///< The id of the operation. 38 char input[64]; ///< The request (input) string. 39 char output[64]; ///< The response (output) string. 40 u64 start; ///< The start timestamp of the operation. 41 u64 end; ///< The end timestamp of the operation. 42 u64 duration; ///< The duration of the operation. 43}; 44 45/** 46 * The output buffer, which will be used to push the latency event data to user space. 47 */ 48BPF_PERF_OUTPUT(operation_event); 49 50/** 51 * @brief Reads the operation response arguments, calculates the latency event data, and writes it to the user output buffer. 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 struct end_data_t end_data = {}; 65 end_data.operation_id = operation_id; 66 bpf_usdt_readarg_p(2, ctx, &end_data.output, sizeof(end_data.output)); 67 end_data.end = bpf_ktime_get_ns(); 68 end_data.start = start_data->start; 69 end_data.duration = end_data.end - end_data.start; 70 __builtin_memcpy(&end_data.input, start_data->input, sizeof(end_data.input)); 71 72 start_hash.delete(&end_data.operation_id); 73 74 operation_event.perf_submit(ctx, &end_data, sizeof(end_data)); 75 return 0; 76} 77""" 78 79bpf_text = bpf_text.replace("FILTER_STRING", this_filter) 80if this_filter: 81 bpf_text = bpf_text.replace("FILTER_STATEMENT", "if (!filter(start_data.input)) { return 0; }") 82else: 83 bpf_text = bpf_text.replace("FILTER_STATEMENT", "") 84 85# Create USDT context 86print("latency.py - Attaching probes to pid: %d; filter: %s" % (this_pid, this_filter)) 87usdt_ctx = USDT(pid=this_pid) 88 89if args.sdt: 90 usdt_ctx.enable_probe(probe="usdt_sample_lib1_sdt:operation_start_sdt", fn_name="trace_operation_start") 91 usdt_ctx.enable_probe(probe="usdt_sample_lib1_sdt:operation_end_sdt", fn_name="trace_operation_end") 92else: 93 usdt_ctx.enable_probe(probe="usdt_sample_lib1:operation_start", fn_name="trace_operation_start") 94 usdt_ctx.enable_probe(probe="usdt_sample_lib1:operation_end", fn_name="trace_operation_end") 95 96# Create BPF context, load BPF program 97bpf_ctx = BPF(text=bpf_text, usdt_contexts=[usdt_ctx], debug=debugLevel) 98 99# Define latency event and print function 100class OperationEventData(ct.Structure): 101 _fields_ = [("operation_id", ct.c_ulonglong), 102 ("input", ct.c_char * 64), 103 ("output", ct.c_char * 64), 104 ("start", ct.c_ulonglong), 105 ("end", ct.c_ulonglong), 106 ("duration", ct.c_ulonglong)] 107 108start = 0 109def print_event(cpu, data, size): 110 global start 111 event = ct.cast(data, ct.POINTER(OperationEventData)).contents 112 if start == 0: 113 start = event.start 114 time_s = (float(event.start - start)) / 1000000000 115 latency = (float(event.duration) / 1000) 116 print("%-18.9f %-10d %-32s %-32s %16d %16d %16d" % (time_s, event.operation_id, event.input, event.output, event.start, event.end, latency)) 117 118# Print header 119print("Tracing... Hit Ctrl-C to end.") 120print("%-18s %-10s %-32s %-32s %16s %16s %16s" % ("time(s)", "id", "input", "output", "start (ns)", "end (ns)", "duration (us)")) 121 122# Output latency events 123bpf_ctx["operation_event"].open_perf_buffer(print_event) 124while 1: 125 bpf_ctx.perf_buffer_poll() 126