state/wire: do not use sync.Pool for single-byte buffers

This package critically depends on reading/writing single bytes. Since
arguments to interface methods io.Reader.Read / io.Writer.Write escape, a naive
implementation would heap-allocate a one-byte array per read/write.

Prior to cl/625167495, the wire package provided custom Reader/Writer
interfaces, implementations of which were required to provide their own
ReadByte/WriteByte methods that did not take any escaping arguments.
cl/625167495 eliminated these interfaces and made the wire package use
sync.Pool to allocate one-byte arrays instead, simplifying the pipelining of
readers and writers but introducing non-trivial overhead. This CL re-introduces
wire.Reader/Writer, but as structs combining an io.Reader/Writer and a one-byte
array; this preserves the relative ease of using arbitrary io.Readers/Writers,
while eliminating sync.Pool overhead by essentially having callers of the wire
package provide the persistent buffer.

PiperOrigin-RevId: 666946511

pkg/sentry/pgalloc/save_restore.go

@@ -30,6 +30,7 @@ import (
 	"gvisor.dev/gvisor/pkg/sentry/usage"
 	"gvisor.dev/gvisor/pkg/state"
 	"gvisor.dev/gvisor/pkg/state/statefile"
+	"gvisor.dev/gvisor/pkg/state/wire"
 	"gvisor.dev/gvisor/pkg/sync"
 )
 
