mirror of
https://github.com/mulgadc/spinifex.git
synced 2026-10-15 17:11:30 +00:00
refactor: introduce vm.Manager (phase 1 of lifecycle extraction) (#176)
This commit is contained in:
+64
-75
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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{
|
||||
|
||||
@@ -177,9 +177,11 @@ func (d *Daemon) updateGuestDeviceNames(instance *vm.VM) {
|
||||
instance.EBSRequests.Mu.Unlock()
|
||||
|
||||
// Update BlockDeviceMappings using the mapping
|
||||
d.Instances.Mu.Lock()
|
||||
if instance.Instance != nil {
|
||||
for _, bdm := range instance.Instance.BlockDeviceMappings {
|
||||
d.vmMgr.Inspect(instance, func(v *vm.VM) {
|
||||
if v.Instance == nil {
|
||||
return
|
||||
}
|
||||
for _, bdm := range v.Instance.BlockDeviceMappings {
|
||||
if bdm.Ebs == nil || bdm.Ebs.VolumeId == nil || bdm.DeviceName == nil {
|
||||
continue
|
||||
}
|
||||
@@ -187,8 +189,7 @@ func (d *Daemon) updateGuestDeviceNames(instance *vm.VM) {
|
||||
bdm.DeviceName = &gd
|
||||
}
|
||||
}
|
||||
}
|
||||
d.Instances.Mu.Unlock()
|
||||
})
|
||||
|
||||
if err := d.WriteState(); err != nil {
|
||||
slog.Error("Failed to persist state after guest device name update", "instanceId", instance.ID, "err", err)
|
||||
|
||||
+52
-46
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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{
|
||||
|
||||
@@ -12,16 +12,22 @@ import (
|
||||
// On validation failure, the instance status is unchanged and an error is returned.
|
||||
// If WriteState fails, the in-memory status retains the new value (the VM has
|
||||
// physically changed state regardless) and an error is returned.
|
||||
// The caller must NOT hold d.Instances.Mu; this method acquires it internally.
|
||||
func (d *Daemon) TransitionState(instance *vm.VM, target vm.InstanceState) error {
|
||||
d.Instances.Mu.Lock()
|
||||
current := instance.Status
|
||||
if !vm.IsValidTransition(current, target) {
|
||||
d.Instances.Mu.Unlock()
|
||||
var (
|
||||
current vm.InstanceState
|
||||
invalid bool
|
||||
)
|
||||
d.vmMgr.Inspect(instance, func(v *vm.VM) {
|
||||
current = v.Status
|
||||
if !vm.IsValidTransition(current, target) {
|
||||
invalid = true
|
||||
return
|
||||
}
|
||||
v.Status = target
|
||||
})
|
||||
if invalid {
|
||||
return fmt.Errorf("invalid state transition: %s -> %s for instance %s", current, target, instance.ID)
|
||||
}
|
||||
instance.Status = target
|
||||
d.Instances.Mu.Unlock()
|
||||
|
||||
slog.Info("Instance state transition", "instanceId", instance.ID, "from", string(current), "to", string(target))
|
||||
|
||||
|
||||
+92
-128
@@ -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
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
package daemon
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
// assertVMPresent asserts that the daemon's VM manager has an entry for id.
|
||||
func assertVMPresent(t *testing.T, d *Daemon, id string, msgAndArgs ...any) {
|
||||
t.Helper()
|
||||
_, ok := d.vmMgr.Get(id)
|
||||
assert.True(t, ok, msgAndArgs...)
|
||||
}
|
||||
|
||||
// assertVMNotPresent asserts that the daemon's VM manager has no entry for id.
|
||||
func assertVMNotPresent(t *testing.T, d *Daemon, id string, msgAndArgs ...any) {
|
||||
t.Helper()
|
||||
_, ok := d.vmMgr.Get(id)
|
||||
assert.False(t, ok, msgAndArgs...)
|
||||
}
|
||||
@@ -231,17 +231,15 @@ type InstanceServiceImpl struct {
|
||||
config *config.Config
|
||||
instanceTypes map[string]*ec2.InstanceTypeInfo
|
||||
natsConn *nats.Conn
|
||||
instances *vm.Instances
|
||||
objectStore objectstore.ObjectStore
|
||||
}
|
||||
|
||||
// NewInstanceServiceImpl creates a new instance service implementation for daemon use
|
||||
func NewInstanceServiceImpl(cfg *config.Config, instanceTypes map[string]*ec2.InstanceTypeInfo, nc *nats.Conn, instances *vm.Instances, store objectstore.ObjectStore) *InstanceServiceImpl {
|
||||
func NewInstanceServiceImpl(cfg *config.Config, instanceTypes map[string]*ec2.InstanceTypeInfo, nc *nats.Conn, store objectstore.ObjectStore) *InstanceServiceImpl {
|
||||
return &InstanceServiceImpl{
|
||||
config: cfg,
|
||||
instanceTypes: instanceTypes,
|
||||
natsConn: nc,
|
||||
instances: instances,
|
||||
objectStore: store,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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"`
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user