1#!/usr/bin/env python
2#
3# stackcount    Count kernel function calls and their stack traces.
4#               For Linux, uses BCC, eBPF.
5#
6# USAGE: stackcount [-h] [-p PID] [-i INTERVAL] [-T] [-r] pattern
7#
8# The pattern is a string with optional '*' wildcards, similar to file
9# globbing. If you'd prefer to use regular expressions, use the -r option.
10#
11# The current implementation uses an unrolled loop for x86_64, and was written
12# as a proof of concept. This implementation should be replaced in the future
13# with an appropriate bpf_ call, when available.
14#
15# Currently limited to a stack trace depth of 11 (maxdepth + 1).
16#
17# Copyright 2016 Netflix, Inc.
18# Licensed under the Apache License, Version 2.0 (the "License")
19#
20# 12-Jan-2016	Brendan Gregg	Created this.
21
22from __future__ import print_function
23from bcc import BPF
24from time import sleep, strftime
25import argparse
26import signal
27
28# arguments
29examples = """examples:
30    ./stackcount submit_bio       # count kernel stack traces for submit_bio
31    ./stackcount ip_output        # count kernel stack traces for ip_output
32    ./stackcount -s ip_output     # show symbol offsets
33    ./stackcount -sv ip_output    # show offsets and raw addresses (verbose)
34    ./stackcount 'tcp_send*'      # count stacks for funcs matching tcp_send*
35    ./stackcount -r '^tcp_send.*' # same as above, using regular expressions
36    ./stackcount -Ti 5 ip_output  # output every 5 seconds, with timestamps
37    ./stackcount -p 185 ip_output # count ip_output stacks for PID 185 only
38"""
39parser = argparse.ArgumentParser(
40    description="Count kernel function calls and their stack traces",
41    formatter_class=argparse.RawDescriptionHelpFormatter,
42    epilog=examples)
43parser.add_argument("-p", "--pid",
44    help="trace this PID only")
45parser.add_argument("-i", "--interval", default=99999999,
46    help="summary interval, seconds")
47parser.add_argument("-T", "--timestamp", action="store_true",
48    help="include timestamp on output")
49parser.add_argument("-r", "--regexp", action="store_true",
50    help="use regular expressions. Default is \"*\" wildcards only.")
51parser.add_argument("-s", "--offset", action="store_true",
52    help="show address offsets")
53parser.add_argument("-v", "--verbose", action="store_true",
54    help="show raw addresses")
55parser.add_argument("pattern",
56    help="search expression for kernel functions")
57args = parser.parse_args()
58pattern = args.pattern
59if not args.regexp:
60    pattern = pattern.replace('*', '.*')
61    pattern = '^' + pattern + '$'
62offset = args.offset
63verbose = args.verbose
64debug = 0
65maxdepth = 10    # and MAXDEPTH
66
67# signal handler
68def signal_ignore(signal, frame):
69    print()
70
71# load BPF program
72bpf_text = """
73#include <uapi/linux/ptrace.h>
74
75#define MAXDEPTH	10
76
77struct key_t {
78    u64 ip;
79    u64 ret[MAXDEPTH];
80};
81BPF_HASH(counts, struct key_t);
82
83static u64 get_frame(u64 *bp) {
84    if (*bp) {
85        // The following stack walker is x86_64 specific
86        u64 ret = 0;
87        if (bpf_probe_read(&ret, sizeof(ret), (void *)(*bp+8)))
88            return 0;
89        if (bpf_probe_read(bp, sizeof(*bp), (void *)*bp))
90            *bp = 0;
91        if (ret < __START_KERNEL_map)
92            return 0;
93        return ret;
94    }
95    return 0;
96}
97
98int trace_count(struct pt_regs *ctx) {
99    FILTER
100    struct key_t key = {};
101    u64 zero = 0, *val, bp = 0;
102    int depth = 0;
103
104    key.ip = ctx->ip;
105    bp = ctx->bp;
106
107    // unrolled loop, 10 (MAXDEPTH) frames deep:
108    if (!(key.ret[depth++] = get_frame(&bp))) goto out;
109    if (!(key.ret[depth++] = get_frame(&bp))) goto out;
110    if (!(key.ret[depth++] = get_frame(&bp))) goto out;
111    if (!(key.ret[depth++] = get_frame(&bp))) goto out;
112    if (!(key.ret[depth++] = get_frame(&bp))) goto out;
113    if (!(key.ret[depth++] = get_frame(&bp))) goto out;
114    if (!(key.ret[depth++] = get_frame(&bp))) goto out;
115    if (!(key.ret[depth++] = get_frame(&bp))) goto out;
116    if (!(key.ret[depth++] = get_frame(&bp))) goto out;
117    if (!(key.ret[depth++] = get_frame(&bp))) goto out;
118
119out:
120    val = counts.lookup_or_init(&key, &zero);
121    (*val)++;
122    return 0;
123}
124"""
125if args.pid:
126    bpf_text = bpf_text.replace('FILTER',
127        ('u32 pid; pid = bpf_get_current_pid_tgid(); ' +
128        'if (pid != %s) { return 0; }') % (args.pid))
129else:
130    bpf_text = bpf_text.replace('FILTER', '')
131if debug:
132    print(bpf_text)
133b = BPF(text=bpf_text)
134b.attach_kprobe(event_re=pattern, fn_name="trace_count")
135matched = b.num_open_kprobes()
136if matched == 0:
137    print("0 functions matched by \"%s\". Exiting." % args.pattern)
138    exit()
139
140# header
141print("Tracing %d functions for \"%s\"... Hit Ctrl-C to end." %
142    (matched, args.pattern))
143
144def print_frame(addr):
145    print("  ", end="")
146    if verbose:
147        print("%-16x " % addr, end="")
148    print(b.ksym(addr, show_offset=offset))
149
150# output
151exiting = 0 if args.interval else 1
152while (1):
153    try:
154        sleep(int(args.interval))
155    except KeyboardInterrupt:
156        exiting = 1
157        # as cleanup can take many seconds, trap Ctrl-C:
158        signal.signal(signal.SIGINT, signal_ignore)
159
160    print()
161    if args.timestamp:
162        print("%-8s\n" % strftime("%H:%M:%S"), end="")
163
164    counts = b.get_table("counts")
165    for k, v in sorted(counts.items(), key=lambda counts: counts[1].value):
166        print_frame(k.ip)
167        for i in range(0, maxdepth):
168            if k.ret[i] == 0:
169                break
170            print_frame(k.ret[i])
171        print("    %d\n" % v.value)
172    counts.clear()
173
174    if exiting:
175        print("Detaching...")
176        exit()
177