@@ -171,6 +172,7 @@ func (f *MemoryFile) SaveTo(ctx context.Context, w io.Writer, pw io.Writer, opts
 	log.Debugf("MemoryFile.SaveTo: saved metadata in %s", time.Since(timeMetadataStart))
 
 	// Dump out committed pages.
+	ww := wire.Writer{Writer: w}
 	timePagesStart := time.Now()
 	savedBytes := uint64(0)
 	for maseg := f.memAcct.FirstSegment(); maseg.Ok(); maseg = maseg.NextSegment() {
@@ -178,7 +180,7 @@ func (f *MemoryFile) SaveTo(ctx context.Context, w io.Writer, pw io.Writer, opts
 			continue
 		}
 		// Write a header to distinguish from objects.
-		if err := state.WriteHeader(w, uint64(maseg.Range().Length()), false); err != nil {
+		if err := state.WriteHeader(&ww, uint64(maseg.Range().Length()), false); err != nil {
 			return err
 		}
 		// Write out data.
@@ -304,6 +306,7 @@ func (f *MemoryFile) LoadFrom(ctx context.Context, r io.Reader, pr *statefile.As
 	defer madviseWG.Wait()
 
 	// Load committed pages.
+	wr := wire.Reader{Reader: r}
 	timePagesStart := time.Now()
 	loadedBytes := uint64(0)
 	for maseg := f.memAcct.FirstSegment(); maseg.Ok(); maseg = maseg.NextSegment() {
@@ -311,7 +314,7 @@ func (f *MemoryFile) LoadFrom(ctx context.Context, r io.Reader, pr *statefile.As
 			continue
 		}
 		// Verify header.
-		length, object, err := state.ReadHeader(r)
+		length, object, err := state.ReadHeader(&wr)
 		if err != nil {
 			return err
 		}

pkg/state/decode.go

@@ -18,7 +18,6 @@ import (
 	"bytes"
 	"context"
 	"fmt"
-	"io"
 	"math"
 	"reflect"
 
@@ -143,7 +142,7 @@ type decodeState struct {
 	ctx context.Context
 
 	// r is the input stream.
-	r io.Reader
+	r wire.Reader
 
 	// types is the type database.
 	types typeDecodeDatabase
@@ -591,7 +590,7 @@ func (ds *decodeState) Load(obj reflect.Value) {
 	ds.pending.PushBack(rootOds)
 
 	// Read the number of objects.
-	numObjects, object, err := ReadHeader(ds.r)
+	numObjects, object, err := ReadHeader(&ds.r)
 	if err != nil {
 		Failf("header error: %w", err)
 	}
@@ -613,7 +612,7 @@ func (ds *decodeState) Load(obj reflect.Value) {
 		// decoding loop in state/pretty/pretty.printer.printStream().
 		for i := uint64(0); i < numObjects; {
 			// Unmarshal either a type object or object ID.
-			encoded = wire.Load(ds.r)
+			encoded = wire.Load(&ds.r)
 			switch we := encoded.(type) {
 			case *wire.Type:
 				ds.types.Register(we)
@@ -624,7 +623,7 @@ func (ds *decodeState) Load(obj reflect.Value) {
 				id = objectID(we)
 				i++
 				// Unmarshal and resolve the actual object.
-				encoded = wire.Load(ds.r)
+				encoded = wire.Load(&ds.r)
 				ods = ds.lookup(id)
 				if ods != nil {
 					// Decode the object.
@@ -718,7 +717,7 @@ func (ds *decodeState) Load(obj reflect.Value) {
 // Each object written to the statefile is prefixed with a header. See
 // WriteHeader for more information; these functions are exported to allow
 // non-state writes to the file to play nice with debugging tools.
-func ReadHeader(r io.Reader) (length uint64, object bool, err error) {
+func ReadHeader(r *wire.Reader) (length uint64, object bool, err error) {
 	// Read the header.
 	err = safely(func() {
 		length = wire.LoadUint(r)

pkg/state/encode.go

@@ -16,7 +16,6 @@ package state
 
 import (
 	"context"
-	"io"
 	"reflect"
 	"sort"
 
@@ -62,7 +61,7 @@ type encodeState struct {
 	ctx context.Context
 
 	// w is the output stream.
-	w io.Writer
+	w wire.Writer
 
 	// types is the type database.
 	types typeEncodeDatabase
@@ -781,7 +780,7 @@ func (es *encodeState) Save(obj reflect.Value) {
 	}
 
 	// Write the header with the number of objects.
-	if err := WriteHeader(es.w, uint64(len(es.pending)), true); err != nil {
+	if err := WriteHeader(&es.w, uint64(len(es.pending)), true); err != nil {
 		Failf("error writing header: %w", err)
 	}
 
@@ -791,7 +790,7 @@ func (es *encodeState) Save(obj reflect.Value) {
 	if err := safely(func() {
 		for _, wt := range es.pendingTypes {
 			// Encode the type.
-			wire.Save(es.w, &wt)
+			wire.Save(&es.w, &wt)
 		}
 		// Emit objects in ID order.
 		ids := make([]objectID, 0, len(es.pending))
@@ -803,10 +802,10 @@ func (es *encodeState) Save(obj reflect.Value) {
 		})
 		for _, id := range ids {
 			// Encode the id.
-			wire.Save(es.w, wire.Uint(id))
+			wire.Save(&es.w, wire.Uint(id))
 			// Marshal the object.
 			oes := es.pending[id]
-			wire.Save(es.w, oes.encoded)
+			wire.Save(&es.w, oes.encoded)
 		}
 	}); err != nil {
 		// Include the object and the error.
@@ -825,7 +824,7 @@ const objectFlag uint64 = 1 << 63
 // order to generate statefiles that play nicely with debugging tools, raw
 // writes should be prefixed with a header with object set to false and the
 // appropriate length. This will allow tools to skip these regions.
-func WriteHeader(w io.Writer, length uint64, object bool) error {
+func WriteHeader(w *wire.Writer, length uint64, object bool) error {
 	// Sanity check the length.
 	if length&objectFlag != 0 {
 		Failf("impossibly huge length: %d", length)

pkg/state/pretty/pretty.go

@@ -196,6 +196,8 @@ func (p *printer) format(graph uint64, depth int, encoded wire.Object) (string,
 
 // printStream is the basic print implementation.
 func (p *printer) printStream(w io.Writer, r io.Reader) (err error) {
+	wr := wire.Reader{Reader: r}
+
 	// current graph ID.
 	var graph uint64
 
@@ -218,7 +220,7 @@ func (p *printer) printStream(w io.Writer, r io.Reader) (err error) {
 
 	for {
 		// Find the first object to begin generation.
-		length, object, err := state.ReadHeader(r)
+		length, object, err := state.ReadHeader(&wr)
 		if err == io.EOF {
 			// Nothing else to do.
 			break
@@ -252,7 +254,7 @@ func (p *printer) printStream(w io.Writer, r io.Reader) (err error) {
 		)
 		for i := uint64(0); i < length; {
 			// Unmarshal either a type object or object ID.
-			encoded := wire.Load(r)
+			encoded := wire.Load(&wr)
 			switch we := encoded.(type) {
 			case *wire.Type:
 				str, _ := p.format(graph, 0, encoded)
@@ -270,7 +272,7 @@ func (p *printer) printStream(w io.Writer, r io.Reader) (err error) {
 				// Unmarshal the actual object.
 				objects = append(objects, objectAndID{
 					id:  uint64(we),
-					obj: wire.Load(r),
+					obj: wire.Load(&wr),
 				})
 				i++
 			default:

pkg/state/state.go

@@ -92,7 +92,7 @@ func Save(ctx context.Context, w io.Writer, rootPtr any) (Stats, error) {
 	// Create the encoding state.
 	es := encodeState{
 		ctx:            ctx,
-		w:              w,
+		w:              wire.Writer{Writer: w},
 		types:          makeTypeEncodeDatabase(),
 		zeroValues:     make(map[reflect.Type]*objectEncodeState),
 		pending:        make(map[objectID]*objectEncodeState),
@@ -111,7 +111,7 @@ func Load(ctx context.Context, r io.Reader, rootPtr any) (Stats, error) {
 	// Create the decoding state.
 	ds := decodeState{
 		ctx:      ctx,
-		r:        r,
+		r:        wire.Reader{Reader: r},
 		types:    makeTypeDecodeDatabase(),
 		deferred: make(map[objectID]wire.Object),
 	}

pkg/state/wire/BUILD

@@ -13,6 +13,5 @@ go_library(
     visibility = ["//:sandbox"],
     deps = [
         "//pkg/gohacks",
-        "//pkg/sync",
     ],
 )