From 905d769f6ffbb37aef439b71fdb90fe7271ba762 Mon Sep 17 00:00:00 2001 From: Ayush Ranjan Date: Mon, 9 Sep 2024 14:41:22 -0700 Subject: [PATCH] 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 | 10 ++++----- pkg/state/tests/bench_test.go | 2 +- pkg/state/tests/tests.go | 12 ++--------- pkg/state/wire/BUILD | 8 +++++++- pkg/state/wire/wire.go | 22 ++++++++------------ pkg/state/wire/wire_test.go | 32 +++++++++++++++++++++++++++++ 6 files changed, 56 insertions(+), 30 deletions(-) create mode 100644 pkg/state/wire/wire_test.go diff --git a/pkg/compressio/nocompressio_test.go b/pkg/compressio/nocompressio_test.go index 0e445fbbb..25c2ba322 100644 --- a/pkg/compressio/nocompressio_test.go +++ b/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) diff --git a/pkg/state/tests/bench_test.go b/pkg/state/tests/bench_test.go index 1fdff4dbc..2a94721d4 100644 --- a/pkg/state/tests/bench_test.go +++ b/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) } } diff --git a/pkg/state/tests/tests.go b/pkg/state/tests/tests.go index c1d5e3849..d09f10ffa 100644 --- a/pkg/state/tests/tests.go +++ b/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) diff --git a/pkg/state/wire/BUILD b/pkg/state/wire/BUILD index 2d522242e..997a714ad 100644 --- a/pkg/state/wire/BUILD +++ b/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", +) diff --git a/pkg/state/wire/wire.go b/pkg/state/wire/wire.go index 657f3dbb5..f89067acd 100644 --- a/pkg/state/wire/wire.go +++ b/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. diff --git a/pkg/state/wire/wire_test.go b/pkg/state/wire/wire_test.go new file mode 100644 index 000000000..eb5089e9d --- /dev/null +++ b/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) + } +}