1import argparse
2from time import sleep
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 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("-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.")
14parser.add_argument("-v", "--verbose", dest="verbose", action="store_true", help="If true, will output verbose logging information.")
15parser.set_defaults(verbose=False)
16args = parser.parse_args()
17this_pid = int(args.pid)
18this_filter = str(args.filterstr)
19
20debugLevel=0
21if args.verbose:
22    debugLevel=4
23
24# BPF program
25bpf_text_shared = "%s/bpf_text_shared.c" % os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
26bpf_text = open(bpf_text_shared, 'r').read()
27bpf_text += """
28
29/**
30 * @brief Contains the latency data w.r.t. the complete operation from request to response.
31 */
32struct 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 */
45BPF_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 */
51int 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
76bpf_text = bpf_text.replace("FILTER_STRING", this_filter)
77if this_filter:
78    bpf_text = bpf_text.replace("FILTER", "if (!filter(start_data.input)) { return 0; }")
79else:
80    bpf_text = bpf_text.replace("FILTER", "")
81
82# Create USDT context
83print("Attaching probes to pid %d" % this_pid)
84usdt_ctx = USDT(pid=this_pid)
85usdt_ctx.enable_probe(probe="operation_start", fn_name="trace_operation_start")
86usdt_ctx.enable_probe(probe="operation_end", fn_name="trace_operation_end")
87
88# Create BPF context, load BPF program
89bpf_ctx = BPF(text=bpf_text, usdt_contexts=[usdt_ctx], debug=debugLevel)
90
91# Define latency event and print function
92class 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
100start = 0
101def 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
111print("Tracing... Hit Ctrl-C to end.")
112print("%-18s %-10s %-32s %-32s %16s %16s %16s" % ("time(s)", "id", "input", "output", "start (ns)", "end (ns)", "duration (us)"))
113
114# Output latency events
115bpf_ctx["operation_event"].open_perf_buffer(print_event)
116while 1:
117    bpf_ctx.perf_buffer_poll()
118