btf: fix slow LoadKernelSpec by making Spec.Copy lazy

LoadKernelSpec is currently very slow, on the order of tens of milliseconds.
The reason for this is that we copy all ~150k types before returning the Spec
to the caller. This is necessary since we don't want independent callers of
LoadKernelSpec to be able to influence each other.

So far we've put up with the slowness in case of features that are not used
that much yet (kfuncs) or worked around it by moving APIs into the btf
package where we can avoid the copy.

I've made multiple attempts at fixing this problem. In no particlar order:

* Have an internal non-copying API and an external copying one. Users have
  to make sure they only call LoadKernelSpec once since it's expensive.
  The problem here is that we'd still like the library internal code to use
  a single copy of the kernel Spec, but this is very difficult due to
  import cycles.

* Decode raw BTF lazily instead of slurping all types into memory. This
  saves heap memory since we keep less inflated types around and only
  needs to do minimal upfront work. However, the upfront work still takes
  10ish ms to complete and the decoding logic becomes much more complicated
  since we need to lazily apply fixups. Both Dylan and I have written
  implementations for this, which go into the 500-600 lines of changes.

Finally, I've decided to dust of a third approach: we still parse kernel
BTF eagerly but make copying a Spec lazy. This makes LoadKernelSpec
basically free, at the cost of more expensive Spec.TypeByID, etc.
Doing CO-RE and reading split BTF also slows down, although it probably
ends up faster overall since we save a lot by not copying types.
Finally, we still pin one copy of the kernel spec in memory, so
FlushKernelSpec() is here to stay. It's hard to estimate how much
of an issue this really is since Go's tooling for heap analysis is
really poor.

Signed-off-by: Lorenz Bauer <lmb@isovalent.com>

btf/btf.go

