Repository: cilium/ebpf
Pull request #1778 (2025-05-07 10:34:54)

collection: share kernel BTF when loading a collection
    We currently create a new copy of kernel BTF when loading each program in a 
    collection, which is quite expensive.

    Introduce a new type btf.Cache which allows re-using kernel BTF as-needed. 
    As an optimisation, the Cache is opportunistically populated from the global 
    cache on creation.

        core: 1
       goos: linux
       goarch: amd64
       pkg: github.com/cilium/ebpf
       cpu: 13th Gen Intel(R) Core(TM) i7-1365U
                           │   base.txt   │             cache.txt              │
                           │    sec/op    │   sec/op     vs base               │
       NewCollectionManyProgs   347.84m ± 1%   14.54m ± 2%  -95.82% (p=0.002
    n=6)

                            │    base.txt    │              cache.txt           
      │
                           │      B/op      │     B/op      vs base             
     │
       NewCollectionManyProgs   108.056Mi ± 0%   3.841Mi ± 0%  -96.44% (p=0.002
    n=6)

                            │   base.txt    │             cache.txt             
    │
                           │   allocs/op   │  allocs/op   vs base              
    │
       NewCollectionManyProgs   1026.83k ± 0%   38.04k ± 0%  -96.30% (p=0.002
    n=6)

    Signed-off-by: Lorenz Bauer <lmb@isovalent.com>


lmb: @patrickpichler could you give this branch a try and let me know what you think?
patrickpichler: Sorry for the super late response. I did run the benchmark again and here are the results:
Without your change
BenchmarkNewCollectionManyProgs-4   	     15	 75442629 ns/op	74477811 B/op	 809763 allocs/op

With your change
BenchmarkNewCollectionManyProgs-4   	    272	  4193884 ns/op	2734829 B/op	  30806 allocs/op

