Repository: cilium/ebpf
Pull request #1069 (2023-06-20 11:20:17)

btf: move Spec.Add to dedicated type
internal: add NewBuffer

    Share a single pool of bytes.Buffer across the whole library.

    Signed-off-by: Lorenz Bauer <lmb@isovalent.com>

btf: add marshalOptions.Order

    Allow specifying which ByteOrder to encode BTF into via a new field in
    marshalOptions.

    Signed-off-by: Lorenz Bauer <lmb@isovalent.com>

btf: use number of types to pre-size string table

    Remove the dependency on the ELF string table to pre-size the string table
    builder for large Specs. Instead we use the heuristic that for vmlinux the
    number of strings is roughly equal to the number of types.

    Signed-off-by: Lorenz Bauer <lmb@isovalent.com>

btf: refactor to Builder

    Spec allows both querying Types as well as constructing new BTF blobs via
    Add. This has some unfortunate consequences: we need to do extra work so
    that Added types can be queried and query code needs to be aware that the
    set of types can change. The query API also allows observing side-effects of
    adding types which we'd rather hide, for example when working around datasec
    bugs. It's also difficult to make the implementation efficient for both use
    cases at once.

    Split the code for constructing new BTF blobs into a separate struct.

    NewHandle now takes a Builder instead of a Spec which is a breaking change.
    Other API changes are to functions which have not appeared in a release yet.

    Signed-off-by: Lorenz Bauer <lmb@isovalent.com>

btf: unexport NewSpec

    NewSpec was added to be able to create BTF from scratch using Spec.Add. This
    usecase was migrated to Builder and Spec.Add doesn't exist anymore.

    Unexport NewSpec so we can still use it for testing but users only
    experience API breakage once.

    Signed-off-by: Lorenz Bauer <lmb@isovalent.com>

btf: generalise datasecWorkaround to Builder.Add

    Extend the datasec validation workaround to all Datasec. This is safe to do
    because there are no methods on Builder which allow going from TypeID to
    Type.

    Signed-off-by: Lorenz Bauer <lmb@isovalent.com>



Patch:
diff --git a/btf/btf.go b/btf/btf.go
index 45d7a35..5589a7e 100644
--- a/btf/btf.go
+++ b/btf/btf.go
@@ -79,8 +79,8 @@ func (h *btfHeader) stringStart() int64 {
 	return int64(h.HdrLen + h.StringOff)
 }
 
