uflow: Trace method execution flow

This tool traces method executions and prints out messages
with a nesting depth indicator that allows relatively easy
visualization of program flow. To reduce overhead, `-C` and
`-M` switches take class and method prefixes to filter the
output by. The filtering takes place in the BPF program, so
the unnecessary data never ends up in user space.

Usage examples:

```
uflow java 185
uflow -M index python 180
uflow -C 'IRB::Irb' ruby 10041
```

tools/uflow.py

@@ -0,0 +1,174 @@
+#!/usr/bin/python
+# @lint-avoid-python-3-compatibility-imports
+#
+# uflow  Trace method execution flow in high-level languages.
+#        For Linux, uses BCC, eBPF.
+#
+# USAGE: uflow [-C CLASS] [-M METHOD] [-v] {java,python,ruby} pid
+#
+# Copyright 2016 Sasha Goldshtein
+# Licensed under the Apache License, Version 2.0 (the "License")
+#
+# 27-Oct-2016   Sasha Goldshtein   Created this.
+
+from __future__ import print_function
+import argparse
+from bcc import BPF, USDT
+import ctypes as ct
+import time
+
+examples = """examples:
+    ./uflow java 185                # trace Java method calls in process 185
+    ./uflow ruby 1344               # trace Ruby method calls in process 1344
+    ./uflow -M indexOf java 185     # trace only 'indexOf'-prefixed methods
+    ./uflow -C '<stdin>' python 180 # trace only REPL-defined methods
+"""
+parser = argparse.ArgumentParser(
+    description="Trace method execution flow in high-level languages.",
+    formatter_class=argparse.RawDescriptionHelpFormatter,
+    epilog=examples)
+parser.add_argument("language", choices=["java", "python", "ruby"],
+    help="language to trace")
+parser.add_argument("pid", type=int, help="process id to attach to")
+parser.add_argument("-M", "--method",
+    help="trace only calls to methods starting with this prefix")
+parser.add_argument("-C", "--class", dest="clazz",
+    help="trace only calls to classes starting with this prefix")
+parser.add_argument("-v", "--verbose", action="store_true",
+    help="verbose mode: print the BPF program (for debugging purposes)")
+args = parser.parse_args()
+
+usdt = USDT(pid=args.pid)
+
+program = """
+struct call_t {
+    u64 depth;                  // first bit is direction (0 entry, 1 return)
+    u64 pid;                    // (tgid << 32) + pid from bpf_get_current...
+    u64 timestamp;              // ns
+    char clazz[80];
+    char method[80];
+};
+
+BPF_PERF_OUTPUT(calls);
+BPF_HASH(entry, u64, u64);
+"""
+
+prefix_template = """
+static inline bool prefix_%s(char *actual) {
+    char expected[] = "%s";
+    for (int i = 0; i < sizeof(expected) - 1; ++i) {
+        if (expected[i] != actual[i]) {
+            return false;
+        }
+    }
+    return true;
+}
+"""
+
+if args.clazz:
+    program += prefix_template % ("class", args.clazz)
+if args.method:
+    program += prefix_template % ("method", args.method)
+
+trace_template = """
+int NAME(struct pt_regs *ctx) {
+    u64 *depth, zero = 0, clazz = 0, method = 0 ;
+    struct call_t data = {};
+
+    READ_CLASS
+    READ_METHOD
+    bpf_probe_read(&data.clazz, sizeof(data.clazz), (void *)clazz);
+    bpf_probe_read(&data.method, sizeof(data.method), (void *)method);
+
+    FILTER_CLASS
+    FILTER_METHOD
+
+    data.pid = bpf_get_current_pid_tgid();
+    data.timestamp = bpf_ktime_get_ns();
+    depth = entry.lookup_or_init(&data.pid, &zero);
+    data.depth = DEPTH;
+    UPDATE
+
+    calls.perf_submit(ctx, &data, sizeof(data));
+    return 0;
+}
+"""
+
+def enable_probe(probe_name, func_name, read_class, read_method, is_return):
+    global program, trace_template, usdt
+    depth = "*depth + 1" if not is_return else "*depth | (1ULL << 63)"
+    update = "++(*depth);" if not is_return else  "if (*depth) --(*depth);"
+    filter_class = "if (!prefix_class(data.clazz)) { return 0; }" \
+                   if args.clazz else ""
+    filter_method = "if (!prefix_method(data.method)) { return 0; }" \
+                   if args.method else ""
+    program += trace_template.replace("NAME", func_name)                \
+                             .replace("READ_CLASS", read_class)         \
+                             .replace("READ_METHOD", read_method)       \
+                             .replace("FILTER_CLASS", filter_class)     \
+                             .replace("FILTER_METHOD", filter_method)   \
+                             .replace("DEPTH", depth)                   \
+                             .replace("UPDATE", update)
+    usdt.enable_probe(probe_name, func_name)
+
+usdt = USDT(pid=args.pid)
+
+if args.language == "java":
+    enable_probe("method__entry", "java_entry",
+                 "bpf_usdt_readarg(2, ctx, &clazz);",
+                 "bpf_usdt_readarg(4, ctx, &method);", is_return=False)
+    enable_probe("method__return", "java_return",
+                 "bpf_usdt_readarg(2, ctx, &clazz);",
+                 "bpf_usdt_readarg(4, ctx, &method);", is_return=True)
+elif args.language == "python":
+    enable_probe("function__entry", "python_entry",
+                 "bpf_usdt_readarg(1, ctx, &clazz);",   # filename really
+                 "bpf_usdt_readarg(2, ctx, &method);", is_return=False)
+    enable_probe("function__return", "python_return",
+                 "bpf_usdt_readarg(1, ctx, &clazz);",   # filename really
+                 "bpf_usdt_readarg(2, ctx, &method);", is_return=True)
+elif args.language == "ruby":
+    enable_probe("method__entry", "ruby_entry",
+                 "bpf_usdt_readarg(1, ctx, &clazz);",
+                 "bpf_usdt_readarg(2, ctx, &method);", is_return=False)
+    enable_probe("method__return", "ruby_return",
+                 "bpf_usdt_readarg(1, ctx, &clazz);",
+                 "bpf_usdt_readarg(2, ctx, &method);", is_return=True)
+    enable_probe("cmethod__entry", "ruby_centry",
+                 "bpf_usdt_readarg(1, ctx, &clazz);",
+                 "bpf_usdt_readarg(2, ctx, &method);", is_return=False)
+    enable_probe("cmethod__return", "ruby_creturn",
+                 "bpf_usdt_readarg(1, ctx, &clazz);",
+                 "bpf_usdt_readarg(2, ctx, &method);", is_return=True)
+
+if args.verbose:
+    print(usdt.get_text())
+    print(program)
+
+bpf = BPF(text=program, usdt_contexts=[usdt])
+print("Tracing method calls in %s process %d... Ctrl-C to quit." %
+      (args.language, args.pid))
+print("%-3s %-6s %-6s %-8s %s" % ("CPU", "PID", "TID", "TIME(us)", "METHOD"))
+
+class CallEvent(ct.Structure):
+    _fields_ = [
+        ("depth", ct.c_ulonglong),
+        ("pid", ct.c_ulonglong),
+        ("timestamp", ct.c_ulonglong),
+        ("clazz", ct.c_char * 80),
+        ("method", ct.c_char * 80)
+        ]
+
+start_ts = time.time()
+
+def print_event(cpu, data, size):
+    event = ct.cast(data, ct.POINTER(CallEvent)).contents
+    depth = event.depth & (~(1 << 63))
+    direction = "<- " if event.depth & (1 << 63) else "-> "
+    print("%-3d %-6d %-6d %-8.3f %-40s" % (cpu, event.pid >> 32,
+        event.pid & 0xFFFFFFFF, time.time() - start_ts,
+        ("  " * depth) + direction + event.clazz + "." + event.method))
+
+bpf["calls"].open_perf_buffer(print_event)
+while 1:
+    bpf.kprobe_poll()