fix(vm): resume drain-stopped VMs after host reboot (#494)

The coordinated host-shutdown DRAIN phase transitioned every running VM to StateStopped and migrated it to the operator-stopped shared KV bucket, evicting it from node-local running state. On the next boot Restore only loads local running state and treated every StateStopped instance as user-stopped, so previously-running VMs never relaunched after a host reboot. Gate on the existing operator-intent signal Attributes.StopInstance (set only by StopInstances, cleared by StartInstance): stopOne keeps drain-stopped VMs in local running state instead of migrating them, and Restore relaunches drain-stopped VMs while operator-stopped ones stay stopped. markUnschedulable now stamps StopInstance so capacity-demoted instances are not relaunched.
2026-10-19 10:52:04 +00:00 · 2026-07-10 10:13:47 +10:00
parent 635915b5b0
commit 293cebee68
4 changed files with 155 additions and 20 deletions
@@ -89,15 +89,24 @@ func (m *Manager) classifyRestoredInstances() []*VM {
 		}

 		if instance.Status == StateStopped {
-			if !m.MigrateStoppedToSharedKV(instance) {
-				// KV write failed — keep in local state so the next restart
-				// retries the migration. Deleting here would create a "void":
-				// the instance disappears from both local state and the
-				// stopped KV, making it invisible to DescribeStoppedInstances.
-				slog.Warn("Stopped instance KV migration failed, will retry on next restart",
-					"instance", instance.ID)
+			if instance.Attributes.StopInstance {
+				if !m.MigrateStoppedToSharedKV(instance) {
+					// KV write failed — keep in local state so the next restart
+					// retries the migration. Deleting here would create a "void":
+					// the instance disappears from both local state and the
+					// stopped KV, making it invisible to DescribeStoppedInstances.
+					slog.Warn("Stopped instance KV migration failed, will retry on next restart",
+						"instance", instance.ID)
+				}
+				continue
 			}
-			continue
+
+			// StateStopped with no operator StopInstance intent means the
+			// host DRAIN sequence stopped it for a coordinated reboot/shutdown,
+			// not the operator. Treat it like a running instance whose QEMU
+			// exited: reset to Pending and relaunch.
+			slog.Info("Instance was drain-stopped (not operator-stopped), relaunching", "instance", instance.ID)
+			instance.Status = StatePending
 		}

 		// Recovery-failed instances stay in StateError until the operator
@@ -180,8 +189,11 @@ func (m *Manager) classifyRestoredInstances() []*VM {

 // markUnschedulable flips an instance to Stopped with InsufficientInstanceCapacity
 // so DescribeInstances reports a useful error when the node can no longer host the type.
+// StopInstance is stamped so a future restore's drain-relaunch path (which
+// resumes StateStopped without operator intent) does not try to relaunch it.
 func markUnschedulable(instance *VM, reason string) {
 	instance.Status = StateStopped
+	instance.Attributes.StopInstance = true
 	if instance.Instance != nil {
 		instance.Instance.StateReason = &ec2.StateReason{}
 		instance.Instance.StateReason.SetCode("Server.InsufficientInstanceCapacity")
@@ -114,19 +114,78 @@ func TestClassifyRestoredInstances_TerminatedMigrates(t *testing.T) {
 	assert.False(t, stillLocal, "terminated must be removed from the local map")
 }

+// TestClassifyRestoredInstances_StoppedMigrates covers the operator-stopped
+// case: Attributes.StopInstance=true is the signal stamped only by the
+// operator StopInstances path, so classify must honour it and never relaunch.
 func TestClassifyRestoredInstances_StoppedMigrates(t *testing.T) {
 	m, store, _ := classifyTestManager(t)
-	v := &VM{ID: "i-stopped", Status: StateStopped, InstanceType: "t3.micro"}
+	v := &VM{
+		ID:           "i-stopped",
+		Status:       StateStopped,
+		InstanceType: "t3.micro",
+		Attributes:   types.EC2CommandAttributes{StopInstance: true},
+	}
 	m.Replace(map[string]*VM{v.ID: v})

 	toLaunch := m.classifyRestoredInstances()

-	assert.Empty(t, toLaunch, "stopped must not be queued for relaunch")
+	assert.Empty(t, toLaunch, "operator-stopped must not be queued for relaunch")
 	assert.NotNil(t, store.stopped[v.ID], "stopped must be migrated to the shared bucket")
 	_, stillLocal := m.Get(v.ID)
 	assert.False(t, stillLocal, "stopped must be removed from the local map after migration")
 }

+// TestClassifyRestoredInstances_DrainStoppedRelaunches covers the
+// coordinated-shutdown DRAIN case: a VM left StateStopped with no operator
+// StopInstance intent (Attributes zero value) must be queued for relaunch,
+// not migrated to the stopped bucket, exactly like a running instance whose
+// QEMU exited.
+func TestClassifyRestoredInstances_DrainStoppedRelaunches(t *testing.T) {
+	m, store, rc := classifyTestManager(t)
+	v := &VM{
+		ID:           "i-drain-stopped",
+		Status:       StateStopped,
+		InstanceType: "t3.micro",
+		Attributes:   types.EC2CommandAttributes{},
+		Instance:     &ec2.Instance{},
+	}
+	m.Replace(map[string]*VM{v.ID: v})
+
+	toLaunch := m.classifyRestoredInstances()
+
+	require.Len(t, toLaunch, 1)
+	assert.Same(t, v, toLaunch[0])
+	assert.Equal(t, StatePending, v.Status,
+		"drain-stopped (no operator StopInstance) must reset to Pending so the relaunch path accepts it")
+	require.NotNil(t, v.Instance.LaunchTime, "LaunchTime must be reset for the pending watchdog")
+	assert.Nil(t, store.stopped[v.ID], "drain-stopped instance must not be migrated to the stopped bucket")
+	assert.Equal(t, 1, rc.allocateCount("t3.micro"),
+		"resources must be re-allocated before queuing for relaunch")
+}
+
+// TestClassifyRestoredInstances_UnschedulableInstanceNotRelaunched covers the
+// markUnschedulable guard: an instance demoted for insufficient
+// capacity/unknown type has StopInstance stamped true so the new
+// drain-relaunch branch does not mistake it for a drained instance and
+// wrongly resume it on a later restore.
+func TestClassifyRestoredInstances_UnschedulableInstanceNotRelaunched(t *testing.T) {
+	m, store, _ := classifyTestManager(t)
+	v := &VM{
+		ID:           "i-unschedulable",
+		InstanceType: "t3.micro",
+		Instance:     &ec2.Instance{},
+	}
+	markUnschedulable(v, "insufficient resources")
+	require.True(t, v.Attributes.StopInstance,
+		"markUnschedulable must stamp StopInstance so a later restore does not resume it")
+	m.Replace(map[string]*VM{v.ID: v})
+
+	toLaunch := m.classifyRestoredInstances()
+
+	assert.Empty(t, toLaunch, "an unschedulable instance must not be relaunched by the drain-stopped path")
+	assert.NotNil(t, store.stopped[v.ID], "must stay migrated to the stopped bucket, not relaunch")
+}
+
 func TestClassifyRestoredInstances_UnknownTypeMarkedUnschedulable(t *testing.T) {
 	m, store, _ := classifyTestManager(t)
 	v := &VM{
@@ -310,8 +369,11 @@ func TestRestore_HappyPath(t *testing.T) {
 	store := &markerStateStore{
 		fakeStateStore: newFakeStateStore(),
 		snapshot: map[string]*VM{
-			"i-term":    {ID: "i-term", Status: StateTerminated, InstanceType: "t3.micro"},
-			"i-stopped": {ID: "i-stopped", Status: StateStopped, InstanceType: "t3.micro"},
+			"i-term": {ID: "i-term", Status: StateTerminated, InstanceType: "t3.micro"},
+			"i-stopped": {
+				ID: "i-stopped", Status: StateStopped, InstanceType: "t3.micro",
+				Attributes: types.EC2CommandAttributes{StopInstance: true},
+			},
 		},
 	}

@@ -64,6 +64,14 @@ func (m *Manager) stopOne(instance *VM) (bool, error) {
 		slog.Error("Failed to transition to stopped", "instanceId", instance.ID, "err", err)
 	}

+	if !instance.Attributes.StopInstance {
+		// Host DRAIN stop (not operator): keep the VM in the local running
+		// map at StateStopped so Restore relaunches it on the next boot. Do
+		// not migrate to the operator-stopped shared bucket or fire
+		// OnInstanceDown; QEMU is already down and resources released.
+		return false, nil
+	}
+
 	if !m.MigrateStoppedToSharedKV(instance) {
 		// Either StateStore unavailable / write failed (instance stays in
 		// local map; restoreInstances retries on next boot) OR a concurrent
@@ -11,6 +11,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/mulgadc/spinifex/spinifex/gpu"
 	"github.com/mulgadc/spinifex/spinifex/qmp"
+	"github.com/mulgadc/spinifex/spinifex/types"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -241,14 +242,12 @@ func TestStopAll_EmptyMap_FastPath(t *testing.T) {
 	assert.Empty(t, cleaner.cleanupMgmt, "empty StopAll must not invoke cleaner")
 }

-// TestStopAll_FiresOnInstanceDownAndMigratesPerVM verifies the DRAIN
-// hook contract: every Running VM transitions through Stopping → Stopped,
-// is migrated to the cluster-shared "stopped" bucket, fires
-// OnInstanceDown once, and is removed from the local running map. The
-// pre-fix StopAll only ran stopCleanup and left every VM stuck in
-// Running, so a daemon restart promoted user-stopped VMs through the
-// failed-recovery path and surfaced them as terminated — data-equivalent
-// loss for the customer.
+// TestStopAll_FiresOnInstanceDownAndMigratesPerVM verifies the
+// operator-stop hook contract: a Running VM stamped with the operator
+// StopInstance intent transitions through Stopping → Stopped, is migrated
+// to the cluster-shared "stopped" bucket, fires OnInstanceDown once, and
+// is removed from the local running map. Without honouring the operator
+// intent a restart would relaunch a deliberately stopped VM.
 func TestStopAll_FiresOnInstanceDownAndMigratesPerVM(t *testing.T) {
 	m, store, _, _, rt := shutdownTestManager(t)
 	var (
@@ -278,6 +277,7 @@ func TestStopAll_FiresOnInstanceDownAndMigratesPerVM(t *testing.T) {
 			Status:       StateRunning,
 			InstanceType: "t3.micro",
 			Instance:     &ec2.Instance{},
+			Attributes:   types.EC2CommandAttributes{StopInstance: true},
 		})
 	}

@@ -298,6 +298,56 @@ func TestStopAll_FiresOnInstanceDownAndMigratesPerVM(t *testing.T) {
 	}
 }

+// TestStopAll_DrainStopKeepsLocalNoMigrate verifies the host-DRAIN
+// contract: a Running VM with no operator StopInstance intent is stopped
+// gracefully (Stopping → Stopped) but stays in the local running map so
+// Restore relaunches it on the next boot. It must NOT be migrated to the
+// operator-stopped shared bucket and must NOT fire OnInstanceDown. This is
+// the reboot-resilience fix — DRAIN previously migrated every VM to the
+// shared bucket, hiding it from the node-local restore path.
+func TestStopAll_DrainStopKeepsLocalNoMigrate(t *testing.T) {
+	m, store, _, _, rt := shutdownTestManager(t)
+	var down atomic.Int64
+	m.SetDeps(Deps{
+		NodeID:          "test-node",
+		StateStore:      store,
+		VolumeMounter:   &fakeVolumeMounter{},
+		InstanceCleaner: &recordingInstanceCleaner{},
+		TransitionState: rt.apply,
+		ShutdownSignal:  func() bool { return false },
+		Hooks: ManagerHooks{
+			OnInstanceDown: func(string) { down.Add(1) },
+		},
+	})
+
+	ids := []string{"i-1", "i-2", "i-3"}
+	for _, id := range ids {
+		m.Insert(&VM{
+			ID:           id,
+			Status:       StateRunning,
+			InstanceType: "t3.micro",
+			Instance:     &ec2.Instance{},
+			Attributes:   types.EC2CommandAttributes{},
+		})
+	}
+
+	require.NoError(t, m.StopAll())
+
+	assert.Zero(t, down.Load(),
+		"drain stop must not fire OnInstanceDown — the VM is being relaunched, not released")
+	assert.Equal(t, len(ids), m.Count(),
+		"drain-stopped VMs must stay in the local running map so Restore relaunches them")
+	for _, id := range ids {
+		v, ok := m.Get(id)
+		require.True(t, ok, "drain-stopped VM %s must remain in the local map", id)
+		assert.Equal(t, StateStopped, v.Status, "drain-stopped VM %s must be at StateStopped", id)
+		stored, err := store.LoadStoppedInstance(id)
+		require.NoError(t, err)
+		assert.Nil(t, stored,
+			"drain-stopped VM %s must NOT migrate to the operator-stopped shared bucket", id)
+	}
+}
+
 // TestStopAll_SkipsNonRunningWithoutFiringHook verifies the precheck
 // skip path: non-Running VMs (Pending, Stopping mid-flight, Stopped,
 // Error, ShuttingDown, Terminated) surface as ErrInvalidTransition from
@@ -381,6 +431,7 @@ func TestStopAll_WriteRunningStateFailure(t *testing.T) {
 			Status:       StateRunning,
 			InstanceType: "t3.micro",
 			Instance:     &ec2.Instance{},
+			Attributes:   types.EC2CommandAttributes{StopInstance: true},
 		})
 	}

@@ -505,6 +556,7 @@ func TestStop_FiresOnInstanceDownExactlyOnce(t *testing.T) {
 		Status:       StateRunning,
 		InstanceType: "t3.micro",
 		Instance:     &ec2.Instance{},
+		Attributes:   types.EC2CommandAttributes{StopInstance: true},
 	}
 	m.Insert(v)

@@ -544,6 +596,7 @@ func TestStop_DoesNotFireOnInstanceDown_OnSlotReclaim(t *testing.T) {
 		Status:       StateRunning,
 		InstanceType: "t3.micro",
 		Instance:     &ec2.Instance{},
+		Attributes:   types.EC2CommandAttributes{StopInstance: true},
 	}
 	m.Insert(v)