Repository: cilium/ebpf
Pull request #637 (2022-04-22 17:46:08)

btf: split string table ahead of time
We currently read the string table into one big byte slice, and
allocate substrings in Lookup. This forces us to hang on to one
large allocation for the string table, and many small allocations
for substrings. Calling Lookup for the same offset also allocates
on each invocation.

Instead, split the string table into substrings at load time. In
Lookup we can then use a binary search to find the correct string.
As a result, both overall memory and allocations drop:

    name            old time/op    new time/op    delta
    ParseVmlinux-4    76.6ms ±12%    73.9ms ± 4%     ~     (p=0.343 n=4+4)

    name            old alloc/op   new alloc/op   delta
    ParseVmlinux-4    92.9MB ± 0%    84.3MB ± 0%   -9.22%  (p=0.029 n=4+4)

    name            old allocs/op  new allocs/op  delta
    ParseVmlinux-4      786k ± 0%      693k ± 0%  -11.87%  (p=0.029 n=4+4)

To keep the speed roughly the same I had to "vendor" sort.Search and
specialised it for uint32, since the function was prominent in profiles.

lmb: cc @paulcacheux, this was the idea I mentioned, please take a look.
lmb: I also tried replacing `offsets, strings` with a `map[uint32]string` but its slower:

```
name            old time/op    new time/op    delta
ParseVmlinux-4    74.3ms ± 5%    82.4ms ± 3%  +10.87%  (p=0.029 n=4+4)

name            old alloc/op   new alloc/op   delta
ParseVmlinux-4    84.3MB ± 0%    87.8MB ± 0%   +4.16%  (p=0.029 n=4+4)

name            old allocs/op  new allocs/op  delta
ParseVmlinux-4      693k ± 0%      693k ± 0%     ~     (p=0.943 n=4+4)
```
paulcacheux: This looks great, thanks for working on this

