From 711d415080be5df94e08a8f3a1018be251b82760 Mon Sep 17 00:00:00 2001
From: Leon Hwang <hffilwlqm@gmail.com>
Date: Tue, 12 Nov 2024 20:24:33 +0800
Subject: [PATCH] bpflbr: Reduce overhead LBR entries caused by ringbuf

Signed-off-by: Leon Hwang <hffilwlqm@gmail.com>
---
 bpf/lbr.c              | 12 +++++++-----
 internal/bpflbr/lbr.go | 24 ++++++++++++------------
 main.go                | 13 ++++++++++++-
 3 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/bpf/lbr.c b/bpf/lbr.c
index c6d816e..3862c31 100644
--- a/bpf/lbr.c
+++ b/bpf/lbr.c
@@ -3,7 +3,6 @@
 #include "vmlinux.h"
 #include "bpf_helpers.h"
 #include "bpf_tracing.h"
-#include "bpf_cleanup.h"
 
 #define MAX_LBR_ENTRIES 32
 
@@ -19,22 +18,25 @@ struct event {
     __u64 func_ip;
 } __attribute__((packed));
 
+struct event lbr_events[1] SEC(".data.lbrs");
+
 SEC("fexit")
 int BPF_PROG(fexit_fn)
 {
     struct event *event;
     __u64 retval;
-    int err = 0;
+    __u32 cpu;
 
-    guard_ringbuf(&events, event, &err);
-    if (!event)
-        return BPF_OK;
+    cpu = bpf_get_smp_processor_id();
+    event = &lbr_events[cpu];
 
     event->nr_bytes = bpf_get_branch_snapshot(event->lbr, sizeof(event->lbr), 0); /* required 5.16 kernel. */
     bpf_get_func_ret(ctx, (void *) &retval); /* required 5.17 kernel. */
     event->func_ret = retval;
     event->func_ip = bpf_get_func_ip(ctx); /* required 5.17 kernel. */
 
+    bpf_ringbuf_output(&events, event, sizeof(*event), 0);
+
     return BPF_OK;
 }
 
diff --git a/internal/bpflbr/lbr.go b/internal/bpflbr/lbr.go
index 33473f6..8af89e7 100644
--- a/internal/bpflbr/lbr.go
+++ b/internal/bpflbr/lbr.go
@@ -14,19 +14,19 @@ import (
 	"golang.org/x/sys/unix"
 )
 
-func Run(reader *ringbuf.Reader, progs *bpfProgs, addr2line *Addr2Line, ksyms *Kallsyms, w io.Writer) error {
-	type LbrEntry struct {
-		From  uintptr
-		To    uintptr
-		Flags uint64
-	}
-	type Event struct {
-		Entries [32]LbrEntry
-		NrBytes int64
-		Retval  int64
-		FuncIP  uintptr
-	}
+type LbrEntry struct {
+	From  uintptr
+	To    uintptr
+	Flags uint64
+}
+type Event struct {
+	Entries [32]LbrEntry
+	NrBytes int64
+	Retval  int64
+	FuncIP  uintptr
+}
 
+func Run(reader *ringbuf.Reader, progs *bpfProgs, addr2line *Addr2Line, ksyms *Kallsyms, w io.Writer) error {
 	stack := newLBRStack()
 
 	var sb strings.Builder
diff --git a/main.go b/main.go
index d11eb8e..1a4fe48 100644
--- a/main.go
+++ b/main.go
@@ -10,6 +10,7 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
+	"unsafe"
 
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/ringbuf"
@@ -77,12 +78,22 @@ func main() {
 	bpfSpec, err := loadLbr()
 	assert.NoErr(err, "Failed to load bpf spec: %v")
 
+	numCPU, err := ebpf.PossibleCPU()
+	assert.NoErr(err, "Failed to get possible cpu: %v")
+
+	lbrsMapSpec := bpfSpec.Maps[".data.lbrs"]
+	lbrsMapSpec.ValueSize = uint32(unsafe.Sizeof(bpflbr.Event{})) * uint32(numCPU)
+	lbrsMapSpec.Contents[0].Value = make([]byte, lbrsMapSpec.ValueSize)
+	lbrs, err := ebpf.NewMap(lbrsMapSpec)
+	assert.NoErr(err, "Failed to create lbrs map: %v")
+
 	events, err := ebpf.NewMap(bpfSpec.Maps["events"])
 	assert.NoErr(err, "Failed to create events map: %v")
 	defer events.Close()
 
 	reusedMaps := map[string]*ebpf.Map{
-		"events": events,
+		"events":     events,
+		".data.lbrs": lbrs,
 	}
 
 	tracings, err := bpflbr.NewBPFTracing(bpfSpec, reusedMaps, tracingTargets)