Added memory leak tracer

diff --git a/tools/memleak.c b/tools/memleak.c
new file mode 100644 (file)
index 0000000..b1b8f1e
--- /dev/null
+++ b/tools/memleak.c
@@ -0,0 +1,88 @@
+#include <uapi/linux/ptrace.h>
+
+#define MAX_STACK_SIZE 10
+
+struct alloc_info_t {
+       u64 size;
+       int num_frames;
+       u64 callstack[MAX_STACK_SIZE];
+};
+
+BPF_HASH(sizes, u64);
+BPF_HASH(allocs, u64, struct alloc_info_t);
+
+// Adapted from https://github.com/iovisor/bcc/tools/offcputime.py
+static u64 get_frame(u64 *bp) {
+       if (*bp) {
+               // The following stack walker is x86_64 specific
+               u64 ret = 0;
+               if (bpf_probe_read(&ret, sizeof(ret), (void *)(*bp+8)))
+                       return 0;
+               if (bpf_probe_read(bp, sizeof(*bp), (void *)*bp))
+                       *bp = 0;
+               return ret;
+       }
+       return 0;
+}
+static int grab_stack(struct pt_regs *ctx, struct alloc_info_t *info)
+{
+       int depth = 0;
+       u64 bp = ctx->bp;
+       if (!(info->callstack[depth++] = get_frame(&bp))) return depth;
+       if (!(info->callstack[depth++] = get_frame(&bp))) return depth;
+        if (!(info->callstack[depth++] = get_frame(&bp))) return depth;
+        if (!(info->callstack[depth++] = get_frame(&bp))) return depth;
+        if (!(info->callstack[depth++] = get_frame(&bp))) return depth;
+        if (!(info->callstack[depth++] = get_frame(&bp))) return depth;
+        if (!(info->callstack[depth++] = get_frame(&bp))) return depth;
+        if (!(info->callstack[depth++] = get_frame(&bp))) return depth;
+        if (!(info->callstack[depth++] = get_frame(&bp))) return depth;
+        if (!(info->callstack[depth++] = get_frame(&bp))) return depth;
+       return depth;
+}
+
+int alloc_enter(struct pt_regs *ctx, size_t size)
+{
+       u64 pid = bpf_get_current_pid_tgid();
+       u64 size64 = size;
+       sizes.update(&pid, &size64);
+
+        if (SHOULD_PRINT)
+                bpf_trace_printk("alloc entered, size = %u\n", size);
+        return 0;
+}
+
+int alloc_exit(struct pt_regs *ctx)
+{
+       u64 address = ctx->ax;
+       u64 pid = bpf_get_current_pid_tgid();
+       u64* size64 = sizes.lookup(&pid);
+       struct alloc_info_t info = {0};
+
+       if (size64 == 0)
+               return 0; // missed alloc entry
+
+       info.size = *size64;
+       sizes.delete(&pid);
+
+       info.num_frames = grab_stack(ctx, &info) - 2;
+       allocs.update(&address, &info);
+       
+        if (SHOULD_PRINT)
+                bpf_trace_printk("alloc exited, size = %lu, result = %lx, frames = %d\n", info.size, address, info.num_frames);
+        return 0;
+}
+
+int free_enter(struct pt_regs *ctx, void *address)
+{
+       u64 addr = (u64)address;
+       struct alloc_info_t *info = allocs.lookup(&addr);
+       if (info == 0)
+               return 0;
+
+       allocs.delete(&addr);
+
+        if (SHOULD_PRINT)
+                bpf_trace_printk("free entered, address = %lx, size = %lu\n", address, info->size);
+        return 0;
+}

