fix(orchestrator): reject snapshot readiness after uffd exits

docs/ARCHITECTURE.md

@@ -323,7 +323,9 @@ Key mechanisms (all under `pkg/sandbox/`):
   `/metrics`.
 - **Lazy memory / UFFD** (`uffd/`): on resume, Firecracker restores the VM without loading
   memory; a userfaultfd handler serves page faults directly from the template's memfile, so only
-  touched pages are read. An optional prefetcher warms known-hot pages.
+  touched pages are read. Snapshot loading checks the handler's recorded exit after its readiness
+  notification, so a startup failure stops the resume with the original cause. An optional
+  prefetcher warms known-hot pages.
 - **Copy-on-write rootfs** (`rootfs/`, `nbd/`, `block/`): the template rootfs stays read-only;
   writes go to a per-sandbox COW cache exposed to Firecracker as an NBD block device served by
   an in-process userspace NBD server. On pause, the dirty blocks are exported as a diff.

packages/orchestrator/pkg/sandbox/fc/client.go

@@ -22,6 +22,7 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/fc/models"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
 const archARM64 = "arm64"
@@ -74,6 +75,7 @@ func (c *apiClient) loadSnapshot(
 	ctx context.Context,
 	uffdSocketPath string,
 	uffdReady chan struct{},
+	uffdExit *utils.ErrorOnce,
 	snapfile template.File,
 	useMemfd bool,
 	useSyncWP bool,
@@ -133,6 +135,13 @@ func (c *apiClient) loadSnapshot(
 		// Wait for the uffd to be ready to serve requests
 	}
 
+	// Ready also closes on startup failure. The exit watcher may not have run.
+	select {
+	case <-uffdExit.Done():
+		return errors.Join(errors.New("uffd exited while loading snapshot"), uffdExit.Wait())
+	default:
+	}
+
 	telemetry.ReportEvent(ctx, "uffd ready")
 
 	return nil

packages/orchestrator/pkg/sandbox/fc/client_test.go

@@ -13,16 +13,12 @@ import (
 
 	"github.com/stretchr/testify/require"
 	"golang.org/x/sys/unix"
+
+	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/block"
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
-// A uffd handler that exits after Firecracker loaded the snapshot ends the
-// wait for it with its exit error as the cause, and the cause must reach the
-// caller: Create tells a node whose pool could not back the guest memory
-// (ENOMEM from the memfd mmap) from a failed create by it. The stub keeps the
-// ready channel open so the wait takes the context branch; in production the
-// handler closes the channel right after recording the exit, so this branch is
-// the rarer ordering and the cause otherwise arrives through the next
-// Firecracker call.
+// Cancellation while waiting for readiness preserves the handler's exit cause.
 func TestLoadSnapshotKeepsUffdExitCause(t *testing.T) {
 	t.Parallel()
 
@@ -53,6 +49,59 @@ func TestLoadSnapshotKeepsUffdExitCause(t *testing.T) {
 		cancel(fmt.Errorf("uffd process exited: failed to wrap memfd: mmap memfd: %w", unix.ENOMEM))
 	}()
 
-	err = newApiClient(socket).loadSnapshot(ctx, dir+"/uffd.sock", make(chan struct{}), stubSnapfile(dir+"/snapfile"), false, false, "", "team-1")
+	err = newApiClient(socket).loadSnapshot(ctx, dir+"/uffd.sock", make(chan struct{}), utils.NewErrorOnce(), stubSnapfile(dir+"/snapfile"), false, false, "", "team-1")
 	require.ErrorIs(t, err, unix.ENOMEM)
 }
+
+// The exit watcher may not have canceled the start context when Ready closes.
+func TestLoadSnapshotChecksUffdExitAfterReady(t *testing.T) {
+	t.Parallel()
+
+	for _, tt := range []struct {
+		name    string
+		exited  bool
+		exitErr error
+	}{
+		{name: "ready"},
+		{name: "failed", exited: true, exitErr: block.GuestMemoryError{Err: unix.ENOMEM}},
+		{name: "stopped", exited: true},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			dir, err := os.MkdirTemp("", "fc") //nolint:usetesting // unix socket path limit
+			require.NoError(t, err)
+			t.Cleanup(func() { _ = os.RemoveAll(dir) })
+			socket := dir + "/fc.sock"
+			listener, err := new(net.ListenConfig).Listen(t.Context(), "unix", socket)
+			require.NoError(t, err)
+
+			exit := utils.NewErrorOnce()
+			ready := make(chan struct{})
+			mux := http.NewServeMux()
+			mux.HandleFunc("PUT /snapshot/load", func(w http.ResponseWriter, _ *http.Request) {
+				if tt.exited {
+					_ = exit.SetError(tt.exitErr)
+				}
+				close(ready)
+				w.WriteHeader(http.StatusNoContent)
+			})
+			server := &http.Server{Handler: mux} //nolint:gosec // private test socket
+			go func() { _ = server.Serve(listener) }()
+			t.Cleanup(func() { _ = server.Close() })
+
+			err = newApiClient(socket).loadSnapshot(t.Context(), dir+"/uffd.sock", ready, exit, stubSnapfile(dir+"/snapfile"), false, false, "", "team-1")
+			if !tt.exited {
+				require.NoError(t, err)
+
+				return
+			}
+			require.ErrorContains(t, err, "uffd exited while loading snapshot")
+			if tt.exitErr != nil {
+				require.ErrorIs(t, err, unix.ENOMEM)
+				var memoryErr block.GuestMemoryError
+				require.ErrorAs(t, err, &memoryErr)
+			}
+		})
+	}
+}

