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>

derp/client_test.go

@@ -6,7 +6,6 @@ package derp
 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

derp/derp.go

@@ -168,6 +168,8 @@ const FastStartHeader = "Derp-Fast-Start"
 
 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
 	}

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