rendezvous: stream registry snapshot to disk (drop full-size save buffer)

server_persist.go

@@ -9,51 +9,56 @@
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
+	"io"
 	"log/slog"
 	"os"
 	"path/filepath"
 	"runtime/debug"
-	"sync"
 	"sync/atomic"
 	"time"
 
-	"github.com/pilot-protocol/common/fsutil"
 	"github.com/pilot-protocol/common/registry/wire"
 	"github.com/pilot-protocol/common/urlvalidate"
 	dashpkg "github.com/pilot-protocol/rendezvous/dashboard"
 	trustpkg "github.com/pilot-protocol/rendezvous/trust"
 )
 
 const (
-	maxPooledSaveBuf    = 128 << 20
 	scavengeIntervalMs  = 180_000
 	scavengeMinSnapshot = 64 << 20
 )
 
 var lastScavengeMs atomic.Int64
 
-// flushSaveBufPool reuses the bytes buffer that backs the snapshot JSON
-// across save ticks. After the first few saves the pool returns a
-// buffer at peak capacity, so subsequent saves do zero allocation in
-// the encode path. Eliminates the ~1 GB live `bytes.growSlice` heap
-// that was driving GC STW pauses in fleet-scale production.
-//
-// Starts at 1 MB; bytes.Buffer grows organically up to whatever the
-// real snapshot needs. Production fleet (~50-100 MB JSON) hits steady
-// state within a few saves; CI runners with empty registries stay at
-// 1 MB. The 128 MB pre-grow that lived here previously caused CI OOM
-// flakes under 4-way parallel integration tests.
-var flushSaveBufPool = sync.Pool{
-	New: func() interface{} {
-		b := make([]byte, 0, 1*1024*1024)
-		return &b
-	},
+// trimTrailingNewlineWriter forwards every byte it is given except the final
+// one. json.Encoder.Encode terminates its output with '\n', but the snapshot
+// checksum is defined over the encoded value WITHOUT that newline (load()
+// strips it before recomputing too). Holding back the last byte yields a file
+// and hash of `{...}` while writing straight to disk — no full-size buffer.
+type trimTrailingNewlineWriter struct {
+	w       io.Writer
+	held    byte
+	hasHeld bool
 }
 
-func putSaveBuf(bp *[]byte) {
-	if cap(*bp) <= maxPooledSaveBuf {
-		flushSaveBufPool.Put(bp)
+func (t *trimTrailingNewlineWriter) Write(p []byte) (int, error) {
+	if len(p) == 0 {
+		return 0, nil
 	}
+	if t.hasHeld {
+		if _, err := t.w.Write([]byte{t.held}); err != nil {
+			return 0, err
+		}
+		t.hasHeld = false
+	}
+	if len(p) > 1 {
+		if _, err := t.w.Write(p[:len(p)-1]); err != nil {
+			return 0, err
+		}
+	}
+	t.held = p[len(p)-1]
+	t.hasHeld = true
+	return len(p), nil
 }
 
 func shouldScavenge(dataLen int, nowMs int64) bool {
@@ -64,6 +69,42 @@
 	return nowMs-last >= scavengeIntervalMs && lastScavengeMs.CompareAndSwap(last, nowMs)
 }
 
+// atomicWriteStream writes path atomically via a temp file + rename, but
+// streams the content: fn receives the open *os.File and may write
+// incrementally (and Seek/WriteAt to backfill a trailer) instead of the
+// whole payload being materialized in a buffer first. Durability protocol
+// matches AtomicWrite: fsync the file, close, rename, best-effort dir fsync.
+// On fn error the temp file is removed and path is left untouched.
+func atomicWriteStream(path string, fn func(f *os.File) error) error {
+	tmp := path + ".tmp"
+	f, err := os.OpenFile(tmp, os.O_CREATE|os.O_RDWR|os.O_TRUNC, 0600)
+	if err != nil {
+		return err
+	}
+	if err := fn(f); err != nil {
+		f.Close()
+		os.Remove(tmp)
+		return err
+	}
+	if err := f.Sync(); err != nil {
+		f.Close()
+		os.Remove(tmp)
+		return err
+	}
+	if err := f.Close(); err != nil {
+		os.Remove(tmp)
+		return err
+	}
+	if err := os.Rename(tmp, path); err != nil {
+		return err
+	}
+	if dir, err := os.Open(filepath.Dir(path)); err == nil {
+		_ = dir.Sync()
+		_ = dir.Close()
+	}
+	return nil
+}
+
 // rawNodeCopy holds raw node fields copied under RLock (no encoding).
 // base64/time.Format happens outside the lock to minimize lock hold time.
 type rawNodeCopy struct {
@@ -510,56 +551,58 @@
 	}
 	s.auditMu.Unlock()
 
-	// Compute checksum: encode once without checksum (omitempty omits it), hash,
-	// then inject the checksum into the JSON without a second encode.
+	// Stream-encode straight to the temp file: no full-size in-memory buffer
+	// and no second copy through AtomicWrite. The body is hashed as it is
+	// written, then the checksum is backfilled as a trailing JSON member.
+	//
+	// At 170k+ nodes the old path materialized a ~50-100 MB pooled buffer and
+	// was the dominant heap allocator (~1 GB live, GC STW pauses causing UDP
+	// drops). Streaming writes the same bytes directly to disk instead.
 	//
-	// 2026-05-14: switched from json.Marshal (one-shot allocate) to json.Encoder
-	// writing into a pooled bytes.Buffer. At 170k+ active nodes the snapshot is
-	// ~50-100 MB and was the dominant heap allocator (~1 GB live, GC pressure
-	// causing kernel UDP drops during STW pauses). The pooled buffer is reused
-	// across save ticks — the underlying slice grows once and stays at peak,
-	// no per-tick allocation thereafter.
+	// Checksum invariant: load() re-encodes the struct with Checksum omitted
+	// and hashes those bytes (newline stripped). So the hash here must cover
+	// exactly the encoder output minus its trailing '\n' — which
+	// trimTrailingNewlineWriter guarantees by withholding the final byte.
 	snap.Checksum = ""
-	bp := flushSaveBufPool.Get().(*[]byte)
-	buf := bytes.NewBuffer((*bp)[:0])
-	enc := json.NewEncoder(buf)
-	enc.SetEscapeHTML(false)
-	if err := enc.Encode(snap); err != nil {
-		*bp = buf.Bytes()[:0]
-		putSaveBuf(bp)
-		slog.Error("registry save encode error", "err", err)
-		return fmt.Errorf("encode snapshot: %w", err)
-	}
-	data := buf.Bytes()
-	// json.Encoder.Encode appends a newline; drop it to match prior Marshal output.
-	if len(data) > 0 && data[len(data)-1] == '\n' {
-		data = data[:len(data)-1]
-	}
-	hash := sha256.Sum256(data)
-	checksum := hex.EncodeToString(hash[:])
-	// Insert "checksum":"<hex>" before the closing brace. json.Encoder of a struct
-	// always produces a JSON object ending with '}' (after newline trim).
-	if len(data) == 0 || data[len(data)-1] != '}' {
-		*bp = buf.Bytes()[:0]
-		putSaveBuf(bp)
-		return fmt.Errorf("encode snapshot: unexpected JSON format (expected trailing '}')")
-	}
-	data = append(data[:len(data)-1], []byte(`,"checksum":"`+checksum+`"}`)...)
-	defer func() {
-		// Return the (possibly grown) underlying buffer to the pool. AtomicWrite
-		// has copied data to disk by this point, so it is safe to release.
-		*bp = data[:0]
-		putSaveBuf(bp)
-	}()
-
-	// Persist to disk atomically
+	var written int64
 	if s.storePath != "" {
-		if err := fsutil.AtomicWrite(s.storePath, data); err != nil {
+		err := atomicWriteStream(s.storePath, func(f *os.File) error {
+			h := sha256.New()
+			enc := json.NewEncoder(&trimTrailingNewlineWriter{w: io.MultiWriter(f, h)})
+			enc.SetEscapeHTML(false)
+			if err := enc.Encode(snap); err != nil {
+				return fmt.Errorf("encode snapshot: %w", err)
+			}
+			// The encoded object ends with '}' (the trailing newline was
+			// withheld). Replace that '}' with the checksum trailer so the
+			// on-disk file is `{...` + `,"checksum":"<hex>"}`.
+			end, err := f.Seek(0, io.SeekEnd)
+			if err != nil {
+				return fmt.Errorf("seek snapshot temp: %w", err)
+			}
+			var last [1]byte
+			if end == 0 {
+				return fmt.Errorf("encode snapshot: empty output")
+			}
+			if _, err := f.ReadAt(last[:], end-1); err != nil {
+				return fmt.Errorf("read snapshot tail: %w", err)
+			}
+			if last[0] != '}' {
+				return fmt.Errorf("encode snapshot: unexpected JSON format (expected trailing '}')")
+			}
+			trailer := `,"checksum":"` + hex.EncodeToString(h.Sum(nil)) + `"}`
+			if _, err := f.WriteAt([]byte(trailer), end-1); err != nil {
+				return fmt.Errorf("write checksum trailer: %w", err)
+			}
+			written = end - 1 + int64(len(trailer))
+			return nil
+		})
+		if err != nil {
 			slog.Error("registry save error", "err", err)
 			return fmt.Errorf("write snapshot: %w", err)
 		}
-		s.lastSnapshotSizeB.Store(int64(len(data)))
-		if shouldScavenge(len(data), time.Now().UnixMilli()) {
+		s.lastSnapshotSizeB.Store(written)
+		if shouldScavenge(int(written), time.Now().UnixMilli()) {
 			go debug.FreeOSMemory()
 		}
 		// Truncate WAL after successful snapshot (compaction).

zz_save_memory_test.go

@@ -2,23 +2,50 @@
 
 package server
 
-import "testing"
+import (
+	"bytes"
+	"testing"
+)
 
-func TestPutSaveBufDropsOversized(t *testing.T) {
-	big := make([]byte, 0, maxPooledSaveBuf+1)
-	putSaveBuf(&big)
-	got := flushSaveBufPool.Get().(*[]byte)
-	if cap(*got) > maxPooledSaveBuf {
-		t.Fatalf("pool retained an oversized buffer: cap=%d", cap(*got))
+// The snapshot encode path streams JSON straight to disk (no pooled buffer).
+// trimTrailingNewlineWriter withholds the encoder's trailing '\n' so the
+// checksum covers `{...}` exactly as load() recomputes it.
+
+func TestTrimTrailingNewlineWriterDropsFinalByte(t *testing.T) {
+	t.Parallel()
+	var buf bytes.Buffer
+	w := &trimTrailingNewlineWriter{w: &buf}
+	if _, err := w.Write([]byte(`{"a":1}` + "\n")); err != nil {
+		t.Fatalf("write: %v", err)
+	}
+	if got, want := buf.String(), `{"a":1}`; got != want {
+		t.Fatalf("body = %q, want %q (trailing newline must be withheld)", got, want)
 	}
 }
 
-func TestPutSaveBufKeepsNormal(t *testing.T) {
-	buf := make([]byte, 0, 1<<20)
-	putSaveBuf(&buf)
-	got := flushSaveBufPool.Get().(*[]byte)
-	if cap(*got) < 1<<20 {
-		t.Fatalf("normal-sized buffer was not pooled: cap=%d", cap(*got))
+func TestTrimTrailingNewlineWriterAcrossChunks(t *testing.T) {
+	t.Parallel()
+	var buf bytes.Buffer
+	w := &trimTrailingNewlineWriter{w: &buf}
+	for _, chunk := range []string{`{"a"`, `:1}`, "\n"} {
+		if _, err := w.Write([]byte(chunk)); err != nil {
+			t.Fatalf("write %q: %v", chunk, err)
+		}
+	}
+	if got, want := buf.String(), `{"a":1}`; got != want {
+		t.Fatalf("body = %q, want %q (streaming across chunks)", got, want)
+	}
+}
+
+func TestTrimTrailingNewlineWriterEmptyWrite(t *testing.T) {
+	t.Parallel()
+	var buf bytes.Buffer
+	w := &trimTrailingNewlineWriter{w: &buf}
+	if n, err := w.Write(nil); n != 0 || err != nil {
+		t.Fatalf("empty write = (%d, %v), want (0, nil)", n, err)
+	}
+	if buf.Len() != 0 {
+		t.Fatalf("buffer = %q, want empty", buf.String())
 	}
 }