packages/orchestrator/pkg/sandbox/fc/process.go

@@ -726,6 +726,7 @@ func (p *Process) Resume(
 	uffdSocketPath string,
 	snapfile template.File,
 	uffdReady chan struct{},
+	uffdExit *utils.ErrorOnce,
 	accessToken *string,
 	cgroupHandle *cgroup.CgroupHandle,
 	useMemfd bool,
@@ -819,6 +820,7 @@ func (p *Process) Resume(
 		ctx,
 		uffdSocketPath,
 		uffdReady,
+		uffdExit,
 		snapfile,
 		useMemfd,
 		useSyncWP,

packages/orchestrator/pkg/sandbox/fc/snapshot_load_failure_metric_test.go

@@ -15,6 +15,7 @@ import (
 	"go.opentelemetry.io/otel/sdk/metric/metricdata"
 
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
 type stubSnapfile string
@@ -56,7 +57,7 @@ func TestLoadSnapshotRecordsRefusal(t *testing.T) {
 	t.Cleanup(func() { _ = server.Close() })
 
 	// uffdReady is never closed: the load fails before it is waited on.
-	err = newApiClient(socket).loadSnapshot(t.Context(), dir+"/uffd.sock", make(chan struct{}), stubSnapfile(dir+"/snapfile"), false, false, "abc123", "team-1")
+	err = newApiClient(socket).loadSnapshot(t.Context(), dir+"/uffd.sock", make(chan struct{}), utils.NewErrorOnce(), stubSnapfile(dir+"/snapfile"), false, false, "abc123", "team-1")
 	require.Error(t, err)
 
 	assert.Equal(t, int64(1), counterValue(t, reader, string(telemetry.SandboxFCSnapshotLoadFailures), "reason", string(snapshotLoadVcpuMSR)))

packages/orchestrator/pkg/sandbox/sandbox.go

@@ -2003,6 +2003,7 @@ func (f *Factory) ResumeSandbox(
 		fcUffdPath,
 		snapfile,
 		fcUffd.Ready(),
+		fcUffd.Exit(),
 		config.Envd.AccessToken,
 		cgroupHandle,
 		useMemfd,