-// NewSpec creates a Spec containing only Void.
-func NewSpec() *Spec {
+// newSpec creates a Spec containing only Void.
+func newSpec() *Spec {
 	return &Spec{
 		[]Type{(*Void)(nil)},
 		map[Type]TypeID{(*Void)(nil): 0},
@@ -493,35 +493,6 @@ func (sw sliceWriter) Write(p []byte) (int, error) {
 	return copy(sw, p), nil
 }
 
-// Add a Type to the Spec, making it queryable via [TypeByName], etc.
-//
-// Adding the identical Type multiple times is valid and will return a stable ID.
-//
-// See [Type] for details on identity.
-func (s *Spec) Add(typ Type) (TypeID, error) {
-	if typ == nil {
-		return 0, fmt.Errorf("can't add nil Type")
-	}
-
-	if id, err := s.TypeID(typ); err == nil {
-		return id, nil
-	}
-
-	id, err := s.nextTypeID()
-	if err != nil {
-		return 0, err
-	}
-
-	s.typeIDs[typ] = id
-	s.types = append(s.types, typ)
-
-	if name := newEssentialName(typ.TypeName()); name != "" {
-		s.namedTypes[name] = append(s.namedTypes[name], typ)
-	}
-
-	return id, nil
-}
-
 // nextTypeID returns the next unallocated type ID or an error if there are no
 // more type IDs.
 func (s *Spec) nextTypeID() (TypeID, error) {
@@ -773,16 +744,19 @@ var haveFuncLinkage = internal.NewFeatureTest("BTF func linkage", "5.6", func()
 })
 
 func probeBTF(typ Type) error {
-	buf := getBuffer()
-	defer putBuffer(buf)
+	b, err := NewBuilder([]Type{typ})
+	if err != nil {
+		return err
+	}
 
-	if err := marshalTypes(buf, []Type{&Void{}, typ}, nil, nil); err != nil {
+	buf, err := b.Marshal(nil, nil)
+	if err != nil {
 		return err
 	}
 
 	fd, err := sys.BtfLoad(&sys.BtfLoadAttr{
-		Btf:     sys.NewSlicePointer(buf.Bytes()),
-		BtfSize: uint32(buf.Len()),
+		Btf:     sys.NewSlicePointer(buf),
+		BtfSize: uint32(len(buf)),
 	})
 
 	if err == nil {
diff --git a/btf/ext_info.go b/btf/ext_info.go
index fe5b2b9..f97ad78 100644
--- a/btf/ext_info.go
+++ b/btf/ext_info.go
@@ -153,8 +153,7 @@ func MarshalExtInfos(insns asm.Instructions) (_ *Handle, funcInfos, lineInfos []
 	return nil, nil, nil, nil
 
 marshal:
-	stb := newStringTableBuilder(0)
-	spec := NewSpec()
+	var b Builder
 	var fiBuf, liBuf bytes.Buffer
 	for {
 		if fn := FuncMetadata(iter.Ins); fn != nil {
@@ -162,7 +161,7 @@ marshal:
 				fn:     fn,
 				offset: iter.Offset,
 			}
-			if err := fi.marshal(&fiBuf, spec); err != nil {
+			if err := fi.marshal(&fiBuf, &b); err != nil {
 				return nil, nil, nil, fmt.Errorf("write func info: %w", err)
 			}
 		}
@@ -172,7 +171,7 @@ marshal:
 				line:   line,
 				offset: iter.Offset,
 			}
-			if err := li.marshal(&liBuf, stb); err != nil {
+			if err := li.marshal(&liBuf, &b); err != nil {
 				return nil, nil, nil, fmt.Errorf("write line info: %w", err)
 			}
 		}
@@ -182,14 +181,7 @@ marshal:
 		}
 	}
 
-	buf := getBuffer()
-	defer putBuffer(buf)
-
-	if err := marshalTypes(buf, spec.types, stb, kernelMarshalOptions()); err != nil {
-		return nil, nil, nil, fmt.Errorf("marshal BTF: %w", err)
-	}
-
-	handle, err := newHandleFromRawBTF(buf.Bytes())
+	handle, err := NewHandle(&b)
 	return handle, fiBuf.Bytes(), liBuf.Bytes(), err
 }
 
@@ -383,8 +375,8 @@ func newFuncInfos(bfis []bpfFuncInfo, spec *Spec) ([]funcInfo, error) {
 }
 
 // marshal into the BTF wire format.
-func (fi *funcInfo) marshal(w *bytes.Buffer, spec *Spec) error {
-	id, err := spec.Add(fi.fn)
+func (fi *funcInfo) marshal(w *bytes.Buffer, b *Builder) error {
+	id, err := b.Add(fi.fn)
 	if err != nil {
 		return err
 	}
@@ -549,7 +541,7 @@ func newLineInfos(blis []bpfLineInfo, strings *stringTable) ([]lineInfo, error)
 }
 
 // marshal writes the binary representation of the LineInfo to w.
-func (li *lineInfo) marshal(w *bytes.Buffer, stb *stringTableBuilder) error {
+func (li *lineInfo) marshal(w *bytes.Buffer, b *Builder) error {
 	line := li.line
 	if line.lineNumber > bpfLineMax {
 		return fmt.Errorf("line %d exceeds %d", line.lineNumber, bpfLineMax)
@@ -559,12 +551,12 @@ func (li *lineInfo) marshal(w *bytes.Buffer, stb *stringTableBuilder) error {
 		return fmt.Errorf("column %d exceeds %d", line.lineColumn, bpfColumnMax)
 	}
 
-	fileNameOff, err := stb.Add(line.fileName)
+	fileNameOff, err := b.addString(line.fileName)
 	if err != nil {
 		return fmt.Errorf("file name %q: %w", line.fileName, err)
 	}
 
-	lineOff, err := stb.Add(line.line)
+	lineOff, err := b.addString(line.line)
 	if err != nil {
 		return fmt.Errorf("line %q: %w", line.line, err)
 	}
diff --git a/btf/handle.go b/btf/handle.go
index 04aab92..b6b3e87 100644
--- a/btf/handle.go
+++ b/btf/handle.go
@@ -22,38 +22,25 @@ type Handle struct {
 	needsKernelBase bool
 }
 
-// NewHandle loads BTF into the kernel.
+// NewHandle loads the contents of a [Builder] into the kernel.
 //
-// Returns ErrNotSupported if BTF is not supported.
-func NewHandle(spec *Spec) (*Handle, error) {
-	if spec.byteOrder != nil && spec.byteOrder != internal.NativeEndian {
-		return nil, fmt.Errorf("can't load %s BTF on %s", spec.byteOrder, internal.NativeEndian)
-	}
-
-	if spec.firstTypeID != 0 {
-		return nil, fmt.Errorf("split BTF can't be loaded into the kernel")
-	}
-
-	buf := getBuffer()
-	defer putBuffer(buf)
+// Returns an error wrapping ErrNotSupported if the kernel doesn't support BTF.
+func NewHandle(b *Builder) (*Handle, error) {
+	small := getByteSlice()
+	defer putByteSlice(small)
 
-	var stb *stringTableBuilder
-	if spec.strings != nil {
-		// Use the ELF string table as an estimate of the final
-		// string table size. We don't use the ELF string
-		// table since the types may have been changed in the meantime.
-		stb = newStringTableBuilder(spec.strings.Num())
-	}
-
-	err := marshalTypes(buf, spec.types, stb, kernelMarshalOptions())
+	buf, err := b.Marshal(*small, KernelMarshalOptions())
 	if err != nil {
 		return nil, fmt.Errorf("marshal BTF: %w", err)
 	}
 
-	return newHandleFromRawBTF(buf.Bytes())
+	return NewHandleFromRawBTF(buf)
 }
 
-func newHandleFromRawBTF(btf []byte) (*Handle, error) {
+// NewHandleFromRawBTF loads raw BTF into the kernel.
+//
+// Returns an error wrapping ErrNotSupported if the kernel doesn't support BTF.
+func NewHandleFromRawBTF(btf []byte) (*Handle, error) {
 	if uint64(len(btf)) > math.MaxUint32 {
 		return nil, errors.New("BTF exceeds the maximum size")
 	}
diff --git a/btf/marshal.go b/btf/marshal.go
index 9d293c2..bfe53b4 100644
--- a/btf/marshal.go
+++ b/btf/marshal.go
@@ -9,113 +9,173 @@ import (
 	"sync"
 
 	"github.com/cilium/ebpf/internal"
+
+	"golang.org/x/exp/slices"
 )
 
-type marshalOptions struct {
+type MarshalOptions struct {
+	// Target byte order. Defaults to the system's native endianness.
+	Order binary.ByteOrder
 	// Remove function linkage information for compatibility with <5.6 kernels.
 	StripFuncLinkage bool
 }
 
-// kernelMarshalOptions will generate BTF suitable for the current kernel.
-func kernelMarshalOptions() *marshalOptions {
-	return &marshalOptions{
+// KernelMarshalOptions will generate BTF suitable for the current kernel.
+func KernelMarshalOptions() *MarshalOptions {
+	return &MarshalOptions{
+		Order:            internal.NativeEndian,
 		StripFuncLinkage: haveFuncLinkage() != nil,
 	}
 }
 
 // encoder turns Types into raw BTF.
 type encoder struct {
-	marshalOptions
-
-	byteOrder binary.ByteOrder
-	pending   internal.Deque[Type]
-	buf       *bytes.Buffer
-	strings   *stringTableBuilder
-	ids       map[Type]TypeID
-	lastID    TypeID
-}
+	MarshalOptions
 
-var emptyBTFHeader = make([]byte, btfHeaderLen)
+	pending internal.Deque[Type]
+	buf     *bytes.Buffer
+	strings *stringTableBuilder
+	ids     map[Type]TypeID
+	lastID  TypeID
+}
 
 var bufferPool = sync.Pool{
 	New: func() any {
-		return bytes.NewBuffer(make([]byte, btfHeaderLen+128))
+		buf := make([]byte, btfHeaderLen+128)
+		return &buf
 	},
 }
 
-func getBuffer() *bytes.Buffer {
-	buf := bufferPool.Get().(*bytes.Buffer)
-	buf.Reset()
-	return buf
+func getByteSlice() *[]byte {
+	return bufferPool.Get().(*[]byte)
 }
 
-func putBuffer(buf *bytes.Buffer) {
+func putByteSlice(buf *[]byte) {
+	*buf = (*buf)[:0]
 	bufferPool.Put(buf)
 }
 
-// marshalTypes encodes a slice of types into BTF wire format.
+// Builder turns Types into raw BTF.
 //
-// types are guaranteed to be written in the order they are passed to this
-// function. The first type must always be Void.
+// The default value may be used and represents an empty BTF blob. Void is
+// added implicitly if necessary.
+type Builder struct {
+	// Explicitly added types.
+	types []Type
+	// IDs for all added types which the user knows about.
+	stableIDs map[Type]TypeID
+	// Explicitly added strings.
+	strings *stringTableBuilder
+}
+
+// NewBuilder creates a Builder from a list of types.
 //
-// Doesn't support encoding split BTF since it's not possible to load
-// that into the kernel and we don't have a use case for writing BTF
-// out again.
+// It is more efficient than calling [Add] individually.
 //
-// w should be retrieved from bufferPool. opts may be nil.
-func marshalTypes(w *bytes.Buffer, types []Type, stb *stringTableBuilder, opts *marshalOptions) error {
-	if len(types) < 1 {
-		return errors.New("types must contain at least Void")
+// Returns an error if adding any of the types fails.
+func NewBuilder(types []Type) (*Builder, error) {
+	b := &Builder{
+		make([]Type, 0, len(types)),
+		make(map[Type]TypeID, len(types)),
+		nil,
 	}
 
-	if _, ok := types[0].(*Void); !ok {
-		return fmt.Errorf("first type is %s, not Void", types[0])
+	for _, typ := range types {
+		_, err := b.Add(typ)
+		if err != nil {
+			return nil, fmt.Errorf("add %s: %w", typ, err)
+		}
 	}
-	types = types[1:]
 
-	if stb == nil {
-		stb = newStringTableBuilder(0)
-	}
+	return b, nil
+}
 
-	e := encoder{
-		byteOrder: internal.NativeEndian,
-		buf:       w,
-		strings:   stb,
-		ids:       make(map[Type]TypeID, len(types)),
+// Add a Type and allocate a stable ID for it.
+//
+// Adding the identical Type multiple times is valid and will return the same ID.
+//
+// See [Type] for details on identity.
+func (b *Builder) Add(typ Type) (TypeID, error) {
+	if b.stableIDs == nil {
+		b.stableIDs = make(map[Type]TypeID)
 	}
 
-	if opts != nil {
-		e.marshalOptions = *opts
+	if _, ok := typ.(*Void); ok {
+		// Equality is weird for void, since it is a zero sized type.
+		return 0, nil
 	}
 
-	// Ensure that passed types are marshaled in the exact order they were
-	// passed.
-	e.pending.Grow(len(types))
-	for _, typ := range types {
-		if err := e.allocateID(typ); err != nil {
-			return err
+	if ds, ok := typ.(*Datasec); ok {
+		if err := datasecResolveWorkaround(b, ds); err != nil {
+			return 0, err
 		}
 	}
 
+	id, ok := b.stableIDs[typ]
+	if ok {
+		return id, nil
+	}
+
+	b.types = append(b.types, typ)
+
+	id = TypeID(len(b.types))
+	if int(id) != len(b.types) {
+		return 0, fmt.Errorf("no more type IDs")
+	}
+
+	b.stableIDs[typ] = id
+	return id, nil
+}
+
+// Marshal encodes all types in the Marshaler into BTF wire format.
+//
+// opts may be nil.
+func (b *Builder) Marshal(buf []byte, opts *MarshalOptions) ([]byte, error) {
+	stb := b.strings
+	if stb == nil {
+		// Assume that most types are named. This makes encoding large BTF like
+		// vmlinux a lot cheaper.
+		stb = newStringTableBuilder(len(b.types))
+	} else {
+		// Avoid modifying the Builder's string table.
+		stb = b.strings.Copy()
+	}
+
+	if opts == nil {
+		opts = &MarshalOptions{Order: internal.NativeEndian}
+	}
+
 	// Reserve space for the BTF header.
-	_, _ = e.buf.Write(emptyBTFHeader)
+	buf = slices.Grow(buf, btfHeaderLen)[:btfHeaderLen]
+
+	w := internal.NewBuffer(buf)
+	defer internal.PutBuffer(w)
+
+	e := encoder{
+		MarshalOptions: *opts,
+		buf:            w,
+		strings:        stb,
+		lastID:         TypeID(len(b.types)),
+		ids:            make(map[Type]TypeID, len(b.types)),
+	}
+
+	// Ensure that types are marshaled in the exact order they were Add()ed.
+	// Otherwise the ID returned from Add() won't match.
+	e.pending.Grow(len(b.types))
+	for _, typ := range b.types {
+		e.pending.Push(typ)
+		e.ids[typ] = b.stableIDs[typ]
+	}
 
 	if err := e.deflatePending(); err != nil {
-		return err
+		return nil, err
 	}
 
 	length := e.buf.Len()
 	typeLen := uint32(length - btfHeaderLen)
 
-	// Reserve space for the string table.
 	stringLen := e.strings.Length()
-	e.buf.Grow(stringLen)
-	buf := e.strings.AppendEncoded(e.buf.Bytes())
-
-	// Add string table to the unread portion of the buffer, otherwise
-	// it isn't return by Bytes().
-	// The copy is optimized out since src == dst.
-	_, _ = e.buf.Write(buf[length:])
+	buf = e.strings.AppendEncoded(e.buf.Bytes())
 
 	// Fill out the header, and write it out.
 	header := &btfHeader{
@@ -129,12 +189,26 @@ func marshalTypes(w *bytes.Buffer, types []Type, stb *stringTableBuilder, opts *
 		StringLen: uint32(stringLen),
 	}
 
-	err := binary.Write(sliceWriter(buf[:btfHeaderLen]), e.byteOrder, header)
+	err := binary.Write(sliceWriter(buf[:btfHeaderLen]), e.Order, header)
 	if err != nil {
-		return fmt.Errorf("write header: %v", err)
+		return nil, fmt.Errorf("write header: %v", err)
 	}
 
-	return nil
+	return buf, nil
+}
+
+// addString adds a string to the resulting BTF.
+//
+// Adding the same string multiple times will return the same result.
+//
+// Returns an identifier into the string table or an error if the string
+// contains invalid characters.
+func (b *Builder) addString(str string) (uint32, error) {
+	if b.strings == nil {
+		b.strings = newStringTableBuilder(0)
+	}
+
+	return b.strings.Add(str)
 }
 
 func (e *encoder) allocateID(typ Type) error {
@@ -170,7 +244,7 @@ func (e *encoder) deflatePending() error {
 		if t == root {
 			// Force descending into the current root type even if it already
 			// has an ID. Otherwise we miss children of types that have their
-			// ID pre-allocated in marshalTypes.
+			// ID pre-allocated via Add.
 			return false
 		}
 
@@ -338,7 +412,7 @@ func (e *encoder) deflateType(typ Type) (err error) {
 		return err
 	}
 
-	return raw.Marshal(e.buf, e.byteOrder)
+	return raw.Marshal(e.buf, e.Order)
 }
 
 func (e *encoder) convertMembers(header *btfType, members []Member) ([]btfMember, error) {
@@ -441,29 +515,23 @@ func (e *encoder) deflateVarSecinfos(vars []VarSecinfo) []btfVarSecinfo {
 // The function is intended for the use of the ebpf package and may be removed
 // at any point in time.
 func MarshalMapKV(key, value Type) (_ *Handle, keyID, valueID TypeID, err error) {
-	spec := NewSpec()
+	var b Builder
 
 	if key != nil {
-		keyID, err = spec.Add(key)
+		keyID, err = b.Add(key)
 		if err != nil {
 			return nil, 0, 0, fmt.Errorf("add key type: %w", err)
 		}
 	}
 
 	if value != nil {
-		if ds, ok := value.(*Datasec); ok {
-			if err := datasecResolveWorkaround(spec, ds); err != nil {
-				return nil, 0, 0, err
-			}
-		}
-
-		valueID, err = spec.Add(value)
+		valueID, err = b.Add(value)
 		if err != nil {
 			return nil, 0, 0, fmt.Errorf("add value type: %w", err)
 		}
 	}
 
-	handle, err := NewHandle(spec)
+	handle, err := NewHandle(&b)
 	if err != nil {
 		// Check for 'full' map BTF support, since kernels between 4.18 and 5.2
 		// already support BTF blobs for maps without Var or Datasec just fine.
diff --git a/btf/strings.go b/btf/strings.go
index d0eb604..bc6aff2 100644
--- a/btf/strings.go
+++ b/btf/strings.go
@@ -7,6 +7,8 @@ import (
 	"fmt"
 	"io"
 	"strings"
+
+	"golang.org/x/exp/maps"
 )
 
 type stringTable struct {
@@ -185,7 +187,6 @@ func (stb *stringTableBuilder) Lookup(str string) (uint32, error) {
 	}
 
 	return offset, nil
-
 }
 
 // Length returns the length in bytes.
@@ -203,3 +204,11 @@ func (stb *stringTableBuilder) AppendEncoded(buf []byte) []byte {
 	}
 	return buf
 }
+
+// Copy the string table builder.
+func (stb *stringTableBuilder) Copy() *stringTableBuilder {
+	return &stringTableBuilder{
+		stb.length,
+		maps.Clone(stb.strings),
+	}
+}
diff --git a/btf/workarounds.go b/btf/workarounds.go
index 1b068ca..12a89b8 100644
--- a/btf/workarounds.go
+++ b/btf/workarounds.go
@@ -4,7 +4,7 @@ package btf
 // to a Spec before the Datasec. This avoids a bug in kernel BTF validation.
 //
 // See https://lore.kernel.org/bpf/20230302123440.1193507-1-lmb@isovalent.com/
-func datasecResolveWorkaround(spec *Spec, ds *Datasec) error {
+func datasecResolveWorkaround(b *Builder, ds *Datasec) error {
 	for _, vsi := range ds.Vars {
 		v, ok := vsi.Type.(*Var)
 		if !ok {
@@ -13,7 +13,9 @@ func datasecResolveWorkaround(spec *Spec, ds *Datasec) error {
 
 		switch v.Type.(type) {
 		case *Typedef, *Volatile, *Const, *Restrict, *typeTag:
-			_, err := spec.Add(v.Type)
+			// NB: We must never call Add on a Datasec, otherwise we risk
+			// infinite recursion.
+			_, err := b.Add(v.Type)
 			if err != nil {
 				return err
 			}
diff --git a/internal/buffer.go b/internal/buffer.go
new file mode 100644
index 0000000..81c6544
--- /dev/null
+++ b/internal/buffer.go
@@ -0,0 +1,31 @@
+package internal
+
+import (
+	"bytes"
+	"sync"
+)
+
+var bytesBufferPool = sync.Pool{
+	New: func() interface{} {
+		return new(bytes.Buffer)
+	},
+}
+
+// NewBuffer retrieves a [bytes.Buffer] from a pool an re-initialises it.
+//
+// The returned buffer should be passed to [PutBuffer].
+func NewBuffer(buf []byte) *bytes.Buffer {
+	wr := bytesBufferPool.Get().(*bytes.Buffer)
+	// Reinitialize the Buffer with a new backing slice since it is returned to
+	// the caller by wr.Bytes() below. Pooling is faster despite calling
+	// NewBuffer. The pooled alloc is still reused, it only needs to be zeroed.
+	*wr = *bytes.NewBuffer(buf)
+	return wr
+}
+
+// PutBuffer releases a buffer to the pool.
+func PutBuffer(buf *bytes.Buffer) {
+	// Release reference to the backing buffer.
+	*buf = *bytes.NewBuffer(nil)
+	bytesBufferPool.Put(buf)
+}
diff --git a/marshalers.go b/marshalers.go
index f369481..a568bff 100644
--- a/marshalers.go
+++ b/marshalers.go
@@ -33,12 +33,6 @@ func marshalPtr(data interface{}, length int) (sys.Pointer, error) {
 	return sys.NewSlicePointer(buf), nil
 }
 
-var bytesBufferPool = sync.Pool{
-	New: func() interface{} {
-		return new(bytes.Buffer)
-	},
-}
-
 // marshalBytes converts an arbitrary value into a byte buffer.
 //
 // Prefer using Map.marshalKey and Map.marshalValue if possible, since
@@ -63,13 +57,8 @@ func marshalBytes(data interface{}, length int) (buf []byte, err error) {
 	case Map, *Map, Program, *Program:
 		err = fmt.Errorf("can't marshal %T", value)
 	default:
-		wr := bytesBufferPool.Get().(*bytes.Buffer)
-		defer bytesBufferPool.Put(wr)
-
-		// Reinitialize the Buffer with a new backing slice since it is returned to
-		// the caller by wr.Bytes() below. Pooling is faster despite calling
-		// NewBuffer. The pooled alloc is still reused, it only needs to be zeroed.
-		*wr = *bytes.NewBuffer(make([]byte, 0, length))
+		wr := internal.NewBuffer(make([]byte, 0, length))
+		defer internal.PutBuffer(wr)
 
 		err = binary.Write(wr, internal.NativeEndian, value)
 		if err != nil {


Test patch:
diff --git a/btf/btf_test.go b/btf/btf_test.go
index 4ffd690..234189b 100644
--- a/btf/btf_test.go
+++ b/btf/btf_test.go
@@ -203,50 +203,6 @@ func TestTypeByName(t *testing.T) {
 	}
 }
 
-func TestSpecAdd(t *testing.T) {
-	i := &Int{
-		Name:     "foo",
-		Size:     2,
-		Encoding: Signed | Char,
-	}
-	pi := &Pointer{i}
-
-	s := NewSpec()
-	id, err := s.Add(pi)
-	qt.Assert(t, err, qt.IsNil)
-	qt.Assert(t, id, qt.Equals, TypeID(1), qt.Commentf("First non-void type doesn't get id 1"))
-
-	got, err := s.TypeByID(id)
-	qt.Assert(t, err, qt.IsNil)
-	qt.Assert(t, got, qt.Equals, pi)
-
-	id, err = s.Add(pi)
-	qt.Assert(t, err, qt.IsNil)
-	qt.Assert(t, id, qt.Equals, TypeID(1))
-
-	_, err = s.TypeID(i)
-	qt.Assert(t, err, qt.IsNotNil, qt.Commentf("Children mustn't be added"))
-
-	id, err = s.Add(i)
-	qt.Assert(t, err, qt.IsNil)
-	qt.Assert(t, id, qt.Equals, TypeID(2), qt.Commentf("Second type doesn't get id 2"))
-
-	id, err = s.Add(i)
-	qt.Assert(t, err, qt.IsNil)
-	qt.Assert(t, id, qt.Equals, TypeID(2), qt.Commentf("Adding a type twice returns different ids"))
-
-	typ, err := s.AnyTypeByName("foo")
-	qt.Assert(t, err, qt.IsNil, qt.Commentf("Add doesn't make named type queryable"))
-	qt.Assert(t, typ, qt.Equals, i)
-
-	id, err = s.Add(&Typedef{"baz", i})
-	qt.Assert(t, err, qt.IsNil)
-	qt.Assert(t, id, qt.Equals, TypeID(3))
-
-	_, err = s.AnyTypeByName("baz")
-	qt.Assert(t, err, qt.IsNil)
-}
-
 func BenchmarkParseVmlinux(b *testing.B) {
 	rd := vmlinuxTestdataReader(b)
 	b.ReportAllocs()
@@ -323,29 +279,11 @@ func TestLoadSpecFromElf(t *testing.T) {
 				t.Error("Expected global linkage:", v)
 			}
 		}
-
-		if spec.byteOrder != internal.NativeEndian {
-			return
-		}
-
-		t.Run("Handle", func(t *testing.T) {
-			testutils.SkipIfNotSupported(t, haveMapBTF())
-			btf, err := NewHandle(spec)
-			if err != nil {
-				t.Fatal("Can't load BTF:", err)
-			}
-			defer btf.Close()
-		})
 	})
 }
 
 func TestVerifierError(t *testing.T) {
-	var buf bytes.Buffer
-	if err := marshalTypes(&buf, []Type{&Void{}}, nil, nil); err != nil {
-		t.Fatal(err)
-	}
-
-	_, err := newHandleFromRawBTF(buf.Bytes())
+	_, err := NewHandle(&Builder{})
 	testutils.SkipIfNotSupported(t, err)
 	var ve *internal.VerifierError
 	if !errors.As(err, &ve) {
@@ -387,10 +325,10 @@ func TestSpecCopy(t *testing.T) {
 }
 
 func TestSpecTypeByID(t *testing.T) {
-	_, err := NewSpec().TypeByID(0)
+	_, err := newSpec().TypeByID(0)
 	qt.Assert(t, err, qt.IsNil)
 
-	_, err = NewSpec().TypeByID(1)
+	_, err = newSpec().TypeByID(1)
 	qt.Assert(t, err, qt.ErrorIs, ErrNotFound)
 }
 
@@ -427,25 +365,31 @@ func ExampleSpec_TypeByName() {
 }
 
 func TestTypesIterator(t *testing.T) {
-	spec := NewSpec()
-
 	types := []Type{(*Void)(nil), &Int{Size: 4}, &Int{Size: 2}}
-	mustSpecAdd(t, spec, types...)
 
-	iter := spec.Iterate()
+	b, err := NewBuilder(types[1:])
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	raw, err := b.Marshal(nil, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	spec, err := LoadSpecFromReader(bytes.NewReader(raw))
+	if err != nil {
+		t.Fatal(err)
+	}
 
-	// Add a type after calling Iterate() to make sure the iteration
-	// below doesn't pick it up.
-	mustSpecAdd(t, spec, &Const{types[0]})
+	iter := spec.Iterate()
 
 	for i, typ := range types {
 		if !iter.Next() {
 			t.Fatal("Iterator ended early at item", i)
 		}
 
-		if iter.Type != typ {
-			t.Fatalf("Expected %p to match %p (%[1]T)", iter.Type, typ)
-		}
+		qt.Assert(t, iter.Type, qt.DeepEquals, typ)
 	}
 
 	if iter.Next() {
@@ -453,17 +397,6 @@ func TestTypesIterator(t *testing.T) {
 	}
 }
 
-func mustSpecAdd(t *testing.T, s *Spec, types ...Type) {
-	t.Helper()
-
-	for _, typ := range types {
-		_, err := s.Add(typ)
-		if err != nil {
-			t.Fatal(err)
-		}
-	}
-}
-
 func TestLoadSplitSpecFromReader(t *testing.T) {
 	spec := vmlinuxTestdataSpec(t)
 
diff --git a/btf/fuzz_test.go b/btf/fuzz_test.go
index 5c4a113..16da22b 100644
--- a/btf/fuzz_test.go
+++ b/btf/fuzz_test.go
@@ -56,7 +56,7 @@ func FuzzExtInfo(f *testing.F) {
 	}
 	f.Add(buf.Bytes(), []byte("\x00foo\x00barfoo\x00"))
 
-	emptySpec := NewSpec()
+	emptySpec := newSpec()
 
 	f.Fuzz(func(t *testing.T, data, strings []byte) {
 		if len(data) < binary.Size(btfExtHeader{}) {
diff --git a/btf/marshal_test.go b/btf/marshal_test.go
index 6e3a5be..4c79fde 100644
--- a/btf/marshal_test.go
+++ b/btf/marshal_test.go
@@ -8,11 +8,12 @@ import (
 
 	"github.com/cilium/ebpf/internal"
 	"github.com/cilium/ebpf/internal/testutils"
+	"github.com/google/go-cmp/cmp"
 
 	qt "github.com/frankban/quicktest"
 )
 
-func TestBuild(t *testing.T) {
+func TestBuilderMarshal(t *testing.T) {
 	typ := &Int{
 		Name:     "foo",
 		Size:     2,
@@ -26,14 +27,49 @@ func TestBuild(t *testing.T) {
 		&Typedef{"baz", typ},
 	}
 
-	var buf bytes.Buffer
-	qt.Assert(t, marshalTypes(&buf, want, nil, nil), qt.IsNil)
+	b, err := NewBuilder(want)
+	qt.Assert(t, err, qt.IsNil)
 
-	have, err := loadRawSpec(bytes.NewReader(buf.Bytes()), internal.NativeEndian, nil)
+	cpy := *b
+	buf, err := b.Marshal(nil, &MarshalOptions{Order: internal.NativeEndian})
+	qt.Assert(t, err, qt.IsNil)
+	qt.Assert(t, b, qt.CmpEquals(cmp.AllowUnexported(*b)), &cpy, qt.Commentf("Marshaling should not change Builder state"))
+
+	have, err := loadRawSpec(bytes.NewReader(buf), internal.NativeEndian, nil)
 	qt.Assert(t, err, qt.IsNil, qt.Commentf("Couldn't parse BTF"))
 	qt.Assert(t, have.types, qt.DeepEquals, want)
 }
 
+func TestBuilderAdd(t *testing.T) {
+	i := &Int{
+		Name:     "foo",
+		Size:     2,
+		Encoding: Signed | Char,
+	}
+	pi := &Pointer{i}
+
+	var b Builder
+	id, err := b.Add(pi)
+	qt.Assert(t, err, qt.IsNil)
+	qt.Assert(t, id, qt.Equals, TypeID(1), qt.Commentf("First non-void type doesn't get id 1"))
+
+	id, err = b.Add(pi)
+	qt.Assert(t, err, qt.IsNil)
+	qt.Assert(t, id, qt.Equals, TypeID(1))
+
+	id, err = b.Add(i)
+	qt.Assert(t, err, qt.IsNil)
+	qt.Assert(t, id, qt.Equals, TypeID(2), qt.Commentf("Second type doesn't get id 2"))
+
+	id, err = b.Add(i)
+	qt.Assert(t, err, qt.IsNil)
+	qt.Assert(t, id, qt.Equals, TypeID(2), qt.Commentf("Adding a type twice returns different ids"))
+
+	id, err = b.Add(&Typedef{"baz", i})
+	qt.Assert(t, err, qt.IsNil)
+	qt.Assert(t, id, qt.Equals, TypeID(3))
+}
+
 func TestRoundtripVMlinux(t *testing.T) {
 	types := vmlinuxSpec(t).types
 
@@ -70,22 +106,37 @@ limitTypes:
 		}
 	}
 
-	var buf bytes.Buffer
-	qt.Assert(t, marshalTypes(&buf, types, nil, nil), qt.IsNil)
+	buf := marshalNativeEndian(t, types)
 
-	rebuilt, err := loadRawSpec(bytes.NewReader(buf.Bytes()), binary.LittleEndian, nil)
+	rebuilt, err := loadRawSpec(bytes.NewReader(buf), binary.LittleEndian, nil)
 	qt.Assert(t, err, qt.IsNil, qt.Commentf("round tripping BTF failed"))
 
 	if n := len(rebuilt.types); n > math.MaxUint16 {
 		t.Logf("Rebuilt BTF contains %d types which exceeds uint16, test may fail on older kernels", n)
 	}
 
-	h, err := NewHandle(rebuilt)
+	h, err := NewHandleFromRawBTF(buf)
 	testutils.SkipIfNotSupported(t, err)
 	qt.Assert(t, err, qt.IsNil, qt.Commentf("loading rebuilt BTF failed"))
 	h.Close()
 }
 
+func BenchmarkMarshaler(b *testing.B) {
+	spec := vmlinuxTestdataSpec(b)
+	types := spec.types[:100]
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		var b Builder
+		for _, typ := range types {
+			_, _ = b.Add(typ)
+		}
+		_, _ = b.Marshal(nil, nil)
+	}
+}
+
 func BenchmarkBuildVmlinux(b *testing.B) {
 	types := vmlinuxTestdataSpec(b).types
 
@@ -93,9 +144,20 @@ func BenchmarkBuildVmlinux(b *testing.B) {
 	b.ResetTimer()
 
 	for i := 0; i < b.N; i++ {
-		var buf bytes.Buffer
-		if err := marshalTypes(&buf, types, nil, nil); err != nil {
-			b.Fatal(err)
+		var b Builder
+		for _, typ := range types {
+			_, _ = b.Add(typ)
 		}
+		_, _ = b.Marshal(nil, nil)
 	}
 }
+
+func marshalNativeEndian(tb testing.TB, types []Type) []byte {
+	tb.Helper()
+
+	b, err := NewBuilder(types)
+	qt.Assert(tb, err, qt.IsNil)
+	buf, err := b.Marshal(nil, nil)
+	qt.Assert(tb, err, qt.IsNil)
+	return buf
+}
diff --git a/btf/types_test.go b/btf/types_test.go
index 0461c5e..ce3bd5c 100644
--- a/btf/types_test.go
+++ b/btf/types_test.go
@@ -205,11 +205,9 @@ func TestTagMarshaling(t *testing.T) {
 		&typeTag{&Int{}, "foo"},
 	} {
 		t.Run(fmt.Sprint(typ), func(t *testing.T) {
-			var buf bytes.Buffer
-			err := marshalTypes(&buf, []Type{&Void{}, typ}, nil, nil)
-			qt.Assert(t, err, qt.IsNil)
+			buf := marshalNativeEndian(t, []Type{typ})
 
-			s, err := loadRawSpec(bytes.NewReader(buf.Bytes()), internal.NativeEndian, nil)
+			s, err := loadRawSpec(bytes.NewReader(buf), internal.NativeEndian, nil)
 			qt.Assert(t, err, qt.IsNil)
 
 			have, err := s.TypeByID(1)
diff --git a/btf/workarounds_test.go b/btf/workarounds_test.go
index 9769063..1553b43 100644
--- a/btf/workarounds_test.go
+++ b/btf/workarounds_test.go
@@ -47,17 +47,12 @@ func TestDatasecResolveWorkaround(t *testing.T) {
 				},
 			}
 
-			spec := NewSpec()
-			if err := datasecResolveWorkaround(spec, ds); err != nil {
-				t.Fatal(err)
-			}
-
-			_, err := spec.Add(ds)
+			b, err := NewBuilder([]Type{ds})
 			if err != nil {
 				t.Fatal(err)
 			}
 
-			h, err := NewHandle(spec)
+			h, err := NewHandle(b)
 			var ve *internal.VerifierError
 			if errors.As(err, &ve) {
 				t.Fatalf("%+v\n", ve)