program: use cached BTF for fixups and target search

Extend the BTF cache to fixups and target search. This means that
each collection load will load each BTF blob at most once.

Unexport fields in btf.Cache since we don't need to short-circuit
BTF loading from package ebpf anymore.

Signed-off-by: Lorenz Bauer <[email protected]>

Author: lmb

btf/kernel.go

@@ -225,9 +225,9 @@ func findVMLinux() (*os.File, error) {
 //
 // It is not safe for concurrent use.
 type Cache struct {
-	KernelTypes   *Spec
-	ModuleTypes   map[string]*Spec
-	LoadedModules []string
+	kernelTypes   *Spec
+	moduleTypes   map[string]*Spec
+	loadedModules []string
 }
 
 // NewCache creates a new Cache.
@@ -262,26 +262,26 @@ func NewCache() *Cache {
 // 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
+	if c.kernelTypes != nil {
+		return c.kernelTypes, nil
 	}
 
 	var err error
-	c.KernelTypes, err = LoadKernelSpec()
-	return c.KernelTypes, err
+	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.ModuleTypes[name]; spec != nil {
+	if spec := c.moduleTypes[name]; spec != nil {
 		return spec, nil
 	}
 
-	if c.ModuleTypes == nil {
-		c.ModuleTypes = make(map[string]*Spec)
+	if c.moduleTypes == nil {
+		c.moduleTypes = make(map[string]*Spec)
 	}
 
 	base, err := c.Kernel()
@@ -302,14 +302,14 @@ func (c *Cache) Module(name string) (*Spec, error) {
 	}
 
 	spec = &Spec{decoder: decoder}
-	c.ModuleTypes[name] = spec
+	c.moduleTypes[name] = spec
 	return spec, err
 }
 
 // Modules returns a sorted list of all loaded modules.
 func (c *Cache) Modules() ([]string, error) {
-	if c.LoadedModules != nil {
-		return c.LoadedModules, nil
+	if c.loadedModules != nil {
+		return c.loadedModules, nil
 	}
 
 	btfDir, err := os.Open("/sys/kernel/btf")
@@ -328,6 +328,6 @@ func (c *Cache) Modules() ([]string, error) {
 	})
 
 	sort.Strings(entries)
-	c.LoadedModules = entries
+	c.loadedModules = entries
 	return entries, nil
 }

btf/kernel_test.go

@@ -43,9 +43,9 @@ func TestCache(t *testing.T) {
 	FlushKernelSpec()
 	c := NewCache()
 
-	qt.Assert(t, qt.IsNil(c.KernelTypes))
-	qt.Assert(t, qt.HasLen(c.ModuleTypes, 0))
-	qt.Assert(t, qt.IsNil(c.LoadedModules))
+	qt.Assert(t, qt.IsNil(c.kernelTypes))
+	qt.Assert(t, qt.HasLen(c.moduleTypes, 0))
+	qt.Assert(t, qt.IsNil(c.loadedModules))
 
 	// Test that Kernel() creates only one copy
 	spec1, err := c.Kernel()
@@ -83,15 +83,15 @@ func TestCache(t *testing.T) {
 
 	// 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)))
+	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.ModuleTypes["bpf_testmod"]))
-		qt.Assert(t, qt.Not(qt.Equals(c.ModuleTypes["bpf_testmod"], testmod)))
+		qt.Assert(t, qt.IsNotNil(c.moduleTypes["bpf_testmod"]))
+		qt.Assert(t, qt.Not(qt.Equals(c.moduleTypes["bpf_testmod"], testmod)))
 	}
 
 	// Test that Modules only reads modules once.
 	_, err = c.Modules()
 	qt.Assert(t, qt.IsNil(err))
-	qt.Assert(t, qt.IsNotNil(c.LoadedModules))
+	qt.Assert(t, qt.IsNotNil(c.loadedModules))
 }

collection.go

@@ -437,7 +437,7 @@ func newCollectionLoader(coll *CollectionSpec, opts *CollectionOptions) (*collec
 		make(map[string]*Map),
 		make(map[string]*Program),
 		make(map[string]*Variable),
-		newBTFCache(&opts.Programs),
+		btf.NewCache(),
 	}, nil
 }
 

linker.go