diff --git a/tools/memleak.py b/tools/memleak.py
new file mode 100755 (executable)
index 0000000..eed6a16
--- /dev/null
+++ b/tools/memleak.py
@@ -0,0 +1,144 @@
+#!/usr/bin/env python
+
+from bcc import BPF
+from time import sleep
+import argparse
+import subprocess
+
+examples = """
+EXAMPLES:
+
+memleak.py -p $(pidof allocs)
+       Trace allocations and display a summary of "leaked" (outstanding)
+       allocations every 5 seconds
+memleak.py -p $(pidof allocs) -t
+       Trace allocations and display each individual call to malloc/free
+memleak.py -p $(pidof allocs) -a -i 10
+       Trace allocations and display allocated addresses, sizes, and stacks
+       every 10 seconds for outstanding allocations
+memleak.py
+       Trace allocations in kernel mode and display a summary of outstanding
+       allocations every 5 seconds
+"""
+
+description = """
+Trace outstanding memory allocations that weren't freed.
+Supports both user-mode allocations made with malloc/free and kernel-mode
+allocations made with kmalloc/kfree.
+"""
+
+parser = argparse.ArgumentParser(description=description, formatter_class=argparse.RawDescriptionHelpFormatter, epilog=examples)
+parser.add_argument("-p", "--pid", help="the PID to trace; if not specified, trace kernel allocs")
+parser.add_argument("-t", "--trace", action="store_true", help="print trace messages for each alloc/free call")
+parser.add_argument("-i", "--interval", default=5, help="interval in seconds to print outstanding allocations")
+parser.add_argument("-a", "--show-allocs", default=False, action="store_true", help="show allocation addresses and sizes as well as call stacks")
+
+args = parser.parse_args()
+
+pid = -1 if args.pid is None else int(args.pid)
+kernel_trace = (pid == -1)
+trace_all = args.trace
+interval = int(args.interval)
+
+bpf_source = open("memleak.c").read()
+bpf_source = bpf_source.replace("SHOULD_PRINT", "1" if trace_all else "0")
+
+bpf_program = BPF(text=bpf_source)
+
+if not kernel_trace:
+       print("Attaching to malloc and free in pid %d, Ctrl+C to quit." % pid)
+       bpf_program.attach_uprobe(name="c", sym="malloc", fn_name="alloc_enter", pid=pid)
+       bpf_program.attach_uretprobe(name="c", sym="malloc", fn_name="alloc_exit", pid=pid)
+       bpf_program.attach_uprobe(name="c", sym="free", fn_name="free_enter", pid=pid)
+else:
+       print("Attaching to kmalloc and kfree, Ctrl+C to quit.")
+       bpf_program.attach_kprobe(event="__kmalloc", fn_name="alloc_enter")
+       bpf_program.attach_kretprobe(event="__kmalloc", fn_name="alloc_exit")
+       bpf_program.attach_kprobe(event="kfree", fn_name="free_enter")
+
+def get_code_ranges(pid):
+       ranges = {}
+       raw_ranges = open("/proc/%d/maps" % pid).readlines()
+       for raw_range in raw_ranges:
+               parts = raw_range.split()
+               if len(parts) < 6 or parts[5][0] == '[' or not 'x' in parts[1]:
+                       continue
+               binary = parts[5]
+               range_parts = parts[0].split('-')
+               addr_range = (int(range_parts[0], 16), int(range_parts[1], 16))
+               ranges[binary] = addr_range
+       return ranges
+
+def run_command(command):
+       p = subprocess.Popen(command.split(), stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+       return iter(p.stdout.readline, b'')
+
+ranges_cache = {}
+
+def get_sym_ranges(binary):
+       if binary in ranges_cache:
+               return ranges_cache[binary]
+       sym_ranges = {}
+       raw_symbols = run_command("objdump -t %s" % binary)
+       for raw_symbol in raw_symbols:
+               parts = raw_symbol.split()
+               if len(parts) < 6 or parts[3] != ".text" or parts[2] != "F":
+                       continue
+               sym_start = int(parts[0], 16)
+               sym_len = int(parts[4], 16)
+               sym_name = parts[5]
+               sym_ranges[sym_name] = (sym_start, sym_len)
+       ranges_cache[binary] = sym_ranges
+       return sym_ranges
+
+def decode_sym(binary, offset):
+       sym_ranges = get_sym_ranges(binary)
+       for name, (start, length) in sym_ranges.items():
+               if offset >= start and offset <= (start + length):
+                       return "%s+0x%x" % (name, offset - start)
+       return "%x" % offset
+
+def decode_addr(code_ranges, addr):
+       for binary, (start, end) in code_ranges.items():
+               if addr >= start and addr <= end:
+                       offset = addr - start if binary.endswith(".so") else addr
+                       return "%s %s" % (binary, decode_sym(binary, offset))
+       return "%x" % addr
+
+def decode_stack(info):
+       stack = ""
+       if info.num_frames <= 0:
+               return "???"
+       for i in range(0, info.num_frames):
+               addr = info.callstack[i]
+               if kernel_trace:
+                       stack += " %s (%x) ;" % (bpf_program.ksym(addr), addr)
+               else:
+                       stack += " %s (%x) ;" % (decode_addr(code_ranges, addr), addr)
+       return stack
+
+def print_outstanding():
+       stacks = {}
+       print("*** Outstanding allocations:")
+       allocs = bpf_program.get_table("allocs")
+       for address, info in sorted(allocs.items(), key=lambda a: -a[1].size):
+               stack = decode_stack(info)
+               if stack in stacks: stacks[stack] += info.size
+               else:               stacks[stack] = info.size
+               if args.show_allocs:
+                       print("\taddr = %x size = %s" % (address.value, info.size))
+       for stack, size in sorted(stacks.items(), key=lambda s: -s[1]):
+               print("\t%d bytes allocated from stack\n\t\t%s" % (size, stack.replace(";", "\n\t\t")))
+
+while True:
+        if trace_all:
+               print bpf_program.trace_fields()
+       else:
+               try:
+                       sleep(interval)
+               except KeyboardInterrupt:
+                       exit()
+               if not kernel_trace:
+                       code_ranges = get_code_ranges(pid)
+               print_outstanding()
+