Optimize wire.Uint.save() to make only 1 Write() call.

Earlier wire.Uint.save() could make up to 10 Write() calls depending on how
large the Uint being marshaled was. Write() is an interface method call, so
this avoids the dynamic dispatch overhead. Furthermore, compressio's Write
implementations themselves do a bunch of fixed work per call and invoke more
interface functions.

Increased the scratch buffer in the wire.Writer to accommodate 10 bytes. The
wire.Uint can be marshaled into at most 10 bytes.

Before:

goos: linux
goarch: amd64
pkg: pkg/state/wire/wire
cpu: AMD EPYC 7B12
BenchmarkUintSave
BenchmarkUintSave-24    	37557860	        31.13 ns/op	       0 B/op	       0 allocs/op

After:

goos: linux
goarch: amd64
pkg: pkg/state/wire/wire
cpu: AMD EPYC 7B12
BenchmarkUintSave
BenchmarkUintSave-24    	129611625	         9.274 ns/op	       0 B/op	       0 allocs/op
PiperOrigin-RevId: 672679703

pkg/compressio/nocompressio_test.go

@@ -69,22 +69,22 @@ func TestNoCompress(t *testing.T) {
 
 // The following benchmarks aims to be representative of how the wire
 // package writes to the image. In practice, wire package only calls Write
-// with a single byte at a time. Because all types boil down to wire.Uint
-// whose implementation invokes Write with one byte at a time.
+// with very small buffers. Because all types boil down to wire.Uint whose
+// implementation invokes Write with at most 10 bytes at a time.
 
 func BenchmarkTinyIO(b *testing.B) {
 	// Use the same chunk size as the statefile package.
 	const blockSize = 1024 * 1024
 	for _, key := range [][]byte{nil, hashKey} {
 		b.Run(benchmarkName(false, true, key != nil, blockSize), func(b *testing.B) {
-			benchmarkNoCompress1ByteWrite(b, key, blockSize)
+			benchmarkNoCompress8ByteWrite(b, key, blockSize)
 		})
 	}
 }
 
-func benchmarkNoCompress1ByteWrite(b *testing.B, key []byte, blockSize uint32) {
+func benchmarkNoCompress8ByteWrite(b *testing.B, key []byte, blockSize uint32) {
 	var (
-		buf [1]byte
+		buf [8]byte
 		out bytes.Buffer
 	)
 	w := NewSimpleWriter(&out, key, blockSize)

pkg/state/tests/bench_test.go

@@ -117,7 +117,7 @@ func BenchmarkEncoding(b *testing.B) {
 				b.ReportAllocs()
 				b.StartTimer()
 				for i := 0; i < n; i++ {
-					if _, err := algoInfo.Save(context.Background(), discard{}, v); err != nil {
+					if _, err := algoInfo.Save(context.Background(), io.Discard, v); err != nil {
 						b.Errorf("save failed: %v", err)
 					}
 				}

pkg/state/tests/tests.go

@@ -19,6 +19,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"io"
 	"math"
 	"reflect"
 	"testing"
@@ -27,15 +28,6 @@ import (
 	"gvisor.dev/gvisor/pkg/state/pretty"
 )
 
-// discard is an implementation of wire.Writer.
-type discard struct{}
-
-// Write implements wire.Writer.Write.
-func (discard) Write(p []byte) (int, error) { return len(p), nil }
-
-// WriteByte implements wire.Writer.WriteByte.
-func (discard) WriteByte(byte) error { return nil }
-
 // checkEqual checks if two objects are equal.
 //
 // N.B. This only handles one level of dereferences for NaN. Otherwise we
@@ -103,7 +95,7 @@ func runTestCases(t *testing.T, shouldFail bool, prefix string, objects []any) {
 					t.Errorf("PrettyPrint(html=false) failed unexpected: %v", err)
 				}
 			}
-			if err := pretty.PrintHTML(discard{}, bytes.NewReader(saveBuffer.Bytes())); err != nil {
+			if err := pretty.PrintHTML(io.Discard, bytes.NewReader(saveBuffer.Bytes())); err != nil {
 				// See above.
 				if !shouldFail {
 					t.Errorf("PrettyPrint(html=true) failed unexpected: %v", err)

pkg/state/wire/BUILD

@@ -1,4 +1,4 @@
-load("//tools:defs.bzl", "go_library")
+load("//tools:defs.bzl", "go_library", "go_test")
 
 package(
     default_applicable_licenses = ["//:license"],
@@ -15,3 +15,9 @@ go_library(
         "//pkg/gohacks",
     ],
 )
+
+go_test(
+    name = "wire_test",
+    srcs = ["wire_test.go"],
+    library = ":wire",
+)

pkg/state/wire/wire.go

@@ -58,17 +58,8 @@ func (r *Reader) readByte() byte {
 type Writer struct {
 	io.Writer
 
-	buf [1]byte
-}
-
-// writeByte writes a single byte to w.Writer without allocation. It panics on
-// error.
-func (w *Writer) writeByte(b byte) {
-	w.buf[0] = b
-	n, err := w.Write(w.buf[:])
-	if n != 1 {
-		panic(err)
-	}
+	// buf is used by Uint as a scratch buffer.
+	buf [10]byte
 }
 
 // readFull is a utility. The equivalent is not needed for Write, but the API
@@ -173,11 +164,16 @@ func loadUint(r *Reader) Uint {
 
 // save implements Object.save.
 func (u Uint) save(w *Writer) {
+	i := 0
 	for u >= 0x80 {
-		w.writeByte(byte(u) | 0x80)
+		w.buf[i] = byte(u) | 0x80
+		i++
 		u >>= 7
 	}
-	w.writeByte(byte(u))
+	w.buf[i] = byte(u)
+	if _, err := w.Write(w.buf[:i+1]); err != nil {
+		panic(err)
+	}
 }
 
 // load implements Object.load.

pkg/state/wire/wire_test.go

@@ -0,0 +1,32 @@
+// Copyright 2024 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package wire
+
+import (
+	"bufio"
+	"io"
+	"testing"
+)
+
+// BenchmarkUintSave benchmarks saving a Uint. This benchmark is important
+// because almost all types in this package boil down to Uints. So
+// Uint.save() performance is critical for checkpoint performance.
+func BenchmarkUintSave(b *testing.B) {
+	w := Writer{Writer: bufio.NewWriter(io.Discard)}
+	n := Uint(0xdeadbeef)
+	for i := 0; i < b.N; i++ {
+		n.save(&w)
+	}
+}