fix: prevent Go runtime memory corruption and sandbox halt after snapshot restore

Three root causes addressed:

1. Go page allocator corruption: allocations between the pre-snapshot GC
   and VM freeze leave the summary tree inconsistent. After restore, GC
   reads corrupted metadata — either panicking (killing PID 1 → kernel
   panic) or silently failing to collect, causing unbounded heap growth
   until OOM. Fix: move GC to after all HTTP allocations in
   PostSnapshotPrepare, then set GOMAXPROCS(1) so any remaining
   allocations run sequentially with no concurrent page allocator access.
   GOMAXPROCS is restored on first health check after restore.

2. PostInit timeout starvation: WaitUntilReady and PostInit shared a
   single 30s context. If WaitUntilReady consumed most of it, PostInit
   failed — RestoreAfterSnapshot never ran, leaving envd with keep-alives
   disabled and zombie connections. Fix: separate timeout contexts.

3. CP HTTP server missing timeouts: no ReadHeaderTimeout or IdleTimeout
   caused goroutine leaks from hung proxy connections. Fix: add both,
   matching host agent values.

Also adds UFFD prefetch to proactively load all guest pages after restore,
eliminating on-demand page fault latency for subsequent RPC calls.

envd/VERSION

@@ -1 +1 @@
-0.1.1
+0.1.2

envd/internal/api/conntracker.go

