bpf模块完成，着手符号地址解析和json格式处理

eBPF_Supermarket/Stack_Analyser/bcc/mem_count.bpf.c

@@ -0,0 +1,129 @@
+// Copyright 2023 The LMP Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// https://github.com/linuxkerneltravel/lmp/blob/develop/LICENSE
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// author: luiyanbing@foxmail.com
+//
+// 内核态ebpf的内存模块代码
+#include <linux/sched.h>
+
+typedef struct {
+    u32 pid;
+    int32_t ksid, usid;
+} psid;
+
+typedef struct {
+    char str[TASK_COMM_LEN];
+} comm;
+
+BPF_HASH(psid_count, psid, u64);
+BPF_STACK_TRACE(stack_trace, STACK_STORAGE_SIZE);
+BPF_HASH(pid_tgid, u32, u32);
+BPF_HASH(pid_comm, u32, comm);
+
+typedef struct
+{
+    __u64 addr;
+    __u32 pid, o;
+} piddr; // mem info key
+
+typedef struct
+{
+    __u64 size;
+    __u32 usid, o;
+} mem_info; // mem info with stack
+
+BPF_HASH(pid_size, u32, u64);
+BPF_HASH(piddr_meminfo, piddr, mem_info);
+
+int malloc_enter(struct pt_regs *ctx)
+{
+    u64 size = PT_REGS_PARM1(ctx) >> 10;
+    if(!size) return -1;
+    // record data
+    u64 pt = bpf_get_current_pid_tgid();
+    u32 pid = pt >> 32;
+    u32 tgid = pt;
+    pid_tgid.update(&pid, &tgid);
+    comm *p = pid_comm.lookup(&pid);
+    if (!p)
+    {
+        comm name;
+        bpf_get_current_comm(&name, TASK_COMM_LEN);
+        pid_comm.update(&pid, &name);
+    }
+
+    // record size
+    return pid_size.update(&pid, &size);
+}
+
+int malloc_exit(struct pt_regs *ctx)
+{
+    // get size
+    u32 pid = bpf_get_current_pid_tgid() >> 32;
+    u64 *size = pid_size.lookup(&pid);
+    if(!size) return -1;
+
+    // record stack count
+    psid apsid = {
+        .pid = pid,
+        .usid = USER_STACK_GET,
+        .ksid = -1,
+    };
+    u64 *count = psid_count.lookup(&apsid);
+    if (!count)
+        psid_count.update(&apsid, size);
+    else (*count) += *size;
+
+    // record pid_addr-info
+    u64 addr = PT_REGS_RC(ctx);
+    piddr a = {
+        .addr = addr, 
+        .pid = pid,
+        .o = 0,
+    };
+    mem_info info = {
+        .size = *size, 
+        .usid = apsid.usid,
+        .o = 0,
+    };
+    piddr_meminfo.update(&a, &info);
+
+    // delete pid-size
+    return pid_size.delete(&pid);
+}
+
+int free_enter(struct pt_regs *ctx)
+{
+    u64 addr = PT_REGS_PARM1(ctx);
+    // get freeing size
+    u32 pid = bpf_get_current_pid_tgid() >> 32;
+    piddr a = {addr, pid};
+    mem_info *info = piddr_meminfo.lookup(&a);
+    if(!info) return -1;
+
+    // get allocated size
+    psid apsid = {
+        .ksid = -1,
+        .pid = pid,
+        .usid = info->usid,
+    };
+    u64 *size = psid_count.lookup(&apsid);
+    if(!size) return -1;
+
+    // sub the freeing size
+    (*size) -= info->size;
+
+    // del freeing addr info
+    return piddr_meminfo.delete(&a);
+}

eBPF_Supermarket/Stack_Analyser/bcc/stack_count.py

@@ -6,11 +6,12 @@
 from argparse import ArgumentTypeError, ArgumentParser, RawDescriptionHelpFormatter, SUPPRESS
 from subprocess import Popen, PIPE
 
-mode = 'on_cpu'
+mode_list = ['on_cpu', 'off_cpu', 'mem']
 
 # arguments
 examples = """examples:
     sudo -E ./stack_count.py             # trace on-CPU stack time until Ctrl-C
+    sudo -E ./stack_count.py -m off_cpu  # trace off-CPU stack time until Ctrl-C
     sudo -E ./stack_count.py 5           # trace for 5 seconds only
     sudo -E ./stack_count.py -f 5        # 5 seconds, and output as stack_count.svg in flame graph format
     sudo -E ./stack_count.py -s 5        # 5 seconds, and show symbol offsets
@@ -25,7 +26,7 @@
 """
 
 
-def positive_int(val):
+def positive_int(val:str):
     try:
         ival = int(val)
     except ValueError:
@@ -34,6 +35,12 @@ def positive_int(val):
         raise ArgumentTypeError("must be positive")
     return ival
 
+def mode_str(val:str):
+    if val in mode_list:
+        return val
+    else:
+        raise ArgumentTypeError("must be 'on_cpu', 'off_cpu' or 'mem'")
+
 
 parser = ArgumentParser(
     description="Summarize on-CPU time by stack trace",
@@ -69,6 +76,9 @@ def positive_int(val):
                     help="output folded format")
 parser.add_argument("-s", "--offset", action="store_true",
                     help="show address offsets")
