examples/usdt_sample/scripts/lat_avg.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 moving average of 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("-i", "--interval", type=int, help="The interval in seconds on which to report the latency distribution.")
  14 parser.add_argument("-c", "--count", type=int, default=16, help="The count of samples over which to calculate the moving average.")
  15 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.")
  16 parser.add_argument("-v", "--verbose", dest="verbose", action="store_true", help="If true, will output verbose logging information.")
  17 parser.set_defaults(verbose=False)
  18 args = parser.parse_args()
  19 this_pid = int(args.pid)
  20 this_interval = int(args.interval)
  21 this_count = int(args.count)
  22 this_filter = str(args.filterstr)
  23
  24 if this_interval < 1:
  25     print("Invalid value for interval, using 1.")
  26     this_interval = 1
  27
  28 if this_count < 1:
  29     print("Invalid value for count, using 1.")
  30     this_count = 1
  31
  32 debugLevel=0
  33 if args.verbose:
  34     debugLevel=4
  35
  36 # BPF program
  37 bpf_text_shared = "%s/bpf_text_shared.c" % os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
  38 bpf_text = open(bpf_text_shared, 'r').read()
  39 bpf_text += """
  40
  41 const u32 MAX_SAMPLES = SAMPLE_COUNT;
  42
  43 struct hash_key_t
  44 {
  45     char input[64];
  46 };
  47
  48 struct hash_leaf_t
  49 {
  50     u32 count;
  51     u64 total;
  52     u64 average;
  53 };
  54
  55 /**
  56  * @brief Contains the averages for the operation latencies by operation input.
  57  */
  58 BPF_HASH(lat_hash, struct hash_key_t, struct hash_leaf_t, 512);
  59
  60 /**
  61  * @brief Reads the operation response arguments, calculates the latency, and stores it in the histogram.
  62  * @param ctx The BPF context.
  63  */
  64 int trace_operation_end(struct pt_regs* ctx)
  65 {
  66     u64 operation_id;
  67     bpf_usdt_readarg(1, ctx, &operation_id);
  68
  69     struct start_data_t* start_data = start_hash.lookup(&operation_id);
  70     if (0 == start_data) {
  71         return 0;
  72     }
  73
  74     u64 duration = bpf_ktime_get_ns() - start_data->start;
  75     struct hash_key_t hash_key = {};
  76     __builtin_memcpy(&hash_key.input, start_data->input, sizeof(hash_key.input));
  77     start_hash.delete(&operation_id);
  78
  79     struct hash_leaf_t zero = {};
  80     struct hash_leaf_t* hash_leaf = lat_hash.lookup_or_init(&hash_key, &zero);
  81     if (0 == hash_leaf) {
  82         return 0;
  83     }
  84
  85     if (hash_leaf->count < MAX_SAMPLES) {
  86         hash_leaf->count++;
  87     } else {
  88         hash_leaf->total -= hash_leaf->average;
  89     }
  90
  91     hash_leaf->total += duration;
  92     hash_leaf->average = hash_leaf->total / hash_leaf->count;
  93
  94     return 0;
  95 }
  96 """
  97
  98 bpf_text = bpf_text.replace("SAMPLE_COUNT", str(this_count))
  99 bpf_text = bpf_text.replace("FILTER_STRING", this_filter)
 100 if this_filter:
 101     bpf_text = bpf_text.replace("FILTER", "if (!filter(start_data.input)) { return 0; }")
 102 else:
 103     bpf_text = bpf_text.replace("FILTER", "")
 104
 105 # Create USDT context
 106 print("Attaching probes to pid %d" % this_pid)
 107 usdt_ctx = USDT(pid=this_pid)
 108 usdt_ctx.enable_probe(probe="operation_start", fn_name="trace_operation_start")
 109 usdt_ctx.enable_probe(probe="operation_end", fn_name="trace_operation_end")
 110
 111 # Create BPF context, load BPF program
 112 bpf_ctx = BPF(text=bpf_text, usdt_contexts=[usdt_ctx], debug=debugLevel)
 113
 114 print("Tracing... Hit Ctrl-C to end.")
 115
 116 lat_hash = bpf_ctx.get_table("lat_hash")
 117 while (1):
 118     try:
 119         sleep(this_interval)
 120     except KeyboardInterrupt:
 121         exit()
 122
 123     print("[%s]" % strftime("%H:%M:%S"))
 124     print("%-64s %8s %16s" % ("input", "count", "latency (us)"))
 125     for k, v in lat_hash.items():
 126         print("%-64s %8d %16d" % (k.input, v.count, v.average / 1000))