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 <mikeo@tailscale.com>
2026-03-24 18:08:01 -04:00
parent 9992b7c817
commit f52c1e3615
3 changed files with 119 additions and 33 deletions
@@ -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