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>

Author: lmb

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,98 +23,112 @@ 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")
 	if err == nil {
 		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
+}

btf/kernel_test.go

@@ -4,6 +4,8 @@ import (
 	"os"
 	"testing"
 
+	"github.com/cilium/ebpf/internal/testutils"
+
 	"github.com/go-quicktest/qt"
 )
 
@@ -26,3 +28,54 @@ func TestLoadKernelModuleSpec(t *testing.T) {
 	_, 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)))
+	}
+}

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 {

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