@@ -51,7 +51,7 @@ func (t *ServerConnTracker) Track(conn net.Conn, state http.ConnState) {
 // (with keep-alives disabled, the connection closes), RestoreAfterSnapshot
 // will close any that survived into the snapshot as zombie TCP sockets.
 //
-// GC cycles are handled by PortSubsystem.Stop() which runs before this.
+// GC is handled by PostSnapshotPrepare after this returns.
 func (t *ServerConnTracker) PrepareForSnapshot() {
 	t.mu.Lock()
 	defer t.mu.Unlock()

envd/internal/api/init.go

@@ -150,15 +150,17 @@ func (a *API) PostInit(w http.ResponseWriter, r *http.Request) {
 		host.PollForMMDSOpts(ctx, a.mmdsChan, a.defaults.EnvVars)
 	}()
 
-	// Close zombie connections from before the snapshot and re-enable
-	// keep-alives. On first boot this is a no-op (no zombie connections).
+	// Safety net: if the health check's postRestoreRecovery didn't run yet
+	// (e.g. PostInit arrived before the first health check), re-enable GC
+	// here. On first boot needsRestore is false so CAS is a no-op.
+	if a.needsRestore.CompareAndSwap(true, false) {
+		a.postRestoreRecovery()
+	}
+	// RestoreAfterSnapshot is idempotent (clears preSnapshot set), and
+	// Start is a no-op if already running.
 	if a.connTracker != nil {
 		a.connTracker.RestoreAfterSnapshot()
 	}
-
-	// Start the port scanner and forwarder if they were stopped by a
-	// pre-snapshot prepare call. Start is a no-op if already running,
-	// so this is safe on first boot and only takes effect after restore.
 	if a.portSubsystem != nil {
 		a.portSubsystem.Start(a.rootCtx)
 	}

envd/internal/api/snapshot.go

@@ -5,6 +5,8 @@ package api
 
 import (
 	"net/http"
+	"runtime"
+	"runtime/debug"
 )
 
 // PostSnapshotPrepare quiesces continuous goroutines (port scanner, forwarder),
@@ -13,6 +15,14 @@ import (
 // TCP state after snapshot restore. Keep-alives are disabled so the current
 // request's connection also closes after the response.
 //
+// To prevent Go page allocator corruption, GOMAXPROCS is set to 1 after the
+// final GC. With a single P, all goroutines (including any that allocate
+// between now and the VM freeze) run sequentially. This eliminates concurrent
+// page allocator access, so even if the freeze lands mid-allocation, the
+// in-flight operation completes atomically on restore before any GC reads
+// the summary tree. GOMAXPROCS is restored on the first health check after
+// restore (see postRestoreRecovery).
+//
 // Called by the host agent as a best-effort signal before vm.Pause().
 func (a *API) PostSnapshotPrepare(w http.ResponseWriter, r *http.Request) {
 	defer r.Body.Close()
@@ -27,6 +37,26 @@ func (a *API) PostSnapshotPrepare(w http.ResponseWriter, r *http.Request) {
 		a.logger.Info().Msg("snapshot/prepare: idle connections closed, keep-alives disabled")
 	}
 
+	// Send the response before the GC so HTTP buffer allocations happen
+	// while GOMAXPROCS is still at its normal value.
 	w.Header().Set("Cache-Control", "no-store")
 	w.WriteHeader(http.StatusNoContent)
+	if f, ok := w.(http.Flusher); ok {
+		f.Flush()
+	}
+
+	// Final GC pass after all major allocations (connection cleanup,
+	// response write) are complete.
+	runtime.GC()
+	runtime.GC()
+	debug.FreeOSMemory()
+
+	// Reduce to a single P so any post-GC allocations (HTTP server
+	// connection teardown) run sequentially — no concurrent page allocator
+	// access that could leave the summary tree inconsistent if the VM
+	// freezes mid-update.
+	a.prevGOMAXPROCS = runtime.GOMAXPROCS(1)
+
+	a.needsRestore.Store(true)
+	a.logger.Info().Msg("snapshot/prepare: GOMAXPROCS=1, ready for freeze")
 }

envd/internal/api/store.go

@@ -7,7 +7,10 @@ import (
 	"context"
 	"encoding/json"
 	"net/http"
+	"runtime"
+	"runtime/debug"
 	"sync"
+	"sync/atomic"
 
 	"github.com/rs/zerolog"
 
@@ -48,6 +51,12 @@ type API struct {
 	rootCtx       context.Context
 	portSubsystem *publicport.PortSubsystem
 	connTracker   *ServerConnTracker
+
+	// needsRestore is set by PostSnapshotPrepare and cleared on the first
+	// health check or PostInit after restore. While set, GOMAXPROCS is 1
+	// to prevent concurrent page allocator access during the freeze window.
+	needsRestore   atomic.Bool
+	prevGOMAXPROCS int // GOMAXPROCS value before PrepareSnapshot reduced it to 1
 }
 
 func New(l *zerolog.Logger, defaults *execcontext.Defaults, mmdsChan chan *host.MMDSOpts, isNotFC bool, rootCtx context.Context, portSubsystem *publicport.PortSubsystem, connTracker *ServerConnTracker, version string) *API {
@@ -69,6 +78,14 @@ func New(l *zerolog.Logger, defaults *execcontext.Defaults, mmdsChan chan *host.
 func (a *API) GetHealth(w http.ResponseWriter, r *http.Request) {
 	defer r.Body.Close()
 
+	// On the first health check after snapshot restore, re-enable GC and
+	// clean up stale state. By this point, any goroutine that was mid-
+	// allocation when the VM was frozen has completed, so the page allocator
+	// summary tree is consistent and safe for GC to read.
+	if a.needsRestore.CompareAndSwap(true, false) {
+		a.postRestoreRecovery()
+	}
+
 	a.logger.Trace().Msg("Health check")
 
 	w.Header().Set("Cache-Control", "no-store")
@@ -79,6 +96,35 @@ func (a *API) GetHealth(w http.ResponseWriter, r *http.Request) {
 	})
 }
 
+// postRestoreRecovery restores GOMAXPROCS, runs a clean GC cycle, closes
+// zombie TCP connections from before the snapshot, re-enables HTTP keep-alives,
+// and restarts the port subsystem. Called exactly once per restore cycle,
+// guarded by a CAS on needsRestore in both GetHealth and PostInit.
+func (a *API) postRestoreRecovery() {
+	// Restore parallelism first — any goroutine that was mid-allocation
+	// when the VM froze has already completed by the time a health check
+	// or PostInit request is being served, so the page allocator summary
+	// tree is consistent and safe for a full GC.
+	prev := a.prevGOMAXPROCS
+	if prev > 0 {
+		runtime.GOMAXPROCS(prev)
+	}
+	runtime.GC()
+	runtime.GC()
+	debug.FreeOSMemory()
+	a.logger.Info().Msg("restore: GOMAXPROCS restored, GC complete")
+
+	if a.connTracker != nil {
+		a.connTracker.RestoreAfterSnapshot()
+		a.logger.Info().Msg("restore: zombie connections closed, keep-alives re-enabled")
+	}
+
+	if a.portSubsystem != nil {
+		a.portSubsystem.Start(a.rootCtx)
+		a.logger.Info().Msg("restore: port subsystem restarted")
+	}
+}
+
 func (a *API) GetMetrics(w http.ResponseWriter, r *http.Request) {
 	defer r.Body.Close()
 

envd/internal/port/subsystem.go

@@ -5,8 +5,6 @@ package port
 
 import (
 	"context"
-	"runtime"
-	"runtime/debug"
 	"sync"
 	"time"
 
@@ -72,9 +70,12 @@ func (p *PortSubsystem) Start(parentCtx context.Context) {
 	}()
 }
 
-// Stop quiesces the scanner and forwarder goroutines and forces a GC cycle
-// to put the Go runtime's page allocator in a consistent state before snapshot.
+// Stop quiesces the scanner and forwarder goroutines.
 // Blocks until both goroutines have exited. Safe to call if already stopped.
+//
+// GC is NOT run here — it is deferred to PostSnapshotPrepare so that the
+// GC happens after all allocations (connection cleanup, HTTP response) are
+// complete, minimizing the window where page allocator corruption can occur.
 func (p *PortSubsystem) Stop() {
 	p.mu.Lock()
 	if !p.running {
@@ -90,12 +91,6 @@ func (p *PortSubsystem) Stop() {
 
 	cancelFn()
 	wg.Wait()
-
-	// Force two GC cycles to ensure all spans are swept and the page
-	// allocator summary tree is fully consistent before the VM is frozen.
-	runtime.GC()
-	runtime.GC()
-	debug.FreeOSMemory()
 }
 
 // Restart stops the subsystem (if running) and starts it again with a fresh