examples/usdt_sample/scripts/lat_dist.py

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