Home | History | Annotate | Download | only in scripts 
      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)")
     99