libbpf-tools: Support memleak trace third-part allocators

Memleak run failed when trace task who use jemalloc, and luckly i see it
fixed in memleak.py: https://github.com/iovisor/bcc/pull/4812/files, so
we can also support on c version.

compile jemalloc with "je_" prefix:
./configure --with-jemalloc-prefix=je_

jemalloc test case:
 #include <stdio.h>
 #include <unistd.h>
 #include <time.h>
 #include <jemalloc/jemalloc.h>

 #define malloc(size) je_malloc(size)
 #define free(ptr) je_free(ptr)
 int main(void)
 {
        for (int i = 0; i < 100; i++) {
                void *p = malloc(1000);
                sleep(1);
        }
        return 0;
 }
gcc -o alloc alloc.c -ljemalloc

test result:
sudo ./memleak -p $(pidof alloc) -O /usr/local/lib/libjemalloc.so -S je_

libbpf: elf: ambiguous match for 'mmap@GLIBC_2.2.5', 'mmap' in '/usr/local/lib/libjemalloc.so'
libbpf: elf: ambiguous match for 'mmap@GLIBC_2.2.5', 'mmap' in '/usr/local/lib/libjemalloc.so'
libbpf: elf: ambiguous match for 'munmap@GLIBC_2.2.5', 'munmap' in '/usr/local/lib/libjemalloc.so'
libbpf: elf: failed to find symbol 'pvalloc' in '/usr/local/lib/libjemalloc.so'
libbpf: elf: failed to find symbol 'pvalloc' in '/usr/local/lib/libjemalloc.so'
Tracing outstanding memory allocs...  Hit Ctrl-C to end
[15:12:34] Top 1 stacks with outstanding allocations:
5000 bytes in 5 allocations from stack
        0 [<00005584f3217188>] main+0x1f
        1 [<00007fe897617d90>] __libc_init_first+0x90
[15:12:39] Top 1 stacks with outstanding allocations:
10000 bytes in 10 allocations from stack
        0 [<00005584f3217188>] main+0x1f
        1 [<00007fe897617d90>] __libc_init_first+0x90

Signed-off-by: Tao Chen <chen.dylane@gmail.com>

1 files changed, 23 insertions, 4 deletions

diff --git a/libbpf-tools/memleak.c b/libbpf-tools/memleak.c
index a2c7d1cd..39a8ccdc 100644
--- a/libbpf-tools/memleak.c
+++ b/libbpf-tools/memleak.c
@@ -50,6 +50,7 @@ static struct env {
 	bool kernel_trace;
 	bool verbose;
 	char command[32];
+	char symbols_prefix[16];
 } env = {
 	.interval = 5, // posarg 1
 	.nr_intervals = -1, // posarg 2
@@ -71,6 +72,7 @@ static struct env {
 	.kernel_trace = true,
 	.verbose = false,
 	.command = {0}, // -c --command
+	.symbols_prefix = {0},
 };
 
 struct allocation_node {
@@ -88,8 +90,10 @@ struct allocation {
 
 #define __ATTACH_UPROBE(skel, sym_name, prog_name, is_retprobe) \
 	do { \
+		char sym[32]; \
+		sprintf(sym, "%s%s", env.symbols_prefix, #sym_name); \
 		LIBBPF_OPTS(bpf_uprobe_opts, uprobe_opts, \
-				.func_name = #sym_name, \
+				.func_name = sym, \
 				.retprobe = is_retprobe); \
 		skel->links.prog_name = bpf_program__attach_uprobe_opts( \
 				skel->progs.prog_name, \
@@ -181,6 +185,8 @@ const char argp_args_doc[] =
 "        allocations that are at least one minute (60 seconds) old\n"
 "./memleak -s 5\n"
 "        Trace roughly every 5th allocation, to reduce overhead\n"
+"./memleak -p $(pidof allocs) -S je_\n"
+"        Trace task who sue jemalloc\n"
 "";
 
 static const struct argp_option argp_options[] = {
@@ -198,6 +204,7 @@ static const struct argp_option argp_options[] = {
 	{"max-size", 'Z', "MAX_SIZE", 0, "capture only allocations smaller than this size"},
 	{"obj", 'O', "OBJECT", 0, "attach to allocator functions in the specified object"},
 	{"percpu", 'P', NULL, 0, "trace percpu allocations"},
+	{"symbols-prefix", 'S', "SYMBOLS_PREFIX", 0, "memory allocator symbols prefix"},
 	{"verbose", 'v', NULL, 0, "verbose debug output" },
 	{},
 };
@@ -527,6 +534,9 @@ error_t argp_parse_arg(int key, char *arg, struct argp_state *state)
 	case 's':
 		env.sample_rate = argp_parse_long(key, arg, state);
 		break;
+	case 'S':
+		strncpy(env.symbols_prefix, arg, sizeof(env.symbols_prefix) - 1);
+		break;
 	case 'T':
 		env.top_stacks = atoi(arg);
 		break;
@@ -1041,8 +1051,14 @@ int attach_uprobes(struct memleak_bpf *skel)
 	ATTACH_UPROBE_CHECKED(skel, realloc, realloc_enter);
 	ATTACH_URETPROBE_CHECKED(skel, realloc, realloc_exit);
 
-	ATTACH_UPROBE_CHECKED(skel, mmap, mmap_enter);
-	ATTACH_URETPROBE_CHECKED(skel, mmap, mmap_exit);
+	/* third party allocator like jemallloc not support mmap, so remove the check. */
+	if (strlen(env.symbols_prefix)) {
+		ATTACH_UPROBE(skel, mmap, mmap_enter);
+		ATTACH_URETPROBE(skel, mmap, mmap_exit);
+	} else {
+		ATTACH_UPROBE_CHECKED(skel, mmap, mmap_enter);
+		ATTACH_URETPROBE_CHECKED(skel, mmap, mmap_exit);
+	}
 
 	ATTACH_UPROBE_CHECKED(skel, posix_memalign, posix_memalign_enter);
 	ATTACH_URETPROBE_CHECKED(skel, posix_memalign, posix_memalign_exit);
@@ -1051,7 +1067,10 @@ int attach_uprobes(struct memleak_bpf *skel)
 	ATTACH_URETPROBE_CHECKED(skel, memalign, memalign_exit);
 
 	ATTACH_UPROBE_CHECKED(skel, free, free_enter);
-	ATTACH_UPROBE_CHECKED(skel, munmap, munmap_enter);
+	if (strlen(env.symbols_prefix))
+		ATTACH_UPROBE(skel, munmap, munmap_enter);
+	else
+		ATTACH_UPROBE_CHECKED(skel, munmap, munmap_enter);
 
 	// the following probes are intentinally allowed to fail attachment