+parser.add_argument("-m", "--mode", default='on_cpu',
+                    type = mode_str,
+                    help="mode of stack counting, 'on_cpu'/'off_cpu'/'mem'")
 parser.add_argument("--stack-storage-size", default=16384,
                     type=positive_int,
                     help="the number of unique stack traces that can be stored and "
@@ -84,6 +94,7 @@ def positive_int(val):
 args = parser.parse_args()
 folded = args.folded
 duration = int(args.duration)
+mode = args.mode
 debug = 0
 
 # set thread filter
@@ -94,7 +105,7 @@ def positive_int(val):
     # thread_filter = '!(curr->tgid == %d)' % args.tgid
     pid = args.tgid
     thread_context = "PID " + str(pid)
-    if mode != 'on_cpu':
+    if mode == 'off_cpu':
         thread_filter = 'curr->tgid == %d' % pid
 elif args.pid is not None:
     thread_context = "TID %d" % args.pid
@@ -106,7 +117,7 @@ def positive_int(val):
     ps.send_signal(19)
     pid = ps.pid
     thread_context = "PID " + str(pid)
-    if mode != 'on_cpu':
+    if mode == 'off_cpu':
         thread_filter = 'curr->tgid == %d' % pid
     # perf default attach children process
 elif args.user_threads_only:
@@ -202,6 +213,16 @@ def positive_int(val):
     case 'off_cpu':
         b.attach_kprobe(
             event_re="^finish_task_switch$|^finish_task_switch\.isra\.\d$", fn_name='do_stack')
+    case 'mem':
+        arch = Popen(args='uname -m', stdout=PIPE, shell=True).stdout.read().decode().split()[0]
+        lib = "/usr/lib/"+arch+"-linux-gnu/libc.so.6"
+        # if pid != -1:
+        #     from psutil import Process
+        #     lib = Process(pid).exe()
+        b.attach_uprobe(name=lib, sym='malloc', fn_name='malloc_enter', pid=pid)
+        b.attach_uretprobe(name=lib, sym='malloc', fn_name='malloc_exit', pid=pid)
+        b.attach_uprobe(name=lib, sym='free', fn_name='free_enter', pid=pid)
+
 
 if pid != -1:
     print("attach %d." % pid)
@@ -217,12 +238,21 @@ def positive_int(val):
 
 def int_handler(sig=None, frame=None):
     print("\b\bquit...")
-    if mode == 'on_cpu':
-        b.detach_perf_event(ev_type=PerfType.SOFTWARE,
-                            ev_config=PerfSWConfig.CPU_CLOCK)
-    elif mode == 'off_cpu':
-        for tp in b.get_kprobe_functions(b"^finish_task_switch$|^finish_task_switch\.isra\.\d$"):
-            b.detach_kprobe(tp)
+    match mode:
+        case 'on_cpu':
+            b.detach_perf_event(ev_type=PerfType.SOFTWARE,
+                                ev_config=PerfSWConfig.CPU_CLOCK)
+        case 'off_cpu':
+            for tp in b.get_kprobe_functions(b"^finish_task_switch$|^finish_task_switch\.isra\.\d$"):
+                b.detach_kprobe(tp)
+        case 'mem':
+            try:
+                b.detach_uprobe(lib, 'malloc')
+                b.detach_uretprobe(lib, 'malloc')
+                b.detach_uprobe(lib, 'free')
+            except:
+                pass
+
     # system("tput rmcup")
     if auto:
         if not ad:

eBPF_Supermarket/Stack_Analyser/libbpf/Makefile

@@ -27,9 +27,9 @@ ARCH ?= $(shell echo $(ORIG_ARCH) | sed 's/x86_64/x86/' \
 APP = stack_analyzer
 bpf = on_cpu_count off_cpu_count mem_count
 
-all: vmlinux.h $(bpf)
+all: vmlinux.h $(bpf) symbol
 	clang -g -O2 -Wall -I. -I./include -c $(APP).cc -o $(APP).o
-	clang -Wall -O2 -g $(APP).o -static -lbpf -lelf -lz -o $(APP)
+	clang -Wall -O2 -g $(APP).o symbol.o elf.o -static -lbpf -lelf -lz -lstdc++ -o $(APP)
 
 vmlinux.h:
 	bpftool btf dump file /sys/kernel/btf/vmlinux format c > vmlinux.h
@@ -38,5 +38,9 @@ $(bpf): %: bpf/%.bpf.c
 	clang -Wunknown-attributes -g -O2 -target bpf -D__TARGET_ARCH_$(ARCH) -I/usr/include/$(ORIG_ARCH)-linux-gnu -I. -I./include -c bpf/$@.bpf.c -o bpf/$@.bpf.o
 	bpftool gen skeleton bpf/$@.bpf.o > bpf/$@.skel.h
 
+symbol: %: include/%.cc
+	clang -g -O2 -Wall -c include/symbol.cc -o symbol.o
+	clang -g -O2 -Wall -c include/elf.cc -o elf.o
+
 clean:
-	rm bpf/*.o bpf/*.skel.h *.o $(APP) vmlinux.h
+	rm bpf/*.o bpf/*.skel.h *.o include/*.o $(APP) vmlinux.h