@@ -29,9 +29,8 @@ var (
 // ID represents the unique ID of a BTF object.
 type ID = sys.BTFID
 
-// Spec allows querying a set of Types and loading the set into the
-// kernel.
-type Spec struct {
+// immutableTypes is a set of types which musn't be changed.
+type immutableTypes struct {
 	// All types contained by the spec, not including types from the base in
 	// case the spec was parsed from split BTF.
 	types []Type
@@ -44,13 +43,132 @@ type Spec struct {
 
 	// Types indexed by essential name.
 	// Includes all struct flavors and types with the same name.
-	namedTypes map[essentialName][]Type
+	namedTypes map[essentialName][]TypeID
+
+	// Byte order of the types. This affects things like struct member order
+	// when using bitfields.
+	byteOrder binary.ByteOrder
+}
+
+func (s *immutableTypes) typeByID(id TypeID) (Type, bool) {
+	if id < s.firstTypeID {
+		return nil, false
+	}
+
+	index := int(id - s.firstTypeID)
+	if index >= len(s.types) {
+		return nil, false
+	}
+
+	return s.types[index], true
+}
+
+// mutableTypes is a set of types which may be changed.
+type mutableTypes struct {
+	imm           immutableTypes
+	copies        map[Type]Type   // map[orig]copy
+	copiedTypeIDs map[Type]TypeID //map[copy]origID
+}
+
+// add a type to the set of mutable types.
+//
+// Copies type and all of its children once. Repeated calls with the same type
+// do not copy again.
+func (mt *mutableTypes) add(typ Type, typeIDs map[Type]TypeID) Type {
+	return modifyGraphPreorder(typ, func(t Type) (Type, bool) {
+		cpy, ok := mt.copies[t]
+		if ok {
+			// This has been copied previously, no need to continue.
+			return cpy, false
+		}
+
+		cpy = t.copy()
+		mt.copies[t] = cpy
+
+		if id, ok := typeIDs[t]; ok {
+			mt.copiedTypeIDs[cpy] = id
+		}
+
+		// This is a new copy, keep copying children.
+		return cpy, true
+	})
+}
+
+// copy a set of mutable types.
+func (mt *mutableTypes) copy() mutableTypes {
+	mtCopy := mutableTypes{
+		mt.imm,
+		make(map[Type]Type, len(mt.copies)),
+		make(map[Type]TypeID, len(mt.copiedTypeIDs)),
+	}
+
+	copies := make(map[Type]Type, len(mt.copies))
+	for orig, copy := range mt.copies {
+		// NB: We make a copy of copy, not orig, so that changes to mutable types
+		// are preserved.
+		copyOfCopy := mtCopy.add(copy, mt.copiedTypeIDs)
+		copies[orig] = copyOfCopy
+	}
+
+	// mtCopy.copies is currently map[copy]copyOfCopy, replace it with
+	// map[orig]copyOfCopy.
+	mtCopy.copies = copies
+	return mtCopy
+}
+
+func (mt *mutableTypes) typeID(typ Type) (TypeID, error) {
+	if _, ok := typ.(*Void); ok {
+		// Equality is weird for void, since it is a zero sized type.
+		return 0, nil
+	}
+
+	id, ok := mt.copiedTypeIDs[typ]
+	if !ok {
+		return 0, fmt.Errorf("no ID for type %s: %w", typ, ErrNotFound)
+	}
+
+	return id, nil
+}
+
+func (mt *mutableTypes) typeByID(id TypeID) (Type, bool) {
+	immT, ok := mt.imm.typeByID(id)
+	if !ok {
+		return nil, false
+	}
+
+	return mt.add(immT, mt.imm.typeIDs), true
+}
+
+func (mt *mutableTypes) anyTypesByName(name string) ([]Type, error) {
+	immTypes := mt.imm.namedTypes[newEssentialName(name)]
+	if len(immTypes) == 0 {
+		return nil, fmt.Errorf("type name %s: %w", name, ErrNotFound)
+	}
+
+	// Return a copy to prevent changes to namedTypes.
+	result := make([]Type, 0, len(immTypes))
+	for _, id := range immTypes {
+		immT, ok := mt.imm.typeByID(id)
+		if !ok {
+			return nil, fmt.Errorf("no type with ID %d", id)
+		}
+
+		// Match against the full name, not just the essential one
+		// in case the type being looked up is a struct flavor.
+		if immT.TypeName() == name {
+			result = append(result, mt.add(immT, mt.imm.typeIDs))
+		}
+	}
+	return result, nil
+}
+
+// Spec allows querying a set of Types and loading the set into the
+// kernel.
+type Spec struct {
+	mutableTypes
 
 	// String table from ELF.
 	strings *stringTable
-
-	// Byte order of the ELF we decoded the spec from, may be nil.
-	byteOrder binary.ByteOrder
 }
 
 // LoadSpec opens file and calls LoadSpecFromReader on it.
@@ -181,7 +299,7 @@ func loadSpecFromELF(file *internal.SafeELFFile) (*Spec, error) {
 		return nil, err
 	}
 
-	err = fixupDatasec(spec.types, sectionSizes, offsets)
+	err = fixupDatasec(spec.imm.types, sectionSizes, offsets)
 	if err != nil {
 		return nil, err
 	}
@@ -197,7 +315,7 @@ func loadRawSpec(btf io.ReaderAt, bo binary.ByteOrder, base *Spec) (*Spec, error
 	)
 
 	if base != nil {
-		if base.firstTypeID != 0 {
+		if base.imm.firstTypeID != 0 {
 			return nil, fmt.Errorf("can't use split BTF as base")
 		}
 
@@ -217,16 +335,22 @@ func loadRawSpec(btf io.ReaderAt, bo binary.ByteOrder, base *Spec) (*Spec, error
 	typeIDs, typesByName := indexTypes(types, firstTypeID)
 
 	return &Spec{
-		namedTypes:  typesByName,
-		typeIDs:     typeIDs,
-		types:       types,
-		firstTypeID: firstTypeID,
-		strings:     rawStrings,
-		byteOrder:   bo,
+		mutableTypes{
+			immutableTypes{
+				types,
+				typeIDs,
+				firstTypeID,
+				typesByName,
+				bo,
+			},
+			make(map[Type]Type),
+			make(map[Type]TypeID),
+		},
+		rawStrings,
 	}, nil
 }
 
-func indexTypes(types []Type, firstTypeID TypeID) (map[Type]TypeID, map[essentialName][]Type) {
+func indexTypes(types []Type, firstTypeID TypeID) (map[Type]TypeID, map[essentialName][]TypeID) {
 	namedTypes := 0
 	for _, typ := range types {
 		if typ.TypeName() != "" {
@@ -238,13 +362,15 @@ func indexTypes(types []Type, firstTypeID TypeID) (map[Type]TypeID, map[essentia
 	}
 
 	typeIDs := make(map[Type]TypeID, len(types))
-	typesByName := make(map[essentialName][]Type, namedTypes)
+	typesByName := make(map[essentialName][]TypeID, namedTypes)
 
 	for i, typ := range types {
+		id := firstTypeID + TypeID(i)
+		typeIDs[typ] = id
+
 		if name := newEssentialName(typ.TypeName()); name != "" {
-			typesByName[name] = append(typesByName[name], typ)
+			typesByName[name] = append(typesByName[name], id)
 		}
-		typeIDs[typ] = firstTypeID + TypeID(i)
 	}
 
 	return typeIDs, typesByName
@@ -492,17 +618,9 @@ func fixupDatasecLayout(ds *Datasec) error {
 
 // Copy creates a copy of Spec.
 func (s *Spec) Copy() *Spec {
-	types := copyTypes(s.types, nil)
-	typeIDs, typesByName := indexTypes(types, s.firstTypeID)
-
-	// NB: Other parts of spec are not copied since they are immutable.
 	return &Spec{
-		types,
-		typeIDs,
-		s.firstTypeID,
-		typesByName,
+		s.mutableTypes.copy(),
 		s.strings,
-		s.byteOrder,
 	}
 }
 
@@ -519,8 +637,8 @@ func (sw sliceWriter) Write(p []byte) (int, error) {
 // nextTypeID returns the next unallocated type ID or an error if there are no
 // more type IDs.
 func (s *Spec) nextTypeID() (TypeID, error) {
-	id := s.firstTypeID + TypeID(len(s.types))
-	if id < s.firstTypeID {
+	id := s.imm.firstTypeID + TypeID(len(s.imm.types))
+	if id < s.imm.firstTypeID {
 		return 0, fmt.Errorf("no more type IDs")
 	}
 	return id, nil
@@ -533,40 +651,17 @@ func (s *Spec) nextTypeID() (TypeID, error) {
 func (s *Spec) TypeByID(id TypeID) (Type, error) {
 	typ, ok := s.typeByID(id)
 	if !ok {
-		return nil, fmt.Errorf("look up type with ID %d (first ID is %d): %w", id, s.firstTypeID, ErrNotFound)
+		return nil, fmt.Errorf("look up type with ID %d (first ID is %d): %w", id, s.imm.firstTypeID, ErrNotFound)
 	}
 
 	return typ, nil
 }
 
-func (s *Spec) typeByID(id TypeID) (Type, bool) {
-	if id < s.firstTypeID {
-		return nil, false
-	}
-
-	index := int(id - s.firstTypeID)
-	if index >= len(s.types) {
-		return nil, false
-	}
-
-	return s.types[index], true
-}
-
 // TypeID returns the ID for a given Type.
 //
 // Returns an error wrapping ErrNoFound if the type isn't part of the Spec.
 func (s *Spec) TypeID(typ Type) (TypeID, error) {
-	if _, ok := typ.(*Void); ok {
-		// Equality is weird for void, since it is a zero sized type.
-		return 0, nil
-	}
-
-	id, ok := s.typeIDs[typ]
-	if !ok {
-		return 0, fmt.Errorf("no ID for type %s: %w", typ, ErrNotFound)
-	}
-
-	return id, nil
+	return s.mutableTypes.typeID(typ)
 }
 
 // AnyTypesByName returns a list of BTF Types with the given name.
@@ -577,21 +672,7 @@ func (s *Spec) TypeID(typ Type) (TypeID, error) {
 //
 // Returns an error wrapping ErrNotFound if no matching Type exists in the Spec.
 func (s *Spec) AnyTypesByName(name string) ([]Type, error) {
-	types := s.namedTypes[newEssentialName(name)]
-	if len(types) == 0 {
-		return nil, fmt.Errorf("type name %s: %w", name, ErrNotFound)
-	}
-
-	// Return a copy to prevent changes to namedTypes.
-	result := make([]Type, 0, len(types))
-	for _, t := range types {
-		// Match against the full name, not just the essential one
-		// in case the type being looked up is a struct flavor.
-		if t.TypeName() == name {
-			result = append(result, t)
-		}
-	}
-	return result, nil
+	return s.mutableTypes.anyTypesByName(name)
 }
 
 // AnyTypeByName returns a Type with the given name.
@@ -689,7 +770,7 @@ type TypesIterator struct {
 
 // Iterate returns the types iterator.
 func (s *Spec) Iterate() *TypesIterator {
-	return &TypesIterator{spec: s, id: s.firstTypeID}
+	return &TypesIterator{spec: s, id: s.imm.firstTypeID}
 }
 
 // Next returns true as long as there are any remaining types.

btf/btf_test.go

@@ -226,7 +226,7 @@ func BenchmarkParseVmlinux(b *testing.B) {
 func TestParseCurrentKernelBTF(t *testing.T) {
 	spec := vmlinuxSpec(t)
 
-	if len(spec.namedTypes) == 0 {
+	if len(spec.imm.namedTypes) == 0 {
 		t.Fatal("Empty kernel BTF")
 	}
 
@@ -257,7 +257,7 @@ func TestFindVMLinux(t *testing.T) {
 		t.Fatal("Can't load BTF:", err)
 	}
 
-	if len(spec.namedTypes) == 0 {
+	if len(spec.imm.namedTypes) == 0 {
 		t.Fatal("Empty kernel BTF")
 	}
 }
@@ -339,10 +339,10 @@ func TestSpecCopy(t *testing.T) {
 	spec := parseELFBTF(t, "../testdata/loader-el.elf")
 	cpy := spec.Copy()
 
-	have := typesFromSpec(t, spec)
-	qt.Assert(t, qt.IsTrue(len(spec.types) > 0))
+	have := typesFromSpec(spec)
+	qt.Assert(t, qt.IsTrue(len(have) > 0))
 
-	want := typesFromSpec(t, cpy)
+	want := typesFromSpec(cpy)
 	qt.Assert(t, qt.HasLen(want, len(have)))
 
 	for i := range want {
@@ -359,6 +359,28 @@ func TestSpecCopy(t *testing.T) {
 	}
 }
 
+func TestSpecCopyModifications(t *testing.T) {
+	spec := specFromTypes(t, []Type{&Int{Name: "a", Size: 4}})
+
+	typ, err := spec.TypeByID(1)
+	qt.Assert(t, qt.IsNil(err))
+
+	i := typ.(*Int)
+	i.Name = "b"
+	i.Size = 2
+
+	cpy := spec.Copy()
+	typ2, err := cpy.TypeByID(1)
+	qt.Assert(t, qt.IsNil(err))
+	i2 := typ2.(*Int)
+
+	qt.Assert(t, qt.Not(qt.Equals(i2, i)), qt.Commentf("Types are distinct"))
+	qt.Assert(t, qt.DeepEquals(i2, i), qt.Commentf("Modifications are preserved"))
+
+	i.Name = "bar"
+	qt.Assert(t, qt.Equals(i2.Name, "b"))
+}
+
 func TestSpecTypeByID(t *testing.T) {
 	spec := specFromTypes(t, nil)
 
@@ -459,10 +481,8 @@ func TestLoadSplitSpecFromReader(t *testing.T) {
 		t.Fatal("'int' is not supposed to be found in the split BTF")
 	}
 
-	if fnProto.Return != intType {
-		t.Fatalf("Return type of 'bpf_testmod_init()' (%s) does not match 'int' type (%s)",
-			fnProto.Return, intType)
-	}
+	qt.Assert(t, qt.Not(qt.Equals(fnProto.Return, intType)),
+		qt.Commentf("types found in base of split spec should be copies"))
 
 	// Check that copied split-BTF's spec has correct type indexing
 	splitSpecCopy := splitSpec.Copy()