refactor: introduce vm.Manager (phase 1 of lifecycle extraction) (#176)

This commit is contained in:
Julian Sommer
2026-05-01 15:15:54 +10:00
committed by GitHub
parent 57a11c0668
commit e1b3623599
25 changed files with 1121 additions and 761 deletions
+64 -75
View File
@@ -131,8 +131,8 @@ type Daemon struct {
cancel context.CancelFunc
shutdownWg sync.WaitGroup
// Local VM Instances
Instances vm.Instances
// vmMgr owns the in-memory map of VMs running on this node.
vmMgr *vm.Manager
// NAT Subscriptions
natsSubscriptions map[string]*nats.Subscription
@@ -412,7 +412,7 @@ func NewDaemon(cfg *config.ClusterConfig) (*Daemon, error) {
resourceMgr: rm,
ctx: ctx,
cancel: cancel,
Instances: vm.Instances{VMS: make(map[string]*vm.VM)},
vmMgr: vm.NewManager(),
natsSubscriptions: make(map[string]*nats.Subscription),
startTime: time.Now(),
detachDelay: 1 * time.Second,
@@ -615,7 +615,7 @@ func (d *Daemon) Start() error {
// Create services before loading/launching instances, since LaunchInstance depends on them
store := objectstore.NewS3ObjectStoreFromConfig(admin.DialTarget(d.config.Predastore.Host), d.config.Predastore.Region, d.config.Predastore.AccessKey, d.config.Predastore.SecretKey)
d.instanceService = handlers_ec2_instance.NewInstanceServiceImpl(d.config, d.resourceMgr.instanceTypes, d.natsConn, &d.Instances, store)
d.instanceService = handlers_ec2_instance.NewInstanceServiceImpl(d.config, d.resourceMgr.instanceTypes, d.natsConn, store)
d.keyService = handlers_ec2_key.NewKeyServiceImpl(d.config)
d.imageService = handlers_ec2_image.NewImageServiceImpl(d.config, d.natsConn)
@@ -856,9 +856,7 @@ func (d *Daemon) Start() error {
// Rebuild mgmt IP allocator from restored VMs so we don't re-allocate IPs
// that are already in use by running system instances.
if d.mgmtIPAllocator != nil {
d.Instances.Mu.Lock()
d.mgmtIPAllocator.Rebuild(d.Instances.VMS)
d.Instances.Mu.Unlock()
d.mgmtIPAllocator.Rebuild(d.vmMgr.SnapshotMap())
slog.Info("Rebuilt mgmt IP allocator from restored instances", "allocated", d.mgmtIPAllocator.AllocatedCount())
}
@@ -1061,7 +1059,11 @@ func (d *Daemon) migrateInstanceToKV(instance *vm.VM, writeFn func(string, *vm.V
"instance", instance.ID, "bucket", label, "err", err)
return false
}
delete(d.Instances.VMS, instance.ID)
if !d.vmMgr.DeleteIf(instance.ID, instance) {
slog.Info("Slot reclaimed by another handler during migration; skipping local delete",
"instance", instance.ID, "bucket", label)
return true
}
slog.Info("Migrated instance to KV", "instance", instance.ID, "bucket", label)
return true
}
@@ -1103,19 +1105,14 @@ func (d *Daemon) restoreInstances() {
return
}
slog.Info("Loaded state", "instance count", len(d.Instances.VMS))
// Ensure mutexes and QMP clients are usable after deserialization
d.Instances.Mu = sync.Mutex{}
slog.Info("Loaded state", "instance count", d.vmMgr.Count())
// Phase 1: Reconnect running QEMU, finalize transitional states, collect VMs to relaunch
var toLaunch []*vm.VM
for i := range d.Instances.VMS {
d.Instances.VMS[i].EBSRequests.Mu = sync.Mutex{}
d.Instances.VMS[i].QMPClient = &qmp.QMPClient{}
instance := d.Instances.VMS[i]
for _, instance := range d.vmMgr.Snapshot() {
instance.EBSRequests.Mu = sync.Mutex{}
instance.QMPClient = &qmp.QMPClient{}
if instance.Status == vm.StateTerminated {
if !d.migrateTerminatedToKV(instance) {
@@ -1280,9 +1277,8 @@ func (d *Daemon) restoreInstances() {
}
}()
// Skip if instance was terminated while waiting for semaphore
d.Instances.Mu.Lock()
status := inst.Status
d.Instances.Mu.Unlock()
var status vm.InstanceState
d.vmMgr.Inspect(inst, func(v *vm.VM) { status = v.Status })
if status != vm.StatePending && status != vm.StateProvisioning {
slog.Info("Instance state changed during recovery, skipping launch",
"instanceId", inst.ID, "status", string(status))
@@ -1526,32 +1522,38 @@ func (d *Daemon) ClusterManager() error {
}
// WriteState writes the instance state to JetStream KV store (required).
// It acquires d.Instances.Mu internally.
// The marshal+put runs under the manager lock so VM fields can't change
// mid-encode. Lock-across-Put is a known limitation; splitting marshal from
// put requires a JetStreamManager API change and is deferred.
func (d *Daemon) WriteState() error {
if d.jsManager == nil {
return fmt.Errorf("JetStream manager not initialized - cannot write state")
}
if err := d.jsManager.WriteState(d.node, &d.Instances); err != nil {
slog.Error("JetStream write failed", "error", err)
return fmt.Errorf("failed to write state to JetStream: %w", err)
var writeErr error
d.vmMgr.View(func(vms map[string]*vm.VM) {
writeErr = d.jsManager.WriteState(d.node, vms)
})
if writeErr != nil {
slog.Error("JetStream write failed", "error", writeErr)
return fmt.Errorf("failed to write state to JetStream: %w", writeErr)
}
return nil
}
// LoadState loads the instance state from JetStream KV store (required)
// LoadState loads the instance state from JetStream KV store (required).
// Replaces the manager's running set with the loaded map.
func (d *Daemon) LoadState() error {
if d.jsManager == nil {
return fmt.Errorf("JetStream manager not initialized - cannot load state")
}
instances, err := d.jsManager.LoadState(d.node)
loaded, err := d.jsManager.LoadState(d.node)
if err != nil {
slog.Error("JetStream load failed", "error", err)
return fmt.Errorf("failed to load state from JetStream: %w", err)
}
// Copy only the VMS map, not the mutex
d.Instances.VMS = instances.VMS
d.vmMgr.Replace(loaded)
return nil
}
@@ -1841,7 +1843,7 @@ func (d *Daemon) setupShutdown() {
} else {
d.shuttingDown.Store(true)
// Pass instances to terminate
if err := d.stopInstance(d.Instances.VMS, false); err != nil {
if err := d.stopInstance(d.vmMgr.SnapshotMap(), false); err != nil {
slog.Error("Failed to stop instances during shutdown", "err", err)
}
}
@@ -1910,9 +1912,8 @@ func (d *Daemon) CreateQMPClient(instance *vm.VM) (err error) {
time.Sleep(30 * time.Second)
// Check if instance is in a terminal or transitional state - exit heartbeat
d.Instances.Mu.Lock()
status := instance.Status
d.Instances.Mu.Unlock()
var status vm.InstanceState
d.vmMgr.Inspect(instance, func(v *vm.VM) { status = v.Status })
if status == vm.StateStopping || status == vm.StateStopped || status == vm.StateShuttingDown || status == vm.StateTerminated || status == vm.StateError {
slog.Info("QMP heartbeat exiting - instance not running", "instance", instance.ID, "status", status)
@@ -1949,9 +1950,8 @@ func (d *Daemon) CreateQMPClient(instance *vm.VM) (err error) {
// effects. Logs the abort reason at Info so it's visible without polluting
// error logs (terminate-during-pending is an expected user-driven race).
func (d *Daemon) launchStillValid(instance *vm.VM) bool {
d.Instances.Mu.Lock()
status := instance.Status
d.Instances.Mu.Unlock()
var status vm.InstanceState
d.vmMgr.Inspect(instance, func(v *vm.VM) { status = v.Status })
if status == vm.StatePending || status == vm.StateStopped || status == vm.StateProvisioning {
return true
}
@@ -2045,9 +2045,7 @@ func (d *Daemon) LaunchInstance(instance *vm.VM) (err error) {
}
// Step 9: Update the instance metadata for running state and volume attached
d.Instances.Mu.Lock()
d.Instances.VMS[instance.ID] = instance
d.Instances.Mu.Unlock()
d.vmMgr.Insert(instance)
// Final race check — QEMU is up, but if a terminate fired during
// StartInstance/CreateQMPClient, the concurrent goroutine has already
@@ -2084,27 +2082,34 @@ func (d *Daemon) markInstanceFailed(instance *vm.VM, reason string) {
// request transitioned it to shutting-down while LaunchInstance was running),
// don't spawn a second finalizeTermination goroutine — the existing cleanup
// handler owns the lifecycle from here.
d.Instances.Mu.Lock()
if instance.Status == vm.StateShuttingDown || instance.Status == vm.StateTerminated {
d.Instances.Mu.Unlock()
skip := false
var observedStatus vm.InstanceState
d.vmMgr.Inspect(instance, func(v *vm.VM) {
observedStatus = v.Status
if v.Status == vm.StateShuttingDown || v.Status == vm.StateTerminated {
skip = true
return
}
// Set state reason before transition
if v.Instance != nil {
v.Instance.StateReason = &ec2.StateReason{}
v.Instance.StateReason.SetCode("Server.InternalError")
v.Instance.StateReason.SetMessage(reason)
}
})
if skip {
slog.Info("markInstanceFailed: instance already in cleanup state, skipping",
"instanceId", instance.ID, "status", string(instance.Status), "reason", reason)
"instanceId", instance.ID, "status", string(observedStatus), "reason", reason)
return
}
// Set state reason before transition
if instance.Instance != nil {
instance.Instance.StateReason = &ec2.StateReason{}
instance.Instance.StateReason.SetCode("Server.InternalError")
instance.Instance.StateReason.SetMessage(reason)
}
d.Instances.Mu.Unlock()
if err := d.TransitionState(instance, vm.StateShuttingDown); err != nil {
slog.Error("markInstanceFailed transition failed", "instanceId", instance.ID, "err", err)
// If the error was a write failure, the in-memory state is already
// shutting-down. Still proceed with finalization to avoid getting stuck.
if instance.Status != vm.StateShuttingDown {
var postErrStatus vm.InstanceState
d.vmMgr.Inspect(instance, func(v *vm.VM) { postErrStatus = v.Status })
if postErrStatus != vm.StateShuttingDown {
return
}
}
@@ -2131,9 +2136,7 @@ func (d *Daemon) finalizeTermination(instance *vm.VM) {
return
}
d.Instances.Mu.Lock()
instance.LastNode = d.node
d.Instances.Mu.Unlock()
d.vmMgr.Inspect(instance, func(v *vm.VM) { v.LastNode = d.node })
if err := d.TransitionState(instance, vm.StateTerminated); err != nil {
slog.Error("Failed to transition failed instance to terminated", "instanceId", instance.ID, "err", err)
@@ -2150,25 +2153,16 @@ func (d *Daemon) finalizeTermination(instance *vm.VM) {
}
// Guard + delete: another handler may have reclaimed this instance.
d.Instances.Mu.Lock()
current, exists := d.Instances.VMS[instance.ID]
if !exists || current != instance {
d.Instances.Mu.Unlock()
if !d.vmMgr.DeleteIf(instance.ID, instance) {
slog.Info("Instance was reclaimed by another handler, skipping local cleanup",
"instanceId", instance.ID, "state", "terminated")
return
}
delete(d.Instances.VMS, instance.ID)
d.Instances.Mu.Unlock()
if err := d.WriteState(); err != nil {
slog.Error("Failed to persist state after terminating failed instance, re-adding to local map",
"instanceId", instance.ID, "err", err)
d.Instances.Mu.Lock()
if _, occupied := d.Instances.VMS[instance.ID]; !occupied {
d.Instances.VMS[instance.ID] = instance
}
d.Instances.Mu.Unlock()
d.vmMgr.InsertIfAbsent(instance)
} else {
slog.Info("Released failed instance ownership to KV",
"instanceId", instance.ID, "lastNode", d.node)
@@ -2189,16 +2183,11 @@ func (d *Daemon) startPendingWatchdog() {
case <-d.ctx.Done():
return
case <-ticker.C:
d.Instances.Mu.Lock()
var stuck []*vm.VM
for _, instance := range d.Instances.VMS {
if (instance.Status == vm.StatePending || instance.Status == vm.StateProvisioning) &&
instance.Instance != nil && instance.Instance.LaunchTime != nil &&
time.Since(*instance.Instance.LaunchTime) > pendingWatchdogTimeout {
stuck = append(stuck, instance)
}
}
d.Instances.Mu.Unlock()
stuck := d.vmMgr.Filter(func(v *vm.VM) bool {
return (v.Status == vm.StatePending || v.Status == vm.StateProvisioning) &&
v.Instance != nil && v.Instance.LaunchTime != nil &&
time.Since(*v.Instance.LaunchTime) > pendingWatchdogTimeout
})
for _, instance := range stuck {
slog.Warn("Instance stuck in pending, marking failed",
+6 -13
View File
@@ -69,10 +69,7 @@ func (d *Daemon) handleEC2Events(msg *nats.Msg) {
slog.Debug("Received message", "subject", msg.Subject, "data", string(msg.Data))
d.Instances.Mu.Lock()
instance, ok := d.Instances.VMS[command.ID]
d.Instances.Mu.Unlock()
instance, ok := d.vmMgr.Get(command.ID)
if !ok {
slog.Warn("Instance is not running on this node", "id", command.ID)
respondWithError(msg, awserrors.ErrorInvalidInstanceIDNotFound)
@@ -158,14 +155,12 @@ func (d *Daemon) daemonIP() string {
func (d *Daemon) handleNodeStatus(msg *nats.Msg) {
totalVCPU, totalMemGB, reservedVCPU, reservedMemGB, allocVCPU, allocMemGB, caps := d.resourceMgr.GetResourceStats()
d.Instances.Mu.Lock()
vmCount := 0
for _, v := range d.Instances.VMS {
d.vmMgr.ForEach(func(v *vm.VM) {
if v.Status == vm.StateRunning {
vmCount++
}
}
d.Instances.Mu.Unlock()
})
resp := types.NodeStatusResponse{
Node: d.node,
@@ -288,9 +283,8 @@ func fetchPredastoreRole(url string, client *http.Client) string {
// handleNodeVMs responds with the list of VMs running on this node.
// Used by the CLI: spx get vms.
func (d *Daemon) handleNodeVMs(msg *nats.Msg) {
d.Instances.Mu.Lock()
vms := make([]types.VMInfo, 0, len(d.Instances.VMS))
for _, v := range d.Instances.VMS {
vms := make([]types.VMInfo, 0, d.vmMgr.Count())
d.vmMgr.ForEach(func(v *vm.VM) {
info := types.VMInfo{
InstanceID: v.ID,
Status: string(v.Status),
@@ -306,8 +300,7 @@ func (d *Daemon) handleNodeVMs(msg *nats.Msg) {
info.LaunchTime = v.Instance.LaunchTime.Unix()
}
vms = append(vms, info)
}
d.Instances.Mu.Unlock()
})
resp := types.NodeVMsResponse{
Node: d.node,
+1 -4
View File
@@ -33,10 +33,7 @@ func (d *Daemon) handleEC2GetConsoleOutput(msg *nats.Msg) {
instanceID := *input.InstanceId
// Find the instance on this node
d.Instances.Mu.Lock()
instance, exists := d.Instances.VMS[instanceID]
d.Instances.Mu.Unlock()
instance, exists := d.vmMgr.Get(instanceID)
if !exists {
respondWithError(msg, awserrors.ErrorInvalidInstanceIDNotFound)
return
+14 -12
View File
@@ -62,25 +62,27 @@ func (d *Daemon) handleEC2CreateImage(msg *nats.Msg) {
instanceID := *input.InstanceId
// Extract all instance context in a single critical section
d.Instances.Mu.Lock()
instance, ok := d.Instances.VMS[instanceID]
var status vm.InstanceState
var rootVolumeID, sourceImageID string
if ok {
status = instance.Status
if instance.Instance != nil {
for _, bdm := range instance.Instance.BlockDeviceMappings {
var (
instance *vm.VM
status vm.InstanceState
rootVolumeID string
sourceImageID string
)
ok := d.vmMgr.UpdateState(instanceID, func(v *vm.VM) {
instance = v
status = v.Status
if v.Instance != nil {
for _, bdm := range v.Instance.BlockDeviceMappings {
if bdm.Ebs != nil && bdm.Ebs.VolumeId != nil {
rootVolumeID = *bdm.Ebs.VolumeId
break
}
}
if instance.Instance.ImageId != nil {
sourceImageID = *instance.Instance.ImageId
if v.Instance.ImageId != nil {
sourceImageID = *v.Instance.ImageId
}
}
}
d.Instances.Mu.Unlock()
})
if !ok {
slog.Warn("CreateImage: instance not found", "instanceId", instanceID)
+87 -110
View File
@@ -320,11 +320,9 @@ func (d *Daemon) handleEC2RunInstances(msg *nats.Msg) {
// Add all instances to state immediately so DescribeInstances can find them
// while volumes are being prepared and VMs are launching
d.Instances.Mu.Lock()
for _, instance := range instances {
d.Instances.VMS[instance.ID] = instance
d.vmMgr.Insert(instance)
}
d.Instances.Mu.Unlock()
if err := d.WriteState(); err != nil {
slog.Error("handleEC2RunInstances failed to write initial state", "err", err)
@@ -350,9 +348,8 @@ func (d *Daemon) handleEC2RunInstances(msg *nats.Msg) {
var successCount int
for _, instance := range instances {
// Skip if instance was terminated by a concurrent request
d.Instances.Mu.Lock()
status := instance.Status
d.Instances.Mu.Unlock()
var status vm.InstanceState
d.vmMgr.Inspect(instance, func(v *vm.VM) { status = v.Status })
if status != vm.StatePending && status != vm.StateProvisioning {
slog.Info("Instance state changed during provisioning, skipping launch",
"instanceId", instance.ID, "status", string(status))
@@ -407,9 +404,8 @@ func (d *Daemon) handleEC2RunInstances(msg *nats.Msg) {
func (d *Daemon) handleRebootInstance(msg *nats.Msg, command types.EC2InstanceCommand, instance *vm.VM) {
slog.Info("Rebooting instance", "id", command.ID)
d.Instances.Mu.Lock()
status := instance.Status
d.Instances.Mu.Unlock()
var status vm.InstanceState
d.vmMgr.Inspect(instance, func(v *vm.VM) { status = v.Status })
if status != vm.StateRunning {
slog.Error("RebootInstance: instance not in running state", "instanceId", command.ID, "status", status)
@@ -435,9 +431,8 @@ func (d *Daemon) handleStartInstance(msg *nats.Msg, command types.EC2InstanceCom
slog.Info("Starting instance", "id", command.ID)
// Validate instance is in stopped state
d.Instances.Mu.Lock()
status := instance.Status
d.Instances.Mu.Unlock()
var status vm.InstanceState
d.vmMgr.Inspect(instance, func(v *vm.VM) { status = v.Status })
if status != vm.StateStopped {
slog.Error("StartInstance: instance not in stopped state", "instanceId", command.ID, "status", status)
@@ -458,9 +453,7 @@ func (d *Daemon) handleStartInstance(msg *nats.Msg, command types.EC2InstanceCom
// Clear stop attribute before launch so WriteState inside LaunchInstance
// persists the correct attributes. Without this, a daemon restart after
// a stop→start cycle would see StopInstance=true and skip reconnecting QEMU.
d.Instances.Mu.Lock()
instance.Attributes = command.Attributes
d.Instances.Mu.Unlock()
d.vmMgr.Inspect(instance, func(v *vm.VM) { v.Attributes = command.Attributes })
// Launch the instance infrastructure (QEMU, QMP, NATS subscriptions)
err := d.LaunchInstance(instance)
@@ -500,9 +493,8 @@ func (d *Daemon) handleStopOrTerminateInstance(msg *nats.Msg, command types.EC2I
// Check state validity before attempting transition — return the correct
// AWS error code when the instance is already stopped/terminated/etc.
d.Instances.Mu.Lock()
currentState := instance.Status
d.Instances.Mu.Unlock()
var currentState vm.InstanceState
d.vmMgr.Inspect(instance, func(v *vm.VM) { currentState = v.Status })
// If instance is already shutting-down and we're asked to terminate, treat
// as idempotent — the finalizeTermination goroutine is already cleaning up.
@@ -544,10 +536,10 @@ func (d *Daemon) handleStopOrTerminateInstance(msg *nats.Msg, command types.EC2I
slog.Error("Failed to transition to error state", "instanceId", inst.ID, "err", err)
}
} else {
d.Instances.Mu.Lock()
inst.Attributes = attrs
inst.LastNode = d.node
d.Instances.Mu.Unlock()
d.vmMgr.Inspect(inst, func(v *vm.VM) {
v.Attributes = attrs
v.LastNode = d.node
})
if err := d.TransitionState(inst, finalState); err != nil {
slog.Error("Failed to transition to final state", "instanceId", inst.ID, "err", err)
@@ -591,16 +583,11 @@ func (d *Daemon) handleStopOrTerminateInstance(msg *nats.Msg, command types.EC2I
// from stopped KV, re-added it to VMS with a new pointer, and
// launched it. Deleting here would destroy the running instance's
// state — creating a "ghost instance" visible nowhere.
d.Instances.Mu.Lock()
current, exists := d.Instances.VMS[inst.ID]
if !exists || current != inst {
d.Instances.Mu.Unlock()
if !d.vmMgr.DeleteIf(inst.ID, inst) {
slog.Info("Instance was reclaimed by another handler, skipping local cleanup",
"instanceId", inst.ID, "state", string(finalState))
return
}
delete(d.Instances.VMS, inst.ID)
d.Instances.Mu.Unlock()
// Unsubscribe from per-instance NATS topic. Safe to do after
// the delete — LaunchInstance already unsubscribes stale entries
@@ -619,11 +606,7 @@ func (d *Daemon) handleStopOrTerminateInstance(msg *nats.Msg, command types.EC2I
slog.Error("Failed to persist state after releasing instance, re-adding to local map for consistency",
"instanceId", inst.ID, "err", err)
// Only re-add if another handler hasn't claimed the slot
d.Instances.Mu.Lock()
if _, occupied := d.Instances.VMS[inst.ID]; !occupied {
d.Instances.VMS[inst.ID] = inst
}
d.Instances.Mu.Unlock()
d.vmMgr.InsertIfAbsent(inst)
} else {
slog.Info("Released instance ownership to KV",
"instanceId", inst.ID, "state", string(finalState), "lastNode", d.node)
@@ -738,9 +721,6 @@ func (d *Daemon) handleEC2DescribeInstances(msg *nats.Msg) {
slog.Info("Processing DescribeInstances request from this node", "accountID", accountID)
d.Instances.Mu.Lock()
defer d.Instances.Mu.Unlock()
// Validate and filter instances if specific instance IDs were requested
instanceIDFilter := make(map[string]bool)
if len(describeInstancesInput.InstanceIds) > 0 {
@@ -766,74 +746,78 @@ func (d *Daemon) handleEC2DescribeInstances(msg *nats.Msg) {
// Group instances by reservation ID (AWS returns instances grouped by reservation)
reservationMap := make(map[string]*ec2.Reservation)
// Iterate through all instances on this node
for _, instance := range d.Instances.VMS {
// Skip instances not owned by the caller's account.
// Pre-Phase4 instances (empty AccountID) are only visible to root.
if !isInstanceVisible(accountID, instance.AccountID) {
continue
}
// Skip if filtering by instance IDs and this instance is not in the filter
if len(instanceIDFilter) > 0 && !instanceIDFilter[instance.ID] {
continue
}
// Use stored reservation metadata if available
if instance.Reservation != nil && instance.Instance != nil {
resID := ""
if instance.Reservation.ReservationId != nil {
resID = *instance.Reservation.ReservationId
}
// Create reservation entry if it doesn't exist
if _, exists := reservationMap[resID]; !exists {
reservation := &ec2.Reservation{}
reservation.SetReservationId(resID)
if instance.Reservation.OwnerId != nil {
reservation.SetOwnerId(*instance.Reservation.OwnerId)
}
reservation.Instances = []*ec2.Instance{}
reservationMap[resID] = reservation
}
// Update the instance state to current state
instanceCopy := *instance.Instance
instanceCopy.State = &ec2.InstanceState{}
// Populate PublicIpAddress from VM if stored
if instance.PublicIP != "" && instanceCopy.PublicIpAddress == nil {
instanceCopy.PublicIpAddress = aws.String(instance.PublicIP)
}
// Map internal status to EC2 state codes using the centralized mapping
if info, ok := vm.EC2StateCodes[instance.Status]; ok {
instanceCopy.State.SetCode(info.Code)
instanceCopy.State.SetName(info.Name)
} else {
slog.Warn("Instance has unmapped status, reporting as pending",
"instanceId", instance.ID, "status", string(instance.Status))
instanceCopy.State.SetCode(0)
instanceCopy.State.SetName("pending")
}
// Populate Placement if instance belongs to a placement group
if instance.PlacementGroupName != "" {
instanceCopy.Placement = &ec2.Placement{
GroupName: aws.String(instance.PlacementGroupName),
AvailabilityZone: aws.String(d.config.AZ),
}
}
// Apply filters against the fully-built instance copy
if len(parsedFilters) > 0 && !instanceMatchesFilters(instance, &instanceCopy, parsedFilters) {
// Iterate under the manager lock — VM fields (Status, Instance, Reservation,
// PublicIP, PlacementGroupName) are mutated through manager-locked
// Inspect/UpdateState elsewhere, so reading them lock-free would race.
d.vmMgr.View(func(vms map[string]*vm.VM) {
for _, instance := range vms {
// Skip instances not owned by the caller's account.
// Pre-Phase4 instances (empty AccountID) are only visible to root.
if !isInstanceVisible(accountID, instance.AccountID) {
continue
}
// Add instance to its reservation
reservationMap[resID].Instances = append(reservationMap[resID].Instances, &instanceCopy)
// Skip if filtering by instance IDs and this instance is not in the filter
if len(instanceIDFilter) > 0 && !instanceIDFilter[instance.ID] {
continue
}
// Use stored reservation metadata if available
if instance.Reservation != nil && instance.Instance != nil {
resID := ""
if instance.Reservation.ReservationId != nil {
resID = *instance.Reservation.ReservationId
}
// Create reservation entry if it doesn't exist
if _, exists := reservationMap[resID]; !exists {
reservation := &ec2.Reservation{}
reservation.SetReservationId(resID)
if instance.Reservation.OwnerId != nil {
reservation.SetOwnerId(*instance.Reservation.OwnerId)
}
reservation.Instances = []*ec2.Instance{}
reservationMap[resID] = reservation
}
// Update the instance state to current state
instanceCopy := *instance.Instance
instanceCopy.State = &ec2.InstanceState{}
// Populate PublicIpAddress from VM if stored
if instance.PublicIP != "" && instanceCopy.PublicIpAddress == nil {
instanceCopy.PublicIpAddress = aws.String(instance.PublicIP)
}
// Map internal status to EC2 state codes using the centralized mapping
if info, ok := vm.EC2StateCodes[instance.Status]; ok {
instanceCopy.State.SetCode(info.Code)
instanceCopy.State.SetName(info.Name)
} else {
slog.Warn("Instance has unmapped status, reporting as pending",
"instanceId", instance.ID, "status", string(instance.Status))
instanceCopy.State.SetCode(0)
instanceCopy.State.SetName("pending")
}
// Populate Placement if instance belongs to a placement group
if instance.PlacementGroupName != "" {
instanceCopy.Placement = &ec2.Placement{
GroupName: aws.String(instance.PlacementGroupName),
AvailabilityZone: aws.String(d.config.AZ),
}
}
// Apply filters against the fully-built instance copy
if len(parsedFilters) > 0 && !instanceMatchesFilters(instance, &instanceCopy, parsedFilters) {
continue
}
// Add instance to its reservation
reservationMap[resID].Instances = append(reservationMap[resID].Instances, &instanceCopy)
}
}
}
})
// Convert map to slice
reservations := make([]*ec2.Reservation, 0, len(reservationMap))
@@ -963,10 +947,8 @@ func (d *Daemon) handleEC2StartStoppedInstance(msg *nats.Msg) {
}
// Add instance to local map and clear stop attribute before launch
d.Instances.Mu.Lock()
d.Instances.VMS[instance.ID] = instance
instance.Attributes = types.EC2CommandAttributes{StartInstance: true}
d.Instances.Mu.Unlock()
d.vmMgr.Insert(instance)
// Launch the instance infrastructure (QEMU, QMP, NATS subscriptions)
err = d.LaunchInstance(instance)
@@ -976,9 +958,7 @@ func (d *Daemon) handleEC2StartStoppedInstance(msg *nats.Msg) {
if ok {
d.resourceMgr.deallocate(instanceType)
}
d.Instances.Mu.Lock()
delete(d.Instances.VMS, instance.ID)
d.Instances.Mu.Unlock()
d.vmMgr.Delete(instance.ID)
respondWithError(msg, awserrors.ErrorServerInternal)
return
}
@@ -1391,12 +1371,9 @@ func (d *Daemon) handleEC2DescribeInstanceAttribute(msg *nats.Msg) {
// Look up instance: running first, then stopped KV.
var instance *vm.VM
d.Instances.Mu.Lock()
if running, ok := d.Instances.VMS[instanceID]; ok {
if running, ok := d.vmMgr.Get(instanceID); ok {
instance = running
}
d.Instances.Mu.Unlock()
if instance == nil {
if d.jsManager == nil {
+50 -85
View File
@@ -510,7 +510,7 @@ func TestHandleEC2RunInstances_ServiceErrorPropagated(t *testing.T) {
// but RunInstance() will fail with ErrorInvalidInstanceType.
emptyTypes := map[string]*ec2.InstanceTypeInfo{}
daemon.instanceService = handlers_ec2_instance.NewInstanceServiceImpl(
daemon.config, emptyTypes, daemon.natsConn, &daemon.Instances,
daemon.config, emptyTypes, daemon.natsConn,
objectstore.NewMemoryObjectStore(),
)
@@ -542,14 +542,14 @@ func TestHandleEC2Events_StopInstance(t *testing.T) {
daemon := createFullTestDaemonWithJetStream(t, natsURL)
instanceID := "i-test-stop-001"
daemon.Instances.VMS[instanceID] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: instanceID,
InstanceType: getTestInstanceType(t),
Status: vm.StateRunning,
Instance: &ec2.Instance{},
QMPClient: &qmp.QMPClient{},
AccountID: testAccountID,
}
})
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.cmd.%s", instanceID),
@@ -576,9 +576,8 @@ func TestHandleEC2Events_StopInstance(t *testing.T) {
assert.Equal(t, `{}`, string(reply.Data))
// State should transition to stopping
daemon.Instances.Mu.Lock()
status := daemon.Instances.VMS[instanceID].Status
daemon.Instances.Mu.Unlock()
var status vm.InstanceState
daemon.vmMgr.UpdateState(instanceID, func(v *vm.VM) { status = v.Status })
assert.Equal(t, vm.StateStopping, status)
}
@@ -588,14 +587,14 @@ func TestHandleEC2Events_TerminateInstance(t *testing.T) {
daemon := createFullTestDaemonWithJetStream(t, natsURL)
instanceID := "i-test-term-001"
daemon.Instances.VMS[instanceID] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: instanceID,
InstanceType: getTestInstanceType(t),
Status: vm.StateRunning,
Instance: &ec2.Instance{},
QMPClient: &qmp.QMPClient{},
AccountID: testAccountID,
}
})
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.cmd.%s", instanceID),
@@ -620,9 +619,8 @@ func TestHandleEC2Events_TerminateInstance(t *testing.T) {
assert.Equal(t, `{}`, string(reply.Data))
daemon.Instances.Mu.Lock()
status := daemon.Instances.VMS[instanceID].Status
daemon.Instances.Mu.Unlock()
var status vm.InstanceState
daemon.vmMgr.UpdateState(instanceID, func(v *vm.VM) { status = v.Status })
assert.Equal(t, vm.StateShuttingDown, status)
}
@@ -632,14 +630,14 @@ func TestHandleEC2Events_RebootRunningInstance(t *testing.T) {
daemon := createTestDaemon(t, natsURL)
instanceID := "i-test-reboot-001"
daemon.Instances.VMS[instanceID] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: instanceID,
InstanceType: getTestInstanceType(t),
Status: vm.StateRunning,
Instance: &ec2.Instance{},
QMPClient: &qmp.QMPClient{},
AccountID: testAccountID,
}
})
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.cmd.%s", instanceID),
@@ -670,9 +668,8 @@ func TestHandleEC2Events_RebootRunningInstance(t *testing.T) {
assert.Contains(t, errResp, "Code")
// Instance should remain in running state (reboot doesn't change state)
daemon.Instances.Mu.Lock()
status := daemon.Instances.VMS[instanceID].Status
daemon.Instances.Mu.Unlock()
var status vm.InstanceState
daemon.vmMgr.UpdateState(instanceID, func(v *vm.VM) { status = v.Status })
assert.Equal(t, vm.StateRunning, status)
}
@@ -682,14 +679,14 @@ func TestHandleEC2Events_RebootStoppedInstance(t *testing.T) {
daemon := createTestDaemon(t, natsURL)
instanceID := "i-test-reboot-stopped"
daemon.Instances.VMS[instanceID] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: instanceID,
InstanceType: getTestInstanceType(t),
Status: vm.StateStopped,
Instance: &ec2.Instance{},
QMPClient: &qmp.QMPClient{},
AccountID: testAccountID,
}
})
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.cmd.%s", instanceID),
@@ -725,14 +722,14 @@ func TestHandleEC2Events_RebootTerminatedInstance(t *testing.T) {
daemon := createTestDaemon(t, natsURL)
instanceID := "i-test-reboot-terminated"
daemon.Instances.VMS[instanceID] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: instanceID,
InstanceType: getTestInstanceType(t),
Status: vm.StateTerminated,
Instance: &ec2.Instance{},
QMPClient: &qmp.QMPClient{},
AccountID: testAccountID,
}
})
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.cmd.%s", instanceID),
@@ -1055,8 +1052,7 @@ func TestHandleEC2CreateImage_InvalidState(t *testing.T) {
daemon := createFullTestDaemon(t, natsURL)
// Add an instance in "pending" state (not running or stopped)
daemon.Instances.Mu.Lock()
daemon.Instances.VMS["i-pending123"] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: "i-pending123",
Status: vm.StatePending,
AccountID: testAccountID,
@@ -1070,9 +1066,7 @@ func TestHandleEC2CreateImage_InvalidState(t *testing.T) {
},
},
},
}
daemon.Instances.Mu.Unlock()
})
sub, err := daemon.natsConn.Subscribe("ec2.CreateImage", daemon.handleEC2CreateImage)
require.NoError(t, err)
defer sub.Unsubscribe()
@@ -1097,8 +1091,7 @@ func TestHandleEC2CreateImage_NoRootVolume(t *testing.T) {
daemon := createFullTestDaemon(t, natsURL)
// Add instance with no block device mappings
daemon.Instances.Mu.Lock()
daemon.Instances.VMS["i-novol123"] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: "i-novol123",
Status: vm.StateRunning,
AccountID: testAccountID,
@@ -1107,9 +1100,7 @@ func TestHandleEC2CreateImage_NoRootVolume(t *testing.T) {
ImageId: aws.String("ami-source"),
BlockDeviceMappings: []*ec2.InstanceBlockDeviceMapping{},
},
}
daemon.Instances.Mu.Unlock()
})
sub, err := daemon.natsConn.Subscribe("ec2.CreateImage", daemon.handleEC2CreateImage)
require.NoError(t, err)
defer sub.Unsubscribe()
@@ -1479,17 +1470,14 @@ func TestHandleEC2GetConsoleOutput(t *testing.T) {
require.NoError(t, os.WriteFile(logPath, []byte("Hello from serial console\nBoot complete."), 0644))
// Add an instance with console log path
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instanceID] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: instanceID,
Status: vm.StateRunning,
AccountID: testAccountID,
Config: vm.Config{
ConsoleLogPath: logPath,
},
}
daemon.Instances.Mu.Unlock()
})
topic := fmt.Sprintf("ec2.%s.GetConsoleOutput", instanceID)
sub, err := daemon.natsConn.Subscribe(topic, daemon.handleEC2GetConsoleOutput)
require.NoError(t, err)
@@ -1523,17 +1511,14 @@ func TestHandleEC2GetConsoleOutput_EmptyLog(t *testing.T) {
instanceID := "i-console-empty-001"
// Instance exists but no log file yet
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instanceID] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: instanceID,
Status: vm.StateRunning,
AccountID: testAccountID,
Config: vm.Config{
ConsoleLogPath: "/nonexistent/console.log",
},
}
daemon.Instances.Mu.Unlock()
})
topic := fmt.Sprintf("ec2.%s.GetConsoleOutput", instanceID)
sub, err := daemon.natsConn.Subscribe(topic, daemon.handleEC2GetConsoleOutput)
require.NoError(t, err)
@@ -1596,7 +1581,7 @@ func TestAttachVolume_ZoneMismatch(t *testing.T) {
Instance: &ec2.Instance{},
QMPClient: &qmp.QMPClient{},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
// Create a volume in a different AZ
wrapper := struct {
@@ -1966,13 +1951,9 @@ func TestHandleEC2DescribeInstanceAttribute_RunningInstance_InstanceType(t *test
},
}
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instanceID] = instance
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(instance)
t.Cleanup(func() {
daemon.Instances.Mu.Lock()
delete(daemon.Instances.VMS, instanceID)
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Delete(instanceID)
})
input := &ec2.DescribeInstanceAttributeInput{
@@ -2170,13 +2151,9 @@ func TestHandleEC2DescribeInstanceAttribute_GroupSet_WithSecurityGroups(t *testi
},
}
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instanceID] = instance
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(instance)
t.Cleanup(func() {
daemon.Instances.Mu.Lock()
delete(daemon.Instances.VMS, instanceID)
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Delete(instanceID)
})
input := &ec2.DescribeInstanceAttributeInput{
@@ -2459,12 +2436,9 @@ func TestHandleNodeStatus(t *testing.T) {
daemon.config.Daemon.Host = "10.0.0.5:4432"
// Add some VMs (2 running, 1 stopped — only running counted)
daemon.Instances.Mu.Lock()
daemon.Instances.VMS["i-run-1"] = &vm.VM{ID: "i-run-1", Status: vm.StateRunning}
daemon.Instances.VMS["i-run-2"] = &vm.VM{ID: "i-run-2", Status: vm.StateRunning}
daemon.Instances.VMS["i-stop-1"] = &vm.VM{ID: "i-stop-1", Status: vm.StateStopped}
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(&vm.VM{ID: "i-run-1", Status: vm.StateRunning})
daemon.vmMgr.Insert(&vm.VM{ID: "i-run-2", Status: vm.StateRunning})
daemon.vmMgr.Insert(&vm.VM{ID: "i-stop-1", Status: vm.StateStopped})
sub, err := daemon.natsConn.Subscribe("spinifex.node.status.test", daemon.handleNodeStatus)
require.NoError(t, err)
defer sub.Unsubscribe()
@@ -2521,23 +2495,20 @@ func TestHandleNodeVMs(t *testing.T) {
instanceType := getTestInstanceType(t)
launchTime := time.Date(2025, 1, 15, 10, 0, 0, 0, time.UTC)
daemon.Instances.Mu.Lock()
daemon.Instances.VMS["i-vm-1"] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: "i-vm-1",
Status: vm.StateRunning,
InstanceType: instanceType,
Instance: &ec2.Instance{
LaunchTime: &launchTime,
},
}
daemon.Instances.VMS["i-vm-2"] = &vm.VM{
})
daemon.vmMgr.Insert(&vm.VM{
ID: "i-vm-2",
Status: vm.StateStopped,
InstanceType: instanceType,
Instance: nil, // no launch time
}
daemon.Instances.Mu.Unlock()
})
sub, err := daemon.natsConn.Subscribe("spinifex.node.vms.test", daemon.handleNodeVMs)
require.NoError(t, err)
defer sub.Unsubscribe()
@@ -2594,14 +2565,11 @@ func TestHandleNodeVMs_UnknownInstanceType(t *testing.T) {
daemon := createTestDaemon(t, sharedNATSURL)
daemon.config.Daemon.Host = "10.0.0.5:4432"
daemon.Instances.Mu.Lock()
daemon.Instances.VMS["i-vm-unknown"] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: "i-vm-unknown",
Status: vm.StateRunning,
InstanceType: "z99.mega", // not in instanceTypes map
}
daemon.Instances.Mu.Unlock()
})
sub, err := daemon.natsConn.Subscribe("spinifex.node.vms.unknown", daemon.handleNodeVMs)
require.NoError(t, err)
defer sub.Unsubscribe()
@@ -3223,8 +3191,7 @@ func TestHandleEC2CreateImage_RunningInstanceReachesService(t *testing.T) {
Body: strings.NewReader(string(volData)),
})
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instanceID] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: instanceID,
Status: vm.StateRunning,
InstanceType: getTestInstanceType(t),
@@ -3241,9 +3208,7 @@ func TestHandleEC2CreateImage_RunningInstanceReachesService(t *testing.T) {
},
},
},
}
daemon.Instances.Mu.Unlock()
})
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.%s.CreateImage", instanceID),
daemon.handleEC2CreateImage,
@@ -3282,7 +3247,7 @@ func TestAttachVolume_MissingVolumeData(t *testing.T) {
Instance: &ec2.Instance{},
QMPClient: &qmp.QMPClient{},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.cmd.%s", instanceID),
@@ -3322,7 +3287,7 @@ func TestAttachVolume_InstanceNotRunning(t *testing.T) {
Instance: &ec2.Instance{},
QMPClient: &qmp.QMPClient{},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.cmd.%s", instanceID),
@@ -3363,7 +3328,7 @@ func TestAttachVolume_VolumeNotFound(t *testing.T) {
Instance: &ec2.Instance{},
QMPClient: &qmp.QMPClient{},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.cmd.%s", instanceID),
@@ -3406,7 +3371,7 @@ func TestAttachVolume_VolumeInUse(t *testing.T) {
Instance: &ec2.Instance{},
QMPClient: &qmp.QMPClient{},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
// Seed a volume that is already in-use
wrapper := struct {
@@ -3469,7 +3434,7 @@ func TestDetachVolume_MissingVolumeData(t *testing.T) {
Instance: &ec2.Instance{},
QMPClient: &qmp.QMPClient{},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.cmd.%s", instanceID),
@@ -3508,7 +3473,7 @@ func TestDetachVolume_InstanceNotRunning(t *testing.T) {
Instance: &ec2.Instance{},
QMPClient: &qmp.QMPClient{},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.cmd.%s", instanceID),
@@ -3549,7 +3514,7 @@ func TestDetachVolume_VolumeNotAttached(t *testing.T) {
Instance: &ec2.Instance{},
QMPClient: &qmp.QMPClient{},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.cmd.%s", instanceID),
@@ -3593,7 +3558,7 @@ func TestDetachVolume_BootVolumeRejected(t *testing.T) {
instance.EBSRequests.Requests = []types.EBSRequest{
{Name: "vol-boot-001", Boot: true, DeviceName: "/dev/sda1"},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.cmd.%s", instanceID),
@@ -3637,7 +3602,7 @@ func TestDetachVolume_DeviceMismatch(t *testing.T) {
instance.EBSRequests.Requests = []types.EBSRequest{
{Name: "vol-mismatch-001", DeviceName: "/dev/sdf"},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.cmd.%s", instanceID),
+19 -21
View File
@@ -36,9 +36,8 @@ func (d *Daemon) handleAttachVolume(msg *nats.Msg, command types.EC2InstanceComm
device := command.AttachVolumeData.Device
// Validate instance is running
d.Instances.Mu.Lock()
status := instance.Status
d.Instances.Mu.Unlock()
var status vm.InstanceState
d.vmMgr.Inspect(instance, func(v *vm.VM) { status = v.Status })
if status != vm.StateRunning {
slog.Error("AttachVolume: instance not running", "instanceId", command.ID, "status", status)
@@ -82,9 +81,7 @@ func (d *Daemon) handleAttachVolume(msg *nats.Msg, command types.EC2InstanceComm
// Determine device name
if device == "" {
d.Instances.Mu.Lock()
device = nextAvailableDevice(instance)
d.Instances.Mu.Unlock()
d.vmMgr.Inspect(instance, func(v *vm.VM) { device = nextAvailableDevice(v) })
if device == "" {
slog.Error("AttachVolume: no available device names")
respondWithError(msg, awserrors.ErrorAttachmentLimitExceeded)
@@ -265,8 +262,10 @@ func (d *Daemon) handleAttachVolume(msg *nats.Msg, command types.EC2InstanceComm
instance.EBSRequests.Mu.Unlock()
// Update BlockDeviceMappings on the ec2.Instance using actual guest device name
d.Instances.Mu.Lock()
if instance.Instance != nil {
d.vmMgr.Inspect(instance, func(v *vm.VM) {
if v.Instance == nil {
return
}
now := time.Now()
mapping := &ec2.InstanceBlockDeviceMapping{}
mapping.SetDeviceName(guestDevice)
@@ -275,9 +274,8 @@ func (d *Daemon) handleAttachVolume(msg *nats.Msg, command types.EC2InstanceComm
mapping.Ebs.SetAttachTime(now)
mapping.Ebs.SetDeleteOnTermination(false)
mapping.Ebs.SetStatus("attached")
instance.Instance.BlockDeviceMappings = append(instance.Instance.BlockDeviceMappings, mapping)
}
d.Instances.Mu.Unlock()
v.Instance.BlockDeviceMappings = append(v.Instance.BlockDeviceMappings, mapping)
})
// Update volume metadata in S3
if err := d.volumeService.UpdateVolumeState(volumeID, "in-use", command.ID, guestDevice); err != nil {
@@ -313,9 +311,8 @@ func (d *Daemon) handleDetachVolume(msg *nats.Msg, command types.EC2InstanceComm
force := command.DetachVolumeData.Force
// Validate instance is running
d.Instances.Mu.Lock()
status := instance.Status
d.Instances.Mu.Unlock()
var status vm.InstanceState
d.vmMgr.Inspect(instance, func(v *vm.VM) { status = v.Status })
if status != vm.StateRunning {
slog.Error("DetachVolume: instance not running", "instanceId", command.ID, "status", status)
@@ -424,18 +421,19 @@ func (d *Daemon) handleDetachVolume(msg *nats.Msg, command types.EC2InstanceComm
instance.EBSRequests.Mu.Unlock()
// Remove from BlockDeviceMappings
d.Instances.Mu.Lock()
if instance.Instance != nil {
filtered := make([]*ec2.InstanceBlockDeviceMapping, 0, len(instance.Instance.BlockDeviceMappings))
for _, bdm := range instance.Instance.BlockDeviceMappings {
d.vmMgr.Inspect(instance, func(v *vm.VM) {
if v.Instance == nil {
return
}
filtered := make([]*ec2.InstanceBlockDeviceMapping, 0, len(v.Instance.BlockDeviceMappings))
for _, bdm := range v.Instance.BlockDeviceMappings {
if bdm.Ebs != nil && bdm.Ebs.VolumeId != nil && *bdm.Ebs.VolumeId == volumeID {
continue
}
filtered = append(filtered, bdm)
}
instance.Instance.BlockDeviceMappings = filtered
}
d.Instances.Mu.Unlock()
v.Instance.BlockDeviceMappings = filtered
})
// Update volume metadata to "available"
if err := d.volumeService.UpdateVolumeState(volumeID, "available", "", ""); err != nil {
+4 -16
View File
@@ -237,9 +237,7 @@ func (d *Daemon) LaunchSystemInstance(input *handlers_elbv2.SystemInstanceInput)
}
// Add to daemon state so LaunchInstance can find it
d.Instances.Mu.Lock()
d.Instances.VMS[instance.ID] = instance
d.Instances.Mu.Unlock()
d.vmMgr.Insert(instance)
if err := d.WriteState(); err != nil {
slog.Warn("LaunchSystemInstance: failed to write state", "instanceId", instance.ID, "err", err)
@@ -347,10 +345,7 @@ func (d *Daemon) LaunchSystemInstance(input *handlers_elbv2.SystemInstanceInput)
// TerminateSystemInstance stops and cleans up a system-managed VM.
func (d *Daemon) TerminateSystemInstance(instanceID string) error {
d.Instances.Mu.Lock()
instance, exists := d.Instances.VMS[instanceID]
d.Instances.Mu.Unlock()
instance, exists := d.vmMgr.Get(instanceID)
if !exists {
return fmt.Errorf("instance %s not found", instanceID)
}
@@ -404,9 +399,7 @@ func (d *Daemon) TerminateSystemInstance(instanceID string) error {
}
// Clean up local state
d.Instances.Mu.Lock()
delete(d.Instances.VMS, instanceID)
d.Instances.Mu.Unlock()
d.vmMgr.Delete(instanceID)
// Unsubscribe from per-instance NATS topic
d.mu.Lock()
@@ -467,13 +460,8 @@ func (d *Daemon) cleanupFailedSystemInstance(instance *vm.VM, instanceType *ec2.
func (d *Daemon) WaitForSystemInstance(instanceID string, timeout time.Duration) error {
deadline := time.Now().Add(timeout)
for time.Now().Before(deadline) {
d.Instances.Mu.Lock()
inst, exists := d.Instances.VMS[instanceID]
var status vm.InstanceState
if exists {
status = inst.Status
}
d.Instances.Mu.Unlock()
exists := d.vmMgr.UpdateState(instanceID, func(v *vm.VM) { status = v.Status })
if !exists {
return fmt.Errorf("instance %s disappeared", instanceID)
+14 -36
View File
@@ -7,14 +7,16 @@ import (
"github.com/mulgadc/spinifex/spinifex/vm"
)
func TestWaitForSystemInstance_AlreadyRunning(t *testing.T) {
d := &Daemon{
Instances: vm.Instances{
VMS: map[string]*vm.VM{
"i-test1": {ID: "i-test1", Status: vm.StateRunning},
},
},
func newDaemonWithVMs(vms ...*vm.VM) *Daemon {
d := &Daemon{vmMgr: vm.NewManager()}
for _, v := range vms {
d.vmMgr.Insert(v)
}
return d
}
func TestWaitForSystemInstance_AlreadyRunning(t *testing.T) {
d := newDaemonWithVMs(&vm.VM{ID: "i-test1", Status: vm.StateRunning})
err := d.WaitForSystemInstance("i-test1", 1*time.Second)
if err != nil {
@@ -23,11 +25,7 @@ func TestWaitForSystemInstance_AlreadyRunning(t *testing.T) {
}
func TestWaitForSystemInstance_NotFound(t *testing.T) {
d := &Daemon{
Instances: vm.Instances{
VMS: map[string]*vm.VM{},
},
}
d := newDaemonWithVMs()
err := d.WaitForSystemInstance("i-nonexistent", 500*time.Millisecond)
if err == nil {
@@ -36,13 +34,7 @@ func TestWaitForSystemInstance_NotFound(t *testing.T) {
}
func TestWaitForSystemInstance_ErrorState(t *testing.T) {
d := &Daemon{
Instances: vm.Instances{
VMS: map[string]*vm.VM{
"i-failed": {ID: "i-failed", Status: vm.StateError},
},
},
}
d := newDaemonWithVMs(&vm.VM{ID: "i-failed", Status: vm.StateError})
err := d.WaitForSystemInstance("i-failed", 1*time.Second)
if err == nil {
@@ -52,20 +44,12 @@ func TestWaitForSystemInstance_ErrorState(t *testing.T) {
func TestWaitForSystemInstance_TransitionsToRunning(t *testing.T) {
inst := &vm.VM{ID: "i-pending", Status: vm.StateProvisioning}
d := &Daemon{
Instances: vm.Instances{
VMS: map[string]*vm.VM{
"i-pending": inst,
},
},
}
d := newDaemonWithVMs(inst)
// Transition to running after a short delay
go func() {
time.Sleep(200 * time.Millisecond)
d.Instances.Mu.Lock()
inst.Status = vm.StateRunning
d.Instances.Mu.Unlock()
d.vmMgr.Inspect(inst, func(v *vm.VM) { v.Status = vm.StateRunning })
}()
err := d.WaitForSystemInstance("i-pending", 2*time.Second)
@@ -75,13 +59,7 @@ func TestWaitForSystemInstance_TransitionsToRunning(t *testing.T) {
}
func TestWaitForSystemInstance_Timeout(t *testing.T) {
d := &Daemon{
Instances: vm.Instances{
VMS: map[string]*vm.VM{
"i-stuck": {ID: "i-stuck", Status: vm.StateProvisioning},
},
},
}
d := newDaemonWithVMs(&vm.VM{ID: "i-stuck", Status: vm.StateProvisioning})
err := d.WaitForSystemInstance("i-stuck", 600*time.Millisecond)
if err == nil {
+25 -51
View File
@@ -101,7 +101,7 @@ func createTestDaemon(t *testing.T, natsURL string) *Daemon {
daemon.detachDelay = 0 // Skip sleep in tests
// Initialize services (needed for handler tests)
daemon.instanceService = handlers_ec2_instance.NewInstanceServiceImpl(cfg, daemon.resourceMgr.instanceTypes, nc, &daemon.Instances, objectstore.NewMemoryObjectStore())
daemon.instanceService = handlers_ec2_instance.NewInstanceServiceImpl(cfg, daemon.resourceMgr.instanceTypes, nc, objectstore.NewMemoryObjectStore())
daemon.volumeService = handlers_ec2_volume.NewVolumeServiceImplWithStore(cfg, objectstore.NewMemoryObjectStore(), nc)
t.Cleanup(func() {
@@ -393,7 +393,7 @@ func TestDaemon_Initialization(t *testing.T) {
assert.NotNil(t, daemon)
assert.NotNil(t, daemon.resourceMgr)
assert.NotNil(t, daemon.Instances.VMS)
assert.NotNil(t, daemon.vmMgr)
assert.Equal(t, cfg, daemon.config)
}
@@ -944,8 +944,7 @@ func TestDaemon_BootAllocation(t *testing.T) {
require.NoError(t, err)
// Pre-populate JetStream with test state
testInstances := &vm.Instances{VMS: vms}
err = daemon.jsManager.WriteState("node-1", testInstances)
err = daemon.jsManager.WriteState("node-1", vms)
require.NoError(t, err)
// Manually trigger the LoadState and allocation logic normally found in Start()
@@ -953,7 +952,7 @@ func TestDaemon_BootAllocation(t *testing.T) {
require.NoError(t, err)
// Simulate the allocation loop in Start()
for _, instance := range daemon.Instances.VMS {
for _, instance := range daemon.vmMgr.Snapshot() {
if instance.Status != vm.StateTerminated && !instance.Attributes.StopInstance {
instanceType, ok := daemon.resourceMgr.instanceTypes[instance.InstanceType]
if ok {
@@ -985,12 +984,12 @@ func TestStopInstance_Deallocation(t *testing.T) {
instanceId := "i-test-stop"
instanceTypeStr := getTestInstanceType(t)
instanceType := daemon.resourceMgr.instanceTypes[instanceTypeStr]
daemon.Instances.VMS[instanceId] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: instanceId,
InstanceType: instanceTypeStr,
Status: vm.StateRunning,
AccountID: testAccountID,
}
})
err = daemon.resourceMgr.allocate(instanceType)
require.NoError(t, err)
@@ -1188,13 +1187,13 @@ func TestDescribeInstances_ReservationGrouping(t *testing.T) {
ec2Instance.SetInstanceId(instanceID)
ec2Instance.SetInstanceType("t3.micro")
daemon.Instances.VMS[instanceID] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: instanceID,
Status: vm.StateRunning,
AccountID: testAccountID,
Reservation: reservation1,
Instance: ec2Instance,
}
})
}
// Create another reservation with 2 instances
@@ -1208,13 +1207,13 @@ func TestDescribeInstances_ReservationGrouping(t *testing.T) {
ec2Instance.SetInstanceId(instanceID)
ec2Instance.SetInstanceType("t3.small")
daemon.Instances.VMS[instanceID] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: instanceID,
Status: vm.StateRunning,
AccountID: testAccountID,
Reservation: reservation2,
Instance: ec2Instance,
}
})
}
// Create a single-instance reservation
@@ -1226,13 +1225,13 @@ func TestDescribeInstances_ReservationGrouping(t *testing.T) {
ec2Instance.SetInstanceId("i-single-001")
ec2Instance.SetInstanceType("t3.large")
daemon.Instances.VMS["i-single-001"] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: "i-single-001",
Status: vm.StateStopped,
AccountID: testAccountID,
Reservation: reservation3,
Instance: ec2Instance,
}
})
// Subscribe to handle DescribeInstances
sub, err := daemon.natsConn.Subscribe("ec2.DescribeInstances", daemon.handleEC2DescribeInstances)
@@ -1866,7 +1865,7 @@ func TestStopInstance_DeleteOnTermination_VolumeDeletion(t *testing.T) {
err = daemon.resourceMgr.allocate(instanceType)
require.NoError(t, err)
daemon.Instances.VMS[instance.ID] = instance
daemon.vmMgr.Insert(instance)
// Call stopInstance with deleteVolume=true (termination)
err = daemon.stopInstance(map[string]*vm.VM{instance.ID: instance}, true)
@@ -1963,7 +1962,7 @@ func TestStopInstance_DeleteOnTermination_False_SkipsVolumeDeletion(t *testing.T
err = daemon.resourceMgr.allocate(instanceType)
require.NoError(t, err)
daemon.Instances.VMS[instance.ID] = instance
daemon.vmMgr.Insert(instance)
// Call stopInstance with deleteVolume=true (termination)
err = daemon.stopInstance(map[string]*vm.VM{instance.ID: instance}, true)
@@ -2046,7 +2045,7 @@ func TestStopInstance_NoDelete_OnStop(t *testing.T) {
err = daemon.resourceMgr.allocate(instanceType)
require.NoError(t, err)
daemon.Instances.VMS[instance.ID] = instance
daemon.vmMgr.Insert(instance)
// Call stopInstance with deleteVolume=false (stop, not terminate)
err = daemon.stopInstance(map[string]*vm.VM{instance.ID: instance}, false)
@@ -2080,7 +2079,7 @@ func TestHandleEC2Events_AttachVolume(t *testing.T) {
Instance: &ec2.Instance{},
QMPClient: &qmp.QMPClient{}, // nil encoder/decoder
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
// Subscribe the handler to the instance's per-instance topic
sub, err := daemon.natsConn.Subscribe(
@@ -2113,10 +2112,7 @@ func TestHandleEC2Events_AttachVolume(t *testing.T) {
t.Run("InstanceNotRunning", func(t *testing.T) {
// Temporarily set status to stopped
daemon.Instances.Mu.Lock()
instance.Status = vm.StateStopped
daemon.Instances.Mu.Unlock()
command := types.EC2InstanceCommand{
ID: instanceID,
Attributes: types.EC2CommandAttributes{
@@ -2137,9 +2133,7 @@ func TestHandleEC2Events_AttachVolume(t *testing.T) {
assert.Contains(t, string(resp.Data), "IncorrectInstanceState")
// Restore running state
daemon.Instances.Mu.Lock()
instance.Status = vm.StateRunning
daemon.Instances.Mu.Unlock()
})
t.Run("VolumeNotFound", func(t *testing.T) {
@@ -2226,7 +2220,7 @@ func TestHandleEC2Events_DetachVolume(t *testing.T) {
},
},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
// Subscribe the handler to the instance's per-instance topic
sub, err := daemon.natsConn.Subscribe(
@@ -2257,10 +2251,7 @@ func TestHandleEC2Events_DetachVolume(t *testing.T) {
t.Run("InstanceNotRunning", func(t *testing.T) {
// Temporarily set status to stopped
daemon.Instances.Mu.Lock()
instance.Status = vm.StateStopped
daemon.Instances.Mu.Unlock()
command := types.EC2InstanceCommand{
ID: instanceID,
Attributes: types.EC2CommandAttributes{
@@ -2281,9 +2272,7 @@ func TestHandleEC2Events_DetachVolume(t *testing.T) {
assert.Contains(t, string(resp.Data), "IncorrectInstanceState")
// Restore running state
daemon.Instances.Mu.Lock()
instance.Status = vm.StateRunning
daemon.Instances.Mu.Unlock()
})
t.Run("VolumeNotAttached", func(t *testing.T) {
@@ -2575,7 +2564,7 @@ func TestDetachVolume_SuccessPath(t *testing.T) {
},
},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
// Subscribe a mock ebs.unmount handler
ebsUnmountCalled := make(chan string, 1)
@@ -2645,7 +2634,6 @@ func TestDetachVolume_SuccessPath(t *testing.T) {
instance.EBSRequests.Mu.Unlock()
// Verify volume removed from BlockDeviceMappings
daemon.Instances.Mu.Lock()
for _, bdm := range instance.Instance.BlockDeviceMappings {
if bdm.Ebs != nil && bdm.Ebs.VolumeId != nil {
assert.NotEqual(t, volumeID, *bdm.Ebs.VolumeId, "Volume should be removed from BlockDeviceMappings")
@@ -2654,7 +2642,6 @@ func TestDetachVolume_SuccessPath(t *testing.T) {
// Root volume should still be present
assert.Len(t, instance.Instance.BlockDeviceMappings, 1)
assert.Equal(t, "vol-root", *instance.Instance.BlockDeviceMappings[0].Ebs.VolumeId)
daemon.Instances.Mu.Unlock()
}
// TestDetachVolume_ForceFlag tests that force=true continues past device_del failure
@@ -2718,7 +2705,7 @@ func TestDetachVolume_ForceFlag(t *testing.T) {
},
},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
// Mock ebs.unmount
ebsSub, err := daemon.natsConn.Subscribe("ebs.node-1.unmount", func(msg *nats.Msg) {
@@ -2832,7 +2819,7 @@ func TestDetachVolume_BlockdevDelFailure(t *testing.T) {
},
},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.cmd.%s", instanceID),
@@ -2873,7 +2860,6 @@ func TestDetachVolume_BlockdevDelFailure(t *testing.T) {
assert.True(t, found, "Volume must remain in EBSRequests when blockdev-del fails")
// Critical: BlockDeviceMappings must NOT be cleaned up
daemon.Instances.Mu.Lock()
bdmFound := false
for _, bdm := range instance.Instance.BlockDeviceMappings {
if bdm.Ebs != nil && bdm.Ebs.VolumeId != nil && *bdm.Ebs.VolumeId == volumeID {
@@ -2881,7 +2867,6 @@ func TestDetachVolume_BlockdevDelFailure(t *testing.T) {
break
}
}
daemon.Instances.Mu.Unlock()
assert.True(t, bdmFound, "Volume must remain in BlockDeviceMappings when blockdev-del fails")
}
@@ -2924,7 +2909,7 @@ func TestDetachVolume_SuccessWithDeviceMatch(t *testing.T) {
},
},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
ebsSub, err := daemon.natsConn.Subscribe("ebs.node-1.unmount", func(msg *nats.Msg) {
resp := types.EBSUnMountResponse{Mounted: false}
@@ -3000,7 +2985,7 @@ func TestAttachVolume_ReplacesStaleEBSRequest(t *testing.T) {
},
},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
// Mock ebs.mount to return success with a new NBDURI
ebsSub, err := daemon.natsConn.Subscribe("ebs.node-1.mount", func(msg *nats.Msg) {
@@ -3284,13 +3269,10 @@ func TestMarkInstanceFailed(t *testing.T) {
Instance: ec2Instance,
QMPClient: &qmp.QMPClient{},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
daemon.markInstanceFailed(instance, "volume_preparation_failed")
daemon.Instances.Mu.Lock()
defer daemon.Instances.Mu.Unlock()
// Verify state transitioned to shutting-down
assert.Equal(t, vm.StateShuttingDown, instance.Status)
@@ -3316,13 +3298,11 @@ func TestMarkInstanceFailed_NilInstance(t *testing.T) {
Instance: nil, // no ec2.Instance
QMPClient: &qmp.QMPClient{},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
// Should not panic
daemon.markInstanceFailed(instance, "test_failure")
daemon.Instances.Mu.Lock()
defer daemon.Instances.Mu.Unlock()
assert.Equal(t, vm.StateShuttingDown, instance.Status)
}
@@ -3468,7 +3448,7 @@ func TestStopTerminate_IncorrectInstanceState(t *testing.T) {
AccountID: testAccountID,
Instance: &ec2.Instance{},
}
daemon.Instances.VMS[instanceID] = instance
daemon.vmMgr.Insert(instance)
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.cmd.%s", instanceID),
@@ -3478,10 +3458,7 @@ func TestStopTerminate_IncorrectInstanceState(t *testing.T) {
defer sub.Unsubscribe()
t.Run("StopAlreadyStoppedInstance", func(t *testing.T) {
daemon.Instances.Mu.Lock()
instance.Status = vm.StateStopped
daemon.Instances.Mu.Unlock()
command := types.EC2InstanceCommand{
ID: instanceID,
Attributes: types.EC2CommandAttributes{
@@ -3501,10 +3478,7 @@ func TestStopTerminate_IncorrectInstanceState(t *testing.T) {
})
t.Run("TerminateAlreadyTerminatedInstance", func(t *testing.T) {
daemon.Instances.Mu.Lock()
instance.Status = vm.StateTerminated
daemon.Instances.Mu.Unlock()
command := types.EC2InstanceCommand{
ID: instanceID,
Attributes: types.EC2CommandAttributes{
+6 -5
View File
@@ -177,9 +177,11 @@ func (d *Daemon) updateGuestDeviceNames(instance *vm.VM) {
instance.EBSRequests.Mu.Unlock()
// Update BlockDeviceMappings using the mapping
d.Instances.Mu.Lock()
if instance.Instance != nil {
for _, bdm := range instance.Instance.BlockDeviceMappings {
d.vmMgr.Inspect(instance, func(v *vm.VM) {
if v.Instance == nil {
return
}
for _, bdm := range v.Instance.BlockDeviceMappings {
if bdm.Ebs == nil || bdm.Ebs.VolumeId == nil || bdm.DeviceName == nil {
continue
}
@@ -187,8 +189,7 @@ func (d *Daemon) updateGuestDeviceNames(instance *vm.VM) {
bdm.DeviceName = &gd
}
}
}
d.Instances.Mu.Unlock()
})
if err := d.WriteState(); err != nil {
slog.Error("Failed to persist state after guest device name update", "instanceId", instance.ID, "err", err)
+52 -46
View File
@@ -78,9 +78,8 @@ func classifyCrashReason(waitErr error) string {
func (d *Daemon) handleInstanceCrash(instance *vm.VM, waitErr error) {
// Guard: if instance is not running, this was an expected exit
// (stopInstance/terminateInstance set status before QEMU exits)
d.Instances.Mu.Lock()
status := instance.Status
d.Instances.Mu.Unlock()
var status vm.InstanceState
d.vmMgr.Inspect(instance, func(v *vm.VM) { status = v.Status })
if status != vm.StateRunning {
slog.Debug("QEMU exited but instance not in running state, skipping crash handler",
@@ -106,16 +105,16 @@ func (d *Daemon) handleInstanceCrash(instance *vm.VM, waitErr error) {
// Update health tracking
now := time.Now()
d.Instances.Mu.Lock()
instance.Health.CrashCount++
instance.Health.LastCrashTime = now
instance.Health.LastCrashReason = reason
if instance.Health.FirstCrashTime.IsZero() {
instance.Health.FirstCrashTime = now
}
instance.Running = false
instance.PID = 0
d.Instances.Mu.Unlock()
d.vmMgr.Inspect(instance, func(v *vm.VM) {
v.Health.CrashCount++
v.Health.LastCrashTime = now
v.Health.LastCrashReason = reason
if v.Health.FirstCrashTime.IsZero() {
v.Health.FirstCrashTime = now
}
v.Running = false
v.PID = 0
})
// Deallocate resources to fix phantom reservation
instanceType, ok := d.resourceMgr.instanceTypes[instance.InstanceType]
@@ -184,31 +183,35 @@ func (d *Daemon) maybeRestartInstance(instance *vm.VM) {
now := time.Now()
d.Instances.Mu.Lock()
health := &instance.Health
// If crashes are outside the restart window, reset counters
if !health.FirstCrashTime.IsZero() && now.Sub(health.FirstCrashTime) > restartWindow {
slog.Info("Crash window expired, resetting counters", "instance", instance.ID)
health.CrashCount = 1
health.FirstCrashTime = now
health.RestartCount = 0
}
// Check if we've exceeded the max restarts in the window
if health.CrashCount > maxRestartsInWindow {
d.Instances.Mu.Unlock()
var (
crashCount int
restartCount int
exceeded bool
)
d.vmMgr.Inspect(instance, func(v *vm.VM) {
health := &v.Health
if !health.FirstCrashTime.IsZero() && now.Sub(health.FirstCrashTime) > restartWindow {
slog.Info("Crash window expired, resetting counters", "instance", v.ID)
health.CrashCount = 1
health.FirstCrashTime = now
health.RestartCount = 0
}
if health.CrashCount > maxRestartsInWindow {
exceeded = true
crashCount = health.CrashCount
return
}
restartCount = health.RestartCount
})
if exceeded {
slog.Error("Instance exceeded max restarts in window, leaving in error state",
"instance", instance.ID,
"crashes", health.CrashCount,
"crashes", crashCount,
"window", restartWindow,
"max", maxRestartsInWindow)
return
}
restartCount := health.RestartCount
d.Instances.Mu.Unlock()
// Check resource availability
instanceType, ok := d.resourceMgr.instanceTypes[instance.InstanceType]
if !ok || instanceType == nil {
@@ -238,26 +241,29 @@ func (d *Daemon) maybeRestartInstance(instance *vm.VM) {
// restartCrashedInstance re-verifies the instance is still in error state
// and relaunches it via LaunchInstance.
func (d *Daemon) restartCrashedInstance(instance *vm.VM) {
d.Instances.Mu.Lock()
if instance.Status != vm.StateError {
d.Instances.Mu.Unlock()
slog.Info("Instance no longer in error state, skipping restart",
"instance", instance.ID, "status", instance.Status)
var skipReason string
var restartCount int
d.vmMgr.Inspect(instance, func(v *vm.VM) {
if v.Status != vm.StateError {
skipReason = fmt.Sprintf("not in error state (%s)", v.Status)
return
}
if d.shuttingDown.Load() {
skipReason = "shutting down"
return
}
v.Health.RestartCount++
restartCount = v.Health.RestartCount
})
if skipReason != "" {
slog.Info("Skipping restart of crashed instance",
"instance", instance.ID, "reason", skipReason)
return
}
if d.shuttingDown.Load() {
d.Instances.Mu.Unlock()
slog.Info("Daemon shutting down, skipping restart", "instance", instance.ID)
return
}
instance.Health.RestartCount++
d.Instances.Mu.Unlock()
slog.Info("Restarting crashed instance",
"instance", instance.ID,
"restartCount", instance.Health.RestartCount)
"restartCount", restartCount)
// Re-allocate resources before relaunch — handleCrashedInstance deallocates
// them, so we must re-reserve before starting the VM again.
+21 -21
View File
@@ -57,14 +57,14 @@ func TestHandleInstanceCrash_SkipsNonRunning(t *testing.T) {
d := &Daemon{
natsConn: nc,
resourceMgr: rm,
Instances: vm.Instances{VMS: make(map[string]*vm.VM)},
vmMgr: vm.NewManager(),
}
instance := &vm.VM{
ID: "i-test-stopped",
Status: vm.StateStopped,
}
d.Instances.VMS[instance.ID] = instance
d.vmMgr.Insert(instance)
// Should return without action (no panic, no state change)
d.handleInstanceCrash(instance, fmt.Errorf("test error"))
@@ -82,7 +82,7 @@ func TestHandleInstanceCrash_SkipsShuttingDown(t *testing.T) {
d := &Daemon{
natsConn: nc,
resourceMgr: rm,
Instances: vm.Instances{VMS: make(map[string]*vm.VM)},
vmMgr: vm.NewManager(),
}
d.shuttingDown.Store(true)
@@ -90,7 +90,7 @@ func TestHandleInstanceCrash_SkipsShuttingDown(t *testing.T) {
ID: "i-test-shutdown",
Status: vm.StateRunning,
}
d.Instances.VMS[instance.ID] = instance
d.vmMgr.Insert(instance)
d.handleInstanceCrash(instance, fmt.Errorf("test error"))
@@ -113,7 +113,7 @@ func TestMaybeRestart_ExceedsMaxInWindow(t *testing.T) {
Services: []string{"daemon"},
},
resourceMgr: rm,
Instances: vm.Instances{VMS: make(map[string]*vm.VM)},
vmMgr: vm.NewManager(),
}
// Find an allocatable instance type
@@ -134,7 +134,7 @@ func TestMaybeRestart_ExceedsMaxInWindow(t *testing.T) {
RestartCount: maxRestartsInWindow,
},
}
d.Instances.VMS[instance.ID] = instance
d.vmMgr.Insert(instance)
// Should not schedule a restart (no panic, instance stays in error)
d.maybeRestartInstance(instance)
@@ -158,7 +158,7 @@ func TestMaybeRestart_ResetsAfterWindow(t *testing.T) {
Services: []string{"daemon"},
},
resourceMgr: rm,
Instances: vm.Instances{VMS: make(map[string]*vm.VM)},
vmMgr: vm.NewManager(),
}
// Find an allocatable instance type
@@ -179,7 +179,7 @@ func TestMaybeRestart_ResetsAfterWindow(t *testing.T) {
RestartCount: 5,
},
}
d.Instances.VMS[instance.ID] = instance
d.vmMgr.Insert(instance)
// maybeRestartInstance should reset counters and schedule a restart
// (restart will fail because there's no real QEMU, but counters should reset)
@@ -228,7 +228,7 @@ func newTestDaemon(t *testing.T) (*Daemon, func()) {
Services: []string{"daemon"},
},
resourceMgr: rm,
Instances: vm.Instances{VMS: make(map[string]*vm.VM)},
vmMgr: vm.NewManager(),
}
return d, func() { nc.Close() }
}
@@ -284,7 +284,7 @@ func TestHandleInstanceCrash_CoreFlow(t *testing.T) {
InstanceType: allocType,
Config: vm.Config{QMPSocket: qmpPath},
}
d.Instances.VMS[instance.ID] = instance
d.vmMgr.Insert(instance)
d.handleInstanceCrash(instance, fmt.Errorf("test crash"))
@@ -325,7 +325,7 @@ func TestHandleInstanceCrash_FirstCrashSetsTime(t *testing.T) {
PID: 111,
InstanceType: allocType,
}
d.Instances.VMS[instance.ID] = instance
d.vmMgr.Insert(instance)
// First crash
d.handleInstanceCrash(instance, fmt.Errorf("crash 1"))
@@ -334,11 +334,11 @@ func TestHandleInstanceCrash_FirstCrashSetsTime(t *testing.T) {
assert.Equal(t, 1, instance.Health.CrashCount)
// Reset to running for a second crash
d.Instances.Mu.Lock()
instance.Status = vm.StateRunning
instance.Running = true
instance.PID = 222
d.Instances.Mu.Unlock()
d.vmMgr.Inspect(instance, func(v *vm.VM) {
v.Status = vm.StateRunning
v.Running = true
v.PID = 222
})
time.Sleep(time.Millisecond) // ensure time advances
d.handleInstanceCrash(instance, fmt.Errorf("crash 2"))
@@ -359,7 +359,7 @@ func TestHandleInstanceCrash_UnknownInstanceType(t *testing.T) {
PID: 333,
InstanceType: "z99.nonexistent",
}
d.Instances.VMS[instance.ID] = instance
d.vmMgr.Insert(instance)
// Should not panic even though instance type is not in resourceMgr
d.handleInstanceCrash(instance, fmt.Errorf("crash"))
@@ -386,7 +386,7 @@ func TestMaybeRestart_SkipsShuttingDown(t *testing.T) {
FirstCrashTime: time.Now(),
},
}
d.Instances.VMS[instance.ID] = instance
d.vmMgr.Insert(instance)
d.maybeRestartInstance(instance)
@@ -407,7 +407,7 @@ func TestMaybeRestart_UnknownInstanceType(t *testing.T) {
FirstCrashTime: time.Now(),
},
}
d.Instances.VMS[instance.ID] = instance
d.vmMgr.Insert(instance)
// Should not panic — just log and return
d.maybeRestartInstance(instance)
@@ -438,7 +438,7 @@ func TestMaybeRestart_InsufficientResources(t *testing.T) {
FirstCrashTime: time.Now(),
},
}
d.Instances.VMS[instance.ID] = instance
d.vmMgr.Insert(instance)
d.maybeRestartInstance(instance)
@@ -462,7 +462,7 @@ func TestMaybeRestart_SchedulesRestart(t *testing.T) {
RestartCount: 0,
},
}
d.Instances.VMS[instance.ID] = instance
d.vmMgr.Insert(instance)
// maybeRestartInstance should pass all guards and call time.AfterFunc.
// The scheduled restartCrashedInstance will increment RestartCount and
+1 -3
View File
@@ -54,9 +54,7 @@ func (d *Daemon) publishHeartbeat() {
func (d *Daemon) buildHeartbeat() *Heartbeat {
totalVCPU, totalMem, reservedVCPU, reservedMem, allocVCPU, allocMem, _ := d.resourceMgr.GetResourceStats()
d.Instances.Mu.Lock()
vmCount := len(d.Instances.VMS)
d.Instances.Mu.Unlock()
vmCount := d.vmMgr.Count()
return &Heartbeat{
Node: d.node,
+4 -6
View File
@@ -25,7 +25,7 @@ func TestBuildHeartbeat(t *testing.T) {
Services: []string{"daemon", "nats", "viperblock"},
},
resourceMgr: rm,
Instances: vm.Instances{VMS: make(map[string]*vm.VM)},
vmMgr: vm.NewManager(),
}
h := d.buildHeartbeat()
@@ -58,7 +58,7 @@ func TestHeartbeatReflectsAllocation(t *testing.T) {
Services: []string{"daemon"},
},
resourceMgr: rm,
Instances: vm.Instances{VMS: make(map[string]*vm.VM)},
vmMgr: vm.NewManager(),
}
// Take a heartbeat before allocation
@@ -79,13 +79,11 @@ func TestHeartbeatReflectsAllocation(t *testing.T) {
require.NoError(t, err)
// Add a VM to the instance map
d.Instances.Mu.Lock()
d.Instances.VMS["i-test-001"] = &vm.VM{
d.vmMgr.Insert(&vm.VM{
ID: "i-test-001",
Status: vm.StateRunning,
InstanceType: allocType,
}
d.Instances.Mu.Unlock()
})
// Take a heartbeat after allocation
after := d.buildHeartbeat()
+18 -20
View File
@@ -390,11 +390,8 @@ func (m *JetStreamManager) WriteServiceManifest(nodeID string, services []string
}
// WriteState writes the instance state to the KV store for the given node.
// It acquires instances.Mu internally.
func (m *JetStreamManager) WriteState(nodeID string, instances *vm.Instances) error {
instances.Mu.Lock()
defer instances.Mu.Unlock()
// vms must be a snapshot owned by the caller — JetStreamManager does not lock.
func (m *JetStreamManager) WriteState(nodeID string, vms map[string]*vm.VM) error {
if m.kv == nil {
return errors.New("KV bucket not initialized")
}
@@ -403,7 +400,7 @@ func (m *JetStreamManager) WriteState(nodeID string, instances *vm.Instances) er
state := struct {
VMS map[string]*vm.VM `json:"vms"`
}{
VMS: instances.VMS,
VMS: vms,
}
jsonData, err := json.Marshal(state)
@@ -423,18 +420,19 @@ func (m *JetStreamManager) WriteState(nodeID string, instances *vm.Instances) er
if _, retryErr := kv.Put(key, jsonData); retryErr != nil {
return retryErr
}
slog.Debug("Wrote state to JetStream KV (after recovery)", "key", key, "instances", len(instances.VMS))
slog.Debug("Wrote state to JetStream KV (after recovery)", "key", key, "instances", len(vms))
return nil
}
return err
}
slog.Debug("Wrote state to JetStream KV", "key", key, "instances", len(instances.VMS))
slog.Debug("Wrote state to JetStream KV", "key", key, "instances", len(vms))
return nil
}
// LoadState loads the instance state from the KV store for the given node
func (m *JetStreamManager) LoadState(nodeID string) (*vm.Instances, error) {
// LoadState loads the instance state from the KV store for the given node.
// Returns an empty (non-nil) map when no state exists for the node.
func (m *JetStreamManager) LoadState(nodeID string) (map[string]*vm.VM, error) {
if m.kv == nil {
return nil, errors.New("KV bucket not initialized")
}
@@ -444,7 +442,7 @@ func (m *JetStreamManager) LoadState(nodeID string) (*vm.Instances, error) {
if err != nil {
if errors.Is(err, nats.ErrKeyNotFound) {
slog.Debug("No existing state in JetStream KV, returning empty state", "key", key)
return &vm.Instances{VMS: make(map[string]*vm.VM)}, nil
return make(map[string]*vm.VM), nil
}
if isStreamUnavailable(err) {
slog.Warn("KV stream unavailable, attempting recovery", "operation", "LoadState", "key", key, "err", err)
@@ -457,7 +455,7 @@ func (m *JetStreamManager) LoadState(nodeID string) (*vm.Instances, error) {
if err != nil {
if errors.Is(err, nats.ErrKeyNotFound) {
slog.Warn("No state found after KV recovery", "key", key)
return &vm.Instances{VMS: make(map[string]*vm.VM)}, nil
return make(map[string]*vm.VM), nil
}
return nil, err
}
@@ -467,18 +465,18 @@ func (m *JetStreamManager) LoadState(nodeID string) (*vm.Instances, error) {
}
}
var instances vm.Instances
if err := json.Unmarshal(entry.Value(), &instances); err != nil {
var state struct {
VMS map[string]*vm.VM `json:"vms"`
}
if err := json.Unmarshal(entry.Value(), &state); err != nil {
return nil, err
}
// Ensure the VMS map is initialized
if instances.VMS == nil {
instances.VMS = make(map[string]*vm.VM)
if state.VMS == nil {
state.VMS = make(map[string]*vm.VM)
}
slog.Debug("Loaded state from JetStream KV", "key", key, "instances", len(instances.VMS))
return &instances, nil
slog.Debug("Loaded state from JetStream KV", "key", key, "instances", len(state.VMS))
return state.VMS, nil
}
// DeleteState removes the instance state from the KV store for the given node
+68 -86
View File
@@ -31,18 +31,16 @@ func TestJetStreamManager_WriteAndLoadState(t *testing.T) {
// Create test instances
testNodeID := "test-node-1"
testInstances := &vm.Instances{
VMS: map[string]*vm.VM{
"i-test-001": {
ID: "i-test-001",
Status: vm.StateRunning,
InstanceType: "t3.micro",
},
"i-test-002": {
ID: "i-test-002",
Status: vm.StateStopped,
InstanceType: "t3.small",
},
testInstances := map[string]*vm.VM{
"i-test-001": {
ID: "i-test-001",
Status: vm.StateRunning,
InstanceType: "t3.micro",
},
"i-test-002": {
ID: "i-test-002",
Status: vm.StateStopped,
InstanceType: "t3.small",
},
}
@@ -56,13 +54,13 @@ func TestJetStreamManager_WriteAndLoadState(t *testing.T) {
require.NotNil(t, loadedInstances, "Loaded instances should not be nil")
// Verify the loaded state matches
assert.Len(t, loadedInstances.VMS, 2, "Should have 2 instances")
assert.NotNil(t, loadedInstances.VMS["i-test-001"], "Should have i-test-001")
assert.NotNil(t, loadedInstances.VMS["i-test-002"], "Should have i-test-002")
assert.Equal(t, vm.StateRunning, loadedInstances.VMS["i-test-001"].Status)
assert.Equal(t, vm.StateStopped, loadedInstances.VMS["i-test-002"].Status)
assert.Equal(t, "t3.micro", loadedInstances.VMS["i-test-001"].InstanceType)
assert.Equal(t, "t3.small", loadedInstances.VMS["i-test-002"].InstanceType)
assert.Len(t, loadedInstances, 2, "Should have 2 instances")
assert.NotNil(t, loadedInstances["i-test-001"], "Should have i-test-001")
assert.NotNil(t, loadedInstances["i-test-002"], "Should have i-test-002")
assert.Equal(t, vm.StateRunning, loadedInstances["i-test-001"].Status)
assert.Equal(t, vm.StateStopped, loadedInstances["i-test-002"].Status)
assert.Equal(t, "t3.micro", loadedInstances["i-test-001"].InstanceType)
assert.Equal(t, "t3.small", loadedInstances["i-test-002"].InstanceType)
}
// TestJetStreamManager_LoadState_KeyNotFound tests that LoadState returns empty state when key doesn't exist
@@ -83,7 +81,7 @@ func TestJetStreamManager_LoadState_KeyNotFound(t *testing.T) {
instances, err := jsm.LoadState("non-existent-node")
require.NoError(t, err, "Should not error when key not found")
require.NotNil(t, instances, "Should return non-nil instances")
assert.Empty(t, instances.VMS, "Should return empty VMS map")
assert.Empty(t, instances, "Should return empty VMS map")
}
// TestJetStreamManager_BucketCreation tests that InitKVBucket creates the bucket when it doesn't exist
@@ -120,12 +118,10 @@ func TestJetStreamManager_BucketReconnection(t *testing.T) {
require.NoError(t, err, "First InitKVBucket should succeed")
// Write some test data
testInstances := &vm.Instances{
VMS: map[string]*vm.VM{
"i-persist": {
ID: "i-persist",
Status: vm.StateRunning,
},
testInstances := map[string]*vm.VM{
"i-persist": {
ID: "i-persist",
Status: vm.StateRunning,
},
}
err = jsm1.WriteState("persist-node", testInstances)
@@ -147,9 +143,9 @@ func TestJetStreamManager_BucketReconnection(t *testing.T) {
// Verify data persisted
loadedInstances, err := jsm2.LoadState("persist-node")
require.NoError(t, err)
assert.NotEmpty(t, loadedInstances.VMS, "Should have persisted instances")
assert.NotNil(t, loadedInstances.VMS["i-persist"], "Should have i-persist")
assert.Equal(t, vm.StateRunning, loadedInstances.VMS["i-persist"].Status)
assert.NotEmpty(t, loadedInstances, "Should have persisted instances")
assert.NotNil(t, loadedInstances["i-persist"], "Should have i-persist")
assert.Equal(t, vm.StateRunning, loadedInstances["i-persist"].Status)
}
// TestJetStreamManager_DeleteState tests deleting state from the KV store
@@ -168,12 +164,10 @@ func TestJetStreamManager_DeleteState(t *testing.T) {
// Write state
testNodeID := "delete-test-node"
testInstances := &vm.Instances{
VMS: map[string]*vm.VM{
"i-delete-me": {
ID: "i-delete-me",
Status: vm.StateRunning,
},
testInstances := map[string]*vm.VM{
"i-delete-me": {
ID: "i-delete-me",
Status: vm.StateRunning,
},
}
err = jsm.WriteState(testNodeID, testInstances)
@@ -182,7 +176,7 @@ func TestJetStreamManager_DeleteState(t *testing.T) {
// Verify state exists
loadedInstances, err := jsm.LoadState(testNodeID)
require.NoError(t, err)
assert.NotEmpty(t, loadedInstances.VMS)
assert.NotEmpty(t, loadedInstances)
// Delete state
err = jsm.DeleteState(testNodeID)
@@ -191,7 +185,7 @@ func TestJetStreamManager_DeleteState(t *testing.T) {
// Verify state is gone (should return empty state)
loadedInstances, err = jsm.LoadState(testNodeID)
require.NoError(t, err)
assert.Empty(t, loadedInstances.VMS, "Should return empty state after deletion")
assert.Empty(t, loadedInstances, "Should return empty state after deletion")
}
// TestJetStreamManager_DeleteState_NonExistent tests deleting state that doesn't exist
@@ -230,28 +224,24 @@ func TestJetStreamManager_WriteState_UpdateExisting(t *testing.T) {
testNodeID := "update-test-node"
// Write initial state
initialInstances := &vm.Instances{
VMS: map[string]*vm.VM{
"i-initial": {
ID: "i-initial",
Status: vm.StateRunning,
},
initialInstances := map[string]*vm.VM{
"i-initial": {
ID: "i-initial",
Status: vm.StateRunning,
},
}
err = jsm.WriteState(testNodeID, initialInstances)
require.NoError(t, err)
// Update state with different instances
updatedInstances := &vm.Instances{
VMS: map[string]*vm.VM{
"i-initial": {
ID: "i-initial",
Status: vm.StateStopped, // Changed status
},
"i-new": { // Added new instance
ID: "i-new",
Status: vm.StateRunning,
},
updatedInstances := map[string]*vm.VM{
"i-initial": {
ID: "i-initial",
Status: vm.StateStopped, // Changed status
},
"i-new": { // Added new instance
ID: "i-new",
Status: vm.StateRunning,
},
}
err = jsm.WriteState(testNodeID, updatedInstances)
@@ -260,9 +250,9 @@ func TestJetStreamManager_WriteState_UpdateExisting(t *testing.T) {
// Load and verify updated state
loadedInstances, err := jsm.LoadState(testNodeID)
require.NoError(t, err)
assert.Len(t, loadedInstances.VMS, 2, "Should have 2 instances")
assert.Equal(t, vm.StateStopped, loadedInstances.VMS["i-initial"].Status, "Status should be updated")
assert.NotNil(t, loadedInstances.VMS["i-new"], "Should have new instance")
assert.Len(t, loadedInstances, 2, "Should have 2 instances")
assert.Equal(t, vm.StateStopped, loadedInstances["i-initial"].Status, "Status should be updated")
assert.NotNil(t, loadedInstances["i-new"], "Should have new instance")
}
// TestJetStreamManager_MultipleNodes tests storing state for multiple nodes
@@ -280,20 +270,16 @@ func TestJetStreamManager_MultipleNodes(t *testing.T) {
require.NoError(t, err)
// Write state for node-1
node1Instances := &vm.Instances{
VMS: map[string]*vm.VM{
"i-node1-001": {ID: "i-node1-001", Status: vm.StateRunning},
},
node1Instances := map[string]*vm.VM{
"i-node1-001": {ID: "i-node1-001", Status: vm.StateRunning},
}
err = jsm.WriteState("node-1", node1Instances)
require.NoError(t, err)
// Write state for node-2
node2Instances := &vm.Instances{
VMS: map[string]*vm.VM{
"i-node2-001": {ID: "i-node2-001", Status: vm.StateStopped},
"i-node2-002": {ID: "i-node2-002", Status: vm.StateRunning},
},
node2Instances := map[string]*vm.VM{
"i-node2-001": {ID: "i-node2-001", Status: vm.StateStopped},
"i-node2-002": {ID: "i-node2-002", Status: vm.StateRunning},
}
err = jsm.WriteState("node-2", node2Instances)
require.NoError(t, err)
@@ -301,18 +287,18 @@ func TestJetStreamManager_MultipleNodes(t *testing.T) {
// Load and verify node-1 state
loadedNode1, err := jsm.LoadState("node-1")
require.NoError(t, err)
assert.Len(t, loadedNode1.VMS, 1)
assert.NotNil(t, loadedNode1.VMS["i-node1-001"])
assert.Len(t, loadedNode1, 1)
assert.NotNil(t, loadedNode1["i-node1-001"])
// Load and verify node-2 state
loadedNode2, err := jsm.LoadState("node-2")
require.NoError(t, err)
assert.Len(t, loadedNode2.VMS, 2)
assert.NotNil(t, loadedNode2.VMS["i-node2-001"])
assert.NotNil(t, loadedNode2.VMS["i-node2-002"])
assert.Len(t, loadedNode2, 2)
assert.NotNil(t, loadedNode2["i-node2-001"])
assert.NotNil(t, loadedNode2["i-node2-002"])
// Verify node isolation - node-1 doesn't have node-2's instances
_, exists := loadedNode1.VMS["i-node2-001"]
_, exists := loadedNode1["i-node2-001"]
assert.False(t, exists, "Node-1 should not have node-2's instances")
}
@@ -328,7 +314,7 @@ func TestJetStreamManager_KVNotInitialized(t *testing.T) {
jsm, err := NewJetStreamManager(nc, 1)
require.NoError(t, err)
testInstances := &vm.Instances{VMS: make(map[string]*vm.VM)}
testInstances := make(map[string]*vm.VM)
err = jsm.WriteState("test-node", testInstances)
assert.Error(t, err, "WriteState should error when KV not initialized")
@@ -522,10 +508,8 @@ func TestJetStreamManager_StoppedInstances_NoInterference(t *testing.T) {
require.NoError(t, err)
// Write per-node state
nodeInstances := &vm.Instances{
VMS: map[string]*vm.VM{
"i-running-001": {ID: "i-running-001", Status: vm.StateRunning},
},
nodeInstances := map[string]*vm.VM{
"i-running-001": {ID: "i-running-001", Status: vm.StateRunning},
}
err = jsm.WriteState("interference-test-node", nodeInstances)
require.NoError(t, err)
@@ -538,8 +522,8 @@ func TestJetStreamManager_StoppedInstances_NoInterference(t *testing.T) {
// Verify per-node state is unaffected
loaded, err := jsm.LoadState("interference-test-node")
require.NoError(t, err)
assert.Len(t, loaded.VMS, 1)
assert.NotNil(t, loaded.VMS["i-running-001"])
assert.Len(t, loaded, 1)
assert.NotNil(t, loaded["i-running-001"])
// Verify stopped instance loads independently
loadedStopped, err := jsm.LoadStoppedInstance(stoppedVM.ID)
@@ -684,10 +668,8 @@ func TestJetStreamManager_WriteState_RecoverAfterStreamLost(t *testing.T) {
deleteInstanceStateBucket(t, nc)
// WriteState should recover by recreating the bucket
testInstances := &vm.Instances{
VMS: map[string]*vm.VM{
"i-recover-001": {ID: "i-recover-001", Status: vm.StateRunning},
},
testInstances := map[string]*vm.VM{
"i-recover-001": {ID: "i-recover-001", Status: vm.StateRunning},
}
err = jsm.WriteState("recovery-node", testInstances)
require.NoError(t, err, "WriteState should recover after stream loss")
@@ -695,8 +677,8 @@ func TestJetStreamManager_WriteState_RecoverAfterStreamLost(t *testing.T) {
// Verify data was written
loaded, err := jsm.LoadState("recovery-node")
require.NoError(t, err)
assert.Len(t, loaded.VMS, 1)
assert.NotNil(t, loaded.VMS["i-recover-001"])
assert.Len(t, loaded, 1)
assert.NotNil(t, loaded["i-recover-001"])
}
func TestJetStreamManager_LoadState_RecoverAfterStreamLost(t *testing.T) {
@@ -714,7 +696,7 @@ func TestJetStreamManager_LoadState_RecoverAfterStreamLost(t *testing.T) {
// LoadState should recover and return empty state
loaded, err := jsm.LoadState("any-node")
require.NoError(t, err, "LoadState should recover after stream loss")
assert.Empty(t, loaded.VMS)
assert.Empty(t, loaded)
}
func TestJetStreamManager_DeleteState_RecoverAfterStreamLost(t *testing.T) {
@@ -843,7 +825,7 @@ func TestJetStreamManager_WriteState_RecoveryFailure(t *testing.T) {
deleteInstanceStateBucket(t, nc)
swapToNonJSContext(t, jsm)
testInstances := &vm.Instances{VMS: map[string]*vm.VM{"i-fail": {ID: "i-fail"}}}
testInstances := map[string]*vm.VM{"i-fail": {ID: "i-fail"}}
err = jsm.WriteState("fail-node", testInstances)
assert.Error(t, err, "WriteState should return error when recovery fails")
}
+5 -8
View File
@@ -112,20 +112,17 @@ func (d *Daemon) handleShutdownDrain(msg *nats.Msg) {
slog.Info("Shutdown DRAIN phase starting", "node", d.node)
// Count total VMs
d.Instances.Mu.Lock()
total := len(d.Instances.VMS)
d.Instances.Mu.Unlock()
// Snapshot VMs to stop. The snapshot decouples iteration from concurrent
// terminate handlers and avoids holding the manager lock across the
// shutdown fan-out.
vms := d.vmMgr.SnapshotMap()
total := len(vms)
// Publish initial progress
d.publishShutdownProgress("drain", total, total)
// Stop all instances (graceful shutdown, no volume deletion)
if total > 0 {
d.Instances.Mu.Lock()
vms := d.Instances.VMS
d.Instances.Mu.Unlock()
if err := d.stopInstance(vms, false); err != nil {
slog.Error("Failed to stop instances during DRAIN", "error", err)
ack := ShutdownACK{
+13 -7
View File
@@ -12,16 +12,22 @@ import (
// On validation failure, the instance status is unchanged and an error is returned.
// If WriteState fails, the in-memory status retains the new value (the VM has
// physically changed state regardless) and an error is returned.
// The caller must NOT hold d.Instances.Mu; this method acquires it internally.
func (d *Daemon) TransitionState(instance *vm.VM, target vm.InstanceState) error {
d.Instances.Mu.Lock()
current := instance.Status
if !vm.IsValidTransition(current, target) {
d.Instances.Mu.Unlock()
var (
current vm.InstanceState
invalid bool
)
d.vmMgr.Inspect(instance, func(v *vm.VM) {
current = v.Status
if !vm.IsValidTransition(current, target) {
invalid = true
return
}
v.Status = target
})
if invalid {
return fmt.Errorf("invalid state transition: %s -> %s for instance %s", current, target, instance.ID)
}
instance.Status = target
d.Instances.Mu.Unlock()
slog.Info("Instance state transition", "instanceId", instance.ID, "from", string(current), "to", string(target))
+92 -128
View File
@@ -80,9 +80,7 @@ func TestTransitionState_ValidTransitions(t *testing.T) {
Status: tt.from,
}
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instance.ID] = instance
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(instance)
err := daemon.TransitionState(instance, tt.to)
require.NoError(t, err)
@@ -120,9 +118,7 @@ func TestTransitionState_InvalidTransitions(t *testing.T) {
Status: tt.from,
}
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instance.ID] = instance
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(instance)
err := daemon.TransitionState(instance, tt.to)
assert.Error(t, err)
@@ -143,9 +139,7 @@ func TestTransitionState_NilEC2Instance(t *testing.T) {
Instance: nil, // no EC2 instance metadata
}
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instance.ID] = instance
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(instance)
err := daemon.TransitionState(instance, vm.StateRunning)
require.NoError(t, err)
@@ -265,9 +259,7 @@ func TestStopLifecycle_CompletesFullChain(t *testing.T) {
daemon := createDaemonWithJetStream(t)
instance := &vm.VM{ID: "i-lifecycle-stop", Status: vm.StateProvisioning}
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instance.ID] = instance
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(instance)
chain := []vm.InstanceState{
vm.StateRunning, // launch
@@ -292,9 +284,7 @@ func TestStopLifecycle_ErrorRecovery(t *testing.T) {
daemon := createDaemonWithJetStream(t)
instance := &vm.VM{ID: "i-lifecycle-err", Status: vm.StateRunning}
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instance.ID] = instance
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(instance)
chain := []vm.InstanceState{
vm.StateStopping,
@@ -316,9 +306,7 @@ func TestShuttingDown_ErrorRecovery(t *testing.T) {
daemon := createDaemonWithJetStream(t)
instance := &vm.VM{ID: "i-shutdown-err", Status: vm.StateShuttingDown}
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instance.ID] = instance
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(instance)
chain := []vm.InstanceState{
vm.StateError, // terminate failed
@@ -350,9 +338,7 @@ func TestLaunchInstance_AbortedByConcurrentTerminate(t *testing.T) {
ID: "i-abort-" + string(status),
Status: status,
}
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instance.ID] = instance
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(instance)
err := daemon.LaunchInstance(instance)
require.NoError(t, err, "expected nil for non-launchable status %s", status)
@@ -368,9 +354,7 @@ func TestTransitionState_ConcurrentTransitions(t *testing.T) {
daemon := createDaemonWithJetStream(t)
instance := &vm.VM{ID: "i-concurrent", Status: vm.StateRunning}
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instance.ID] = instance
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(instance)
// Multiple goroutines race to transition from running.
// After one succeeds, the others may chain from the new state
@@ -430,9 +414,7 @@ func TestTransitionState_DoubleTransitionSameTarget(t *testing.T) {
for _, state := range allStates {
t.Run(string(state), func(t *testing.T) {
instance := &vm.VM{ID: fmt.Sprintf("i-double-%s", state), Status: state}
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instance.ID] = instance
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(instance)
err := daemon.TransitionState(instance, state)
assert.Error(t, err, "self-transition %s -> %s should be rejected", state, state)
@@ -449,9 +431,7 @@ func TestTransitionState_InvalidFromTerminated(t *testing.T) {
for _, target := range allStates {
t.Run(fmt.Sprintf("terminated->%s", target), func(t *testing.T) {
instance := &vm.VM{ID: fmt.Sprintf("i-term-%s", target), Status: vm.StateTerminated}
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instance.ID] = instance
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(instance)
err := daemon.TransitionState(instance, target)
assert.Error(t, err, "transition from terminated to %s should be rejected", target)
@@ -468,9 +448,7 @@ func TestTransitionState_StoppingMustReachStoppedOrError(t *testing.T) {
t.Run("happy_path", func(t *testing.T) {
instance := &vm.VM{ID: "i-stop-happy", Status: vm.StateStopping}
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instance.ID] = instance
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(instance)
err := daemon.TransitionState(instance, vm.StateStopped)
require.NoError(t, err)
@@ -479,9 +457,7 @@ func TestTransitionState_StoppingMustReachStoppedOrError(t *testing.T) {
t.Run("error_path", func(t *testing.T) {
instance := &vm.VM{ID: "i-stop-error", Status: vm.StateStopping}
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instance.ID] = instance
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(instance)
err := daemon.TransitionState(instance, vm.StateError)
require.NoError(t, err)
@@ -490,9 +466,7 @@ func TestTransitionState_StoppingMustReachStoppedOrError(t *testing.T) {
t.Run("force_terminate_path", func(t *testing.T) {
instance := &vm.VM{ID: "i-stop-term", Status: vm.StateStopping}
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instance.ID] = instance
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(instance)
err := daemon.TransitionState(instance, vm.StateShuttingDown)
require.NoError(t, err)
@@ -511,9 +485,7 @@ func TestTransitionState_ShuttingDownMustReachTerminatedOrError(t *testing.T) {
t.Run("happy_path", func(t *testing.T) {
instance := &vm.VM{ID: "i-shut-happy", Status: vm.StateShuttingDown}
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instance.ID] = instance
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(instance)
err := daemon.TransitionState(instance, vm.StateTerminated)
require.NoError(t, err)
@@ -522,9 +494,7 @@ func TestTransitionState_ShuttingDownMustReachTerminatedOrError(t *testing.T) {
t.Run("error_path", func(t *testing.T) {
instance := &vm.VM{ID: "i-shut-error", Status: vm.StateShuttingDown}
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[instance.ID] = instance
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(instance)
err := daemon.TransitionState(instance, vm.StateError)
require.NoError(t, err)
@@ -575,11 +545,8 @@ func TestTransitionState_MultipleInstancesIndependent(t *testing.T) {
inst1 := &vm.VM{ID: "i-multi-1", Status: vm.StateRunning}
inst2 := &vm.VM{ID: "i-multi-2", Status: vm.StateRunning}
daemon.Instances.Mu.Lock()
daemon.Instances.VMS[inst1.ID] = inst1
daemon.Instances.VMS[inst2.ID] = inst2
daemon.Instances.Mu.Unlock()
daemon.vmMgr.Insert(inst1)
daemon.vmMgr.Insert(inst2)
err := daemon.TransitionState(inst1, vm.StateStopping)
require.NoError(t, err)
@@ -610,20 +577,20 @@ func TestRestoreInstances_StoppingFinalizedToStopped(t *testing.T) {
daemon := createDaemonWithJetStream(t)
// Simulate: daemon was running, instance was in StateStopping, daemon crashed.
daemon.Instances.VMS["i-restore-stop"] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: "i-restore-stop",
Status: vm.StateStopping,
}
})
require.NoError(t, daemon.WriteState())
// Clear in-memory state to simulate fresh daemon startup.
daemon.Instances.VMS = make(map[string]*vm.VM)
daemon.vmMgr.Replace(map[string]*vm.VM{})
// restoreInstances loads from JetStream and resolves transitional states.
simulateCleanRestore(t, daemon)
// Stopped instances are migrated to shared KV and removed from local map
_, ok := daemon.Instances.VMS["i-restore-stop"]
_, ok := daemon.vmMgr.Get("i-restore-stop")
assert.False(t, ok, "stopping→stopped instance should be migrated to shared KV, not in local map")
stoppedInst, err := daemon.jsManager.LoadStoppedInstance("i-restore-stop")
@@ -638,17 +605,17 @@ func TestRestoreInstances_StoppingFinalizedToStopped(t *testing.T) {
func TestRestoreInstances_ShuttingDownFinalizedToTerminated(t *testing.T) {
daemon := createDaemonWithJetStream(t)
daemon.Instances.VMS["i-restore-shut"] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: "i-restore-shut",
Status: vm.StateShuttingDown,
}
})
require.NoError(t, daemon.WriteState())
daemon.Instances.VMS = make(map[string]*vm.VM)
daemon.vmMgr.Replace(map[string]*vm.VM{})
simulateCleanRestore(t, daemon)
// shutting-down→terminated should be migrated to terminated KV
_, ok := daemon.Instances.VMS["i-restore-shut"]
_, ok := daemon.vmMgr.Get("i-restore-shut")
assert.False(t, ok, "terminated instance should be removed from local map")
terminatedInst, err := daemon.jsManager.LoadTerminatedInstance("i-restore-shut")
@@ -663,16 +630,16 @@ func TestRestoreInstances_ShuttingDownFinalizedToTerminated(t *testing.T) {
func TestRestoreInstances_TerminatedRemoved(t *testing.T) {
daemon := createDaemonWithJetStream(t)
daemon.Instances.VMS["i-restore-term"] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: "i-restore-term",
Status: vm.StateTerminated,
}
})
require.NoError(t, daemon.WriteState())
daemon.Instances.VMS = make(map[string]*vm.VM)
daemon.vmMgr.Replace(map[string]*vm.VM{})
simulateCleanRestore(t, daemon)
_, ok := daemon.Instances.VMS["i-restore-term"]
_, ok := daemon.vmMgr.Get("i-restore-term")
assert.False(t, ok, "terminated instance should be removed from local map")
}
@@ -681,20 +648,20 @@ func TestRestoreInstances_TerminatedRemoved(t *testing.T) {
func TestRestoreInstances_UserStoppedMigratedToSharedKV(t *testing.T) {
daemon := createDaemonWithJetStream(t)
daemon.Instances.VMS["i-restore-userstop"] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: "i-restore-userstop",
Status: vm.StateStopped,
Attributes: types.EC2CommandAttributes{
StopInstance: true,
},
}
})
require.NoError(t, daemon.WriteState())
daemon.Instances.VMS = make(map[string]*vm.VM)
daemon.vmMgr.Replace(map[string]*vm.VM{})
simulateCleanRestore(t, daemon)
// Stopped instances should be migrated to shared KV
_, ok := daemon.Instances.VMS["i-restore-userstop"]
_, ok := daemon.vmMgr.Get("i-restore-userstop")
assert.False(t, ok, "user-stopped instance should be migrated to shared KV, not in local map")
stoppedInst, err := daemon.jsManager.LoadStoppedInstance("i-restore-userstop")
@@ -711,16 +678,16 @@ func TestRestoreInstances_UserStoppedMigratedToSharedKV(t *testing.T) {
func TestRestoreInstances_RunningResetToPending(t *testing.T) {
daemon := createDaemonWithJetStream(t)
daemon.Instances.VMS["i-restore-run"] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: "i-restore-run",
Status: vm.StateRunning,
}
})
require.NoError(t, daemon.WriteState())
daemon.Instances.VMS = make(map[string]*vm.VM)
daemon.vmMgr.Replace(map[string]*vm.VM{})
simulateCleanRestore(t, daemon)
instance, ok := daemon.Instances.VMS["i-restore-run"]
instance, ok := daemon.vmMgr.Get("i-restore-run")
require.True(t, ok, "running instance should still be loaded")
// LaunchInstance fails (no QEMU), so it stays pending or gets marked failed.
// Either way it should NOT remain "running" — that would be a lie.
@@ -736,20 +703,20 @@ func TestRestoreInstances_RunningResetsLaunchTime(t *testing.T) {
daemon := createDaemonWithJetStream(t)
staleTime := time.Now().Add(-30 * time.Minute)
daemon.Instances.VMS["i-stale-launch"] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: "i-stale-launch",
Status: vm.StateRunning,
Instance: &ec2.Instance{
LaunchTime: &staleTime,
},
}
})
require.NoError(t, daemon.WriteState())
daemon.Instances.VMS = make(map[string]*vm.VM)
daemon.vmMgr.Replace(map[string]*vm.VM{})
before := time.Now()
simulateCleanRestore(t, daemon)
instance, ok := daemon.Instances.VMS["i-stale-launch"]
instance, ok := daemon.vmMgr.Get("i-stale-launch")
require.True(t, ok, "instance should be loaded from state")
require.NotNil(t, instance.Instance)
require.NotNil(t, instance.Instance.LaunchTime)
@@ -767,26 +734,26 @@ func TestRestoreInstances_RunningResetsLaunchTime(t *testing.T) {
func TestRestoreInstances_MixedStates(t *testing.T) {
daemon := createDaemonWithJetStream(t)
daemon.Instances.VMS["i-mix-stopping"] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: "i-mix-stopping", Status: vm.StateStopping,
}
daemon.Instances.VMS["i-mix-shutting"] = &vm.VM{
})
daemon.vmMgr.Insert(&vm.VM{
ID: "i-mix-shutting", Status: vm.StateShuttingDown,
}
daemon.Instances.VMS["i-mix-terminated"] = &vm.VM{
})
daemon.vmMgr.Insert(&vm.VM{
ID: "i-mix-terminated", Status: vm.StateTerminated,
}
daemon.Instances.VMS["i-mix-stopped"] = &vm.VM{
})
daemon.vmMgr.Insert(&vm.VM{
ID: "i-mix-stopped", Status: vm.StateStopped,
Attributes: types.EC2CommandAttributes{StopInstance: true},
}
})
require.NoError(t, daemon.WriteState())
daemon.Instances.VMS = make(map[string]*vm.VM)
daemon.vmMgr.Replace(map[string]*vm.VM{})
simulateCleanRestore(t, daemon)
// Stopped instances are migrated to shared KV and removed from local map
assert.Nil(t, daemon.Instances.VMS["i-mix-stopping"],
assertVMNotPresent(t, daemon, "i-mix-stopping",
"stopping→stopped should be migrated to shared KV")
stoppedFromStopping, err := daemon.jsManager.LoadStoppedInstance("i-mix-stopping")
require.NoError(t, err)
@@ -794,7 +761,7 @@ func TestRestoreInstances_MixedStates(t *testing.T) {
assert.Equal(t, vm.StateStopped, stoppedFromStopping.Status)
// shutting-down should finalize to terminated and migrate to terminated KV
assert.Nil(t, daemon.Instances.VMS["i-mix-shutting"],
assertVMNotPresent(t, daemon, "i-mix-shutting",
"shutting-down→terminated should be migrated to terminated KV")
terminatedFromShutting, err2 := daemon.jsManager.LoadTerminatedInstance("i-mix-shutting")
require.NoError(t, err2)
@@ -802,10 +769,10 @@ func TestRestoreInstances_MixedStates(t *testing.T) {
assert.Equal(t, vm.StateTerminated, terminatedFromShutting.Status)
// already-terminated should be removed from local map (no migration, just cleanup)
assert.Nil(t, daemon.Instances.VMS["i-mix-terminated"],
assertVMNotPresent(t, daemon, "i-mix-terminated",
"terminated should be removed from local map")
assert.Nil(t, daemon.Instances.VMS["i-mix-stopped"],
assertVMNotPresent(t, daemon, "i-mix-stopped",
"user-stopped should be migrated to shared KV")
stoppedFromUser, err := daemon.jsManager.LoadStoppedInstance("i-mix-stopped")
require.NoError(t, err)
@@ -819,20 +786,20 @@ func TestRestoreInstances_MixedStates(t *testing.T) {
func TestRestoreInstances_StatePersistsAfterRecovery(t *testing.T) {
daemon := createDaemonWithJetStream(t)
daemon.Instances.VMS["i-persist-stop"] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: "i-persist-stop", Status: vm.StateStopping,
}
daemon.Instances.VMS["i-persist-shut"] = &vm.VM{
})
daemon.vmMgr.Insert(&vm.VM{
ID: "i-persist-shut", Status: vm.StateShuttingDown,
}
})
require.NoError(t, daemon.WriteState())
// First restart: finalize transitional states.
daemon.Instances.VMS = make(map[string]*vm.VM)
daemon.vmMgr.Replace(map[string]*vm.VM{})
simulateCleanRestore(t, daemon)
// stopping→stopped migrates to shared KV
assert.Nil(t, daemon.Instances.VMS["i-persist-stop"],
assertVMNotPresent(t, daemon, "i-persist-stop",
"stopping→stopped should be migrated to shared KV")
stoppedInst, err := daemon.jsManager.LoadStoppedInstance("i-persist-stop")
require.NoError(t, err)
@@ -840,7 +807,7 @@ func TestRestoreInstances_StatePersistsAfterRecovery(t *testing.T) {
assert.Equal(t, vm.StateStopped, stoppedInst.Status)
// shutting-down→terminated should be migrated to terminated KV
assert.Nil(t, daemon.Instances.VMS["i-persist-shut"],
assertVMNotPresent(t, daemon, "i-persist-shut",
"terminated instance should be removed from local map")
terminatedInst, err := daemon.jsManager.LoadTerminatedInstance("i-persist-shut")
require.NoError(t, err)
@@ -849,7 +816,7 @@ func TestRestoreInstances_StatePersistsAfterRecovery(t *testing.T) {
// Second restart: stopped instance is in shared KV (not per-node state),
// terminated instance is in terminated KV (not per-node state).
daemon.Instances.VMS = make(map[string]*vm.VM)
daemon.vmMgr.Replace(map[string]*vm.VM{})
simulateCleanRestore(t, daemon)
// Stopped instance should still be in shared KV
@@ -871,9 +838,9 @@ func TestMigrateInstanceToKV(t *testing.T) {
// Test stopped migration
stoppedVM := &vm.VM{ID: "i-migrate-stop", Status: vm.StateStopped}
daemon.Instances.VMS[stoppedVM.ID] = stoppedVM
daemon.vmMgr.Insert(stoppedVM)
assert.True(t, daemon.migrateStoppedToSharedKV(stoppedVM))
assert.Nil(t, daemon.Instances.VMS[stoppedVM.ID], "should be removed from local map")
assertVMNotPresent(t, daemon, stoppedVM.ID, "should be removed from local map")
loaded, err := daemon.jsManager.LoadStoppedInstance(stoppedVM.ID)
require.NoError(t, err)
require.NotNil(t, loaded)
@@ -881,9 +848,9 @@ func TestMigrateInstanceToKV(t *testing.T) {
// Test terminated migration
terminatedVM := &vm.VM{ID: "i-migrate-term", Status: vm.StateTerminated}
daemon.Instances.VMS[terminatedVM.ID] = terminatedVM
daemon.vmMgr.Insert(terminatedVM)
assert.True(t, daemon.migrateTerminatedToKV(terminatedVM))
assert.Nil(t, daemon.Instances.VMS[terminatedVM.ID], "should be removed from local map")
assertVMNotPresent(t, daemon, terminatedVM.ID, "should be removed from local map")
loadedTerm, err := daemon.jsManager.LoadTerminatedInstance(terminatedVM.ID)
require.NoError(t, err)
require.NotNil(t, loadedTerm)
@@ -896,11 +863,11 @@ func TestMigrateInstanceToKV_NoJetStream(t *testing.T) {
daemon.jsManager = nil
vm1 := &vm.VM{ID: "i-no-js", Status: vm.StateStopped}
daemon.Instances.VMS[vm1.ID] = vm1
daemon.vmMgr.Insert(vm1)
assert.False(t, daemon.migrateStoppedToSharedKV(vm1))
assert.False(t, daemon.migrateTerminatedToKV(vm1))
// Instance should still be in local map since migration failed
assert.NotNil(t, daemon.Instances.VMS[vm1.ID])
assertVMPresent(t, daemon, vm1.ID)
}
// TestStatePersistence_RoundTrip verifies that all instance fields survive
@@ -917,14 +884,14 @@ func TestStatePersistence_RoundTrip(t *testing.T) {
TerminateInstance: false,
},
}
daemon.Instances.VMS[original.ID] = original
daemon.vmMgr.Insert(original)
require.NoError(t, daemon.WriteState())
// Simulate restart: clear and reload.
daemon.Instances.VMS = make(map[string]*vm.VM)
daemon.vmMgr.Replace(map[string]*vm.VM{})
require.NoError(t, daemon.LoadState())
loaded := daemon.Instances.VMS["i-roundtrip"]
loaded, _ := daemon.vmMgr.Get("i-roundtrip")
require.NotNil(t, loaded)
assert.Equal(t, original.ID, loaded.ID)
assert.Equal(t, original.Status, loaded.Status)
@@ -939,20 +906,20 @@ func TestRestoreInstances_StoppedInstanceMigratedToSharedKV(t *testing.T) {
daemon := createDaemonWithJetStream(t)
instanceID := "i-restore-start"
daemon.Instances.VMS[instanceID] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: instanceID,
InstanceType: "t3.micro",
Status: vm.StateStopped,
Attributes: types.EC2CommandAttributes{StopInstance: true},
}
})
require.NoError(t, daemon.WriteState())
// Simulate daemon restart: clear in-memory state and restore from JetStream.
daemon.Instances.VMS = make(map[string]*vm.VM)
daemon.vmMgr.Replace(map[string]*vm.VM{})
simulateCleanRestore(t, daemon)
// Instance should not be in local map (migrated to shared KV).
_, ok := daemon.Instances.VMS[instanceID]
_, ok := daemon.vmMgr.Get(instanceID)
assert.False(t, ok, "stopped instance should be migrated to shared KV, not in local map")
// Instance should be in shared KV
@@ -971,20 +938,20 @@ func TestRestoreInstances_UnknownTypeMigratedToStoppedWithReason(t *testing.T) {
daemon := createDaemonWithJetStream(t)
instanceID := "i-restore-notype"
daemon.Instances.VMS[instanceID] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: instanceID,
InstanceType: "m7i.small", // type that won't exist on this host
Status: vm.StateRunning,
Instance: &ec2.Instance{},
}
})
require.NoError(t, daemon.WriteState())
// Simulate daemon restart
daemon.Instances.VMS = make(map[string]*vm.VM)
daemon.vmMgr.Replace(map[string]*vm.VM{})
simulateCleanRestore(t, daemon)
// Instance should not be in local map (migrated to shared KV)
_, ok := daemon.Instances.VMS[instanceID]
_, ok := daemon.vmMgr.Get(instanceID)
assert.False(t, ok, "instance with unknown type should be migrated to shared KV")
// Instance should be stopped in shared KV with capacity reason
@@ -1005,46 +972,43 @@ func TestPendingWatchdog_MarksStuckInstanceFailed(t *testing.T) {
staleTime := time.Now().Add(-6 * time.Minute)
recentTime := time.Now()
daemon.Instances.VMS["i-stuck"] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: "i-stuck",
Status: vm.StatePending,
Instance: &ec2.Instance{
LaunchTime: &staleTime,
},
}
daemon.Instances.VMS["i-fresh"] = &vm.VM{
})
daemon.vmMgr.Insert(&vm.VM{
ID: "i-fresh",
Status: vm.StatePending,
Instance: &ec2.Instance{
LaunchTime: &recentTime,
},
}
})
// Run one watchdog tick manually instead of waiting for the ticker
daemon.Instances.Mu.Lock()
var stuck []*vm.VM
for _, instance := range daemon.Instances.VMS {
for _, instance := range daemon.vmMgr.Snapshot() {
if (instance.Status == vm.StatePending || instance.Status == vm.StateProvisioning) &&
instance.Instance != nil && instance.Instance.LaunchTime != nil &&
time.Since(*instance.Instance.LaunchTime) > 5*time.Minute {
stuck = append(stuck, instance)
}
}
daemon.Instances.Mu.Unlock()
for _, instance := range stuck {
daemon.markInstanceFailed(instance, "launch_timeout")
}
// Stuck instance should have transitioned to shutting-down
stuckInst := daemon.Instances.VMS["i-stuck"]
stuckInst, _ := daemon.vmMgr.Get("i-stuck")
assert.Equal(t, vm.StateShuttingDown, stuckInst.Status)
require.NotNil(t, stuckInst.Instance.StateReason)
assert.Equal(t, "Server.InternalError", *stuckInst.Instance.StateReason.Code)
assert.Equal(t, "launch_timeout", *stuckInst.Instance.StateReason.Message)
// Fresh instance should still be pending
freshInst := daemon.Instances.VMS["i-fresh"]
freshInst, _ := daemon.vmMgr.Get("i-fresh")
assert.Equal(t, vm.StatePending, freshInst.Status)
}
@@ -1137,7 +1101,7 @@ func TestMarkInstanceFailed_AlreadyShuttingDown(t *testing.T) {
Status: vm.StateShuttingDown,
Instance: &ec2.Instance{},
}
daemon.Instances.VMS[instance.ID] = instance
daemon.vmMgr.Insert(instance)
// Should be a no-op — instance is already being cleaned up
daemon.markInstanceFailed(instance, "test_reason")
@@ -1154,7 +1118,7 @@ func TestMarkInstanceFailed_AlreadyTerminated(t *testing.T) {
Status: vm.StateTerminated,
Instance: &ec2.Instance{},
}
daemon.Instances.VMS[instance.ID] = instance
daemon.vmMgr.Insert(instance)
daemon.markInstanceFailed(instance, "test_reason")
@@ -1169,14 +1133,14 @@ func TestRestoreInstances_ProvisioningInstanceLaunchTimeReset(t *testing.T) {
// their LaunchTime reset so the pending watchdog doesn't immediately kill
// them after a prolonged outage.
staleTime := time.Now().Add(-30 * time.Minute)
daemon.Instances.VMS["i-provisioning"] = &vm.VM{
daemon.vmMgr.Insert(&vm.VM{
ID: "i-provisioning",
Status: vm.StateProvisioning,
Instance: &ec2.Instance{LaunchTime: &staleTime},
}
})
require.NoError(t, daemon.WriteState())
daemon.Instances.VMS = make(map[string]*vm.VM)
daemon.vmMgr.Replace(map[string]*vm.VM{})
before := time.Now()
simulateCleanRestore(t, daemon)
@@ -1184,7 +1148,7 @@ func TestRestoreInstances_ProvisioningInstanceLaunchTimeReset(t *testing.T) {
// from the map. The instance may have been removed by markInstanceFailed
// (launch fails in test env), so load directly from JetStream state.
require.NoError(t, daemon.LoadState())
instance := daemon.Instances.VMS["i-provisioning"]
instance, _ := daemon.vmMgr.Get("i-provisioning")
if instance == nil {
// Instance was cleaned up by markInstanceFailed — that's fine.
// The important thing is that the LaunchTime WAS reset before
+21
View File
@@ -0,0 +1,21 @@
package daemon
import (
"testing"
"github.com/stretchr/testify/assert"
)
// assertVMPresent asserts that the daemon's VM manager has an entry for id.
func assertVMPresent(t *testing.T, d *Daemon, id string, msgAndArgs ...any) {
t.Helper()
_, ok := d.vmMgr.Get(id)
assert.True(t, ok, msgAndArgs...)
}
// assertVMNotPresent asserts that the daemon's VM manager has no entry for id.
func assertVMNotPresent(t *testing.T, d *Daemon, id string, msgAndArgs ...any) {
t.Helper()
_, ok := d.vmMgr.Get(id)
assert.False(t, ok, msgAndArgs...)
}
@@ -231,17 +231,15 @@ type InstanceServiceImpl struct {
config *config.Config
instanceTypes map[string]*ec2.InstanceTypeInfo
natsConn *nats.Conn
instances *vm.Instances
objectStore objectstore.ObjectStore
}
// NewInstanceServiceImpl creates a new instance service implementation for daemon use
func NewInstanceServiceImpl(cfg *config.Config, instanceTypes map[string]*ec2.InstanceTypeInfo, nc *nats.Conn, instances *vm.Instances, store objectstore.ObjectStore) *InstanceServiceImpl {
func NewInstanceServiceImpl(cfg *config.Config, instanceTypes map[string]*ec2.InstanceTypeInfo, nc *nats.Conn, store objectstore.ObjectStore) *InstanceServiceImpl {
return &InstanceServiceImpl{
config: cfg,
instanceTypes: instanceTypes,
natsConn: nc,
instances: instances,
objectStore: store,
}
}
+162
View File
@@ -0,0 +1,162 @@
package vm
import (
"maps"
"sync"
)
// Manager owns the in-memory map of running VMs on this node.
type Manager struct {
mu sync.Mutex
vms map[string]*VM
}
func NewManager() *Manager {
return &Manager{vms: make(map[string]*VM)}
}
// Get returns the VM for id (and true) or (nil, false).
func (m *Manager) Get(id string) (*VM, bool) {
m.mu.Lock()
defer m.mu.Unlock()
v, ok := m.vms[id]
return v, ok
}
// Insert unconditionally stores v under v.ID, overwriting any existing entry.
func (m *Manager) Insert(v *VM) {
m.mu.Lock()
defer m.mu.Unlock()
m.vms[v.ID] = v
}
// InsertIfAbsent stores v under v.ID only if no entry currently exists.
// Returns true if the insert happened, false if the slot was already occupied.
func (m *Manager) InsertIfAbsent(v *VM) bool {
m.mu.Lock()
defer m.mu.Unlock()
if _, exists := m.vms[v.ID]; exists {
return false
}
m.vms[v.ID] = v
return true
}
// Delete removes the entry for id. No-op if absent.
func (m *Manager) Delete(id string) {
m.mu.Lock()
defer m.mu.Unlock()
delete(m.vms, id)
}
// DeleteIf deletes the entry for id only if the stored pointer matches want.
// Returns true if the delete happened. Used by stop/terminate handlers to guard
// against the slot being reclaimed by a concurrent start handler.
func (m *Manager) DeleteIf(id string, want *VM) bool {
m.mu.Lock()
defer m.mu.Unlock()
current, exists := m.vms[id]
if !exists || current != want {
return false
}
delete(m.vms, id)
return true
}
// Count returns the current number of VMs.
func (m *Manager) Count() int {
m.mu.Lock()
defer m.mu.Unlock()
return len(m.vms)
}
// ForEach calls fn for each VM under the manager lock. fn must not call back
// into Manager methods that take the lock — that will deadlock. For lock-free
// iteration, use Snapshot.
func (m *Manager) ForEach(fn func(*VM)) {
m.mu.Lock()
defer m.mu.Unlock()
for _, v := range m.vms {
fn(v)
}
}
// Snapshot returns the current set of VMs as a slice. Callers must not mutate
// the slice's *VM entries' map-related fields; the underlying VM pointers are
// shared with the manager.
func (m *Manager) Snapshot() []*VM {
m.mu.Lock()
defer m.mu.Unlock()
out := make([]*VM, 0, len(m.vms))
for _, v := range m.vms {
out = append(out, v)
}
return out
}
// SnapshotMap returns a copy of the id→VM map. Used by the state persistence
// adapter so serialization can happen without holding the manager lock.
func (m *Manager) SnapshotMap() map[string]*VM {
m.mu.Lock()
defer m.mu.Unlock()
out := make(map[string]*VM, len(m.vms))
maps.Copy(out, m.vms)
return out
}
// Filter returns the VMs for which pred returns true. pred runs under lock.
func (m *Manager) Filter(pred func(*VM) bool) []*VM {
m.mu.Lock()
defer m.mu.Unlock()
var out []*VM
for _, v := range m.vms {
if pred(v) {
out = append(out, v)
}
}
return out
}
// UpdateState looks up id and, if found, runs fn(v) under lock. Used for
// atomic check-then-mutate or read-under-lock on a single VM. Returns true
// if the VM was found and fn was invoked.
func (m *Manager) UpdateState(id string, fn func(*VM)) bool {
m.mu.Lock()
defer m.mu.Unlock()
v, ok := m.vms[id]
if !ok {
return false
}
fn(v)
return true
}
// Inspect runs fn(v) under the manager lock for an already-resolved VM
// pointer. Used by call sites that hold a *VM (e.g. from a NATS handler
// dispatch) and need to read or mutate its fields with the same memory-
// ordering guarantee as map-keyed access.
func (m *Manager) Inspect(v *VM, fn func(*VM)) {
m.mu.Lock()
defer m.mu.Unlock()
fn(v)
}
// View runs fn with the live id→VM map under the manager lock. fn must not
// mutate the map nor retain references after returning. Used by serialization
// paths that need every VM's fields to be stable for the duration of the
// encode (e.g. JSON marshal during state persistence).
func (m *Manager) View(fn func(map[string]*VM)) {
m.mu.Lock()
defer m.mu.Unlock()
fn(m.vms)
}
// Replace bulk-resets the manager to the given set of VMs. The supplied map
// is copied; callers may mutate it after the call returns. Used by Restore
// after loading persisted state.
func (m *Manager) Replace(vms map[string]*VM) {
m.mu.Lock()
defer m.mu.Unlock()
m.vms = make(map[string]*VM, len(vms))
maps.Copy(m.vms, vms)
}
+373
View File
@@ -0,0 +1,373 @@
package vm
import (
"sync"
"sync/atomic"
"testing"
)
func mkVM(id string) *VM { return &VM{ID: id} }
func TestManager_GetInsertDelete(t *testing.T) {
m := NewManager()
if v, ok := m.Get("missing"); ok || v != nil {
t.Fatalf("Get on empty manager: got (%v, %v), want (nil, false)", v, ok)
}
a := mkVM("i-a")
m.Insert(a)
got, ok := m.Get("i-a")
if !ok || got != a {
t.Fatalf("Get after Insert: got (%v, %v), want (%v, true)", got, ok, a)
}
if m.Count() != 1 {
t.Fatalf("Count after Insert: got %d, want 1", m.Count())
}
m.Delete("i-a")
if _, ok := m.Get("i-a"); ok {
t.Fatalf("Get after Delete: got ok=true, want false")
}
m.Delete("i-a") // delete is idempotent
}
func TestManager_InsertIfAbsent(t *testing.T) {
m := NewManager()
a := mkVM("i-a")
a2 := mkVM("i-a")
if !m.InsertIfAbsent(a) {
t.Fatal("InsertIfAbsent on empty: want true")
}
if m.InsertIfAbsent(a2) {
t.Fatal("InsertIfAbsent on occupied: want false")
}
got, _ := m.Get("i-a")
if got != a {
t.Fatalf("Get after collision: got %v, want %v (original)", got, a)
}
}
func TestManager_DeleteIf(t *testing.T) {
m := NewManager()
a := mkVM("i-a")
b := mkVM("i-a") // different pointer, same ID
m.Insert(a)
if m.DeleteIf("i-a", b) {
t.Fatal("DeleteIf with non-matching pointer: want false")
}
if got, _ := m.Get("i-a"); got != a {
t.Fatalf("DeleteIf preserved entry: got %v, want %v", got, a)
}
if !m.DeleteIf("i-a", a) {
t.Fatal("DeleteIf with matching pointer: want true")
}
if _, ok := m.Get("i-a"); ok {
t.Fatal("entry still present after successful DeleteIf")
}
if m.DeleteIf("missing", a) {
t.Fatal("DeleteIf on missing id: want false")
}
}
func TestManager_UpdateState(t *testing.T) {
m := NewManager()
a := mkVM("i-a")
a.Status = StatePending
m.Insert(a)
called := false
if !m.UpdateState("i-a", func(v *VM) {
called = true
v.Status = StateRunning
}) {
t.Fatal("UpdateState on existing id: want true")
}
if !called {
t.Fatal("UpdateState fn not invoked")
}
if a.Status != StateRunning {
t.Fatalf("UpdateState mutation: got %s, want %s", a.Status, StateRunning)
}
if m.UpdateState("missing", func(*VM) { t.Fatal("fn called on missing id") }) {
t.Fatal("UpdateState on missing id: want false")
}
}
func TestManager_Inspect(t *testing.T) {
m := NewManager()
a := mkVM("i-a")
a.Status = StatePending
m.Insert(a)
var seen InstanceState
m.Inspect(a, func(v *VM) { seen = v.Status })
if seen != StatePending {
t.Fatalf("Inspect: got %s, want %s", seen, StatePending)
}
// Inspect works on VMs not in the map (e.g., during build before insert).
loose := mkVM("i-loose")
loose.Status = StateRunning
m.Inspect(loose, func(v *VM) { seen = v.Status })
if seen != StateRunning {
t.Fatalf("Inspect on loose VM: got %s, want %s", seen, StateRunning)
}
}
func TestManager_ForEachAndSnapshot(t *testing.T) {
m := NewManager()
want := map[string]bool{"i-1": true, "i-2": true, "i-3": true}
for id := range want {
m.Insert(mkVM(id))
}
seen := map[string]bool{}
m.ForEach(func(v *VM) { seen[v.ID] = true })
if len(seen) != len(want) {
t.Fatalf("ForEach saw %d, want %d", len(seen), len(want))
}
for id := range want {
if !seen[id] {
t.Fatalf("ForEach missed %s", id)
}
}
snap := m.Snapshot()
if len(snap) != len(want) {
t.Fatalf("Snapshot len=%d, want %d", len(snap), len(want))
}
snapMap := m.SnapshotMap()
if len(snapMap) != len(want) {
t.Fatalf("SnapshotMap len=%d, want %d", len(snapMap), len(want))
}
// SnapshotMap is decoupled from the live map: mutating the returned map
// must not affect the manager.
delete(snapMap, "i-1")
if _, ok := m.Get("i-1"); !ok {
t.Fatal("mutating SnapshotMap result leaked into manager")
}
}
func TestManager_Filter(t *testing.T) {
m := NewManager()
a, b, c := mkVM("i-a"), mkVM("i-b"), mkVM("i-c")
a.Status = StateRunning
b.Status = StateStopped
c.Status = StateRunning
m.Insert(a)
m.Insert(b)
m.Insert(c)
running := m.Filter(func(v *VM) bool { return v.Status == StateRunning })
if len(running) != 2 {
t.Fatalf("Filter running: got %d, want 2", len(running))
}
}
func TestManager_Replace(t *testing.T) {
m := NewManager()
m.Insert(mkVM("i-old"))
src := map[string]*VM{"i-new1": mkVM("i-new1"), "i-new2": mkVM("i-new2")}
m.Replace(src)
if _, ok := m.Get("i-old"); ok {
t.Fatal("old entry survived Replace")
}
if m.Count() != 2 {
t.Fatalf("Count after Replace: got %d, want 2", m.Count())
}
// Mutating the source after Replace must not affect the manager.
delete(src, "i-new1")
if _, ok := m.Get("i-new1"); !ok {
t.Fatal("Replace did not copy the source map")
}
}
// TestManager_ConcurrentSoak runs many goroutines doing Insert/Delete/ForEach/
// UpdateState concurrently. Run with `go test -race ./vm/...` to validate.
func TestManager_ConcurrentSoak(t *testing.T) {
if testing.Short() {
t.Skip("skipping soak in short mode")
}
m := NewManager()
const writers = 8
const iters = 2000
var writersWG sync.WaitGroup
var readersWG sync.WaitGroup
stop := make(chan struct{})
var foreachCount atomic.Int64
readersWG.Add(2)
go func() {
defer readersWG.Done()
for {
select {
case <-stop:
return
default:
}
m.ForEach(func(v *VM) { _ = v.ID })
foreachCount.Add(1)
}
}()
go func() {
defer readersWG.Done()
for {
select {
case <-stop:
return
default:
}
_ = m.Snapshot()
}
}()
for w := range writers {
writersWG.Add(1)
go func(seed int) {
defer writersWG.Done()
for i := range iters {
id := mkID(seed, i)
m.Insert(mkVM(id))
m.UpdateState(id, func(v *VM) { v.Status = StateRunning })
m.Delete(id)
}
}(w)
}
writersWG.Wait()
close(stop)
readersWG.Wait()
if foreachCount.Load() == 0 {
t.Fatal("readers never observed any state — possible reader starvation")
}
}
// TestManager_SlotOccupancyRace simulates the production race the plan calls
// out at daemon_handlers_instance.go:594-626 and daemon.go:2140/2149: a
// terminate handler running DeleteIf(id, original) concurrently with a start
// handler running InsertIfAbsent(replacement) for the same slot, plus the
// rollback InsertIfAbsent(original) on a write failure. The invariant is that
// the slot never contains a torn or unexpected pointer — at any moment Get(id)
// returns either the original, the replacement, or nothing.
func TestManager_SlotOccupancyRace(t *testing.T) {
if testing.Short() {
t.Skip("skipping race stress in short mode")
}
m := NewManager()
const iters = 5000
const id = "i-slot"
original := &VM{ID: id}
replacement := &VM{ID: id}
// Pre-seed: original is in the slot.
m.Insert(original)
var wg sync.WaitGroup
wg.Add(3)
// Terminator: DeleteIf(original); on the next round, re-Insert original to
// reset the slot for the next iteration.
go func() {
defer wg.Done()
for range iters {
m.DeleteIf(id, original)
m.InsertIfAbsent(original)
}
}()
// Starter: tries to occupy the slot with replacement after a delete; rolls
// back via DeleteIf to keep the test loop in steady state.
go func() {
defer wg.Done()
for range iters {
if m.InsertIfAbsent(replacement) {
m.DeleteIf(id, replacement)
}
}
}()
// Observer: at every observation, Get must return original, replacement,
// or nothing. Anything else means the slot is corrupted.
go func() {
defer wg.Done()
for range iters {
v, ok := m.Get(id)
if !ok {
continue
}
if v != original && v != replacement {
t.Errorf("slot held an unexpected pointer: %p", v)
return
}
}
}()
wg.Wait()
}
// TestManager_ForEachIsRaceSafeForMutableFields documents the safe contract:
// ForEach holds the lock for the entire iteration, so concurrent mutators
// using Inspect/UpdateState are serialized against the iteration. This is the
// race-clean alternative when callers need to read fields like Status that
// can change concurrently. (Snapshot, by contrast, returns shared *VM
// pointers — reading mutable fields after release is unsynchronized.)
func TestManager_ForEachIsRaceSafeForMutableFields(t *testing.T) {
if testing.Short() {
t.Skip("skipping race stress in short mode")
}
m := NewManager()
v := &VM{ID: "i-1", Status: StatePending}
m.Insert(v)
stop := make(chan struct{})
var wg sync.WaitGroup
wg.Go(func() {
toggle := true
for {
select {
case <-stop:
return
default:
}
if toggle {
m.UpdateState("i-1", func(v *VM) { v.Status = StateRunning })
} else {
m.UpdateState("i-1", func(v *VM) { v.Status = StatePending })
}
toggle = !toggle
}
})
for range 5000 {
m.ForEach(func(v *VM) {
_ = v.Status
})
}
close(stop)
wg.Wait()
}
func mkID(w, i int) string {
const hexDigits = "0123456789abcdef"
buf := []byte("i-")
x := uint64(w)<<32 | uint64(i)
for shift := 60; shift >= 0; shift -= 4 {
buf = append(buf, hexDigits[(x>>uint(shift))&0xF])
}
return string(buf)
}
-5
View File
@@ -130,11 +130,6 @@ func (v *VM) ResetNodeLocalState() {
v.EBSRequests.Mu = sync.Mutex{}
}
type Instances struct {
VMS map[string]*VM `json:"vms"`
Mu sync.Mutex `json:"-"`
}
type NetDev struct {
Value string `json:"value"`
}