From 27b061d3024572c66e646e29afa4cf270216bc34 Mon Sep 17 00:00:00 2001 From: Leon Hwang Date: Tue, 24 Oct 2023 20:04:54 +0800 Subject: [PATCH 1/3] Accelerate attaching/detaching kprobes In order to accelerate attaching/detaching kprobes, do attach/detach kprobes concurrently. By concurrent way, it is a little faster than original way. On my 6 CPU cores VM, run by concurrent way: ``` 2023/10/25 14:16:03 Attaching kprobes (via kprobe)... 1462 / 1462 [----------------------------------------------------------------------------------------------------] 100.00% 342 p/s 2023/10/25 14:16:07 Attached (ignored 0) 2023/10/25 14:16:07 Listening for events.. SKB CPU PROCESS FUNC ^C2023/10/25 14:16:08 Received signal, exiting program.. 2023/10/25 14:16:08 Detaching kprobes... 1462 / 1462 [-----------------------------------------------------------------------------------------------------] 100.00% 35 p/s ``` run by original way: ``` 2023/10/25 14:17:27 Attaching kprobes (via kprobe)... 1462 / 1462 [----------------------------------------------------------------------------------------------------] 100.00% 282 p/s 2023/10/25 14:17:32 Attached (ignored 0) 2023/10/25 14:17:32 Listening for events.. SKB CPU PROCESS FUNC ^C2023/10/25 14:17:33 Received signal, exiting program.. 2023/10/25 14:17:33 Detaching kprobes... 1462 / 1462 [-----------------------------------------------------------------------------------------------------] 100.00% 21 p/s ``` Signed-off-by: Leon Hwang --- README.md | 3 +- internal/pwru/kprobe.go | 198 ++++++++++++++++++++++++++++++++++++++++ internal/pwru/types.go | 16 ++-- main.go | 91 ++++++------------ 4 files changed, 235 insertions(+), 73 deletions(-) create mode 100644 internal/pwru/kprobe.go diff --git a/README.md b/README.md index c85b2461..6594bc77 100644 --- a/README.md +++ b/README.md @@ -56,6 +56,7 @@ Usage: pwru [options] [pcap-filter] --backend string Tracing backend('kprobe', 'kprobe-multi'). Will auto-detect if not specified. --filter-func string filter kernel functions to be probed by name (exact match, supports RE2 regular expression) --filter-ifname string filter skb ifname in --filter-netns (if not specified, use current netns) + --filter-kprobe-batch uint batch size for kprobe attaching/detaching (default 10) --filter-mark uint32 filter skb mark --filter-netns string filter netns ("/proc//ns/net", "inode:") --filter-trace-tc trace TC bpf progs @@ -138,7 +139,7 @@ See [docs/vagrant.md](docs/vagrant.md) * Go >= 1.16 * LLVM/clang >= 1.12 -* Bison +* Bison * Lex/Flex >= 2.5.31 ### Building diff --git a/internal/pwru/kprobe.go b/internal/pwru/kprobe.go new file mode 100644 index 00000000..4bac1072 --- /dev/null +++ b/internal/pwru/kprobe.go @@ -0,0 +1,198 @@ +// SPDX-License-Identifier: Apache-2.0 +/* Copyright 2024 Authors of Cilium */ + +package pwru + +import ( + "context" + "errors" + "fmt" + "log" + "os" + "sync" + "syscall" + + "github.com/cheggaaa/pb/v3" + "github.com/cilium/ebpf" + "github.com/cilium/ebpf/link" + "golang.org/x/sync/errgroup" +) + +type Kprobe struct { + hookFunc string // internal use + HookFuncs []string + Prog *ebpf.Program +} + +func attachKprobes(ctx context.Context, bar *pb.ProgressBar, kprobes []Kprobe) (links []link.Link, ignored int, err error) { + links = make([]link.Link, 0, len(kprobes)) + for _, kprobe := range kprobes { + select { + case <-ctx.Done(): + return + + default: + } + + var kp link.Link + kp, err = link.Kprobe(kprobe.hookFunc, kprobe.Prog, nil) + if err != nil { + if !errors.Is(err, os.ErrNotExist) && !errors.Is(err, syscall.EADDRNOTAVAIL) { + err = fmt.Errorf("opening kprobe %s: %w", kprobe.hookFunc, err) + return + } else { + err = nil + ignored++ + } + } else { + links = append(links, kp) + } + + bar.Increment() + } + + return +} + +// AttachKprobes attaches kprobes concurrently. +func AttachKprobes(ctx context.Context, bar *pb.ProgressBar, kps []Kprobe, batch uint) (links []link.Link, ignored int) { + if batch == 0 { + log.Fatal("--filter-kprobe-batch must be greater than 0") + } + + var kprobes []Kprobe + for _, kp := range kps { + for _, fn := range kp.HookFuncs { + kprobes = append(kprobes, Kprobe{ + hookFunc: fn, + Prog: kp.Prog, + }) + } + } + + if len(kprobes) == 0 { + return + } + + errg, ctx := errgroup.WithContext(ctx) + + var mu sync.Mutex + links = make([]link.Link, 0, len(kprobes)) + + attaching := func(kprobes []Kprobe) error { + l, i, e := attachKprobes(ctx, bar, kprobes) + if e != nil { + return e + } + + mu.Lock() + links = append(links, l...) + ignored += i + mu.Unlock() + + return nil + } + + var i uint + for i = 0; i+batch < uint(len(kprobes)); i += batch { + kps := kprobes[i : i+batch] + errg.Go(func() error { + return attaching(kps) + }) + } + if i < uint(len(kprobes)) { + kps := kprobes[i:] + errg.Go(func() error { + return attaching(kps) + }) + } + + if err := errg.Wait(); err != nil { + log.Fatalf("Attaching kprobes: %v\n", err) + } + + return +} + +// DetachKprobes detaches kprobes concurrently. +func DetachKprobes(links []link.Link, showProgressBar bool, batch uint) { + log.Println("Detaching kprobes...") + + if batch < 2 { + for _, l := range links { + _ = l.Close() + } + + return + } + + var errg errgroup.Group + var bar *pb.ProgressBar + + if showProgressBar { + bar = pb.StartNew(len(links)) + defer bar.Finish() + } + increment := func() { + if showProgressBar { + bar.Increment() + } + } + + var i uint + for i = 0; i+batch < uint(len(links)); i += batch { + l := links[i : i+batch] + errg.Go(func() error { + for _, l := range l { + _ = l.Close() + increment() + } + return nil + }) + } + for ; i < uint(len(links)); i++ { + _ = links[i].Close() + increment() + } + + _ = errg.Wait() +} + +// AttachKprobeMulti attaches kprobe-multi serially. +func AttachKprobeMulti(ctx context.Context, bar *pb.ProgressBar, kprobes []Kprobe, a2n Addr2Name) (links []link.Link, ignored int) { + links = make([]link.Link, 0, len(kprobes)) + + for _, kp := range kprobes { + select { + case <-ctx.Done(): + return + default: + } + + addrs := make([]uintptr, 0, len(kp.HookFuncs)) + for _, fn := range kp.HookFuncs { + if addr, ok := a2n.Name2AddrMap[fn]; ok { + addrs = append(addrs, addr...) + } else { + ignored += 1 + bar.Increment() + continue + } + } + + if len(addrs) == 0 { + continue + } + + opts := link.KprobeMultiOptions{Addresses: addrs} + l, err := link.KprobeMulti(kp.Prog, opts) + bar.Add(len(kp.HookFuncs)) + if err != nil { + log.Fatalf("Opening kprobe-multi: %s\n", err) + } + + links = append(links, l) + } + + return +} diff --git a/internal/pwru/types.go b/internal/pwru/types.go index fb7a7c9a..a5c6bd43 100644 --- a/internal/pwru/types.go +++ b/internal/pwru/types.go @@ -25,13 +25,14 @@ type Flags struct { KernelBTF string - FilterNetns string - FilterMark uint32 - FilterFunc string - FilterTrackSkb bool - FilterTraceTc bool - FilterIfname string - FilterPcap string + FilterNetns string + FilterMark uint32 + FilterFunc string + FilterTrackSkb bool + FilterTraceTc bool + FilterIfname string + FilterPcap string + FilterKprobeBatch uint OutputTS string OutputMeta bool @@ -61,6 +62,7 @@ func (f *Flags) SetFlags() { flag.BoolVar(&f.FilterTrackSkb, "filter-track-skb", false, "trace a packet even if it does not match given filters (e.g., after NAT or tunnel decapsulation)") flag.BoolVar(&f.FilterTraceTc, "filter-trace-tc", false, "trace TC bpf progs") flag.StringVar(&f.FilterIfname, "filter-ifname", "", "filter skb ifname in --filter-netns (if not specified, use current netns)") + flag.UintVar(&f.FilterKprobeBatch, "filter-kprobe-batch", 10, "batch size for kprobe attaching/detaching") flag.StringVar(&f.OutputTS, "timestamp", "none", "print timestamp per skb (\"current\", \"relative\", \"absolute\", \"none\")") flag.BoolVar(&f.OutputMeta, "output-meta", false, "print skb metadata") flag.BoolVar(&f.OutputTuple, "output-tuple", false, "print L4 tuple") diff --git a/main.go b/main.go index 233def09..a0d08988 100644 --- a/main.go +++ b/main.go @@ -195,21 +195,19 @@ func main() { var kprobes []link.Link defer func() { + var showProgressBar bool select { case <-ctx.Done(): - log.Println("Detaching kprobes...") - bar := pb.StartNew(len(kprobes)) - for _, kp := range kprobes { - _ = kp.Close() - bar.Increment() - } - bar.Finish() + showProgressBar = true default: - for _, kp := range kprobes { - _ = kp.Close() - } } + + batch := uint(1) + if !useKprobeMulti { + batch = flags.FilterKprobeBatch + } + pwru.DetachKprobes(kprobes, showProgressBar, batch) }() msg := "kprobe" @@ -257,69 +255,32 @@ func main() { } } + pwruKprobes := make([]pwru.Kprobe, 0, len(funcs)) funcsByPos := pwru.GetFuncsByPos(funcs) for pos, fns := range funcsByPos { fn, ok := coll.Programs[fmt.Sprintf("kprobe_skb_%d", pos)] - if !ok { - ignored += len(fns) - bar.Add(len(fns)) - continue - } - - if !useKprobeMulti { - for _, name := range fns { - select { - case <-ctx.Done(): - bar.Finish() - return - default: - } - - kp, err := link.Kprobe(name, fn, nil) - bar.Increment() - if err != nil { - if !errors.Is(err, os.ErrNotExist) && !errors.Is(err, syscall.EADDRNOTAVAIL) { - log.Fatalf("Opening kprobe %s: %s\n", name, err) - } else { - ignored += 1 - } - } else { - kprobes = append(kprobes, kp) - } - } + if ok { + pwruKprobes = append(pwruKprobes, pwru.Kprobe{HookFuncs: fns, Prog: fn}) } else { - select { - case <-ctx.Done(): - bar.Finish() - return - default: - } - - addrs := make([]uintptr, 0, len(fns)) - for _, fn := range fns { - if addr, ok := addr2name.Name2AddrMap[fn]; ok { - addrs = append(addrs, addr...) - } else { - ignored += 1 - bar.Increment() - continue - } - } - - if len(addrs) == 0 { - continue - } - - opts := link.KprobeMultiOptions{Addresses: addrs} - kp, err := link.KprobeMulti(fn, opts) + ignored += len(fns) bar.Add(len(fns)) - if err != nil { - log.Fatalf("Opening kprobe-multi for pos %d: %s\n", pos, err) - } - kprobes = append(kprobes, kp) } } + if !useKprobeMulti { + l, i := pwru.AttachKprobes(ctx, bar, pwruKprobes, flags.FilterKprobeBatch) + kprobes = append(kprobes, l...) + ignored += i + } else { + l, i := pwru.AttachKprobeMulti(ctx, bar, pwruKprobes, addr2name) + kprobes = append(kprobes, l...) + ignored += i + } bar.Finish() + select { + case <-ctx.Done(): + return + default: + } log.Printf("Attached (ignored %d)\n", ignored) log.Println("Listening for events..") From ac519f7befce975741c36996cc17886283cb39b4 Mon Sep 17 00:00:00 2001 From: Leon Hwang Date: Thu, 26 Oct 2023 22:13:01 +0800 Subject: [PATCH 2/3] Flush output file Flush output file before closing it. Signed-off-by: Leon Hwang --- internal/pwru/output.go | 10 ++++++++-- main.go | 1 + 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/internal/pwru/output.go b/internal/pwru/output.go index dd58c4e2..6ff53a70 100644 --- a/internal/pwru/output.go +++ b/internal/pwru/output.go @@ -7,7 +7,6 @@ package pwru import ( "errors" "fmt" - "io" "log" "net" "os" @@ -36,7 +35,7 @@ type output struct { printSkbMap *ebpf.Map printStackMap *ebpf.Map addr2name Addr2Name - writer io.Writer + writer *os.File kprobeMulti bool kfreeReasons map[uint64]string ifaceCache map[uint64]map[uint32]string @@ -81,6 +80,13 @@ func NewOutput(flags *Flags, printSkbMap *ebpf.Map, printStackMap *ebpf.Map, }, nil } +func (o *output) Close() { + if o.writer != os.Stdout { + _ = o.writer.Sync() + _ = o.writer.Close() + } +} + func (o *output) PrintHeader() { if o.flags.OutputTS == "absolute" { fmt.Fprintf(o.writer, "%12s ", "TIME") diff --git a/main.go b/main.go index a0d08988..77e0c04f 100644 --- a/main.go +++ b/main.go @@ -299,6 +299,7 @@ func main() { if err != nil { log.Fatalf("Failed to create outputer: %s", err) } + defer output.Close() output.PrintHeader() defer func() { From 4ebecc711a60df6aed52f30e7ae1324dc436409e Mon Sep 17 00:00:00 2001 From: Leon Hwang Date: Sun, 18 Feb 2024 14:14:34 +0800 Subject: [PATCH 3/3] Fix go.mod Signed-off-by: Leon Hwang --- go.mod | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/go.mod b/go.mod index 776d36c5..28f338d1 100644 --- a/go.mod +++ b/go.mod @@ -13,6 +13,7 @@ require ( github.com/tklauser/ps v0.0.2 github.com/vishvananda/netns v0.0.4 golang.org/x/net v0.21.0 + golang.org/x/sync v0.3.0 golang.org/x/sys v0.17.0 ) @@ -29,5 +30,4 @@ require ( github.com/pkg/errors v0.9.1 // indirect github.com/rivo/uniseg v0.2.0 // indirect golang.org/x/exp v0.0.0-20240213143201-ec583247a57a // indirect - golang.org/x/sync v0.3.0 // indirect )