Only read the BTF when new types are requested

btf/kernel.go

@@ -3,6 +3,7 @@ package btf
 import (
 	"errors"
 	"fmt"
+	"maps"
 	"os"
 	"path/filepath"
 	"sync"
@@ -15,7 +16,6 @@ var kernelBTF = struct {
 	sync.RWMutex
 	kernel   *Spec
 	fullSpec bool
-	opts     *SpecOptions
 	modules  map[string]*Spec
 }{
 	modules: make(map[string]*Spec),
@@ -70,34 +70,44 @@ func LoadKernelSpecWithOptions(opts *SpecOptions) (*Spec, error) {
 		return LoadKernelSpec()
 	}
 
+	// Work on a copy of the paramter
+	newOpts := &SpecOptions{
+		TypeNames: make(map[string]struct{}, len(opts.TypeNames)),
+	}
+	maps.Copy(newOpts.TypeNames, opts.TypeNames)
+
 	// TODO: minimize critical section
 	kernelBTF.Lock()
 	defer kernelBTF.Unlock()
 
-	if kernelBTF.opts == nil {
-		newSpec, _, err := loadKernelSpec(opts)
-		if err != nil {
-			return nil, err
+	needsUpdate := true
+	if kernelBTF.kernel != nil {
+		// Check if we already have all requested types
+		for _, t := range kernelBTF.kernel.imm.types {
+			newOpts.TypeNames[t.TypeName()] = struct{}{}
 		}
-		kernelBTF.opts = opts
-		kernelBTF.kernel = newSpec
-		return newSpec.Copy(), nil
+
+		needsUpdate = len(newOpts.TypeNames) > len(kernelBTF.kernel.imm.types)
 	}
 
-	for typeName := range opts.TypeNames {
-		kernelBTF.opts.TypeNames[typeName] = struct{}{}
+	if !needsUpdate {
+		return kernelBTF.kernel.Copy(), nil
 	}
 
-	newSpec, _, err := loadKernelSpec(kernelBTF.opts)
+	newSpec, _, err := loadKernelSpec(newOpts)
 	if err != nil {
 		return nil, err
 	}
 
-	kernelBTF.kernel.mu.Lock()
-	kernelBTF.kernel.imm.types = newSpec.imm.types
-	kernelBTF.kernel.imm.typeIDs = newSpec.imm.typeIDs
-	kernelBTF.kernel.imm.namedTypes = newSpec.imm.namedTypes
-	kernelBTF.kernel.mu.Unlock()
+	if kernelBTF.kernel == nil {
+		kernelBTF.kernel = newSpec
+	} else {
+		kernelBTF.kernel.mu.Lock()
+		kernelBTF.kernel.imm.types = newSpec.imm.types
+		kernelBTF.kernel.imm.typeIDs = newSpec.imm.typeIDs
+		kernelBTF.kernel.imm.namedTypes = newSpec.imm.namedTypes
+		kernelBTF.kernel.mu.Unlock()
+	}
 
 	return kernelBTF.kernel.Copy(), nil
 }

btf/kernel_test.go

@@ -32,6 +32,9 @@ func TestLoadKernelSpecWithOpts(t *testing.T) {
 		t.Skip("/sys/kernel/btf/vmlinux not present")
 	}
 
+	kernelBTF.fullSpec = false
+	kernelBTF.kernel = nil
+
 	opts := &SpecOptions{
 		TypeNames: map[string]struct{}{
 			"task_struct": {},
@@ -54,3 +57,75 @@ func TestLoadKernelSpecWithOpts(t *testing.T) {
 	qt.Assert(t, qt.Equals(len(spec1.imm.typeIDs), len(spec2.imm.typeIDs)))
 	qt.Assert(t, qt.DeepEquals(spec1.imm.namedTypes, spec2.imm.namedTypes))
 }
+
+func TestLoadKernelSpecWithMoreOpts(t *testing.T) {
+	if _, err := os.Stat("/sys/kernel/btf/vmlinux"); os.IsNotExist(err) {
+		t.Skip("/sys/kernel/btf/vmlinux not present")
+	}
+
+	opts1 := &SpecOptions{
+		TypeNames: map[string]struct{}{
+			"pt_regs": {},
+		},
+	}
+
+	kernelBTF.fullSpec = false
+	kernelBTF.kernel = nil
+
+	spec1, err := LoadKernelSpecWithOptions(opts1)
+	if err != nil {
+		t.Fatal("Can't load kernel spec:", err)
+	}
+
+	opts2 := &SpecOptions{
+		TypeNames: map[string]struct{}{
+			"task_struct": {},
+			"socket":      {},
+		},
+	}
+
+	spec2, err := LoadKernelSpecWithOptions(opts2)
+	if err != nil {
+		t.Fatal("Can't load kernel spec:", err)
+	}
+
+	for typeName := range opts1.TypeNames {
+		contains := false
+		for _, t := range spec1.imm.types {
+			if t.TypeName() == typeName {
+				contains = true
+				break
+			}
+		}
+		qt.Assert(t, qt.IsTrue(contains))
+
+		contains = false
+		for _, t := range spec2.imm.types {
+			if t.TypeName() == typeName {
+				contains = true
+				break
+			}
+		}
+		qt.Assert(t, qt.IsTrue(contains))
+	}
+
+	for typeName := range opts2.TypeNames {
+		contains := false
+		for _, t := range spec1.imm.types {
+			if t.TypeName() == typeName {
+				contains = true
+				break
+			}
+		}
+		qt.Assert(t, qt.IsFalse(contains))
+
+		contains = false
+		for _, t := range spec2.imm.types {
+			if t.TypeName() == typeName {
+				contains = true
+				break
+			}
+		}
+		qt.Assert(t, qt.IsTrue(contains))
+	}
+}