examples/usdt_sample/scripts/latency.py

   1 import argparse
   2 from time import sleep
   3 from sys import argv
   4 import ctypes as ct
   5 from bcc import BPF, USDT
   6 import inspect
   7 import os
   8
   9 # Parse command line arguments
  10 parser = argparse.ArgumentParser(description="Trace the latency of an operation using usdt probes.",
  11     formatter_class=argparse.RawDescriptionHelpFormatter)
  12 parser.add_argument("-p", "--pid", type=int, help="The id of the process to trace.")
  13 parser.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.")
  14 parser.add_argument("-v", "--verbose", dest="verbose", action="store_true", help="If true, will output verbose logging information.")
  15 parser.set_defaults(verbose=False)
  16 args = parser.parse_args()
  17 this_pid = int(args.pid)
  18 this_filter = str(args.filterstr)
  19
  20 debugLevel=0
  21 if args.verbose:
  22     debugLevel=4
  23
  24 # BPF program
  25 bpf_text_shared = "%s/bpf_text_shared.c" % os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
  26 bpf_text = open(bpf_text_shared, 'r').read()
  27 bpf_text += """
  28
  29 /**
  30  * @brief Contains the latency data w.r.t. the complete operation from request to response.
  31  */
  32 struct end_data_t
  33 {
  34     u64 operation_id; ///< The id of the operation.
  35     char input[64];   ///< The request (input) string.
  36     char output[64];  ///< The response (output) string.
  37     u64 start;        ///< The start timestamp of the operation.
  38     u64 end;          ///< The end timestamp of the operation.
  39     u64 duration;     ///< The duration of the operation.
  40 };
  41
  42 /**
  43  * The output buffer, which will be used to push the latency event data to user space.
  44  */
  45 BPF_PERF_OUTPUT(operation_event);
  46
  47 /**
  48  * @brief Reads the operation response arguments, calculates the latency event data, and writes it to the user output buffer.
  49  * @param ctx The BPF context.
  50  */
  51 int trace_operation_end(struct pt_regs* ctx)
  52 {
  53     u64 operation_id;
  54     bpf_usdt_readarg(1, ctx, &operation_id);
  55
  56     struct start_data_t* start_data = start_hash.lookup(&operation_id);
  57     if (0 == start_data) {
  58         return 0;
  59     }
  60
  61     struct end_data_t end_data = {};
  62     end_data.operation_id = operation_id;
  63     bpf_usdt_readarg_p(2, ctx, &end_data.output, sizeof(end_data.output));
  64     end_data.end = bpf_ktime_get_ns();
  65     end_data.start = start_data->start;
  66     end_data.duration = end_data.end - end_data.start;
  67     __builtin_memcpy(&end_data.input, start_data->input, sizeof(end_data.input));
  68
  69     start_hash.delete(&end_data.operation_id);
  70
  71     operation_event.perf_submit(ctx, &end_data, sizeof(end_data));
  72     return 0;
  73 }
  74 """
  75
  76 bpf_text = bpf_text.replace("FILTER_STRING", this_filter)
  77 if this_filter:
  78     bpf_text = bpf_text.replace("FILTER", "if (!filter(start_data.input)) { return 0; }")
  79 else:
  80     bpf_text = bpf_text.replace("FILTER", "")
  81
  82 # Create USDT context
  83 print("Attaching probes to pid %d" % this_pid)
  84 usdt_ctx = USDT(pid=this_pid)
  85 usdt_ctx.enable_probe(probe="operation_start", fn_name="trace_operation_start")
  86 usdt_ctx.enable_probe(probe="operation_end", fn_name="trace_operation_end")
  87
  88 # Create BPF context, load BPF program
  89 bpf_ctx = BPF(text=bpf_text, usdt_contexts=[usdt_ctx], debug=debugLevel)
  90
  91 # Define latency event and print function
  92 class OperationEventData(ct.Structure):
  93   _fields_ = [("operation_id", ct.c_ulonglong),
  94               ("input", ct.c_char * 64),
  95               ("output", ct.c_char * 64),
  96               ("start", ct.c_ulonglong),
  97               ("end", ct.c_ulonglong),
  98               ("duration", ct.c_ulonglong)]
  99
 100 start = 0
 101 def print_event(cpu, data, size):
 102     global start
 103     event = ct.cast(data, ct.POINTER(OperationEventData)).contents
 104     if start == 0:
 105         start = event.start
 106     time_s = (float(event.start - start)) / 1000000000
 107     latency = (float(event.duration) / 1000)
 108     print("%-18.9f %-10d %-32s %-32s %16d %16d %16d" % (time_s, event.operation_id, event.input, event.output, event.start, event.end, latency))
 109
 110 # Print header
 111 print("Tracing... Hit Ctrl-C to end.")
 112 print("%-18s %-10s %-32s %-32s %16s %16s %16s" % ("time(s)", "id", "input", "output", "start (ns)", "end (ns)", "duration (us)"))
 113
 114 # Output latency events
 115 bpf_ctx["operation_event"].open_perf_buffer(print_event)
 116 while 1:
 117     bpf_ctx.kprobe_poll()