• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python
2import argparse
3from time import sleep
4from sys import argv
5import ctypes as ct
6from bcc import BPF, USDT
7import inspect
8import os
9
10# Parse command line arguments
11parser = argparse.ArgumentParser(description="Trace the latency of an operation using usdt probes.",
12    formatter_class=argparse.RawDescriptionHelpFormatter)
13parser.add_argument("-p", "--pid", type=int, help="The id of the process to trace.")
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_filter = str(args.filterstr)
20
21debugLevel=0
22if args.verbose:
23    debugLevel=4
24
25# BPF program
26bpf_text_shared = "%s/bpf_text_shared.c" % os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
27bpf_text = open(bpf_text_shared, 'r').read()
28bpf_text += """
29
30/**
31 * @brief Contains the latency data w.r.t. the complete operation from request to response.
32 */
33struct end_data_t
34{
35    u64 operation_id; ///< The id of the operation.
36    char input[64];   ///< The request (input) string.
37    char output[64];  ///< The response (output) string.
38    u64 start;        ///< The start timestamp of the operation.
39    u64 end;          ///< The end timestamp of the operation.
40    u64 duration;     ///< The duration of the operation.
41};
42
43/**
44 * The output buffer, which will be used to push the latency event data to user space.
45 */
46BPF_PERF_OUTPUT(operation_event);
47
48/**
49 * @brief Reads the operation response arguments, calculates the latency event data, and writes it to the user output buffer.
50 * @param ctx The BPF context.
51 */
52int trace_operation_end(struct pt_regs* ctx)
53{
54    u64 operation_id;
55    bpf_usdt_readarg(1, ctx, &operation_id);
56
57    struct start_data_t* start_data = start_hash.lookup(&operation_id);
58    if (0 == start_data) {
59        return 0;
60    }
61
62    struct end_data_t end_data = {};
63    end_data.operation_id = operation_id;
64    bpf_usdt_readarg_p(2, ctx, &end_data.output, sizeof(end_data.output));
65    end_data.end = bpf_ktime_get_ns();
66    end_data.start = start_data->start;
67    end_data.duration = end_data.end - end_data.start;
68    __builtin_memcpy(&end_data.input, start_data->input, sizeof(end_data.input));
69
70    start_hash.delete(&end_data.operation_id);
71
72    operation_event.perf_submit(ctx, &end_data, sizeof(end_data));
73    return 0;
74}
75"""
76
77bpf_text = bpf_text.replace("FILTER_STRING", this_filter)
78if this_filter:
79    bpf_text = bpf_text.replace("FILTER", "if (!filter(start_data.input)) { return 0; }")
80else:
81    bpf_text = bpf_text.replace("FILTER", "")
82
83# Create USDT context
84print("Attaching probes to pid %d" % this_pid)
85usdt_ctx = USDT(pid=this_pid)
86usdt_ctx.enable_probe(probe="operation_start", fn_name="trace_operation_start")
87usdt_ctx.enable_probe(probe="operation_end", fn_name="trace_operation_end")
88
89# Create BPF context, load BPF program
90bpf_ctx = BPF(text=bpf_text, usdt_contexts=[usdt_ctx], debug=debugLevel)
91
92# Define latency event and print function
93class OperationEventData(ct.Structure):
94  _fields_ = [("operation_id", ct.c_ulonglong),
95              ("input", ct.c_char * 64),
96              ("output", ct.c_char * 64),
97              ("start", ct.c_ulonglong),
98              ("end", ct.c_ulonglong),
99              ("duration", ct.c_ulonglong)]
100
101start = 0
102def print_event(cpu, data, size):
103    global start
104    event = ct.cast(data, ct.POINTER(OperationEventData)).contents
105    if start == 0:
106        start = event.start
107    time_s = (float(event.start - start)) / 1000000000
108    latency = (float(event.duration) / 1000)
109    print("%-18.9f %-10d %-32s %-32s %16d %16d %16d" % (time_s, event.operation_id, event.input, event.output, event.start, event.end, latency))
110
111# Print header
112print("Tracing... Hit Ctrl-C to end.")
113print("%-18s %-10s %-32s %-32s %16s %16s %16s" % ("time(s)", "id", "input", "output", "start (ns)", "end (ns)", "duration (us)"))
114
115# Output latency events
116bpf_ctx["operation_event"].open_perf_buffer(print_event)
117while 1:
118    bpf_ctx.perf_buffer_poll()
119