From f52c1e3615629dd1ac8a2d73a5d9887245d4a2da Mon Sep 17 00:00:00 2001 From: Mike O'Driscoll Date: Tue, 24 Mar 2026 18:08:01 -0400 Subject: [PATCH] derp: use AvailableBuffer for WriteFrameHeader, consolidate tests (#19101) Use bufio.Writer.AvailableBuffer to write the frame header directly into bufio's internal buffer as a single append+Write, avoiding 5 separate WriteByte calls. Fall back to the existing writeUint32 byte-at-a-time path when the buffer has insufficient space. ``` name old ns/op new ns/op speedup WriteFrameHeader-8 18.8 7.8 ~2.4x (0 allocs/op in both) ``` Add TestWriteFrameHeader with correctness checks, allocation assertions, and coverage of both fast and slow write paths. Move BenchmarkReadFrameHeader from client_test.go to derp_test.go alongside BenchmarkWriteFrameHeader, co-located with the functions under test. Updates tailscale/corp#38509 Signed-off-by: Mike O'Driscoll --- derp/client_test.go | 28 ------------- derp/derp.go | 24 ++++++++--- derp/derp_test.go | 100 ++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 119 insertions(+), 33 deletions(-) diff --git a/derp/client_test.go b/derp/client_test.go index b3bacdb85..697aa4961 100644 --- a/derp/client_test.go +++ b/derp/client_test.go @@ -6,7 +6,6 @@ import ( "bufio" "bytes" - "io" "net" "reflect" "sync" @@ -126,33 +125,6 @@ func TestClientSendPong(t *testing.T) { } } -func BenchmarkWriteUint32(b *testing.B) { - w := bufio.NewWriter(io.Discard) - b.ReportAllocs() - b.ResetTimer() - for range b.N { - writeUint32(w, 0x0ba3a) - } -} - -type nopRead struct{} - -func (r nopRead) Read(p []byte) (int, error) { - return len(p), nil -} - -func BenchmarkReadFrameHeader(b *testing.B) { - r := bufio.NewReader(nopRead{}) - b.ReportAllocs() - b.ResetTimer() - for range b.N { - _, _, err := ReadFrameHeader(r) - if err != nil { - b.Fatal(err) - } - } -} - type countWriter struct { mu sync.Mutex writes int diff --git a/derp/derp.go b/derp/derp.go index d4a09415a..46d1ab453 100644 --- a/derp/derp.go +++ b/derp/derp.go @@ -168,6 +168,8 @@ var bin = binary.BigEndian +// writeUint32 writes v to bw one byte at a time +// as a big-endian uint32. func writeUint32(bw *bufio.Writer, v uint32) error { var b [4]byte bin.PutUint32(b[:], v) @@ -243,14 +245,26 @@ func readFrame(br *bufio.Reader, maxSize uint32, b []byte) (t FrameType, frameLe return t, frameLen, err } -// WriteFrameHeader writes a frame header to bw. -// -// The frame header is 5 bytes: a one byte frame type -// followed by a big-endian uint32 length of the -// remaining frame (not including the 5 byte header). +// WriteFrameHeader writes a DERP frame header to bw: a one-byte frame +// type followed by a big-endian uint32 frame length. // +// It uses AvailableBuffer to append the header directly into bufio's +// internal buffer without allocation, falling back to WriteByte when +// the buffer has insufficient space. // It does not flush bw. func WriteFrameHeader(bw *bufio.Writer, t FrameType, frameLen uint32) error { + // Fast path: enough space in the buffer to append the header + // directly without allocation via AvailableBuffer. + if bw.Available() >= FrameHeaderLen { + buf := bw.AvailableBuffer() + buf = append(buf, byte(t)) + buf = bin.AppendUint32(buf, frameLen) + _, 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. if err := bw.WriteByte(byte(t)); err != nil { return err } diff --git a/derp/derp_test.go b/derp/derp_test.go index 0df49969e..a5b5aa4b8 100644 --- a/derp/derp_test.go +++ b/derp/derp_test.go @@ -93,6 +93,106 @@ func TestReadFrameHeader(t *testing.T) { } } +func TestWriteFrameHeader(t *testing.T) { + tests := []struct { + name string + typ derp.FrameType + frameLen uint32 + want [derp.FrameHeaderLen]byte + }{ + { + name: "SendPacket", + typ: derp.FrameSendPacket, + frameLen: 1024, + want: [derp.FrameHeaderLen]byte{byte(derp.FrameSendPacket), 0x00, 0x00, 0x04, 0x00}, + }, + { + name: "KeepAlive", + typ: derp.FrameKeepAlive, + frameLen: 0, + want: [derp.FrameHeaderLen]byte{byte(derp.FrameKeepAlive), 0x00, 0x00, 0x00, 0x00}, + }, + { + name: "MaxLen", + typ: derp.FrameRecvPacket, + frameLen: 0xffffffff, + want: [derp.FrameHeaderLen]byte{byte(derp.FrameRecvPacket), 0xff, 0xff, 0xff, 0xff}, + }, + } + for _, tt := range tests { + // Test fast path (empty buffer, plenty of space). + t.Run(tt.name+"/fast", func(t *testing.T) { + var buf bytes.Buffer + bw := bufio.NewWriter(&buf) + if err := derp.WriteFrameHeader(bw, tt.typ, tt.frameLen); err != nil { + t.Fatalf("WriteFrameHeader: %v", err) + } + bw.Flush() + if got := buf.Bytes(); !bytes.Equal(got, tt.want[:]) { + t.Errorf("wrote % 02x, want % 02x", got, tt.want) + } + }) + + // Test slow path (buffer nearly full, less than FrameHeaderLen available). + t.Run(tt.name+"/slow", func(t *testing.T) { + var buf bytes.Buffer + const smallBuf = 8 // small enough to force slow path + bw := bufio.NewWriterSize(&buf, smallBuf) + // Fill buffer to leave less than FrameHeaderLen bytes available. + padding := make([]byte, smallBuf-derp.FrameHeaderLen+1) + if _, err := bw.Write(padding); err != nil { + t.Fatalf("Write padding: %v", err) + } + if err := derp.WriteFrameHeader(bw, tt.typ, tt.frameLen); err != nil { + t.Fatalf("WriteFrameHeader: %v", err) + } + bw.Flush() + got := buf.Bytes() + // The header is after the padding bytes. + got = got[len(padding):] + if !bytes.Equal(got, tt.want[:]) { + t.Errorf("wrote % 02x, want % 02x", got, tt.want) + } + }) + } + + // Verify zero allocations on fast path. + bw := bufio.NewWriter(io.Discard) + got := testing.AllocsPerRun(1000, func() { + if err := derp.WriteFrameHeader(bw, derp.FrameSendPacket, 1024); err != nil { + t.Fatalf("WriteFrameHeader: %v", err) + } + }) + if got != 0 { + t.Fatalf("WriteFrameHeader allocs = %f, want 0", got) + } +} + +type nopRead struct{} + +func (nopRead) Read(p []byte) (int, error) { return len(p), nil } + +func BenchmarkReadFrameHeader(b *testing.B) { + r := bufio.NewReader(nopRead{}) + b.ReportAllocs() + for b.Loop() { + _, _, err := derp.ReadFrameHeader(r) + if err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkWriteFrameHeader(b *testing.B) { + bw := bufio.NewWriter(io.Discard) + b.ReportAllocs() + for b.Loop() { + if err := derp.WriteFrameHeader(bw, derp.FrameSendPacket, 1024); err != nil { + b.Fatal(err) + } + } +} + func TestClientInfoUnmarshal(t *testing.T) { for i, in := range map[string]struct { json string