btf: avoid constructing strings for named type index

The most common use case of a Spec is to look up a type by its name.
For this purpose we maintain a map[essentialName][]TypeID. This
requires allocating a string for each named type, which causes a
very large overhead when parsing BTF.

In reality, only a very small number of the
named types will ever be looked up. The intuition here is that a
couple of structs in the kernel contain most of the interesting
information, for example struct sk_buff.

Move as much of the cost of looking up a type by name to the actual
lookup. Instead of spending a lot of time constructing an index up
front we only maintaing an index going from the hash of a name to
a type ID.

1. We can compute the hash on a byte slice and therefore avoid
   allocating a string.
2. Storing the index as a (hash, id) tuple allows us to store it
   in a slice. Lookups are just a binary search into the index.
3. Hash collisions do not introduce additional complexity because
   types can already share the same name. At the same time the
   common case of a 1:1 mapping from name to type is fast.

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

Author: lmb

btf/btf.go

@@ -12,6 +12,7 @@ import (
 	"math"
 	"os"
 	"reflect"
+	"slices"
 
 	"github.com/cilium/ebpf/internal"
 	"github.com/cilium/ebpf/internal/sys"
@@ -399,28 +400,20 @@ 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) {
-	typeIDs := s.TypeIDsByName(newEssentialName(name))
-	if len(typeIDs) == 0 {
-		return nil, fmt.Errorf("type name %s: %w", name, ErrNotFound)
-	}
-
-	// Return a copy to prevent changes to namedTypes.
-	result := make([]Type, 0, len(typeIDs))
-	for _, id := range typeIDs {
-		typ, err := s.TypeByID(id)
-		if errors.Is(err, ErrNotFound) {
-			return nil, fmt.Errorf("no type with ID %d", id)
-		} else if err != nil {
-			return nil, err
-		}
+	types, err := s.decoder.TypesByName(newEssentialName(name))
+	if err != nil {
+		return nil, err
+	}
 
+	for i := 0; i < len(types); i++ {
 		// Match against the full name, not just the essential one
 		// in case the type being looked up is a struct flavor.
-		if typ.TypeName() == name {
-			result = append(result, typ)
+		if types[i].TypeName() != name {
+			types = slices.Delete(types, i, i+1)
 		}
 	}
-	return result, nil
+
+	return types, nil
 }
 
 // AnyTypeByName returns a Type with the given name.

btf/btf_test.go

@@ -249,7 +249,7 @@ func BenchmarkIterateVmlinux(b *testing.B) {
 func TestParseCurrentKernelBTF(t *testing.T) {
 	spec := vmlinuxSpec(t)
 
-	if len(spec.namedTypes) == 0 {
+	if len(spec.offsets) == 0 {
 		t.Fatal("Empty kernel BTF")
 	}
 }
@@ -267,7 +267,7 @@ func TestFindVMLinux(t *testing.T) {
 		t.Fatal("Can't load BTF:", err)
 	}
 
-	if len(spec.namedTypes) == 0 {
+	if len(spec.offsets) == 0 {
 		t.Fatal("Empty kernel BTF")
 	}
 }

btf/core.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"math"
 	"reflect"
-	"slices"
 	"strconv"
 	"strings"
 
@@ -265,16 +264,12 @@ func CORERelocate(relos []*CORERelocation, targets []*Spec, bo binary.ByteOrder,
 
 		var targetTypes []Type
 		for _, target := range targets {
-			namedTypeIDs := target.TypeIDsByName(essentialName)
-			targetTypes = slices.Grow(targetTypes, len(namedTypeIDs))
-			for _, id := range namedTypeIDs {
-				typ, err := target.TypeByID(id)
-				if err != nil {
-					return nil, err
-				}
-
-				targetTypes = append(targetTypes, typ)
+			namedTypes, err := target.TypesByName(essentialName)
+			if err != nil {
+				return nil, err
 			}
+
+			targetTypes = append(targetTypes, namedTypes...)
 		}
 
 		fixups, err := coreCalculateFixups(group.relos, targetTypes, bo, resolveTargetTypeID)

btf/strings.go

@@ -54,10 +54,20 @@ func (st *stringTable) Lookup(offset uint32) (string, error) {
 		return "", nil
 	}
 
+	b, err := st.lookupSlow(offset)
+	return string(b), err
+}
+
+func (st *stringTable) LookupBytes(offset uint32) ([]byte, error) {
+	// Fast path: zero offset is the empty string, looked up frequently.
+	if offset == 0 {
+		return nil, nil
+	}
+
 	return st.lookupSlow(offset)
 }
 
-func (st *stringTable) lookupSlow(offset uint32) (string, error) {
+func (st *stringTable) lookupSlow(offset uint32) ([]byte, error) {
 	if st.base != nil {
 		n := uint32(len(st.base.bytes))
 		if offset < n {
@@ -67,15 +77,15 @@ func (st *stringTable) lookupSlow(offset uint32) (string, error) {
 	}
 
 	if offset > uint32(len(st.bytes)) {
-		return "", fmt.Errorf("offset %d is out of bounds of string table", offset)
+		return nil, fmt.Errorf("offset %d is out of bounds of string table", offset)
 	}
 
 	if offset > 0 && st.bytes[offset-1] != 0 {
-		return "", fmt.Errorf("offset %d is not the beginning of a string", offset)
+		return nil, fmt.Errorf("offset %d is not the beginning of a string", offset)
 	}
 
 	i := bytes.IndexByte(st.bytes[offset:], 0)
-	return string(st.bytes[offset : offset+uint32(i)]), nil
+	return st.bytes[offset : offset+uint32(i)], nil
 }
 
 // stringTableBuilder builds BTF string tables.

btf/unmarshal.go

@@ -1,12 +1,15 @@
 package btf
 
 import (
+	"bytes"
 	"encoding/binary"
 	"fmt"
+	"hash/maphash"
 	"io"
 	"iter"
 	"maps"
 	"math"
+	"slices"
 	"sync"
 )
 
@@ -21,9 +24,10 @@ type decoder struct {
 	firstTypeID TypeID
 	// Map from TypeID to offset of the marshaled data in raw. Contains an entry
 	// for each TypeID, including 0 aka Void. The offset for Void is invalid.
-	offsets    []int
-	declTags   map[TypeID][]TypeID
-	namedTypes map[essentialName][]TypeID
+	offsets  []int
+	declTags map[TypeID][]TypeID
+	// An index from essentialName to TypeID.
+	namedTypes *fuzzyStringIndex
 
 	// Protection for mutable fields below.
 	mu              sync.Mutex
@@ -74,7 +78,7 @@ func newDecoder(raw []byte, bo binary.ByteOrder, strings *stringTable, base *dec
 
 	offsets := make([]int, 0, numTypes)
 	declTags := make(map[TypeID][]TypeID, numDeclTags)
-	namedTypes := make(map[essentialName][]TypeID, numNamedTypes)
+	namedTypes := newFuzzyStringIndex(numNamedTypes)
 
 	if firstTypeID == 0 {
 		// Add a sentinel for Void.
@@ -94,27 +98,27 @@ func newDecoder(raw []byte, bo binary.ByteOrder, strings *stringTable, base *dec
 		}
 
 		// Build named type index.
-		name, err := strings.Lookup(header.NameOff)
+		name, err := strings.LookupBytes(header.NameOff)
 		if err != nil {
 			return nil, fmt.Errorf("lookup type name for id %v: %w", id, err)
 		}
 
-		if name := newEssentialName(name); name != "" {
-			ids := namedTypes[name]
-			if ids == nil {
-				// Almost all names will only have a single name to them.
-				// Explicitly allocate a slice of capacity 1 instead of relying
-				// on append behaviour.
-				ids = []TypeID{id}
-			} else {
-				ids = append(ids, id)
+		if len(name) > 0 {
+			if i := bytes.Index(name, []byte("___")); i != -1 {
+				// Flavours are rare. It's cheaper to find the first index for some
+				// reason.
+				i = bytes.LastIndex(name, []byte("___"))
+				name = name[:i]
 			}
-			namedTypes[name] = ids
+
+			namedTypes.Add(name, id)
 		}
 
 		id++
 	}
 
+	namedTypes.Build()
+
 	return &decoder{
 		base,
 		bo,
@@ -211,11 +215,28 @@ func (d *decoder) TypeID(typ Type) (TypeID, error) {
 	return id, nil
 }
 
-// TypeIDsByName returns all type IDs which have the given essential name.
+// TypesByName returns all types which have the given essential name.
 //
-// The returned slice must not be modified.
-func (d *decoder) TypeIDsByName(name essentialName) []TypeID {
-	return d.namedTypes[name]
+// Returns ErrNotFound if no matching Type exists.
+func (d *decoder) TypesByName(name essentialName) ([]Type, error) {
+	var types []Type
+	for id := range d.namedTypes.Find(string(name)) {
+		typ, err := d.TypeByID(id)
+		if err != nil {
+			return nil, err
+		}
+
+		if newEssentialName(typ.TypeName()) == name {
+			// Deal with hash collisions by checking against the name.
+			types = append(types, typ)
+		}
+	}
+
+	if len(types) == 0 {
+		return nil, fmt.Errorf("type with name %s: %w", name, ErrNotFound)
+	}
+
+	return types, nil
 }
 
 // TypeByID decodes a type and any of its descendants.
@@ -649,3 +670,73 @@ func (d *decoder) inflateType(id TypeID) (typ Type, err error) {
 
 	return typ, nil
 }
+
+// An index from string to TypeID.
+//
+// Fuzzy because it may return false positive matches.
+type fuzzyStringIndex struct {
+	seed    maphash.Seed
+	entries []fuzzyStringIndexEntry
+}
+
+func newFuzzyStringIndex(capacity int) *fuzzyStringIndex {
+	return &fuzzyStringIndex{
+		maphash.MakeSeed(),
+		make([]fuzzyStringIndexEntry, 0, capacity),
+	}
+}
+
+// Add a string to the index.
+//
+// Calling the method with identical arguments will create duplicate entries.
+func (idx *fuzzyStringIndex) Add(name []byte, id TypeID) {
+	hash := uint32(maphash.Bytes(idx.seed, name))
+	idx.entries = append(idx.entries, newFuzzyStringIndexEntry(hash, id))
+}
+
+// Build the index.
+//
+// Must be called after [Add] and before [Match].
+func (idx *fuzzyStringIndex) Build() {
+	slices.Sort(idx.entries)
+}
+
+// Find TypeIDs which may match the name.
+//
+// May return false positives, but is guaranteed to not have false negatives.
+//
+// You must call [Build] at least once before calling this method.
+func (idx *fuzzyStringIndex) Find(name string) iter.Seq[TypeID] {
+	return func(yield func(TypeID) bool) {
+		hash := uint32(maphash.String(idx.seed, name))
+
+		// We match only on the first 32 bits here, so ignore found.
+		i, _ := slices.BinarySearch(idx.entries, fuzzyStringIndexEntry(hash)<<32)
+		for i := i; i < len(idx.entries); i++ {
+			if idx.entries[i].hash() != hash {
+				break
+			}
+
+			if !yield(idx.entries[i].id()) {
+				return
+			}
+		}
+	}
+}
+
+// Tuple mapping the hash of an essential name to a type.
+//
+// Encoded in an uint64 so that it implements cmp.Ordered.
+type fuzzyStringIndexEntry uint64
+
+func newFuzzyStringIndexEntry(hash uint32, id TypeID) fuzzyStringIndexEntry {
+	return fuzzyStringIndexEntry(hash)<<32 | fuzzyStringIndexEntry(id)
+}
+
+func (e fuzzyStringIndexEntry) hash() uint32 {
+	return uint32(e >> 32)
+}
+
+func (e fuzzyStringIndexEntry) id() TypeID {
+	return TypeID(e)
+}

btf/unmarshal_test.go

@@ -0,0 +1,36 @@
+package btf
+
+import (
+	"iter"
+	"math"
+	"testing"
+
+	"github.com/go-quicktest/qt"
+)
+
+func TestFuzzyStringIndex(t *testing.T) {
+	idx := newFuzzyStringIndex(10)
+	count := testing.AllocsPerRun(1, func() {
+		idx.Add([]byte("foo"), 1)
+	})
+	qt.Assert(t, qt.Equals(count, 0))
+
+	idx.entries = idx.entries[:0]
+	idx.Add([]byte("foo"), 1)
+	idx.Add([]byte("bar"), 2)
+	idx.Add([]byte("baz"), 3)
+	idx.Build()
+
+	all := func(it iter.Seq[TypeID]) (ids []TypeID) {
+		for id := range it {
+			ids = append(ids, id)
+		}
+		return
+	}
+
+	qt.Assert(t, qt.SliceContains(all(idx.Find("foo")), 1))
+	qt.Assert(t, qt.SliceContains(all(idx.Find("bar")), 2))
+	qt.Assert(t, qt.SliceContains(all(idx.Find("baz")), 3))
+
+	qt.Assert(t, qt.IsTrue(newFuzzyStringIndexEntry(0, math.MaxUint32) < newFuzzyStringIndexEntry(1, 0)))
+}