With the optimization I did in the draft PR (#1690)
BenchmarkNewCollectionManyProgs-4 184 5807742 ns/op 4134444 B/op 17325 allocs/op

When performing a single run, the total number of allocated memory according to pprof went down from `~660MB` to `~26MB`, so this is definitely a win. 

I guess there are a few more ways on how to reduce the number of allocations, but for my use case, this should be good enough. 

Thanks a lot!

Patch:
diff --git a/btf/kernel.go b/btf/kernel.go
index 11cb4a9..5dd1473 100644
--- a/btf/kernel.go
+++ b/btf/kernel.go
@@ -12,7 +12,8 @@ import (
 	"github.com/cilium/ebpf/internal/platform"
 )
 
-var kernelBTF = struct {
+// globalCache amortises decoding BTF across all users of the library.
+var globalCache = struct {
 	sync.RWMutex
 	kernel  *Spec
 	modules map[string]*Spec
@@ -22,80 +23,94 @@ var kernelBTF = struct {
 
 // FlushKernelSpec removes any cached kernel type information.
 func FlushKernelSpec() {
-	kernelBTF.Lock()
-	defer kernelBTF.Unlock()
+	globalCache.Lock()
+	defer globalCache.Unlock()
 
-	kernelBTF.kernel = nil
-	kernelBTF.modules = make(map[string]*Spec)
+	globalCache.kernel = nil
+	globalCache.modules = make(map[string]*Spec)
 }
 
 // LoadKernelSpec returns the current kernel's BTF information.
 //
 // Defaults to /sys/kernel/btf/vmlinux and falls back to scanning the file system
 // for vmlinux ELFs. Returns an error wrapping ErrNotSupported if BTF is not enabled.
+//
+// Consider using [Cache] instead.
 func LoadKernelSpec() (*Spec, error) {
-	kernelBTF.RLock()
-	spec := kernelBTF.kernel
-	kernelBTF.RUnlock()
-
-	if spec == nil {
-		kernelBTF.Lock()
-		defer kernelBTF.Unlock()
+	spec, err := loadCachedKernelSpec()
+	return spec.Copy(), err
+}
 
-		spec = kernelBTF.kernel
-	}
+// load (and cache) the kernel spec.
+//
+// Does not copy Spec.
+func loadCachedKernelSpec() (*Spec, error) {
+	globalCache.RLock()
+	spec := globalCache.kernel
+	globalCache.RUnlock()
 
 	if spec != nil {
-		return spec.Copy(), nil
+		return spec, nil
 	}
 
-	spec, _, err := loadKernelSpec()
+	globalCache.Lock()
+	defer globalCache.Unlock()
+
+	spec, err := loadKernelSpec()
 	if err != nil {
 		return nil, err
 	}
 
-	kernelBTF.kernel = spec
-	return spec.Copy(), nil
+	globalCache.kernel = spec
+	return spec, nil
 }
 
 // LoadKernelModuleSpec returns the BTF information for the named kernel module.
 //
+// Using [Cache.Module] is faster when loading BTF for more than one module.
+//
 // Defaults to /sys/kernel/btf/<module>.
 // Returns an error wrapping ErrNotSupported if BTF is not enabled.
 // Returns an error wrapping fs.ErrNotExist if BTF for the specific module doesn't exist.
 func LoadKernelModuleSpec(module string) (*Spec, error) {
-	kernelBTF.RLock()
-	spec := kernelBTF.modules[module]
-	kernelBTF.RUnlock()
+	spec, err := loadCachedKernelModuleSpec(module)
+	return spec.Copy(), err
+}
+
+// load (and cache) a module spec.
+//
+// Does not copy Spec.
+func loadCachedKernelModuleSpec(module string) (*Spec, error) {
+	globalCache.RLock()
+	spec := globalCache.modules[module]
+	globalCache.RUnlock()
 
 	if spec != nil {
-		return spec.Copy(), nil
+		return spec, nil
 	}
 
-	base, err := LoadKernelSpec()
+	base, err := loadCachedKernelSpec()
 	if err != nil {
-		return nil, fmt.Errorf("load kernel spec: %w", err)
+		return nil, err
 	}
 
-	kernelBTF.Lock()
-	defer kernelBTF.Unlock()
-
-	if spec = kernelBTF.modules[module]; spec != nil {
-		return spec.Copy(), nil
-	}
+	// NB: This only allows a single module to be parsed at a time. Not sure
+	// it makes a difference.
+	globalCache.Lock()
+	defer globalCache.Unlock()
 
 	spec, err = loadKernelModuleSpec(module, base)
 	if err != nil {
-		return nil, fmt.Errorf("load kernel module: %w", err)
+		return nil, err
 	}
 
-	kernelBTF.modules[module] = spec
-	return spec.Copy(), nil
+	globalCache.modules[module] = spec
+	return spec, nil
 }
 
-func loadKernelSpec() (_ *Spec, fallback bool, _ error) {
+func loadKernelSpec() (_ *Spec, _ error) {
 	if platform.IsWindows {
-		return nil, false, internal.ErrNotSupportedOnOS
+		return nil, internal.ErrNotSupportedOnOS
 	}
 
 	fh, err := os.Open("/sys/kernel/btf/vmlinux")
@@ -103,17 +118,17 @@ func loadKernelSpec() (_ *Spec, fallback bool, _ error) {
 		defer fh.Close()
 
 		spec, err := loadRawSpec(fh, internal.NativeEndian, nil)
-		return spec, false, err
+		return spec, err
 	}
 
 	file, err := findVMLinux()
 	if err != nil {
-		return nil, false, err
+		return nil, err
 	}
 	defer file.Close()
 
 	spec, err := LoadSpecFromReader(file)
-	return spec, true, err
+	return spec, err
 }
 
 func loadKernelModuleSpec(module string, base *Spec) (*Spec, error) {
@@ -168,3 +183,83 @@ func findVMLinux() (*os.File, error) {
 
 	return nil, fmt.Errorf("no BTF found for kernel version %s: %w", release, internal.ErrNotSupported)
 }
+
+// Cache allows to amortise the cost of decoding BTF across multiple call-sites.
+//
+// It is not safe for concurrent use.
+type Cache struct {
+	KernelTypes   *Spec
+	KernelModules map[string]*Spec
+}
+
+// NewCache creates a new Cache.
+//
+// Opportunistically reuses a global cache if possible.
+func NewCache() *Cache {
+	globalCache.RLock()
+	defer globalCache.RUnlock()
+
+	// This copy is either a no-op or very cheap, since the spec won't contain
+	// any inflated types.
+	kernel := globalCache.kernel.Copy()
+	if kernel == nil {
+		return &Cache{}
+	}
+
+	modules := make(map[string]*Spec, len(globalCache.modules))
+	for name, spec := range globalCache.modules {
+		decoder, _ := rebaseDecoder(spec.decoder, kernel.decoder)
+		// NB: Kernel module BTF can't contain ELF fixups because it is always
+		// read from sysfs.
+		modules[name] = &Spec{decoder: decoder}
+	}
+
+	return &Cache{kernel, modules}
+}
+
+// Kernel is equivalent to [LoadKernelSpec], except that repeated calls do
+// not copy the Spec.
+func (c *Cache) Kernel() (*Spec, error) {
+	if c.KernelTypes != nil {
+		return c.KernelTypes, nil
+	}
+
+	var err error
+	c.KernelTypes, err = LoadKernelSpec()
+	return c.KernelTypes, err
+}
+
+// Module is equivalent to [LoadKernelModuleSpec], except that repeated calls do
+// not copy the spec.
+//
+// All modules also share the return value of [Kernel] as their base.
+func (c *Cache) Module(name string) (*Spec, error) {
+	if spec := c.KernelModules[name]; spec != nil {
+		return spec, nil
+	}
+
+	if c.KernelModules == nil {
+		c.KernelModules = make(map[string]*Spec)
+	}
+
+	base, err := c.Kernel()
+	if err != nil {
+		return nil, err
+	}
+
+	spec, err := loadCachedKernelModuleSpec(name)
+	if err != nil {
+		return nil, err
+	}
+
+	// Important: base is shared between modules. This allows inflating common
+	// types only once.
+	decoder, err := rebaseDecoder(spec.decoder, base.decoder)
+	if err != nil {
+		return nil, err
+	}
+
+	spec = &Spec{decoder: decoder}
+	c.KernelModules[name] = spec
+	return spec, err
+}
diff --git a/btf/unmarshal.go b/btf/unmarshal.go
index 8704140..4dbfd41 100644
--- a/btf/unmarshal.go
+++ b/btf/unmarshal.go
@@ -165,6 +165,31 @@ func allBtfTypeOffsets(buf []byte, bo binary.ByteOrder, header *btfType) iter.Se
 	}
 }
 
+func rebaseDecoder(d *decoder, base *decoder) (*decoder, error) {
+	if d.base == nil {
+		return nil, fmt.Errorf("rebase split spec: not a split spec")
+	}
+
+	if &d.base.raw[0] != &base.raw[0] || len(d.base.raw) != len(base.raw) {
+		return nil, fmt.Errorf("rebase split spec: raw BTF differs")
+	}
+
+	return &decoder{
+		base,
+		d.byteOrder,
+		d.raw,
+		d.strings,
+		d.firstTypeID,
+		d.offsets,
+		d.declTags,
+		d.namedTypes,
+		sync.Mutex{},
+		make(map[TypeID]Type),
+		make(map[Type]TypeID),
+		make(map[TypeID][2]Bits),
+	}, nil
+}
+
 // Copy performs a deep copy of a decoder and its base.
 func (d *decoder) Copy() *decoder {
 	if d == nil {
diff --git a/collection.go b/collection.go
index 089c7d3..77476bd 100644
--- a/collection.go
+++ b/collection.go
@@ -412,6 +412,7 @@ type collectionLoader struct {
 	maps     map[string]*Map
 	programs map[string]*Program
 	vars     map[string]*Variable
+	types    *btf.Cache
 }
 
 func newCollectionLoader(coll *CollectionSpec, opts *CollectionOptions) (*collectionLoader, error) {
@@ -436,6 +437,7 @@ func newCollectionLoader(coll *CollectionSpec, opts *CollectionOptions) (*collec
 		make(map[string]*Map),
 		make(map[string]*Program),
 		make(map[string]*Variable),
+		newBTFCache(&opts.Programs),
 	}, nil
 }
 
@@ -587,7 +589,7 @@ func (cl *collectionLoader) loadProgram(progName string) (*Program, error) {
 		}
 	}
 
-	prog, err := newProgramWithOptions(progSpec, cl.opts.Programs)
+	prog, err := newProgramWithOptions(progSpec, cl.opts.Programs, cl.types)
 	if err != nil {
 		return nil, fmt.Errorf("program %s: %w", progName, err)
 	}
diff --git a/linker.go b/linker.go
index a09bb9c..d5ba281 100644
--- a/linker.go
+++ b/linker.go
@@ -124,7 +124,7 @@ func hasFunctionReferences(insns asm.Instructions) bool {
 //
 // Passing a nil target will relocate against the running kernel. insns are
 // modified in place.
-func applyRelocations(insns asm.Instructions, targets []*btf.Spec, kmodName string, bo binary.ByteOrder, b *btf.Builder) error {
+func applyRelocations(insns asm.Instructions, kmodName string, bo binary.ByteOrder, b *btf.Builder, c *btf.Cache) error {
 	var relos []*btf.CORERelocation
 	var reloInsns []*asm.Instruction
 	iter := insns.Iterate()
@@ -143,22 +143,21 @@ func applyRelocations(insns asm.Instructions, targets []*btf.Spec, kmodName stri
 		bo = internal.NativeEndian
 	}
 
-	if len(targets) == 0 {
-		kernelTarget, err := btf.LoadKernelSpec()
-		if err != nil {
-			return fmt.Errorf("load kernel spec: %w", err)
-		}
-		targets = append(targets, kernelTarget)
+	var targets []*btf.Spec
+	kernelTarget, err := c.Kernel()
+	if err != nil {
+		return fmt.Errorf("load kernel spec: %w", err)
+	}
+	targets = append(targets, kernelTarget)
 
-		if kmodName != "" {
-			kmodTarget, err := btf.LoadKernelModuleSpec(kmodName)
-			// Ignore ErrNotExists to cater to kernels which have CONFIG_DEBUG_INFO_BTF_MODULES disabled.
-			if err != nil && !errors.Is(err, fs.ErrNotExist) {
-				return fmt.Errorf("load kernel module spec: %w", err)
-			}
-			if err == nil {
-				targets = append(targets, kmodTarget)
-			}
+	if kmodName != "" {
+		kmodTarget, err := c.Module(kmodName)
+		// Ignore ErrNotExists to cater to kernels which have CONFIG_DEBUG_INFO_BTF_MODULES disabled.
+		if err != nil && !errors.Is(err, fs.ErrNotExist) {
+			return fmt.Errorf("load kernel module spec: %w", err)
+		}
+		if err == nil {
+			targets = append(targets, kmodTarget)
 		}
 	}
 
diff --git a/prog.go b/prog.go
index b36e741..a9c2a99 100644
--- a/prog.go
+++ b/prog.go
@@ -244,7 +244,7 @@ func NewProgramWithOptions(spec *ProgramSpec, opts ProgramOptions) (*Program, er
 		return nil, errors.New("can't load a program from a nil spec")
 	}
 
-	prog, err := newProgramWithOptions(spec, opts)
+	prog, err := newProgramWithOptions(spec, opts, newBTFCache(&opts))
 	if errors.Is(err, asm.ErrUnsatisfiedMapReference) {
 		return nil, fmt.Errorf("cannot load program without loading its whole collection: %w", err)
 	}
@@ -260,7 +260,7 @@ var (
 	kfuncBadCall = []byte(fmt.Sprintf("invalid func unknown#%d\n", kfuncCallPoisonBase))
 )
 
-func newProgramWithOptions(spec *ProgramSpec, opts ProgramOptions) (*Program, error) {
+func newProgramWithOptions(spec *ProgramSpec, opts ProgramOptions, c *btf.Cache) (*Program, error) {
 	if len(spec.Instructions) == 0 {
 		return nil, errors.New("instructions cannot be empty")
 	}
@@ -308,18 +308,8 @@ func newProgramWithOptions(spec *ProgramSpec, opts ProgramOptions) (*Program, er
 		return nil, fmt.Errorf("kernel module search: %w", err)
 	}
 
-	var targets []*btf.Spec
-	if opts.KernelTypes != nil {
-		targets = append(targets, opts.KernelTypes)
-	}
-	if kmodName != "" && opts.KernelModuleTypes != nil {
-		if modBTF, ok := opts.KernelModuleTypes[kmodName]; ok {
-			targets = append(targets, modBTF)
-		}
-	}
-
 	var b btf.Builder
-	if err := applyRelocations(insns, targets, kmodName, spec.ByteOrder, &b); err != nil {
+	if err := applyRelocations(insns, kmodName, spec.ByteOrder, &b, c); err != nil {
 		return nil, fmt.Errorf("apply CO-RE relocations: %w", err)
 	}
 
@@ -1208,3 +1198,12 @@ func findTargetInProgram(prog *Program, name string, progType ProgramType, attac
 
 	return spec.TypeID(targetFunc)
 }
+
+func newBTFCache(opts *ProgramOptions) *btf.Cache {
+	c := btf.NewCache()
+	if opts.KernelTypes != nil {
+		c.KernelTypes = opts.KernelTypes
+		c.KernelModules = opts.KernelModuleTypes
+	}
+	return c
+}


Test patch:
diff --git a/btf/kernel_test.go b/btf/kernel_test.go
index 0558cb7..5297bbc 100644
--- a/btf/kernel_test.go
+++ b/btf/kernel_test.go
@@ -4,6 +4,8 @@ import (
 	"os"
 	"testing"
 
+	"github.com/cilium/ebpf/internal/testutils"
+
 	"github.com/go-quicktest/qt"
 )
 
@@ -19,10 +21,61 @@ func TestLoadKernelSpec(t *testing.T) {
 }
 
 func TestLoadKernelModuleSpec(t *testing.T) {
-	if _, err := os.Stat("/sys/kernel/btf/btf_testmod"); os.IsNotExist(err) {
-		t.Skip("/sys/kernel/btf/btf_testmod not present")
+	if _, err := os.Stat("/sys/kernel/btf/bpf_testmod"); os.IsNotExist(err) {
+		t.Skip("/sys/kernel/btf/bpf_testmod not present")
 	}
 
-	_, err := LoadKernelModuleSpec("btf_testmod")
+	_, err := LoadKernelModuleSpec("bpf_testmod")
+	qt.Assert(t, qt.IsNil(err))
+}
+
+func TestCache(t *testing.T) {
+	FlushKernelSpec()
+	c := NewCache()
+
+	qt.Assert(t, qt.IsNil(c.KernelTypes))
+	qt.Assert(t, qt.HasLen(c.KernelModules, 0))
+
+	// Test that Kernel() creates only one copy
+	spec1, err := c.Kernel()
+	testutils.SkipIfNotSupported(t, err)
+	qt.Assert(t, qt.IsNil(err))
+	qt.Assert(t, qt.IsNotNil(spec1))
+
+	spec2, err := c.Kernel()
 	qt.Assert(t, qt.IsNil(err))
+	qt.Assert(t, qt.IsNotNil(spec2))
+
+	qt.Assert(t, qt.Equals(spec1, spec2))
+
+	// Test that Module() creates only one copy
+	mod1, err := c.Module("bpf_testmod")
+	if !os.IsNotExist(err) {
+		qt.Assert(t, qt.IsNil(err))
+		qt.Assert(t, qt.IsNotNil(mod1))
+
+		mod2, err := c.Module("bpf_testmod")
+		qt.Assert(t, qt.IsNil(err))
+		qt.Assert(t, qt.IsNotNil(mod2))
+
+		qt.Assert(t, qt.Equals(mod1, mod2))
+	}
+
+	// Pre-populate global cache
+	vmlinux, err := LoadKernelSpec()
+	qt.Assert(t, qt.IsNil(err))
+
+	testmod, err := LoadKernelModuleSpec("bpf_testmod")
+	if !os.IsNotExist(err) {
+		qt.Assert(t, qt.IsNil(err))
+	}
+
+	// Test that NewCache populates from global cache
+	c = NewCache()
+	qt.Assert(t, qt.IsNotNil(c.KernelTypes))
+	qt.Assert(t, qt.Not(qt.Equals(c.KernelTypes, vmlinux)))
+	if testmod != nil {
+		qt.Assert(t, qt.IsNotNil(c.KernelModules["bpf_testmod"]))
+		qt.Assert(t, qt.Not(qt.Equals(c.KernelModules["bpf_testmod"], testmod)))
+	}
 }