fix(vm): keep a stopped VM restorable after a transient restore failure

hypervisor/cloudhypervisor/restore.go

@@ -46,12 +46,6 @@ func (ch *CloudHypervisor) terminateVMM(ctx context.Context, rec *hypervisor.VMR
 func (ch *CloudHypervisor) restoreAfterExtract(ctx context.Context, vmID string, vmCfg *types.VMConfig, rec *hypervisor.VMRecord, directBoot bool) (_ *types.VM, err error) {
 	logger := log.WithFunc("cloudhypervisor.Restore")
 
-	defer func() {
-		if err != nil {
-			ch.MarkError(ctx, vmID)
-		}
-	}()
-
 	chConfigPath := filepath.Join(rec.RunDir, configJSONName)
 	// rec may have trailing cidata absent from the snapshot (cloudimg post-first-boot); slice to sidecar length.
 	meta, metaErr := hypervisor.LoadAndValidateMeta(rec.RunDir, ch.conf.RootDir, ch.conf.Config.RunDir)

hypervisor/firecracker/restore.go

@@ -54,12 +54,6 @@ func (fc *Firecracker) terminateVMM(ctx context.Context, rec *hypervisor.VMRecor
 func (fc *Firecracker) restoreAfterExtract(ctx context.Context, vmID string, vmCfg *types.VMConfig, rec *hypervisor.VMRecord, cowPath string) (_ *types.VM, err error) {
 	logger := log.WithFunc("firecracker.Restore")
 
-	defer func() {
-		if err != nil {
-			fc.MarkError(ctx, vmID)
-		}
-	}()
-
 	snapshotCOW := filepath.Join(rec.RunDir, cowFileName)
 	if snapshotCOW != cowPath {
 		if _, statErr := os.Stat(snapshotCOW); statErr == nil {

hypervisor/restore.go

@@ -25,6 +25,17 @@ func (b *Backend) KillForRestore(ctx context.Context, vmID string, rec *VMRecord
 	return nil
 }
 
+// FailRestore marks the VM error after a post-kill restore failure; a stopped
+// origin is spared so hibernate wake stays retryable. Run-dir-mutating steps
+// (staged merge, direct populate) quarantine unconditionally at their own
+// site; origin is the pre-kill state — the DB may read stopped after the kill.
+func (b *Backend) FailRestore(ctx context.Context, vmID string, origin types.VMState) {
+	if origin == types.VMStateStopped {
+		return
+	}
+	b.MarkError(ctx, vmID)
+}
+
 func (b *Backend) ResolveForRestore(ctx context.Context, vmRef string) (string, *VMRecord, error) {
 	vmID, err := b.ResolveRef(ctx, vmRef)
 	if err != nil {
@@ -101,6 +112,8 @@ func (b *Backend) RestoreSequence(ctx context.Context, vmRef string, spec Restor
 			}
 		}
 		if mergeErr := MergeDirInto(stagingDir, rec.RunDir); mergeErr != nil {
+			// A partial merge leaves mixed-vintage files in the run dir;
+			// quarantine regardless of origin so vm start cannot boot them.
 			b.MarkError(ctx, vmID)
 			return fmt.Errorf("apply staged snapshot: %w", mergeErr)
 		}
@@ -109,6 +122,7 @@ func (b *Backend) RestoreSequence(ctx context.Context, vmRef string, spec Restor
 		return afterErr
 	}
 	if err := runWrapped(rec, spec.Wrap, inner); err != nil {
+		b.FailRestore(ctx, vmID, rec.State)
 		return nil, err
 	}
 	b.emitRestoreSuccess(ctx, result, oldShape, spec.SourceSnapshotID)
@@ -137,6 +151,8 @@ func (b *Backend) DirectRestoreSequence(ctx context.Context, vmRef string, spec
 	var result *types.VM
 	inner := func() error {
 		if populateErr := spec.Populate(rec, spec.SrcDir); populateErr != nil {
+			// Populate cleans then clones with no rollback; a partial run
+			// dir must quarantine regardless of origin, like the merge.
 			b.MarkError(ctx, vmID)
 			return populateErr
 		}
@@ -145,6 +161,7 @@ func (b *Backend) DirectRestoreSequence(ctx context.Context, vmRef string, spec
 		return afterErr
 	}
 	if err := runWrapped(rec, spec.Wrap, inner); err != nil {
+		b.FailRestore(ctx, vmID, rec.State)
 		return nil, err
 	}
 	b.emitRestoreSuccess(ctx, result, oldShape, spec.SourceSnapshotID)

hypervisor/restore_test.go

@@ -1,12 +1,168 @@
 package hypervisor
 
 import (
+	"archive/tar"
+	"bytes"
+	"context"
+	"errors"
+	"io"
+	"os"
+	"path/filepath"
 	"strings"
 	"testing"
 
 	"github.com/cocoonstack/cocoon/types"
 )
 
+func TestFailRestoreSparesStoppedOrigin(t *testing.T) {
+	b, _ := newMeteringTestBackend(t)
+	tests := []struct {
+		name   string
+		origin types.VMState
+		want   types.VMState
+	}{
+		{"stopped origin stays stopped (wake retryable)", types.VMStateStopped, types.VMStateStopped},
+		{"running origin marks error", types.VMStateRunning, types.VMStateError},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			id := "vm-fail-" + string(tt.origin)
+			seedVMRecord(t, b, id, 1, 512, 1024, true)
+			if err := b.DB.Update(t.Context(), func(idx *VMIndex) error {
+				idx.VMs[id].State = tt.origin
+				return nil
+			}); err != nil {
+				t.Fatalf("seed state: %v", err)
+			}
+
+			b.FailRestore(t.Context(), id, tt.origin)
+
+			rec, err := b.LoadRecord(t.Context(), id)
+			if err != nil {
+				t.Fatalf("load record: %v", err)
+			}
+			if rec.State != tt.want {
+				t.Errorf("state %s, want %s", rec.State, tt.want)
+			}
+		})
+	}
+}
+
+func TestDirectRestoreFailureKeepsOriginContract(t *testing.T) {
+	failAfterExtract := func(spec *DirectRestoreSpec) {
+		spec.AfterExtract = func(context.Context, string, *types.VMConfig, *VMRecord) (*types.VM, error) {
+			return nil, errors.New("launch boom")
+		}
+	}
+	failPopulate := func(spec *DirectRestoreSpec) {
+		spec.Populate = func(*VMRecord, string) error { return errors.New("populate boom") }
+	}
+	tests := []struct {
+		name   string
+		origin types.VMState
+		fail   func(*DirectRestoreSpec)
+		want   types.VMState
+	}{
+		{"stopped origin survives a post-populate failure", types.VMStateStopped, failAfterExtract, types.VMStateStopped},
+		{"running origin quarantines on a post-populate failure", types.VMStateRunning, failAfterExtract, types.VMStateError},
+		{"partial populate quarantines even a stopped origin", types.VMStateStopped, failPopulate, types.VMStateError},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			b, _ := newMeteringTestBackend(t)
+			const id = "vm-direct-fail"
+			seedVMRecord(t, b, id, 1, 512, 1024, true)
+			if err := b.DB.Update(t.Context(), func(idx *VMIndex) error {
+				idx.VMs[id].State = tt.origin
+				return nil
+			}); err != nil {
+				t.Fatalf("seed state: %v", err)
+			}
+
+			spec := DirectRestoreSpec{
+				VMCfg:     &types.VMConfig{Config: types.Config{CPU: 1, Memory: 512, Storage: 1024}},
+				SrcDir:    t.TempDir(),
+				Preflight: func(string, *VMRecord) error { return nil },
+				Kill:      func(context.Context, string, *VMRecord) error { return nil },
+				Populate:  func(*VMRecord, string) error { return nil },
+			}
+			tt.fail(&spec)
+			if _, err := b.DirectRestoreSequence(t.Context(), id, spec); err == nil {
+				t.Fatal("expected restore failure")
+			}
+			rec, err := b.LoadRecord(t.Context(), id)
+			if err != nil {
+				t.Fatalf("load record: %v", err)
+			}
+			if rec.State != tt.want {
+				t.Errorf("state %s, want %s", rec.State, tt.want)
+			}
+		})
+	}
+}
+
+func TestRestorePartialMergeQuarantinesEvenStoppedOrigin(t *testing.T) {
+	b, _ := newMeteringTestBackend(t)
+	const id = "vm-merge-fail"
+	runDir := t.TempDir()
+	seedVMRecord(t, b, id, 1, 512, 1024, true)
+	if err := b.DB.Update(t.Context(), func(idx *VMIndex) error {
+		idx.VMs[id].State = types.VMStateStopped
+		idx.VMs[id].RunDir = runDir
+		return nil
+	}); err != nil {
+		t.Fatalf("seed state: %v", err)
+	}
+
+	// Staged "a" merges, then staged file "b" cannot rename over the run
+	// dir's directory "b": the merge fails halfway through, which must
+	// quarantine even a stopped origin.
+	if err := os.MkdirAll(filepath.Join(runDir, "b"), 0o750); err != nil {
+		t.Fatalf("setup: %v", err)
+	}
+	spec := RestoreSpec{
+		VMCfg:     &types.VMConfig{Config: types.Config{CPU: 1, Memory: 512, Storage: 1024}},
+		Snapshot:  tarWithFiles(t, "a", "b"),
+		Preflight: func(string, *VMRecord) error { return nil },
+		Kill:      func(context.Context, string, *VMRecord) error { return nil },
+		AfterExtract: func(context.Context, string, *types.VMConfig, *VMRecord) (*types.VM, error) {
+			t.Fatal("AfterExtract must not run after a failed merge")
+			return nil, nil
+		},
+	}
+	if _, err := b.RestoreSequence(t.Context(), id, spec); err == nil {
+		t.Fatal("expected merge failure")
+	}
+	rec, err := b.LoadRecord(t.Context(), id)
+	if err != nil {
+		t.Fatalf("load record: %v", err)
+	}
+	if rec.State != types.VMStateError {
+		t.Errorf("state %s, want error (mixed run dir must quarantine)", rec.State)
+	}
+	if _, err := os.Stat(filepath.Join(runDir, "a")); err != nil {
+		t.Errorf("merge was not partial: %v", err)
+	}
+}
+
+func tarWithFiles(t *testing.T, names ...string) io.Reader {
+	t.Helper()
+	var buf bytes.Buffer
+	tw := tar.NewWriter(&buf)
+	for _, name := range names {
+		if err := tw.WriteHeader(&tar.Header{Name: name, Mode: 0o600, Size: 1}); err != nil {
+			t.Fatalf("tar hdr %s: %v", name, err)
+		}
+		if _, err := tw.Write([]byte("y")); err != nil {
+			t.Fatalf("tar body %s: %v", name, err)
+		}
+	}
+	if err := tw.Close(); err != nil {
+		t.Fatalf("tar close: %v", err)
+	}
+	return &buf
+}
+
 func TestResolveForRestoreStates(t *testing.T) {
 	b, _ := newMeteringTestBackend(t)
 	tests := []struct {