From e1b3623599a293eff9f28e1b4e9b838287d26b60 Mon Sep 17 00:00:00 2001 From: Julian Sommer <115516918+juliansommer@users.noreply.github.com> Date: Fri, 1 May 2026 15:15:54 +1000 Subject: [PATCH] refactor: introduce vm.Manager (phase 1 of lifecycle extraction) (#176) --- spinifex/daemon/daemon.go | 139 +++---- spinifex/daemon/daemon_handlers.go | 19 +- spinifex/daemon/daemon_handlers_console.go | 5 +- spinifex/daemon/daemon_handlers_image.go | 26 +- spinifex/daemon/daemon_handlers_instance.go | 197 ++++----- spinifex/daemon/daemon_handlers_test.go | 135 +++---- spinifex/daemon/daemon_handlers_volume.go | 40 +- spinifex/daemon/daemon_system_instance.go | 20 +- .../daemon/daemon_system_instance_test.go | 50 +-- spinifex/daemon/daemon_test.go | 76 ++-- spinifex/daemon/device_map.go | 11 +- spinifex/daemon/health.go | 98 ++--- spinifex/daemon/health_test.go | 42 +- spinifex/daemon/heartbeat.go | 4 +- spinifex/daemon/heartbeat_test.go | 10 +- spinifex/daemon/jetstream.go | 38 +- spinifex/daemon/jetstream_test.go | 154 ++++---- spinifex/daemon/shutdown.go | 13 +- spinifex/daemon/state.go | 20 +- spinifex/daemon/state_test.go | 220 +++++------ spinifex/daemon/test_helpers_test.go | 21 + .../handlers/ec2/instance/service_impl.go | 4 +- spinifex/vm/manager.go | 162 ++++++++ spinifex/vm/manager_test.go | 373 ++++++++++++++++++ spinifex/vm/vm.go | 5 - 25 files changed, 1121 insertions(+), 761 deletions(-) create mode 100644 spinifex/daemon/test_helpers_test.go create mode 100644 spinifex/vm/manager.go create mode 100644 spinifex/vm/manager_test.go diff --git a/spinifex/daemon/daemon.go b/spinifex/daemon/daemon.go index 7dca45df0..345651d61 100644 --- a/spinifex/daemon/daemon.go +++ b/spinifex/daemon/daemon.go @@ -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", diff --git a/spinifex/daemon/daemon_handlers.go b/spinifex/daemon/daemon_handlers.go index 4a42d97b2..f3905a094 100644 --- a/spinifex/daemon/daemon_handlers.go +++ b/spinifex/daemon/daemon_handlers.go @@ -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, diff --git a/spinifex/daemon/daemon_handlers_console.go b/spinifex/daemon/daemon_handlers_console.go index d3f5ef645..fb38185dd 100644 --- a/spinifex/daemon/daemon_handlers_console.go +++ b/spinifex/daemon/daemon_handlers_console.go @@ -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 diff --git a/spinifex/daemon/daemon_handlers_image.go b/spinifex/daemon/daemon_handlers_image.go index f15704c58..eb8614085 100644 --- a/spinifex/daemon/daemon_handlers_image.go +++ b/spinifex/daemon/daemon_handlers_image.go @@ -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) diff --git a/spinifex/daemon/daemon_handlers_instance.go b/spinifex/daemon/daemon_handlers_instance.go index f2a90ec37..ec6fc9e2f 100644 --- a/spinifex/daemon/daemon_handlers_instance.go +++ b/spinifex/daemon/daemon_handlers_instance.go @@ -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 { diff --git a/spinifex/daemon/daemon_handlers_test.go b/spinifex/daemon/daemon_handlers_test.go index 507d54b44..898055cbd 100644 --- a/spinifex/daemon/daemon_handlers_test.go +++ b/spinifex/daemon/daemon_handlers_test.go @@ -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), diff --git a/spinifex/daemon/daemon_handlers_volume.go b/spinifex/daemon/daemon_handlers_volume.go index 6961d7d2b..b70a17f74 100644 --- a/spinifex/daemon/daemon_handlers_volume.go +++ b/spinifex/daemon/daemon_handlers_volume.go @@ -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 { diff --git a/spinifex/daemon/daemon_system_instance.go b/spinifex/daemon/daemon_system_instance.go index 67a136449..6572f6018 100644 --- a/spinifex/daemon/daemon_system_instance.go +++ b/spinifex/daemon/daemon_system_instance.go @@ -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) diff --git a/spinifex/daemon/daemon_system_instance_test.go b/spinifex/daemon/daemon_system_instance_test.go index c29bd00c5..2038bd731 100644 --- a/spinifex/daemon/daemon_system_instance_test.go +++ b/spinifex/daemon/daemon_system_instance_test.go @@ -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 { diff --git a/spinifex/daemon/daemon_test.go b/spinifex/daemon/daemon_test.go index 7d7e1d8ae..e5b380c1e 100644 --- a/spinifex/daemon/daemon_test.go +++ b/spinifex/daemon/daemon_test.go @@ -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{ diff --git a/spinifex/daemon/device_map.go b/spinifex/daemon/device_map.go index b5e4d8b40..1ed97b9d6 100644 --- a/spinifex/daemon/device_map.go +++ b/spinifex/daemon/device_map.go @@ -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) diff --git a/spinifex/daemon/health.go b/spinifex/daemon/health.go index 00b01be63..17b62bbd0 100644 --- a/spinifex/daemon/health.go +++ b/spinifex/daemon/health.go @@ -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. diff --git a/spinifex/daemon/health_test.go b/spinifex/daemon/health_test.go index f4f2c3fc6..cccb5bb2f 100644 --- a/spinifex/daemon/health_test.go +++ b/spinifex/daemon/health_test.go @@ -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 diff --git a/spinifex/daemon/heartbeat.go b/spinifex/daemon/heartbeat.go index 7d7bb7945..d43467a1e 100644 --- a/spinifex/daemon/heartbeat.go +++ b/spinifex/daemon/heartbeat.go @@ -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, diff --git a/spinifex/daemon/heartbeat_test.go b/spinifex/daemon/heartbeat_test.go index 3d0cbfa2c..838ccb869 100644 --- a/spinifex/daemon/heartbeat_test.go +++ b/spinifex/daemon/heartbeat_test.go @@ -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() diff --git a/spinifex/daemon/jetstream.go b/spinifex/daemon/jetstream.go index b7c3d4095..83a1d7b5e 100644 --- a/spinifex/daemon/jetstream.go +++ b/spinifex/daemon/jetstream.go @@ -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 diff --git a/spinifex/daemon/jetstream_test.go b/spinifex/daemon/jetstream_test.go index e9486e46b..44b38bea4 100644 --- a/spinifex/daemon/jetstream_test.go +++ b/spinifex/daemon/jetstream_test.go @@ -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") } diff --git a/spinifex/daemon/shutdown.go b/spinifex/daemon/shutdown.go index 7ed8c281e..1ba782a2d 100644 --- a/spinifex/daemon/shutdown.go +++ b/spinifex/daemon/shutdown.go @@ -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{ diff --git a/spinifex/daemon/state.go b/spinifex/daemon/state.go index 15175624d..e2529bc4c 100644 --- a/spinifex/daemon/state.go +++ b/spinifex/daemon/state.go @@ -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)) diff --git a/spinifex/daemon/state_test.go b/spinifex/daemon/state_test.go index 8f2e04eaa..ec9571de2 100644 --- a/spinifex/daemon/state_test.go +++ b/spinifex/daemon/state_test.go @@ -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 diff --git a/spinifex/daemon/test_helpers_test.go b/spinifex/daemon/test_helpers_test.go new file mode 100644 index 000000000..3b69c5831 --- /dev/null +++ b/spinifex/daemon/test_helpers_test.go @@ -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...) +} diff --git a/spinifex/handlers/ec2/instance/service_impl.go b/spinifex/handlers/ec2/instance/service_impl.go index 4c7c097ce..b5e2e2bfa 100644 --- a/spinifex/handlers/ec2/instance/service_impl.go +++ b/spinifex/handlers/ec2/instance/service_impl.go @@ -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, } } diff --git a/spinifex/vm/manager.go b/spinifex/vm/manager.go new file mode 100644 index 000000000..06f150521 --- /dev/null +++ b/spinifex/vm/manager.go @@ -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) +} diff --git a/spinifex/vm/manager_test.go b/spinifex/vm/manager_test.go new file mode 100644 index 000000000..c98db5acb --- /dev/null +++ b/spinifex/vm/manager_test.go @@ -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) +} diff --git a/spinifex/vm/vm.go b/spinifex/vm/vm.go index b53b7a4be..812c8e00b 100644 --- a/spinifex/vm/vm.go +++ b/spinifex/vm/vm.go @@ -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"` }