From 293cebee68c5215c5ca7e54fa5a2cba0bf388833 Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Fri, 10 Jul 2026 10:13:47 +1000 Subject: [PATCH] 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. --- spinifex/vm/restore.go | 28 ++++++++++----- spinifex/vm/restore_test.go | 70 +++++++++++++++++++++++++++++++++--- spinifex/vm/shutdown.go | 8 +++++ spinifex/vm/shutdown_test.go | 69 ++++++++++++++++++++++++++++++----- 4 files changed, 155 insertions(+), 20 deletions(-) diff --git a/spinifex/vm/restore.go b/spinifex/vm/restore.go index 7f1d49ad9..909a75e80 100644 --- a/spinifex/vm/restore.go +++ b/spinifex/vm/restore.go @@ -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") diff --git a/spinifex/vm/restore_test.go b/spinifex/vm/restore_test.go index 897f99ee6..7ffbb3684 100644 --- a/spinifex/vm/restore_test.go +++ b/spinifex/vm/restore_test.go @@ -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}, + }, }, } diff --git a/spinifex/vm/shutdown.go b/spinifex/vm/shutdown.go index 30d9b3575..71155c281 100644 --- a/spinifex/vm/shutdown.go +++ b/spinifex/vm/shutdown.go @@ -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 diff --git a/spinifex/vm/shutdown_test.go b/spinifex/vm/shutdown_test.go index ded834aa4..d28801cca 100644 --- a/spinifex/vm/shutdown_test.go +++ b/spinifex/vm/shutdown_test.go @@ -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)