diff --git a/types/key/node.go b/types/key/node.go index 1b0232b11..98f72c719 100644 --- a/types/key/node.go +++ b/types/key/node.go @@ -253,18 +253,24 @@ func (k *NodePublic) ReadRawWithoutAllocating(br *bufio.Reader) error { return err } -// WriteRawWithoutAllocating writes out k as 32 bytes to bw. -// The writing is done ~3x slower than bw.Write, but in exchange is -// allocation-free. +// WriteRawWithoutAllocating writes out k as 32 big-endian bytes to bw. +// +// It uses AvailableBuffer to append directly into bufio's internal +// buffer without allocation, falling back to WriteByte when the +// buffer has insufficient space. func (k NodePublic) WriteRawWithoutAllocating(bw *bufio.Writer) error { - // Equivalent to bw.Write(k.k[:]), but without causing an - // escape-related alloc. - // - // Dear future: if bw.Write(k.k[:]) stops causing stuff to escape, - // you should switch back to that. + // Fast path: enough space in the buffer to append directly. + if bw.Available() >= len(k.k) { + buf := bw.AvailableBuffer() + buf = append(buf, k.k[:]...) + _, err := bw.Write(buf) + return err + } + // Slow path: buffer nearly full. Write byte-at-a-time to let + // bufio flush as needed, avoiding a heap allocation from append + // growing past AvailableBuffer's capacity. for _, b := range k.k { - err := bw.WriteByte(b) - if err != nil { + if err := bw.WriteByte(b); err != nil { return err } } diff --git a/types/key/node_test.go b/types/key/node_test.go index 0bcda0465..020ddd1f1 100644 --- a/types/key/node_test.go +++ b/types/key/node_test.go @@ -7,6 +7,7 @@ import ( "bufio" "bytes" "encoding/json" + "io" "strings" "testing" ) @@ -133,8 +134,7 @@ func BenchmarkNodeReadRawWithoutAllocating(b *testing.B) { r := bytes.NewReader(buf) br := bufio.NewReader(r) b.ReportAllocs() - b.ResetTimer() - for range b.N { + for b.Loop() { r.Reset(buf) br.Reset(r) var k NodePublic @@ -145,19 +145,72 @@ func BenchmarkNodeReadRawWithoutAllocating(b *testing.B) { } func TestNodeWriteRawWithoutAllocating(t *testing.T) { - buf := make([]byte, 0, 32) - w := bytes.NewBuffer(buf) - bw := bufio.NewWriter(w) - got := testing.AllocsPerRun(1000, func() { - w.Reset() - bw.Reset(w) - var k NodePublic + var k NodePublic + for i := range k.k { + k.k[i] = byte(i) + } + + // Test fast path (empty buffer, plenty of space). + t.Run("fast", func(t *testing.T) { + var buf bytes.Buffer + bw := bufio.NewWriter(&buf) if err := k.WriteRawWithoutAllocating(bw); err != nil { t.Fatalf("WriteRawWithoutAllocating: %v", err) } + bw.Flush() + if got := buf.Bytes(); !bytes.Equal(got, k.k[:]) { + t.Errorf("wrote % 02x, want % 02x", got, k.k) + } }) - if want := 0.0; got != want { - t.Fatalf("WriteRawWithoutAllocating got %f allocs, want %f", got, want) + + // Test slow path (buffer nearly full, less than 32 bytes available). + t.Run("slow", func(t *testing.T) { + var buf bytes.Buffer + const smallBuf = 40 + bw := bufio.NewWriterSize(&buf, smallBuf) + // Fill buffer to leave less than 32 bytes available. + padding := make([]byte, smallBuf-len(k.k)+1) + if _, err := bw.Write(padding); err != nil { + t.Fatalf("Write padding: %v", err) + } + if err := k.WriteRawWithoutAllocating(bw); err != nil { + t.Fatalf("WriteRawWithoutAllocating: %v", err) + } + bw.Flush() + got := buf.Bytes()[len(padding):] + if !bytes.Equal(got, k.k[:]) { + t.Errorf("wrote % 02x, want % 02x", got, k.k) + } + }) + + // Verify zero allocations on fast path. + t.Run("allocs", func(t *testing.T) { + w := bytes.NewBuffer(make([]byte, 0, 32)) + bw := bufio.NewWriter(w) + got := testing.AllocsPerRun(1000, func() { + w.Reset() + bw.Reset(w) + if err := k.WriteRawWithoutAllocating(bw); err != nil { + t.Fatalf("WriteRawWithoutAllocating: %v", err) + } + }) + if got != 0 { + t.Fatalf("WriteRawWithoutAllocating allocs = %f, want 0", got) + } + }) +} + +func BenchmarkNodeWriteRawWithoutAllocating(b *testing.B) { + bw := bufio.NewWriter(io.Discard) + var k NodePublic + for i := range k.k { + k.k[i] = 0x42 + } + b.ReportAllocs() + for b.Loop() { + if err := k.WriteRawWithoutAllocating(bw); err != nil { + b.Fatal(err) + } } }