Accelerate attaching/detaching kprobes

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/<pid>/ns/net", "inode:<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

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
 )

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
+}

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")

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")