1import argparse
2from time import sleep, strftime
3from sys import argv
4import ctypes as ct
5from bcc import BPF, USDT
6import inspect
7import os
8
9# Parse command line arguments
10parser = argparse.ArgumentParser(description="Trace the latency distribution of an operation using usdt probes.",
11    formatter_class=argparse.RawDescriptionHelpFormatter)
12parser.add_argument("-p", "--pid", type=int, help="The id of the process to trace.")
13parser.add_argument("-i", "--interval", type=int, help="The interval in seconds on which to report the latency distribution.")
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 verbose logging information.")
16parser.set_defaults(verbose=False)
17args = parser.parse_args()
18this_pid = int(args.pid)
19this_interval = int(args.interval)
20this_filter = str(args.filterstr)
21
22if this_interval < 1:
23    print("Invalid value for interval, using 1.")
24    this_interval = 1
25
26debugLevel=0
27if args.verbose:
28    debugLevel=4
29
30# BPF program
31bpf_text_shared = "%s/bpf_text_shared.c" % os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
32bpf_text = open(bpf_text_shared, 'r').read()
33bpf_text += """
34
35/**
36 * @brief The key to use for the latency histogram.
37 */
38struct 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 */
47BPF_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 */
53int 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
74bpf_text = bpf_text.replace("FILTER_STRING", this_filter)
75if this_filter:
76    bpf_text = bpf_text.replace("FILTER", "if (!filter(start_data.input)) { return 0; }")
77else:
78    bpf_text = bpf_text.replace("FILTER", "")
79
80# Create USDT context
81print("Attaching probes to pid %d" % this_pid)
82usdt_ctx = USDT(pid=this_pid)
83usdt_ctx.enable_probe(probe="operation_start", fn_name="trace_operation_start")
84usdt_ctx.enable_probe(probe="operation_end", fn_name="trace_operation_end")
85
86# Create BPF context, load BPF program
87bpf_ctx = BPF(text=bpf_text, usdt_contexts=[usdt_ctx], debug=debugLevel)
88
89start = 0
90dist = bpf_ctx.get_table("dist")
91while (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