1#!/usr/bin/python
2# @lint-avoid-python-3-compatibility-imports
3#
4# ucalls  Summarize method calls in high-level languages and/or system calls.
5#         For Linux, uses BCC, eBPF.
6#
7# USAGE: ucalls [-l {java,perl,php,python,ruby,tcl}] [-h] [-T TOP] [-L] [-S] [-v] [-m]
8#        pid [interval]
9#
10# Copyright 2016 Sasha Goldshtein
11# Licensed under the Apache License, Version 2.0 (the "License")
12#
13# 19-Oct-2016   Sasha Goldshtein   Created this.
14
15from __future__ import print_function
16import argparse
17from bcc import BPF, USDT, utils
18from time import sleep
19import os
20
21languages = ["java", "perl", "php", "python", "ruby", "tcl"]
22
23examples = """examples:
24    ./ucalls -l java 185        # trace Java calls and print statistics on ^C
25    ./ucalls -l python 2020 1   # trace Python calls and print every second
26    ./ucalls -l java 185 -S     # trace Java calls and syscalls
27    ./ucalls 6712 -S            # trace only syscall counts
28    ./ucalls -l ruby 1344 -T 10 # trace top 10 Ruby method calls
29    ./ucalls -l ruby 1344 -L    # trace Ruby calls including latency
30    ./ucalls -l php 443 -LS     # trace PHP calls and syscalls with latency
31    ./ucalls -l python 2020 -mL # trace Python calls including latency in ms
32"""
33parser = argparse.ArgumentParser(
34    description="Summarize method calls in high-level languages.",
35    formatter_class=argparse.RawDescriptionHelpFormatter,
36    epilog=examples)
37parser.add_argument("pid", type=int, help="process id to attach to")
38parser.add_argument("interval", type=int, nargs='?',
39    help="print every specified number of seconds")
40parser.add_argument("-l", "--language", choices=languages + ["none"],
41    help="language to trace (if none, trace syscalls only)")
42parser.add_argument("-T", "--top", type=int,
43    help="number of most frequent/slow calls to print")
44parser.add_argument("-L", "--latency", action="store_true",
45    help="record method latency from enter to exit (except recursive calls)")
46parser.add_argument("-S", "--syscalls", action="store_true",
47    help="record syscall latency (adds overhead)")
48parser.add_argument("-v", "--verbose", action="store_true",
49    help="verbose mode: print the BPF program (for debugging purposes)")
50parser.add_argument("-m", "--milliseconds", action="store_true",
51    help="report times in milliseconds (default is microseconds)")
52parser.add_argument("--ebpf", action="store_true",
53    help=argparse.SUPPRESS)
54args = parser.parse_args()
55
56language = args.language
57if not language:
58    language = utils.detect_language(languages, args.pid)
59
60# We assume that the entry and return probes have the same arguments. This is
61# the case for Java, Python, Ruby, and PHP. If there's a language where it's
62# not the case, we will need to build a custom correlator from entry to exit.
63extra_message = ""
64if language == "java":
65    # TODO for JVM entries, we actually have the real length of the class
66    #      and method strings in arg3 and arg5 respectively, so we can insert
67    #      the null terminator in its proper position.
68    entry_probe = "method__entry"
69    return_probe = "method__return"
70    read_class = "bpf_usdt_readarg(2, ctx, &clazz);"
71    read_method = "bpf_usdt_readarg(4, ctx, &method);"
72    extra_message = ("If you do not see any results, make sure you ran java"
73                     " with option -XX:+ExtendedDTraceProbes")
74elif language == "perl":
75    entry_probe = "sub__entry"
76    return_probe = "sub__return"
77    read_class = "bpf_usdt_readarg(2, ctx, &clazz);"    # filename really
78    read_method = "bpf_usdt_readarg(1, ctx, &method);"
79elif language == "php":
80    entry_probe = "function__entry"
81    return_probe = "function__return"
82    read_class = "bpf_usdt_readarg(4, ctx, &clazz);"
83    read_method = "bpf_usdt_readarg(1, ctx, &method);"
84    extra_message = ("If you do not see any results, make sure the environment"
85                     " variable USE_ZEND_DTRACE is set to 1")
86elif language == "python":
87    entry_probe = "function__entry"
88    return_probe = "function__return"
89    read_class = "bpf_usdt_readarg(1, ctx, &clazz);"    # filename really
90    read_method = "bpf_usdt_readarg(2, ctx, &method);"
91elif language == "ruby":
92    # TODO Also probe cmethod__entry and cmethod__return with same arguments
93    entry_probe = "method__entry"
94    return_probe = "method__return"
95    read_class = "bpf_usdt_readarg(1, ctx, &clazz);"
96    read_method = "bpf_usdt_readarg(2, ctx, &method);"
97elif language == "tcl":
98    # TODO Also consider probe cmd__entry and cmd__return with same arguments
99    entry_probe = "proc__entry"
100    return_probe = "proc__return"
101    read_class = ""  # no class/file info available
102    read_method = "bpf_usdt_readarg(1, ctx, &method);"
103elif not language or language == "none":
104    if not args.syscalls:
105        print("Nothing to do; use -S to trace syscalls.")
106        exit(1)
107    entry_probe, return_probe, read_class, read_method = ("", "", "", "")
108    if language:
109        language = None
110
111program = """
112#include <linux/ptrace.h>
113
114#define MAX_STRING_LENGTH 80
115DEFINE_NOLANG
116DEFINE_LATENCY
117DEFINE_SYSCALLS
118
119struct method_t {
120    char clazz[MAX_STRING_LENGTH];
121    char method[MAX_STRING_LENGTH];
122};
123struct entry_t {
124    u64 pid;
125    struct method_t method;
126};
127struct info_t {
128    u64 num_calls;
129    u64 total_ns;
130};
131struct syscall_entry_t {
132    u64 timestamp;
133    u64 ip;
134};
135
136#ifndef LATENCY
137  BPF_HASH(counts, struct method_t, u64);            // number of calls
138  #ifdef SYSCALLS
139    BPF_HASH(syscounts, u64, u64);                   // number of calls per IP
140  #endif  // SYSCALLS
141#else
142  BPF_HASH(times, struct method_t, struct info_t);
143  BPF_HASH(entry, struct entry_t, u64);              // timestamp at entry
144  #ifdef SYSCALLS
145    BPF_HASH(systimes, u64, struct info_t);          // latency per IP
146    BPF_HASH(sysentry, u64, struct syscall_entry_t); // ts + IP at entry
147  #endif  // SYSCALLS
148#endif
149
150#ifndef NOLANG
151int trace_entry(struct pt_regs *ctx) {
152    u64 clazz = 0, method = 0, val = 0;
153    u64 *valp;
154    struct entry_t data = {0};
155#ifdef LATENCY
156    u64 timestamp = bpf_ktime_get_ns();
157    data.pid = bpf_get_current_pid_tgid();
158#endif
159    READ_CLASS
160    READ_METHOD
161    bpf_probe_read(&data.method.clazz, sizeof(data.method.clazz),
162                   (void *)clazz);
163    bpf_probe_read(&data.method.method, sizeof(data.method.method),
164                   (void *)method);
165#ifndef LATENCY
166    valp = counts.lookup_or_init(&data.method, &val);
167    ++(*valp);
168#endif
169#ifdef LATENCY
170    entry.update(&data, &timestamp);
171#endif
172    return 0;
173}
174
175#ifdef LATENCY
176int trace_return(struct pt_regs *ctx) {
177    u64 *entry_timestamp, clazz = 0, method = 0;
178    struct info_t *info, zero = {};
179    struct entry_t data = {};
180    data.pid = bpf_get_current_pid_tgid();
181    READ_CLASS
182    READ_METHOD
183    bpf_probe_read(&data.method.clazz, sizeof(data.method.clazz),
184                   (void *)clazz);
185    bpf_probe_read(&data.method.method, sizeof(data.method.method),
186                   (void *)method);
187    entry_timestamp = entry.lookup(&data);
188    if (!entry_timestamp) {
189        return 0;   // missed the entry event
190    }
191    info = times.lookup_or_init(&data.method, &zero);
192    info->num_calls += 1;
193    info->total_ns += bpf_ktime_get_ns() - *entry_timestamp;
194    entry.delete(&data);
195    return 0;
196}
197#endif  // LATENCY
198#endif  // NOLANG
199
200#ifdef SYSCALLS
201int syscall_entry(struct pt_regs *ctx) {
202    u64 pid = bpf_get_current_pid_tgid();
203    u64 *valp, ip = PT_REGS_IP(ctx), val = 0;
204    PID_FILTER
205#ifdef LATENCY
206    struct syscall_entry_t data = {};
207    data.timestamp = bpf_ktime_get_ns();
208    data.ip = ip;
209#endif
210#ifndef LATENCY
211    valp = syscounts.lookup_or_init(&ip, &val);
212    ++(*valp);
213#endif
214#ifdef LATENCY
215    sysentry.update(&pid, &data);
216#endif
217    return 0;
218}
219
220#ifdef LATENCY
221int syscall_return(struct pt_regs *ctx) {
222    struct syscall_entry_t *e;
223    struct info_t *info, zero = {};
224    u64 pid = bpf_get_current_pid_tgid(), ip;
225    PID_FILTER
226    e = sysentry.lookup(&pid);
227    if (!e) {
228        return 0;   // missed the entry event
229    }
230    ip = e->ip;
231    info = systimes.lookup_or_init(&ip, &zero);
232    info->num_calls += 1;
233    info->total_ns += bpf_ktime_get_ns() - e->timestamp;
234    sysentry.delete(&pid);
235    return 0;
236}
237#endif  // LATENCY
238#endif  // SYSCALLS
239""".replace("READ_CLASS", read_class) \
240   .replace("READ_METHOD", read_method) \
241   .replace("PID_FILTER", "if ((pid >> 32) != %d) { return 0; }" % args.pid) \
242   .replace("DEFINE_NOLANG", "#define NOLANG" if not language else "") \
243   .replace("DEFINE_LATENCY", "#define LATENCY" if args.latency else "") \
244   .replace("DEFINE_SYSCALLS", "#define SYSCALLS" if args.syscalls else "")
245
246if language:
247    usdt = USDT(pid=args.pid)
248    usdt.enable_probe_or_bail(entry_probe, "trace_entry")
249    if args.latency:
250        usdt.enable_probe_or_bail(return_probe, "trace_return")
251else:
252    usdt = None
253
254if args.ebpf or args.verbose:
255    if args.verbose and usdt:
256        print(usdt.get_text())
257    print(program)
258    if args.ebpf:
259        exit()
260
261bpf = BPF(text=program, usdt_contexts=[usdt] if usdt else [])
262if args.syscalls:
263    syscall_regex = "^[Ss]y[Ss]_.*"
264    bpf.attach_kprobe(event_re=syscall_regex, fn_name="syscall_entry")
265    if args.latency:
266        bpf.attach_kretprobe(event_re=syscall_regex, fn_name="syscall_return")
267    print("Attached %d kernel probes for syscall tracing." %
268          bpf.num_open_kprobes())
269
270def get_data():
271    # Will be empty when no language was specified for tracing
272    if args.latency:
273        data = list(map(lambda kv: (kv[0].clazz.decode('utf-8', 'replace') \
274                                    + "." + \
275                                    kv[0].method.decode('utf-8', 'replace'),
276                                   (kv[1].num_calls, kv[1].total_ns)),
277                   bpf["times"].items()))
278    else:
279        data = list(map(lambda kv: (kv[0].clazz.decode('utf-8', 'replace') \
280                                    + "." + \
281                                    kv[0].method.decode('utf-8', 'replace'),
282                                   (kv[1].value, 0)),
283                   bpf["counts"].items()))
284
285    if args.syscalls:
286        if args.latency:
287            syscalls = map(lambda kv: (bpf.ksym(kv[0].value),
288                                           (kv[1].num_calls, kv[1].total_ns)),
289                           bpf["systimes"].items())
290            data.extend(syscalls)
291        else:
292            syscalls = map(lambda kv: (bpf.ksym(kv[0].value),
293                                       (kv[1].value, 0)),
294                           bpf["syscounts"].items())
295            data.extend(syscalls)
296
297    return sorted(data, key=lambda kv: kv[1][1 if args.latency else 0])
298
299def clear_data():
300    if args.latency:
301        bpf["times"].clear()
302    else:
303        bpf["counts"].clear()
304
305    if args.syscalls:
306        if args.latency:
307            bpf["systimes"].clear()
308        else:
309            bpf["syscounts"].clear()
310
311exit_signaled = False
312print("Tracing calls in process %d (language: %s)... Ctrl-C to quit." %
313      (args.pid, language or "none"))
314if extra_message:
315    print(extra_message)
316while True:
317    try:
318        sleep(args.interval or 99999999)
319    except KeyboardInterrupt:
320        exit_signaled = True
321    print()
322    data = get_data()   # [(function, (num calls, latency in ns))]
323    if args.latency:
324        time_col = "TIME (ms)" if args.milliseconds else "TIME (us)"
325        print("%-50s %8s %8s" % ("METHOD", "# CALLS", time_col))
326    else:
327        print("%-50s %8s" % ("METHOD", "# CALLS"))
328    if args.top:
329        data = data[-args.top:]
330    for key, value in data:
331        if args.latency:
332            time = value[1] / 1000000.0 if args.milliseconds else \
333                   value[1] / 1000.0
334            print("%-50s %8d %6.2f" % (key, value[0], time))
335        else:
336            print("%-50s %8d" % (key, value[0]))
337    if args.interval and not exit_signaled:
338        clear_data()
339    else:
340        if args.syscalls:
341            print("Detaching kernel probes, please wait...")
342        exit()
343