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