@@ -122,9 +122,8 @@ func hasFunctionReferences(insns asm.Instructions) bool {
 
 // applyRelocations collects and applies any CO-RE relocations in insns.
 //
-// Passing a nil target will relocate against the running kernel. insns are
-// modified in place.
-func applyRelocations(insns asm.Instructions, bo binary.ByteOrder, b *btf.Builder, c *btf.Cache, extraTargets []*btf.Spec) error {
+// insns are modified in place.
+func applyRelocations(insns asm.Instructions, bo binary.ByteOrder, b *btf.Builder, c *btf.Cache, kernelOverride *btf.Spec, extraTargets []*btf.Spec) error {
 	var relos []*btf.CORERelocation
 	var reloInsns []*asm.Instruction
 	iter := insns.Iterate()
@@ -143,28 +142,35 @@ func applyRelocations(insns asm.Instructions, bo binary.ByteOrder, b *btf.Builde
 		bo = internal.NativeEndian
 	}
 
-	kernelTarget, err := c.Kernel()
-	if err != nil {
-		return fmt.Errorf("load kernel spec: %w", err)
-	}
+	var targets []*btf.Spec
+	if kernelOverride == nil {
+		kernel, err := c.Kernel()
+		if err != nil {
+			return fmt.Errorf("load kernel spec: %w", err)
+		}
 
-	modules, err := c.Modules()
-	// Ignore ErrNotExists to cater to kernels which have CONFIG_DEBUG_INFO_BTF_MODULES
-	// or CONFIG_DEBUG_INFO_BTF disabled.
-	if err != nil && !errors.Is(err, fs.ErrNotExist) {
-		return err
-	}
+		modules, err := c.Modules()
+		// Ignore ErrNotExists to cater to kernels which have CONFIG_DEBUG_INFO_BTF_MODULES
+		// or CONFIG_DEBUG_INFO_BTF disabled.
+		if err != nil && !errors.Is(err, fs.ErrNotExist) {
+			return err
+		}
 
-	targets := make([]*btf.Spec, 0, 1+len(modules)+len(extraTargets))
-	targets = append(targets, kernelTarget)
+		targets = make([]*btf.Spec, 0, 1+len(modules)+len(extraTargets))
+		targets = append(targets, kernel)
 
-	for _, kmod := range modules {
-		spec, err := c.Module(kmod)
-		if err != nil {
-			return fmt.Errorf("load BTF for kmod %s: %w", kmod, err)
-		}
+		for _, kmod := range modules {
+			spec, err := c.Module(kmod)
+			if err != nil {
+				return fmt.Errorf("load BTF for kmod %s: %w", kmod, err)
+			}
 
-		targets = append(targets, spec)
+			targets = append(targets, spec)
+		}
+	} else {
+		// We expect kernelOverride to contain the merged types
+		// of vmlinux and kernel modules, as distributed by btfhub.
+		targets = []*btf.Spec{kernelOverride}
 	}
 
 	targets = append(targets, extraTargets...)
@@ -289,7 +295,7 @@ const kfuncCallPoisonBase = 0xdedc0de
 // fixupKfuncs loops over all instructions in search for kfunc calls.
 // If at least one is found, the current kernels BTF and module BTFis are searched to set Instruction.Constant
 // and Instruction.Offset to the correct values.
-func fixupKfuncs(insns asm.Instructions) (_ handles, err error) {
+func fixupKfuncs(insns asm.Instructions, cache *btf.Cache) (_ handles, err error) {
 	closeOnError := func(c io.Closer) {
 		if err != nil {
 			c.Close()
@@ -310,7 +316,7 @@ fixups:
 	// Only load kernel BTF if we found at least one kfunc call. kernelSpec can be
 	// nil if the kernel does not have BTF, in which case we poison all kfunc
 	// calls.
-	kernelSpec, err := btf.LoadKernelSpec()
+	_, err = cache.Kernel()
 	// ErrNotSupportedOnOS wraps ErrNotSupported, check for it first.
 	if errors.Is(err, internal.ErrNotSupportedOnOS) {
 		return nil, fmt.Errorf("kfuncs are not supported on this platform: %w", err)
@@ -342,7 +348,7 @@ fixups:
 
 		// findTargetInKernel returns btf.ErrNotFound if the input btf.Spec is nil.
 		target := btf.Type((*btf.Func)(nil))
-		spec, module, err := findTargetInKernel(kernelSpec, kfm.Func.Name, &target)
+		spec, module, err := findTargetInKernel(kfm.Func.Name, &target, cache)
 		if errors.Is(err, btf.ErrNotFound) {
 			if kfm.Binding == elf.STB_WEAK {
 				if ins.IsKfuncCall() {

prog.go

@@ -236,7 +236,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, newBTFCache(&opts))
+	prog, err := newProgramWithOptions(spec, opts, btf.NewCache())
 	if errors.Is(err, asm.ErrUnsatisfiedMapReference) {
 		return nil, fmt.Errorf("cannot load program without loading its whole collection: %w", err)
 	}
@@ -297,7 +297,7 @@ func newProgramWithOptions(spec *ProgramSpec, opts ProgramOptions, c *btf.Cache)
 	copy(insns, spec.Instructions)
 
 	var b btf.Builder
-	if err := applyRelocations(insns, spec.ByteOrder, &b, c, opts.ExtraRelocationTargets); err != nil {
+	if err := applyRelocations(insns, spec.ByteOrder, &b, c, opts.KernelTypes, opts.ExtraRelocationTargets); err != nil {
 		return nil, fmt.Errorf("apply CO-RE relocations: %w", err)
 	}
 
@@ -350,7 +350,7 @@ func newProgramWithOptions(spec *ProgramSpec, opts ProgramOptions, c *btf.Cache)
 		return nil, err
 	}
 
-	handles, err := fixupKfuncs(insns)
+	handles, err := fixupKfuncs(insns, c)
 	if err != nil {
 		return nil, fmt.Errorf("fixing up kfuncs: %w", err)
 	}
@@ -381,7 +381,7 @@ func newProgramWithOptions(spec *ProgramSpec, opts ProgramOptions, c *btf.Cache)
 		attr.AttachBtfObjFd = uint32(spec.AttachTarget.FD())
 		defer runtime.KeepAlive(spec.AttachTarget)
 	} else if spec.AttachTo != "" {
-		module, targetID, err := findProgramTargetInKernel(spec.AttachTo, spec.Type, spec.AttachType)
+		module, targetID, err := findProgramTargetInKernel(spec.AttachTo, spec.Type, spec.AttachType, c)
 		if err != nil && !errors.Is(err, errUnrecognizedAttachType) {
 			// We ignore errUnrecognizedAttachType since AttachTo may be non-empty
 			// for programs that don't attach anywhere.
@@ -1027,7 +1027,7 @@ var errUnrecognizedAttachType = errors.New("unrecognized attach type")
 //
 // Returns errUnrecognizedAttachType if the combination of progType and attachType
 // is not recognised.
-func findProgramTargetInKernel(name string, progType ProgramType, attachType AttachType) (*btf.Handle, btf.TypeID, error) {
+func findProgramTargetInKernel(name string, progType ProgramType, attachType AttachType, cache *btf.Cache) (*btf.Handle, btf.TypeID, error) {
 	type match struct {
 		p ProgramType
 		a AttachType
@@ -1067,12 +1067,7 @@ func findProgramTargetInKernel(name string, progType ProgramType, attachType Att
 		return nil, 0, errUnrecognizedAttachType
 	}
 
-	spec, err := btf.LoadKernelSpec()
-	if err != nil {
-		return nil, 0, fmt.Errorf("load kernel spec: %w", err)
-	}
-
-	spec, module, err := findTargetInKernel(spec, typeName, &target)
+	spec, module, err := findTargetInKernel(typeName, &target, cache)
 	if errors.Is(err, btf.ErrNotFound) {
 		return nil, 0, &internal.UnsupportedFeatureError{Name: featureName}
 	}
@@ -1102,14 +1097,15 @@ func findProgramTargetInKernel(name string, progType ProgramType, attachType Att
 //
 // Returns a non-nil handle if the type was found in a module, or btf.ErrNotFound
 // if the type wasn't found at all.
-func findTargetInKernel(kernelSpec *btf.Spec, typeName string, target *btf.Type) (*btf.Spec, *btf.Handle, error) {
-	if kernelSpec == nil {
-		return nil, nil, fmt.Errorf("nil kernelSpec: %w", btf.ErrNotFound)
+func findTargetInKernel(typeName string, target *btf.Type, cache *btf.Cache) (*btf.Spec, *btf.Handle, error) {
+	kernelSpec, err := cache.Kernel()
+	if err != nil {
+		return nil, nil, err
 	}
 
-	err := kernelSpec.TypeByName(typeName, target)
+	err = kernelSpec.TypeByName(typeName, target)
 	if errors.Is(err, btf.ErrNotFound) {
-		spec, module, err := findTargetInModule(kernelSpec, typeName, target)
+		spec, module, err := findTargetInModule(typeName, target, cache)
 		if err != nil {
 			return nil, nil, fmt.Errorf("find target in modules: %w", err)
 		}
@@ -1127,7 +1123,7 @@ func findTargetInKernel(kernelSpec *btf.Spec, typeName string, target *btf.Type)
 // are searched in the order they were loaded.
 //
 // Returns btf.ErrNotFound if the target can't be found in any module.
-func findTargetInModule(base *btf.Spec, typeName string, target *btf.Type) (*btf.Spec, *btf.Handle, error) {
+func findTargetInModule(typeName string, target *btf.Type, cache *btf.Cache) (*btf.Spec, *btf.Handle, error) {
 	it := new(btf.HandleIterator)
 	defer it.Handle.Close()
 
@@ -1141,7 +1137,7 @@ func findTargetInModule(base *btf.Spec, typeName string, target *btf.Type) (*btf
 			continue
 		}
 
-		spec, err := it.Handle.Spec(base)
+		spec, err := cache.Module(info.Name)
 		if err != nil {
 			return nil, nil, fmt.Errorf("parse types for module %s: %w", info.Name, err)
 		}
@@ -1201,12 +1197,3 @@ 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.LoadedModules = []string{}
-	}
-	return c
-}