Patch:
diff --git a/internal/btf/btf.go b/internal/btf/btf.go
index 8523754..59ac88f 100644
--- a/internal/btf/btf.go
+++ b/internal/btf/btf.go
@@ -35,7 +35,7 @@ type ID uint32
 type Spec struct {
 	// Data from .BTF.
 	rawTypes []rawType
-	strings  stringTable
+	strings  *stringTable
 
 	// Inflated Types.
 	types []Type
@@ -447,15 +447,14 @@ func guessRawBTFByteOrder(r io.ReaderAt) binary.ByteOrder {
 
 // parseBTF reads a .BTF section into memory and parses it into a list of
 // raw types and a string table.
-func parseBTF(btf io.ReaderAt, bo binary.ByteOrder) ([]rawType, stringTable, error) {
+func parseBTF(btf io.ReaderAt, bo binary.ByteOrder) ([]rawType, *stringTable, error) {
 	buf := internal.NewBufferedSectionReader(btf, 0, math.MaxInt64)
 	header, err := parseBTFHeader(buf, bo)
 	if err != nil {
 		return nil, nil, fmt.Errorf("parsing .BTF header: %v", err)
 	}
 
-	buf.Reset(io.NewSectionReader(btf, header.stringStart(), int64(header.StringLen)))
-	rawStrings, err := readStringTable(buf)
+	rawStrings, err := readStringTable(io.NewSectionReader(btf, header.stringStart(), int64(header.StringLen)))
 	if err != nil {
 		return nil, nil, fmt.Errorf("can't read type names: %w", err)
 	}
@@ -474,7 +473,7 @@ type variable struct {
 	name    string
 }
 
-func fixupDatasec(rawTypes []rawType, rawStrings stringTable, sectionSizes map[string]uint32, variableOffsets map[variable]uint32) error {
+func fixupDatasec(rawTypes []rawType, rawStrings *stringTable, sectionSizes map[string]uint32, variableOffsets map[variable]uint32) error {
 	for i, rawType := range rawTypes {
 		if rawType.Kind() != kindDatasec {
 			continue
@@ -580,7 +579,11 @@ func (s *Spec) marshal(opts marshalOpts) ([]byte, error) {
 	typeLen := uint32(buf.Len() - headerLen)
 
 	// Write string section after type section.
-	_, _ = buf.Write(s.strings)
+	stringsLen := s.strings.Length()
+	buf.Grow(stringsLen)
+	if err := s.strings.Marshal(&buf); err != nil {
+		return nil, err
+	}
 
 	// Fill out the header, and write it out.
 	header = &btfHeader{
@@ -591,7 +594,7 @@ func (s *Spec) marshal(opts marshalOpts) ([]byte, error) {
 		TypeOff:   0,
 		TypeLen:   typeLen,
 		StringOff: typeLen,
-		StringLen: uint32(len(s.strings)),
+		StringLen: uint32(stringsLen),
 	}
 
 	raw := buf.Bytes()
diff --git a/internal/btf/ext_info.go b/internal/btf/ext_info.go
index 39ce7b0..11e4057 100644
--- a/internal/btf/ext_info.go
+++ b/internal/btf/ext_info.go
@@ -21,7 +21,7 @@ type extInfo struct {
 }
 
 // loadExtInfos parses the .BTF.ext section into its constituent parts.
-func loadExtInfos(r io.ReaderAt, bo binary.ByteOrder, strings stringTable) (*extInfo, error) {
+func loadExtInfos(r io.ReaderAt, bo binary.ByteOrder, strings *stringTable) (*extInfo, error) {
 	// Open unbuffered section reader. binary.Read() calls io.ReadFull on
 	// the header structs, resulting in one syscall per header.
 	headerRd := io.NewSectionReader(r, 0, math.MaxInt64)
@@ -154,7 +154,7 @@ type btfExtInfoSec struct {
 // appearing within func_info and line_info sub-sections.
 // These headers appear once for each program section in the ELF and are
 // followed by one or more func/line_info records for the section.
-func parseExtInfoSec(r io.Reader, bo binary.ByteOrder, strings stringTable) (string, *btfExtInfoSec, error) {
+func parseExtInfoSec(r io.Reader, bo binary.ByteOrder, strings *stringTable) (string, *btfExtInfoSec, error) {
 	var infoHeader btfExtInfoSec
 	if err := binary.Read(r, bo, &infoHeader); err != nil {
 		return "", nil, fmt.Errorf("read ext info header: %w", err)
@@ -227,7 +227,7 @@ func (fi *FuncInfo) Marshal(w io.Writer, offset uint64) error {
 
 // parseLineInfos parses a func_info sub-section within .BTF.ext ito a map of
 // func infos indexed by section name.
-func parseFuncInfos(r io.Reader, bo binary.ByteOrder, strings stringTable) (map[string][]bpfFuncInfo, error) {
+func parseFuncInfos(r io.Reader, bo binary.ByteOrder, strings *stringTable) (map[string][]bpfFuncInfo, error) {
 	recordSize, err := parseExtInfoRecordSize(r, bo)
 	if err != nil {
 		return nil, err
@@ -374,7 +374,7 @@ func (li LineInfos) Marshal(w io.Writer, offset uint64) error {
 
 // parseLineInfos parses a line_info sub-section within .BTF.ext ito a map of
 // line infos indexed by section name.
-func parseLineInfos(r io.Reader, bo binary.ByteOrder, strings stringTable) (map[string][]bpfLineInfo, error) {
+func parseLineInfos(r io.Reader, bo binary.ByteOrder, strings *stringTable) (map[string][]bpfLineInfo, error) {
 	recordSize, err := parseExtInfoRecordSize(r, bo)
 	if err != nil {
 		return nil, err
@@ -462,7 +462,7 @@ var extInfoReloSize = binary.Size(bpfCORERelo{})
 
 // parseCORERelos parses a core_relos sub-section within .BTF.ext ito a map of
 // CO-RE relocations indexed by section name.
-func parseCORERelos(r io.Reader, bo binary.ByteOrder, strings stringTable) (map[string]CORERelos, error) {
+func parseCORERelos(r io.Reader, bo binary.ByteOrder, strings *stringTable) (map[string]CORERelos, error) {
 	recordSize, err := parseExtInfoRecordSize(r, bo)
 	if err != nil {
 		return nil, err
@@ -494,7 +494,7 @@ func parseCORERelos(r io.Reader, bo binary.ByteOrder, strings stringTable) (map[
 // parseCOREReloRecords parses a stream of CO-RE relocation entries into a
 // coreRelos. These records appear after a btf_ext_info_sec header in the
 // core_relos sub-section of .BTF.ext.
-func parseCOREReloRecords(r io.Reader, bo binary.ByteOrder, recordSize uint32, recordNum uint32, strings stringTable) (CORERelos, error) {
+func parseCOREReloRecords(r io.Reader, bo binary.ByteOrder, recordSize uint32, recordNum uint32, strings *stringTable) (CORERelos, error) {
 	var out CORERelos
 
 	var relo bpfCORERelo
diff --git a/internal/btf/strings.go b/internal/btf/strings.go
index 9876aa2..2f3bf83 100644
--- a/internal/btf/strings.go
+++ b/internal/btf/strings.go
@@ -1,54 +1,112 @@
 package btf
 
 import (
+	"bufio"
 	"bytes"
 	"errors"
 	"fmt"
 	"io"
 )
 
-type stringTable []byte
+type stringTable struct {
+	offsets []uint32
+	strings []string
+}
+
+// sizedReader is implemented by bytes.Reader, io.SectionReader, strings.Reader, etc.
+type sizedReader interface {
+	io.Reader
+	Size() int64
+}
+
+func readStringTable(r sizedReader) (*stringTable, error) {
+	// Derived from vmlinux BTF.
+	const averageStringLength = 16
 
-func readStringTable(r io.Reader) (stringTable, error) {
-	contents, err := io.ReadAll(r)
-	if err != nil {
-		return nil, fmt.Errorf("can't read string table: %v", err)
+	n := int(r.Size() / averageStringLength)
+	offsets := make([]uint32, 0, n)
+	strings := make([]string, 0, n)
+
+	offset := uint32(0)
+	scanner := bufio.NewScanner(r)
+	scanner.Split(splitNull)
+	for scanner.Scan() {
+		str := scanner.Text()
+		offsets = append(offsets, offset)
+		strings = append(strings, str)
+		offset += uint32(len(str)) + 1
+	}
+	if err := scanner.Err(); err != nil {
+		return nil, err
 	}
 
-	if len(contents) < 1 {
+	if len(strings) == 0 {
 		return nil, errors.New("string table is empty")
 	}
 
-	if contents[0] != '\x00' {
+	if strings[0] != "" {
 		return nil, errors.New("first item in string table is non-empty")
 	}
 
-	if contents[len(contents)-1] != '\x00' {
-		return nil, errors.New("string table isn't null terminated")
-	}
-
-	return stringTable(contents), nil
+	return &stringTable{offsets, strings}, nil
 }
 
-func (st stringTable) Lookup(offset uint32) (string, error) {
-	if int64(offset) > int64(^uint(0)>>1) {
-		return "", fmt.Errorf("offset %d overflows int", offset)
+func splitNull(data []byte, atEOF bool) (advance int, token []byte, err error) {
+	i := bytes.IndexByte(data, 0)
+	if i == -1 {
+		if atEOF && len(data) > 0 {
+			return 0, nil, errors.New("string table isn't null terminated")
+		}
+		return 0, nil, nil
 	}
 
-	pos := int(offset)
-	if pos >= len(st) {
-		return "", fmt.Errorf("offset %d is out of bounds", offset)
-	}
+	return i + 1, data[:i], nil
+}
 
-	if pos > 0 && st[pos-1] != '\x00' {
+func (st *stringTable) Lookup(offset uint32) (string, error) {
+	i := search(st.offsets, offset)
+	if i == len(st.offsets) || st.offsets[i] != offset {
 		return "", fmt.Errorf("offset %d isn't start of a string", offset)
 	}
 
-	str := st[pos:]
-	end := bytes.IndexByte(str, '\x00')
-	if end == -1 {
-		return "", fmt.Errorf("offset %d isn't null terminated", offset)
+	return st.strings[i], nil
+}
+
+func (st *stringTable) Length() int {
+	last := len(st.offsets) - 1
+	return int(st.offsets[last]) + len(st.strings[last]) + 1
+}
+
+func (st *stringTable) Marshal(w io.Writer) error {
+	for _, str := range st.strings {
+		_, err := io.WriteString(w, str)
+		if err != nil {
+			return err
+		}
+		_, err = w.Write([]byte{0})
+		if err != nil {
+			return err
+		}
 	}
+	return nil
+}
 
-	return string(str[:end]), nil
+// search is a copy of sort.Search specialised for uint32.
+//
+// Licensed under https://go.dev/LICENSE
+func search(ints []uint32, needle uint32) int {
+	// Define f(-1) == false and f(n) == true.
+	// Invariant: f(i-1) == false, f(j) == true.
+	i, j := 0, len(ints)
+	for i < j {
+		h := int(uint(i+j) >> 1) // avoid overflow when computing h
+		// i ≤ h < j
+		if !(ints[h] >= needle) {
+			i = h + 1 // preserves f(i-1) == false
+		} else {
+			j = h // preserves f(j) == true
+		}
+	}
+	// i == j, f(i-1) == false, and f(j) (= f(i)) == true  =>  answer is i.
+	return i
 }
diff --git a/internal/btf/types.go b/internal/btf/types.go
index 8175c08..8762592 100644
--- a/internal/btf/types.go
+++ b/internal/btf/types.go
@@ -807,7 +807,7 @@ func countFixups(rawTypes []rawType) int {
 // Returns a map of named types (so, where NameOff is non-zero) and a slice of types
 // indexed by TypeID. Since BTF ignores compilation units, multiple types may share
 // the same name. A Type may form a cyclic graph by pointing at itself.
-func inflateRawTypes(rawTypes []rawType, rawStrings stringTable) ([]Type, map[essentialName][]Type, error) {
+func inflateRawTypes(rawTypes []rawType, rawStrings *stringTable) ([]Type, map[essentialName][]Type, error) {
 	type fixupDef struct {
 		id           TypeID
 		expectedKind btfKind


Test patch:
diff --git a/internal/btf/btf_test.go b/internal/btf/btf_test.go
index 2ca82af..41d1e34 100644
--- a/internal/btf/btf_test.go
+++ b/internal/btf/btf_test.go
@@ -185,6 +185,19 @@ func TestParseCurrentKernelBTF(t *testing.T) {
 	if len(spec.namedTypes) == 0 {
 		t.Fatal("Empty kernel BTF")
 	}
+
+	totalBytes := 0
+	distinct := 0
+	seen := make(map[string]bool)
+	for _, str := range spec.strings.strings {
+		totalBytes += len(str)
+		if !seen[str] {
+			distinct++
+			seen[str] = true
+		}
+	}
+	t.Logf("%d strings total, %d distinct", len(spec.strings.strings), distinct)
+	t.Logf("Average string size: %.0f", float64(totalBytes)/float64(len(spec.strings.strings)))
 }
 
 func TestFindVMLinux(t *testing.T) {
diff --git a/internal/btf/ext_info_test.go b/internal/btf/ext_info_test.go
index fa8e641..eb1756e 100644
--- a/internal/btf/ext_info_test.go
+++ b/internal/btf/ext_info_test.go
@@ -1,6 +1,7 @@
 package btf
 
 import (
+	"bytes"
 	"strings"
 	"testing"
 
@@ -9,7 +10,10 @@ import (
 
 func TestParseExtInfoBigRecordSize(t *testing.T) {
 	rd := strings.NewReader("\xff\xff\xff\xff\x00\x00\x00\x000709171295166016")
-	table := stringTable("\x00")
+	table, err := readStringTable(bytes.NewReader([]byte{0}))
+	if err != nil {
+		t.Fatal(err)
+	}
 
 	if _, err := parseFuncInfos(rd, internal.NativeEndian, table); err == nil {
 		t.Error("Parsing func info with large record size doesn't return an error")
diff --git a/internal/btf/fuzz_test.go b/internal/btf/fuzz_test.go
index 7c979ff..89393b4 100644
--- a/internal/btf/fuzz_test.go
+++ b/internal/btf/fuzz_test.go
@@ -55,7 +55,11 @@ func FuzzExtInfo(f *testing.F) {
 			t.Skip("data is too short")
 		}
 
-		table := stringTable(strings)
+		table, err := readStringTable(bytes.NewReader(strings))
+		if err != nil {
+			t.Skip("invalid string table")
+		}
+
 		info, err := loadExtInfos(bytes.NewReader(data), internal.NativeEndian, table)
 		if err != nil {
 			if info != nil {
diff --git a/internal/btf/strings_test.go b/internal/btf/strings_test.go
index 259e27d..9a15ab0 100644
--- a/internal/btf/strings_test.go
+++ b/internal/btf/strings_test.go
@@ -14,7 +14,12 @@ func TestStringTable(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	if !bytes.Equal([]byte(in), []byte(st)) {
+	var buf bytes.Buffer
+	if err := st.Marshal(&buf); err != nil {
+		t.Fatal("Can't marshal string table:", err)
+	}
+
+	if !bytes.Equal([]byte(in), buf.Bytes()) {
 		t.Error("String table doesn't match input")
 	}