types/key: use AvailableBuffer for WriteRawWithoutAllocating (#19102)

Use bufio.Writer.AvailableBuffer to write the 32-byte public key
directly into bufio's internal buffer as a single append+Write,
avoiding 32 separate WriteByte calls. Fall back to the existing
byte-at-a-time path when the buffer has insufficient space.

```
name                                old ns/op  new ns/op  speedup
NodeWriteRawWithoutAllocating-8     121        12.5       ~9.7x
(0 allocs/op in both)
```

Add BenchmarkNodeWriteRawWithoutAllocating and expand
TestNodeWriteRawWithoutAllocating to cover both fast (AvailableBuffer)
and slow (WriteByte fallback) paths with correctness and allocation
checks.

Updates tailscale/corp#38509

Signed-off-by: Mike O'Driscoll <mikeo@tailscale.com>

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
 		}
 	}

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)
+		}
 	}
 }