From 8774fb1e953fc94f5150e563edf01f6b21215e60 Mon Sep 17 00:00:00 2001 From: Julian Sommer <115516918+juliansommer@users.noreply.github.com> Date: Thu, 12 Feb 2026 15:15:05 +1100 Subject: [PATCH] feat: distributed start instances (#37) --- .../start-instance-distributed-testing.md | 369 ++++++++++++++++++ .../feature/start-instance-distributed.md | 141 +++++++ hive/daemon/daemon.go | 41 +- hive/daemon/daemon_handlers.go | 354 +++++++++++++++++ hive/daemon/daemon_handlers_test.go | 301 ++++++++++++++ hive/daemon/jetstream.go | 107 ++++- hive/daemon/jetstream_test.go | 193 +++++++++ hive/daemon/state_test.go | 116 +++--- .../gateway/ec2/instance/DescribeInstances.go | 19 + hive/gateway/ec2/instance/StartInstances.go | 48 +-- .../ec2/instance/StartInstances_test.go | 52 +-- .../ec2/instance/TerminateInstances.go | 22 ++ .../ec2/instance/TerminateInstances_test.go | 73 ++++ hive/gateway/ec2/volume/AttachVolume.go | 38 ++ hive/vm/vm.go | 16 + 15 files changed, 1779 insertions(+), 111 deletions(-) create mode 100644 docs/development/feature/start-instance-distributed-testing.md create mode 100644 docs/development/feature/start-instance-distributed.md diff --git a/docs/development/feature/start-instance-distributed-testing.md b/docs/development/feature/start-instance-distributed-testing.md new file mode 100644 index 000000000..1f7b751ce --- /dev/null +++ b/docs/development/feature/start-instance-distributed-testing.md @@ -0,0 +1,369 @@ +# Distributed StartInstances - Manual Testing Guide + +Manual test plan for verifying stop/start instance lifecycle across daemon nodes, including daemon restarts. Requires a multi-node cluster (see INSTALL.md "Multi-node Configuration"). + +## Prerequisites + +```bash +# Build +make build + +# Set up 3-node cluster (if not already running) +./scripts/create-multi-node.sh +./scripts/start-dev.sh ~/node1/ +./scripts/start-dev.sh ~/node2/ +./scripts/start-dev.sh ~/node3/ + +# AWS CLI setup +export AWS_PROFILE=hive +``` + +Pick a gateway endpoint for all commands (any node works): + +```bash +export EP="--endpoint-url https://10.11.12.1:9999" +``` + +Import an SSH key and AMI if not already done (see INSTALL.md). + +--- + +## Test 1: Basic Stop and Start (Same Cluster) + +Verify an instance can be stopped and restarted by any daemon. + +```bash +# 1. Launch an instance +aws $EP ec2 run-instances \ + --image-id $HIVE_AMI \ + --instance-type $HIVE_INSTANCE \ + --key-name hive-key \ + --count 1 + +export INSTANCE_ID="i-XXX" # from output +``` + +```bash +# 2. Confirm running +aws $EP ec2 describe-instances --instance-ids $INSTANCE_ID +# Expect: State.Name = "running", State.Code = 16 +``` + +```bash +# 3. Stop the instance +aws $EP ec2 stop-instances --instance-ids $INSTANCE_ID +# Expect: PreviousState.Name = "running", CurrentState.Name = "stopping" +``` + +```bash +# 4. Wait for stop to complete, then verify stopped +aws $EP ec2 describe-instances --instance-ids $INSTANCE_ID +# Expect: State.Name = "stopped", State.Code = 80 +``` + +```bash +# 5. Check daemon logs on the original node — should show KV release +grep -i "Released stopped instance ownership" ~/node*/logs/hive.log +# Expect: "Released stopped instance ownership to shared KV" with instanceId and lastNode +``` + +```bash +# 6. Start the instance +aws $EP ec2 start-instances --instance-ids $INSTANCE_ID +# Expect: PreviousState.Name = "stopped", CurrentState.Name = "pending" +``` + +```bash +# 7. Confirm running again +aws $EP ec2 describe-instances --instance-ids $INSTANCE_ID +# Expect: State.Name = "running", State.Code = 16 +``` + +```bash +# 8. Check which node picked up the start +grep -i "Started stopped instance from shared KV" ~/node*/logs/hive.log +# Note the node — it may differ from the original +``` + +--- + +## Test 2: Stop, Restart Daemon, Then Start + +Verify a stopped instance survives a daemon restart and can be started after. + +```bash +# 1. Launch and stop an instance (repeat Test 1 steps 1-5) +export INSTANCE_ID="i-XXX" +``` + +```bash +# 2. Identify which node originally ran the instance +grep "Released stopped instance" ~/node*/logs/hive.log | grep $INSTANCE_ID +# Example: ~/node1/logs/hive.log — node1 was the original +``` + +```bash +# 3. Restart that daemon +./scripts/stop-dev.sh ~/node1/ +./scripts/start-dev.sh ~/node1/ +``` + +```bash +# 4. Verify the instance is still visible as stopped +aws $EP ec2 describe-instances --instance-ids $INSTANCE_ID +# Expect: State.Name = "stopped" +# The instance is in shared KV, so any running daemon can report it +``` + +```bash +# 5. Start the instance +aws $EP ec2 start-instances --instance-ids $INSTANCE_ID +# Expect: PreviousState.Name = "stopped", CurrentState.Name = "pending" +``` + +```bash +# 6. Confirm running and check which node picked it up +aws $EP ec2 describe-instances --instance-ids $INSTANCE_ID +grep "Started stopped instance from shared KV" ~/node*/logs/hive.log | grep $INSTANCE_ID +``` + +--- + +## Test 3: Stop on Node A, Kill Node A, Start on Node B + +Verify a stopped instance can be started when its original daemon is down. + +```bash +# 1. Launch instance (it lands on some node) +aws $EP ec2 run-instances \ + --image-id $HIVE_AMI \ + --instance-type $HIVE_INSTANCE \ + --key-name hive-key \ + --count 1 + +export INSTANCE_ID="i-XXX" +``` + +```bash +# 2. Stop the instance +aws $EP ec2 stop-instances --instance-ids $INSTANCE_ID +``` + +```bash +# 3. Wait for stopped state +aws $EP ec2 describe-instances --instance-ids $INSTANCE_ID +# Expect: State.Name = "stopped" +``` + +```bash +# 4. Identify the original node and shut it down +grep "Released stopped instance" ~/node*/logs/hive.log | grep $INSTANCE_ID +# If node1 — stop node1 +./scripts/stop-dev.sh ~/node1/ +``` + +```bash +# 5. Point CLI to a surviving node and start the instance +export EP="--endpoint-url https://10.11.12.2:9999" +aws $EP ec2 start-instances --instance-ids $INSTANCE_ID +# Expect: CurrentState.Name = "pending" — a different daemon picked it up +``` + +```bash +# 6. Confirm running on the new node +aws $EP ec2 describe-instances --instance-ids $INSTANCE_ID +# Expect: State.Name = "running" +grep "Started stopped instance from shared KV" ~/node2/logs/hive.log ~/node3/logs/hive.log +``` + +```bash +# 7. Bring node1 back up (cleanup) +./scripts/start-dev.sh ~/node1/ +``` + +--- + +## Test 4: Repeated Stop/Start Cycles + +Verify the instance can be stopped and started multiple times without state corruption. + +```bash +# 1. Launch an instance +aws $EP ec2 run-instances \ + --image-id $HIVE_AMI \ + --instance-type $HIVE_INSTANCE \ + --key-name hive-key \ + --count 1 + +export INSTANCE_ID="i-XXX" +``` + +```bash +# 2. Cycle 3 times +for i in 1 2 3; do + echo "=== Cycle $i ===" + + # Wait for running + sleep 5 + aws $EP ec2 describe-instances --instance-ids $INSTANCE_ID \ + --query 'Reservations[0].Instances[0].State.Name' --output text + # Expect: running + + # Stop + aws $EP ec2 stop-instances --instance-ids $INSTANCE_ID + sleep 10 + + aws $EP ec2 describe-instances --instance-ids $INSTANCE_ID \ + --query 'Reservations[0].Instances[0].State.Name' --output text + # Expect: stopped + + # Start + aws $EP ec2 start-instances --instance-ids $INSTANCE_ID + sleep 10 +done +``` + +```bash +# 3. Final state check +aws $EP ec2 describe-instances --instance-ids $INSTANCE_ID +# Expect: State.Name = "running" +``` + +--- + +## Test 5: DescribeInstances Includes Stopped Instances + +Verify stopped instances appear alongside running ones in describe output. + +```bash +# 1. Launch two instances +aws $EP ec2 run-instances \ + --image-id $HIVE_AMI \ + --instance-type $HIVE_INSTANCE \ + --key-name hive-key \ + --count 1 +export RUNNING_ID="i-XXX" + +aws $EP ec2 run-instances \ + --image-id $HIVE_AMI \ + --instance-type $HIVE_INSTANCE \ + --key-name hive-key \ + --count 1 +export STOPPED_ID="i-YYY" +``` + +```bash +# 2. Stop one, keep the other running +aws $EP ec2 stop-instances --instance-ids $STOPPED_ID +sleep 10 +``` + +```bash +# 3. Describe all instances — both should appear +aws $EP ec2 describe-instances +# Expect: $RUNNING_ID with State.Name = "running" +# Expect: $STOPPED_ID with State.Name = "stopped" +``` + +```bash +# 4. Describe with instance ID filter +aws $EP ec2 describe-instances --instance-ids $STOPPED_ID +# Expect: only the stopped instance +``` + +--- + +## Test 6: Start a Non-Existent Instance + +Verify proper error handling. + +```bash +aws $EP ec2 start-instances --instance-ids i-doesnotexist +# Expect: instance remains reported as stopped (state unchanged) +# Check daemon logs for "instance not found in shared KV" +``` + +--- + +## Test 7: Start an Already Running Instance + +Verify that starting a running instance does not cause issues. + +```bash +# 1. Launch and confirm running +aws $EP ec2 run-instances \ + --image-id $HIVE_AMI \ + --instance-type $HIVE_INSTANCE \ + --key-name hive-key \ + --count 1 +export INSTANCE_ID="i-XXX" +sleep 5 +``` + +```bash +# 2. Try to start it again +aws $EP ec2 start-instances --instance-ids $INSTANCE_ID +# Expect: instance not found in shared KV (running instances are not in shared KV) +# The response should report the instance as still "stopped" (no-op from gateway perspective) +``` + +--- + +## Test 8: Daemon Crash During Stop (Migration Recovery) + +Verify that a daemon crash between KV write and local cleanup is handled on restart. + +```bash +# 1. Launch and stop an instance +aws $EP ec2 run-instances \ + --image-id $HIVE_AMI \ + --instance-type $HIVE_INSTANCE \ + --key-name hive-key \ + --count 1 +export INSTANCE_ID="i-XXX" + +aws $EP ec2 stop-instances --instance-ids $INSTANCE_ID +sleep 10 +``` + +```bash +# 2. Force kill the daemon process (simulate crash — skips cleanup) +# Identify the daemon PID on the node that ran the instance +grep "Released stopped instance" ~/node*/logs/hive.log | grep $INSTANCE_ID +# If node1: +kill -9 $(pgrep -f "hive.*node1") +``` + +```bash +# 3. Restart the daemon +./scripts/start-dev.sh ~/node1/ +``` + +```bash +# 4. Check logs for migration handling +grep -i "Migrated stopped instance" ~/node1/logs/hive.log +# If the instance was still in per-node state, restoreInstances migrates it to shared KV +``` + +```bash +# 5. Verify instance is still startable +aws $EP ec2 start-instances --instance-ids $INSTANCE_ID +aws $EP ec2 describe-instances --instance-ids $INSTANCE_ID +# Expect: State.Name = "running" +``` + +--- + +## What to Look For in Logs + +Key log messages and where to find them (`~/node*/logs/hive.log`): + +| Log Message | When | +|-------------|------| +| `Released stopped instance ownership to shared KV` | After successful stop + KV write | +| `Started stopped instance from shared KV` | After a daemon picks up and launches a stopped instance | +| `Migrated stopped instance to shared KV` | During `restoreInstances` on daemon startup | +| `Failed to write stopped instance to shared KV, keeping local ownership` | KV write failed, instance stays local | +| `instance not found in shared KV` | Start request for non-existent instance | +| `instance not in stopped state` | Start request for instance not in stopped state | +| `DescribeInstances: Collected stopped instance reservations` | Gateway merged stopped instances into describe output | diff --git a/docs/development/feature/start-instance-distributed.md b/docs/development/feature/start-instance-distributed.md new file mode 100644 index 000000000..b0d797479 --- /dev/null +++ b/docs/development/feature/start-instance-distributed.md @@ -0,0 +1,141 @@ +# Distributed StartInstances + +Decouple stopped EC2 instances from their original daemon node, allowing any available daemon to handle `start-instances` requests. + +## Problem + +Previously, stopped instances remained bound to the daemon that ran them. The daemon kept an active `ec2.cmd.{instanceID}` NATS subscription after stop, so `start-instances` always routed back to the same node. In multi-node deployments, this meant a stopped instance could not be started if its original daemon was down or at capacity. + +## Request Flow + +``` +AWS SDK (start-instances) + -> Gateway (port 9999): parse AWS query params + -> NATS ec2.start (queue group: hive-workers): any daemon can pick up + -> Daemon handleEC2StartStoppedInstance: + 1. Load instance from shared NATS KV (key: instance.{instanceID}) + 2. Validate state is "stopped" + 3. Re-initialize non-serialized fields (QMPClient, EBSRequests.Mu) + 4. Allocate resources (CPU, memory) + 5. LaunchInstance (mount volumes, start QEMU, subscribe to ec2.cmd.{instanceID}) + 6. Delete from shared KV + <- NATS response: {"status":"running","instanceId":"..."} + <- Gateway: build InstanceStateChange XML + <- AWS SDK: StartInstancesOutput +``` + +## How Ownership Release Works + +When an instance is stopped (not terminated), the daemon releases ownership in this order: + +1. Set `inst.LastNode = d.node` (record which daemon last ran it) +2. Write instance to shared KV at key `instance.{instanceID}` +3. Unsubscribe from `ec2.cmd.{instanceID}` NATS topic +4. Remove instance from local `d.Instances.VMS` map +5. Persist local state via `d.WriteState()` (now without the stopped instance) + +The order is deliberate: shared KV is written **first**, then local state is cleaned up. If the daemon crashes between steps 2 and 5, `restoreInstances` on restart detects the overlap and completes the migration. + +## Shared KV Storage + +Stopped instances are stored in the existing `hive-instance-state` NATS JetStream KV bucket using the key prefix `instance.` (distinct from the per-node prefix `node.`). + +Four methods on `JetStreamManager`: +- `WriteStoppedInstance(instanceID, *vm.VM)` — serialize and write to `instance.{instanceID}` +- `LoadStoppedInstance(instanceID)` — returns `nil, nil` if not found (not an error) +- `DeleteStoppedInstance(instanceID)` — idempotent (ignores `ErrKeyNotFound`) +- `ListStoppedInstances()` — iterates all `instance.*` keys, returns `[]*vm.VM` + +## DescribeInstances Integration + +`DescribeInstances` now includes stopped instances from shared KV. After the existing fan-out to all daemons (which returns running/pending instances from each node's local map), the gateway sends an additional NATS request to `ec2.DescribeStoppedInstances` (queue group, 3s timeout). The response contains reservations for all stopped instances in shared KV. These are merged into the final output. + +If the stopped-instance query fails (timeout, no responders), it is logged and skipped — running instances are still returned. This matches the existing pattern where individual daemon fan-out failures are non-fatal. + +## Migration on Restart + +`restoreInstances` migrates pre-existing stopped instances from per-node state to shared KV: + +- Instances in `StateStopped`: set `LastNode`, write to shared KV, remove from local map +- Instances in `StateStopping`: transition to `StateStopped`, then migrate +- `WriteState()` is called at the end to persist the removals + +If migration fails for an instance (KV write error), the instance stays in local state as a fallback and will be retried on next restart. + +## Files Changed + +### VM struct: `hive/vm/vm.go` + +Added `LastNode string` field (json: `last_node,omitempty`). Records which daemon node last ran the instance. Set during ownership release on stop. Used for diagnostics and future WAL recovery. + +### KV methods: `hive/daemon/jetstream.go` + +Added `StoppedInstancePrefix = "instance."` constant and four CRUD methods for stopped instance state. These reuse the `hive-instance-state` bucket — no new bucket required. + +### Stop flow: `hive/daemon/daemon_handlers.go` + +Extended the background goroutine in `handleStopOrTerminateInstance`. After `TransitionState(inst, StateStopped)`, the stop-only path (not terminate) releases ownership to shared KV. The terminate path is unchanged. + +### New handlers: `hive/daemon/daemon_handlers.go` + +`handleEC2StartStoppedInstance` — loads from shared KV, validates stopped state, re-initializes non-serialized fields, allocates resources, launches instance, removes from shared KV. On failure, rolls back resource allocation and local map insertion. + +`handleEC2DescribeStoppedInstances` — lists all stopped instances from shared KV, applies instance ID filters from request payload, groups by reservation ID, returns `DescribeInstancesOutput`. + +### Subscriptions: `hive/daemon/daemon.go` + +Added to `subscribeAll()`: +- `ec2.start` with queue group `hive-workers` — load-balanced start requests +- `ec2.DescribeStoppedInstances` with queue group `hive-workers` — only one daemon needs to read shared KV + +### Gateway: `hive/gateway/ec2/instance/StartInstances.go` + +Rewritten to send `{"instance_id":"..."}` to `ec2.start` topic instead of constructing QMP commands and routing to `ec2.cmd.{instanceID}`. Removed `qmp` import. + +### Gateway: `hive/gateway/ec2/instance/DescribeInstances.go` + +Added a NATS request to `ec2.DescribeStoppedInstances` after the fan-out loop. Merges stopped instance reservations into the response. + +### Restore: `hive/daemon/daemon.go` + +`restoreInstances` now migrates stopped instances from per-node state to shared KV instead of re-subscribing to per-instance NATS topics. + +## Design Decisions + +**Why queue group for `ec2.start`?** NATS queue groups deliver each message to exactly one subscriber. This prevents two daemons from simultaneously trying to start the same instance, without needing distributed locks. + +**Why shared KV instead of a database?** The `hive-instance-state` KV bucket already exists for per-node state. Reusing it avoids new infrastructure. NATS JetStream KV provides consistency guarantees sufficient for this use case (single writer per key during normal operation). + +**Why write shared KV before removing local state?** Crash safety. If the daemon crashes after the KV write but before local cleanup, the instance exists in both locations. `restoreInstances` detects this overlap and completes the migration. The reverse order (remove local first, then write KV) would risk losing the instance entirely on crash. + +**Why `LastNode` instead of re-routing to the original daemon?** The original daemon may be down, at capacity, or decommissioned. Recording which node last ran an instance is useful for debugging and future WAL recovery (Phase 2), but the start request goes to whichever daemon is available. + +**Why not a dedicated KV bucket for stopped instances?** The `instance.` prefix provides sufficient namespace separation from `node.` prefix entries. A separate bucket would require additional JetStream configuration for no practical benefit. + +## Error Codes + +| Condition | AWS Error Code | +|-----------|---------------| +| No instance IDs provided | `InvalidParameterValue` (gateway) | +| Instance not found in shared KV | `InvalidInstanceID.NotFound` | +| Instance not in stopped state | `IncorrectInstanceState` | +| Resource allocation failure | `InsufficientInstanceCapacity` | +| LaunchInstance failure | `InternalError` | +| NATS timeout (no daemons available) | State reported as still stopped | + +## Testing + +**JetStream KV methods** (`jetstream_test.go`): 6 tests covering write/load round-trip, missing key returns nil, delete (including non-existent), list multiple instances, no interference with per-node state, nil KV bucket handling. + +**Handler error paths** (`daemon_handlers_test.go`): 5 tests covering missing instance ID, instance not found in KV, non-stopped state, describe returns stopped instances, describe with instance ID filter. + +**Gateway** (`StartInstances_test.go`): 6 tests covering single instance success, multiple instances, empty input, nil IDs skipped, NATS failure (reports stopped state), mixed success and failure. + +**State restore** (`state_test.go`): 5 tests updated to verify stopped instances are migrated to shared KV during `restoreInstances` instead of kept in local state. + +## Known Limitations + +- **All daemons must be updated together.** Old daemons don't subscribe to `ec2.start` and won't release stopped instances to shared KV. +- **No WAL flush before ownership release.** Phase 1 relies on `ebs.unmount` during stop to flush volume data. Phase 2 will add explicit WAL flush to Predastore before releasing ownership. +- **No capacity-aware scheduling.** `ec2.start` uses simple NATS queue group round-robin. A daemon might accept a start request even if it lacks resources, failing at allocation time. The gateway reports this as "still stopped" to the caller. +- **TOCTOU on shared KV.** Between loading an instance from KV and deleting it after successful start, another daemon could theoretically load the same key. In practice this doesn't happen because NATS queue group delivers each `ec2.start` message to exactly one subscriber, and the instance ID in the message maps to a single KV key. diff --git a/hive/daemon/daemon.go b/hive/daemon/daemon.go index 45a9bd056..45c05b3e6 100644 --- a/hive/daemon/daemon.go +++ b/hive/daemon/daemon.go @@ -464,6 +464,9 @@ func (d *Daemon) subscribeAll() error { {"ec2.CreateEgressOnlyInternetGateway", d.handleEC2CreateEgressOnlyInternetGateway, "hive-workers"}, {"ec2.DeleteEgressOnlyInternetGateway", d.handleEC2DeleteEgressOnlyInternetGateway, "hive-workers"}, {"ec2.DescribeEgressOnlyInternetGateways", d.handleEC2DescribeEgressOnlyInternetGateways, "hive-workers"}, + {"ec2.start", d.handleEC2StartStoppedInstance, "hive-workers"}, + {"ec2.terminate", d.handleEC2TerminateStoppedInstance, "hive-workers"}, + {"ec2.DescribeStoppedInstances", d.handleEC2DescribeStoppedInstances, "hive-workers"}, {"ec2.DescribeInstances", d.handleEC2DescribeInstances, ""}, {"ec2.DescribeInstanceTypes", d.handleEC2DescribeInstanceTypes, ""}, {"ec2.EnableEbsEncryptionByDefault", d.handleEC2EnableEbsEncryptionByDefault, "hive-workers"}, @@ -596,6 +599,23 @@ func (d *Daemon) initJetStream() error { return nil } +// migrateStoppedToSharedKV writes a stopped instance to shared KV and removes +// it from the local instance map. Returns true if migration succeeded. +func (d *Daemon) migrateStoppedToSharedKV(instance *vm.VM) bool { + if d.jsManager == nil { + return false + } + instance.LastNode = d.node + if err := d.jsManager.WriteStoppedInstance(instance.ID, instance); err != nil { + slog.Error("Failed to migrate stopped instance to shared KV", + "instance", instance.ID, "err", err) + return false + } + delete(d.Instances.VMS, instance.ID) + slog.Info("Migrated stopped instance to shared KV", "instance", instance.ID) + return true +} + // restoreInstances loads persisted VM state and re-launches instances that are // neither terminated nor flagged as user-stopped. func (d *Daemon) restoreInstances() { @@ -622,16 +642,7 @@ func (d *Daemon) restoreInstances() { } if instance.Status == vm.StateStopped { - slog.Info("Instance is stopped, subscribing to NATS for start commands", "instance", instance.ID) - d.mu.Lock() - sub, err := d.natsConn.Subscribe(fmt.Sprintf("ec2.cmd.%s", instance.ID), d.handleEC2Events) - if err != nil { - d.mu.Unlock() - slog.Error("Failed to subscribe stopped instance to NATS", "instance", instance.ID, "err", err) - continue - } - d.natsSubscriptions[instance.ID] = sub - d.mu.Unlock() + d.migrateStoppedToSharedKV(instance) continue } @@ -663,6 +674,11 @@ func (d *Daemon) restoreInstances() { } slog.Info("QEMU exited during transition, finalizing state", "instance", instance.ID, "from", prevStatus, "to", instance.Status) + + if instance.Status == vm.StateStopped && d.migrateStoppedToSharedKV(instance) { + continue + } + if err := d.WriteState(); err != nil { slog.Error("Failed to persist state, will retry on next restart", "instance", instance.ID, "error", err) @@ -680,6 +696,11 @@ func (d *Daemon) restoreInstances() { slog.Error("Failed to launch instance", "instanceId", instance.ID, "err", err) } } + + // Persist state after any migrations/removals during restore + if err := d.WriteState(); err != nil { + slog.Error("Failed to persist state after restore", "error", err) + } } // isInstanceProcessRunning checks if the QEMU process for an instance is still alive. diff --git a/hive/daemon/daemon_handlers.go b/hive/daemon/daemon_handlers.go index 7196d4720..b80cd7128 100644 --- a/hive/daemon/daemon_handlers.go +++ b/hive/daemon/daemon_handlers.go @@ -539,12 +539,54 @@ func (d *Daemon) handleStopOrTerminateInstance(msg *nats.Msg, command qmp.Comman } else { d.Instances.Mu.Lock() inst.Attributes = attrs + inst.LastNode = d.node d.Instances.Mu.Unlock() if err := d.TransitionState(inst, finalState); err != nil { slog.Error("Failed to transition to final state", "instanceId", inst.ID, "err", err) } slog.Info("Instance "+string(finalState), "id", inst.ID) + + // For stop (not terminate): release ownership to shared KV so any + // daemon can pick up the next StartInstance request. + if !isTerminate && d.jsManager != nil { + + // Write to shared KV first — if daemon crashes after this but + // before local cleanup, restoreInstances handles the overlap. + if err := d.jsManager.WriteStoppedInstance(inst.ID, inst); err != nil { + slog.Error("Failed to write stopped instance to shared KV, keeping local ownership", + "instanceId", inst.ID, "err", err) + } else { + // Unsubscribe from per-instance NATS topic + d.mu.Lock() + if sub, ok := d.natsSubscriptions[inst.ID]; ok { + if err := sub.Unsubscribe(); err != nil { + slog.Error("Failed to unsubscribe stopped instance", "instanceId", inst.ID, "err", err) + } + delete(d.natsSubscriptions, inst.ID) + } + d.mu.Unlock() + + // Remove from local instance map + d.Instances.Mu.Lock() + delete(d.Instances.VMS, inst.ID) + d.Instances.Mu.Unlock() + + // Persist local state without the stopped instance + if err := d.WriteState(); err != nil { + slog.Error("Failed to persist state after releasing stopped instance, re-adding to local map for consistency", + "instanceId", inst.ID, "err", err) + // Re-add to local map so in-memory state matches on-disk state. + // On next restart, restoreInstances will retry the migration. + d.Instances.Mu.Lock() + d.Instances.VMS[inst.ID] = inst + d.Instances.Mu.Unlock() + } else { + slog.Info("Released stopped instance ownership to shared KV", + "instanceId", inst.ID, "lastNode", d.node) + } + } + } } }(instance, command.Attributes) } @@ -1284,3 +1326,315 @@ func (d *Daemon) handleEC2EnableSerialConsoleAccess(msg *nats.Msg) { func (d *Daemon) handleEC2DisableSerialConsoleAccess(msg *nats.Msg) { handleNATSRequest(msg, d.accountService.DisableSerialConsoleAccess) } + +// startStoppedInstanceRequest is the payload for ec2.start topic +type startStoppedInstanceRequest struct { + InstanceID string `json:"instance_id"` +} + +// handleEC2StartStoppedInstance picks up a stopped instance from shared KV, +// re-launches it on this daemon node, and removes it from shared KV. +func (d *Daemon) handleEC2StartStoppedInstance(msg *nats.Msg) { + respondWithError := func(errCode string) { + if err := msg.Respond(utils.GenerateErrorPayload(errCode)); err != nil { + slog.Error("Failed to respond to NATS request", "err", err) + } + } + + var req startStoppedInstanceRequest + if err := json.Unmarshal(msg.Data, &req); err != nil { + slog.Error("handleEC2StartStoppedInstance: failed to unmarshal request", "err", err) + respondWithError(awserrors.ErrorServerInternal) + return + } + + if req.InstanceID == "" { + slog.Error("handleEC2StartStoppedInstance: missing instance_id") + respondWithError(awserrors.ErrorMissingParameter) + return + } + + if d.jsManager == nil { + slog.Error("handleEC2StartStoppedInstance: JetStream not available") + respondWithError(awserrors.ErrorServerInternal) + return + } + + // Load instance from shared KV + instance, err := d.jsManager.LoadStoppedInstance(req.InstanceID) + if err != nil { + slog.Error("handleEC2StartStoppedInstance: failed to load stopped instance", "instanceId", req.InstanceID, "err", err) + respondWithError(awserrors.ErrorServerInternal) + return + } + if instance == nil { + slog.Warn("handleEC2StartStoppedInstance: instance not found in shared KV", "instanceId", req.InstanceID) + respondWithError(awserrors.ErrorInvalidInstanceIDNotFound) + return + } + + if instance.Status != vm.StateStopped { + slog.Error("handleEC2StartStoppedInstance: instance not in stopped state", "instanceId", req.InstanceID, "status", instance.Status) + respondWithError(awserrors.ErrorIncorrectInstanceState) + return + } + + // Reset node-local fields that are stale after cross-node migration + instance.ResetNodeLocalState() + + // Allocate resources + instanceType, ok := d.resourceMgr.instanceTypes[instance.InstanceType] + if !ok { + slog.Error("handleEC2StartStoppedInstance: unknown instance type", "instanceId", req.InstanceID, "instanceType", instance.InstanceType) + respondWithError(awserrors.ErrorInvalidInstanceType) + return + } + if err := d.resourceMgr.allocate(instanceType); err != nil { + slog.Error("handleEC2StartStoppedInstance: failed to allocate resources", "instanceId", req.InstanceID, "err", err) + respondWithError(awserrors.ErrorInsufficientInstanceCapacity) + return + } + + // Add instance to local map and clear stop attribute before launch + d.Instances.Mu.Lock() + d.Instances.VMS[instance.ID] = instance + instance.Attributes = qmp.Attributes{StartInstance: true} + d.Instances.Mu.Unlock() + + // Launch the instance infrastructure (QEMU, QMP, NATS subscriptions) + err = d.LaunchInstance(instance) + if err != nil { + slog.Error("handleEC2StartStoppedInstance: LaunchInstance failed", "instanceId", req.InstanceID, "err", err) + // Rollback: deallocate resources and remove from local map + if ok { + d.resourceMgr.deallocate(instanceType) + } + d.Instances.Mu.Lock() + delete(d.Instances.VMS, instance.ID) + d.Instances.Mu.Unlock() + respondWithError(awserrors.ErrorServerInternal) + return + } + + // Remove from shared KV now that it's running locally. + // Retry once on failure — a stale KV entry risks duplicate starts. + if err := d.jsManager.DeleteStoppedInstance(req.InstanceID); err != nil { + slog.Warn("handleEC2StartStoppedInstance: first KV delete failed, retrying", + "instanceId", req.InstanceID, "err", err) + if retryErr := d.jsManager.DeleteStoppedInstance(req.InstanceID); retryErr != nil { + slog.Error("handleEC2StartStoppedInstance: KV delete failed after retry, instance is running locally but stale entry remains in shared KV", + "instanceId", req.InstanceID, "err", retryErr) + } + } + + slog.Info("Started stopped instance from shared KV", "instanceId", instance.ID, "node", d.node) + + if err := msg.Respond(fmt.Appendf(nil, `{"status":"running","instanceId":"%s"}`, instance.ID)); err != nil { + slog.Error("Failed to respond to NATS request", "err", err) + } +} + +// terminateStoppedInstanceRequest is the payload for ec2.terminate topic +type terminateStoppedInstanceRequest struct { + InstanceID string `json:"instance_id"` +} + +// handleEC2TerminateStoppedInstance picks up a stopped instance from shared KV, +// deletes its volumes, and removes it from shared KV. +func (d *Daemon) handleEC2TerminateStoppedInstance(msg *nats.Msg) { + respondWithError := func(errCode string) { + if err := msg.Respond(utils.GenerateErrorPayload(errCode)); err != nil { + slog.Error("Failed to respond to NATS request", "err", err) + } + } + + var req terminateStoppedInstanceRequest + if err := json.Unmarshal(msg.Data, &req); err != nil { + slog.Error("handleEC2TerminateStoppedInstance: failed to unmarshal request", "err", err) + respondWithError(awserrors.ErrorServerInternal) + return + } + + if req.InstanceID == "" { + slog.Error("handleEC2TerminateStoppedInstance: missing instance_id") + respondWithError(awserrors.ErrorMissingParameter) + return + } + + if d.jsManager == nil { + slog.Error("handleEC2TerminateStoppedInstance: JetStream not available") + respondWithError(awserrors.ErrorServerInternal) + return + } + + // Load instance from shared KV + instance, err := d.jsManager.LoadStoppedInstance(req.InstanceID) + if err != nil { + slog.Error("handleEC2TerminateStoppedInstance: failed to load stopped instance", "instanceId", req.InstanceID, "err", err) + respondWithError(awserrors.ErrorServerInternal) + return + } + if instance == nil { + slog.Warn("handleEC2TerminateStoppedInstance: instance not found in shared KV", "instanceId", req.InstanceID) + respondWithError(awserrors.ErrorInvalidInstanceIDNotFound) + return + } + + if instance.Status != vm.StateStopped { + slog.Error("handleEC2TerminateStoppedInstance: instance not in stopped state", "instanceId", req.InstanceID, "status", instance.Status) + respondWithError(awserrors.ErrorIncorrectInstanceState) + return + } + + // Delete volumes — no QEMU shutdown or unmount needed (already done during stop) + instance.EBSRequests.Mu.Lock() + for _, ebsRequest := range instance.EBSRequests.Requests { + // Internal volumes (EFI, cloud-init) are always cleaned up via ebs.delete + if ebsRequest.EFI || ebsRequest.CloudInit { + ebsDeleteData, err := json.Marshal(config.EBSDeleteRequest{Volume: ebsRequest.Name}) + if err != nil { + slog.Error("handleEC2TerminateStoppedInstance: failed to marshal ebs.delete request", "name", ebsRequest.Name, "err", err) + continue + } + deleteMsg, err := d.natsConn.Request("ebs.delete", ebsDeleteData, 30*time.Second) + if err != nil { + slog.Warn("handleEC2TerminateStoppedInstance: ebs.delete failed for internal volume", "name", ebsRequest.Name, "err", err) + } else { + slog.Info("handleEC2TerminateStoppedInstance: ebs.delete sent for internal volume", "name", ebsRequest.Name, "data", string(deleteMsg.Data)) + } + continue + } + + // User-visible volumes: respect DeleteOnTermination flag + if !ebsRequest.DeleteOnTermination { + slog.Info("handleEC2TerminateStoppedInstance: volume has DeleteOnTermination=false, skipping", "name", ebsRequest.Name) + continue + } + + slog.Info("handleEC2TerminateStoppedInstance: deleting volume with DeleteOnTermination=true", "name", ebsRequest.Name) + _, err := d.volumeService.DeleteVolume(&ec2.DeleteVolumeInput{ + VolumeId: &ebsRequest.Name, + }) + if err != nil { + slog.Error("handleEC2TerminateStoppedInstance: failed to delete volume", "name", ebsRequest.Name, "err", err) + } + } + instance.EBSRequests.Mu.Unlock() + + // Remove from shared KV + if err := d.jsManager.DeleteStoppedInstance(req.InstanceID); err != nil { + slog.Error("handleEC2TerminateStoppedInstance: failed to delete from shared KV", "instanceId", req.InstanceID, "err", err) + respondWithError(awserrors.ErrorServerInternal) + return + } + + slog.Info("Terminated stopped instance from shared KV", "instanceId", req.InstanceID) + + if err := msg.Respond(fmt.Appendf(nil, `{"status":"terminated","instanceId":"%s"}`, req.InstanceID)); err != nil { + slog.Error("Failed to respond to NATS request", "err", err) + } +} + +// handleEC2DescribeStoppedInstances returns stopped instances from shared KV. +func (d *Daemon) handleEC2DescribeStoppedInstances(msg *nats.Msg) { + respondWithError := func(errCode string) { + if err := msg.Respond(utils.GenerateErrorPayload(errCode)); err != nil { + slog.Error("Failed to respond to NATS request", "err", err) + } + } + + if d.jsManager == nil { + respondWithError(awserrors.ErrorServerInternal) + return + } + + // Parse optional filters from the request + describeInput := &ec2.DescribeInstancesInput{} + if len(msg.Data) > 0 { + if errResp := utils.UnmarshalJsonPayload(describeInput, msg.Data); errResp != nil { + if err := msg.Respond(errResp); err != nil { + slog.Error("Failed to respond to NATS request", "err", err) + } + return + } + } + + // Build instance ID filter + instanceIDFilter := make(map[string]bool) + if len(describeInput.InstanceIds) > 0 { + for _, id := range describeInput.InstanceIds { + if id != nil { + instanceIDFilter[*id] = true + } + } + } + + instances, err := d.jsManager.ListStoppedInstances() + if err != nil { + slog.Error("handleEC2DescribeStoppedInstances: failed to list stopped instances", "err", err) + respondWithError(awserrors.ErrorServerInternal) + return + } + + // Group instances by reservation ID + reservationMap := make(map[string]*ec2.Reservation) + + for _, instance := range instances { + // Apply instance ID filter + if len(instanceIDFilter) > 0 && !instanceIDFilter[instance.ID] { + continue + } + + if instance.Reservation != nil && instance.Instance != nil { + resID := "" + if instance.Reservation.ReservationId != nil { + resID = *instance.Reservation.ReservationId + } + + 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 + instanceCopy := *instance.Instance + instanceCopy.State = &ec2.InstanceState{} + if info, ok := vm.EC2StateCodes[instance.Status]; ok { + instanceCopy.State.SetCode(info.Code) + instanceCopy.State.SetName(info.Name) + } else { + instanceCopy.State.SetCode(80) + instanceCopy.State.SetName("stopped") + } + + reservationMap[resID].Instances = append(reservationMap[resID].Instances, &instanceCopy) + } + } + + reservations := make([]*ec2.Reservation, 0, len(reservationMap)) + for _, reservation := range reservationMap { + reservations = append(reservations, reservation) + } + + output := &ec2.DescribeInstancesOutput{ + Reservations: reservations, + } + + jsonResponse, err := json.Marshal(output) + if err != nil { + slog.Error("handleEC2DescribeStoppedInstances: failed to marshal output", "err", err) + respondWithError(awserrors.ErrorServerInternal) + return + } + + if err := msg.Respond(jsonResponse); err != nil { + slog.Error("Failed to respond to NATS request", "err", err) + } + + slog.Info("handleEC2DescribeStoppedInstances completed", "count", len(reservations)) +} diff --git a/hive/daemon/daemon_handlers_test.go b/hive/daemon/daemon_handlers_test.go index 6516413b1..464d01714 100644 --- a/hive/daemon/daemon_handlers_test.go +++ b/hive/daemon/daemon_handlers_test.go @@ -795,3 +795,304 @@ func TestSetConfigPath(t *testing.T) { daemon.SetConfigPath("/etc/hive/config.toml") assert.Equal(t, "/etc/hive/config.toml", daemon.configPath) } + +// --- handleEC2StartStoppedInstance tests --- + +func TestHandleEC2StartStoppedInstance_MissingInstance(t *testing.T) { + natsURL := sharedJSNATSURL + + daemon := createFullTestDaemonWithJetStream(t, natsURL) + + sub, err := daemon.natsConn.QueueSubscribe("ec2.start", "hive-workers", daemon.handleEC2StartStoppedInstance) + require.NoError(t, err) + defer sub.Unsubscribe() + + // Request to start a non-existent instance + reqData, _ := json.Marshal(map[string]string{"instance_id": "i-nonexistent"}) + reply, err := daemon.natsConn.Request("ec2.start", reqData, 5*time.Second) + require.NoError(t, err) + + var errResp map[string]any + err = json.Unmarshal(reply.Data, &errResp) + require.NoError(t, err) + assert.Contains(t, errResp, "Code") +} + +func TestHandleEC2StartStoppedInstance_MissingInstanceID(t *testing.T) { + natsURL := sharedJSNATSURL + + daemon := createFullTestDaemonWithJetStream(t, natsURL) + + sub, err := daemon.natsConn.QueueSubscribe("ec2.start", "hive-workers", daemon.handleEC2StartStoppedInstance) + require.NoError(t, err) + defer sub.Unsubscribe() + + // Request with empty instance_id + reqData, _ := json.Marshal(map[string]string{"instance_id": ""}) + reply, err := daemon.natsConn.Request("ec2.start", reqData, 5*time.Second) + require.NoError(t, err) + + var errResp map[string]any + err = json.Unmarshal(reply.Data, &errResp) + require.NoError(t, err) + assert.Contains(t, errResp, "Code") +} + +func TestHandleEC2StartStoppedInstance_NotStoppedState(t *testing.T) { + natsURL := sharedJSNATSURL + + daemon := createFullTestDaemonWithJetStream(t, natsURL) + + // Write an instance in running state to shared KV + runningVM := &vm.VM{ + ID: "i-running-shared", + Status: vm.StateRunning, + InstanceType: getTestInstanceType(), + } + err := daemon.jsManager.WriteStoppedInstance(runningVM.ID, runningVM) + require.NoError(t, err) + + sub, err := daemon.natsConn.QueueSubscribe("ec2.start", "hive-workers", daemon.handleEC2StartStoppedInstance) + require.NoError(t, err) + defer sub.Unsubscribe() + + reqData, _ := json.Marshal(map[string]string{"instance_id": "i-running-shared"}) + reply, err := daemon.natsConn.Request("ec2.start", reqData, 5*time.Second) + require.NoError(t, err) + + var errResp map[string]any + err = json.Unmarshal(reply.Data, &errResp) + require.NoError(t, err) + assert.Contains(t, errResp, "Code") + + // Cleanup + _ = daemon.jsManager.DeleteStoppedInstance(runningVM.ID) +} + +// --- handleEC2DescribeStoppedInstances tests --- + +func TestHandleEC2DescribeStoppedInstances_ReturnsStoppedInstances(t *testing.T) { + natsURL := sharedJSNATSURL + + daemon := createFullTestDaemonWithJetStream(t, natsURL) + + // Write stopped instances to shared KV with full reservation/instance metadata + stoppedVM := &vm.VM{ + ID: "i-describe-stopped-001", + Status: vm.StateStopped, + InstanceType: getTestInstanceType(), + LastNode: "node-1", + Reservation: &ec2.Reservation{ + ReservationId: aws.String("r-test-001"), + OwnerId: aws.String("123456789012"), + }, + Instance: &ec2.Instance{ + InstanceId: aws.String("i-describe-stopped-001"), + InstanceType: aws.String(getTestInstanceType()), + }, + } + err := daemon.jsManager.WriteStoppedInstance(stoppedVM.ID, stoppedVM) + require.NoError(t, err) + + sub, err := daemon.natsConn.QueueSubscribe("ec2.DescribeStoppedInstances", "hive-workers", daemon.handleEC2DescribeStoppedInstances) + require.NoError(t, err) + defer sub.Unsubscribe() + + input := &ec2.DescribeInstancesInput{} + reqData, _ := json.Marshal(input) + reply, err := daemon.natsConn.Request("ec2.DescribeStoppedInstances", reqData, 5*time.Second) + require.NoError(t, err) + + var output ec2.DescribeInstancesOutput + err = json.Unmarshal(reply.Data, &output) + require.NoError(t, err) + + // Find our stopped instance in the output + found := false + for _, res := range output.Reservations { + for _, inst := range res.Instances { + if inst.InstanceId != nil && *inst.InstanceId == "i-describe-stopped-001" { + found = true + assert.Equal(t, "stopped", *inst.State.Name) + assert.Equal(t, int64(80), *inst.State.Code) + } + } + } + assert.True(t, found, "Should find stopped instance in DescribeStoppedInstances output") + + // Cleanup + _ = daemon.jsManager.DeleteStoppedInstance(stoppedVM.ID) +} + +func TestHandleEC2DescribeStoppedInstances_WithFilter(t *testing.T) { + natsURL := sharedJSNATSURL + + daemon := createFullTestDaemonWithJetStream(t, natsURL) + + // Write two stopped instances + for _, id := range []string{"i-filter-001", "i-filter-002"} { + v := &vm.VM{ + ID: id, + Status: vm.StateStopped, + LastNode: "node-1", + Reservation: &ec2.Reservation{ + ReservationId: aws.String("r-filter"), + OwnerId: aws.String("123456789012"), + }, + Instance: &ec2.Instance{ + InstanceId: aws.String(id), + }, + } + err := daemon.jsManager.WriteStoppedInstance(id, v) + require.NoError(t, err) + } + + sub, err := daemon.natsConn.QueueSubscribe("ec2.DescribeStoppedInstances", "hive-workers", daemon.handleEC2DescribeStoppedInstances) + require.NoError(t, err) + defer sub.Unsubscribe() + + // Filter for only one instance + input := &ec2.DescribeInstancesInput{ + InstanceIds: []*string{aws.String("i-filter-001")}, + } + reqData, _ := json.Marshal(input) + reply, err := daemon.natsConn.Request("ec2.DescribeStoppedInstances", reqData, 5*time.Second) + require.NoError(t, err) + + var output ec2.DescribeInstancesOutput + err = json.Unmarshal(reply.Data, &output) + require.NoError(t, err) + + // Count matching instances + var matchCount int + for _, res := range output.Reservations { + for _, inst := range res.Instances { + if inst.InstanceId != nil && *inst.InstanceId == "i-filter-001" { + matchCount++ + } + // Should NOT contain i-filter-002 + if inst.InstanceId != nil && *inst.InstanceId == "i-filter-002" { + t.Error("Should not contain i-filter-002 when filtering for i-filter-001") + } + } + } + assert.Equal(t, 1, matchCount, "Should find exactly one filtered instance") + + // Cleanup + _ = daemon.jsManager.DeleteStoppedInstance("i-filter-001") + _ = daemon.jsManager.DeleteStoppedInstance("i-filter-002") +} + +// --- handleEC2TerminateStoppedInstance tests --- + +func TestHandleEC2TerminateStoppedInstance_MissingInstanceID(t *testing.T) { + natsURL := sharedJSNATSURL + + daemon := createFullTestDaemonWithJetStream(t, natsURL) + + sub, err := daemon.natsConn.QueueSubscribe("ec2.terminate", "hive-workers", daemon.handleEC2TerminateStoppedInstance) + require.NoError(t, err) + defer sub.Unsubscribe() + + reqData, _ := json.Marshal(map[string]string{"instance_id": ""}) + reply, err := daemon.natsConn.Request("ec2.terminate", reqData, 5*time.Second) + require.NoError(t, err) + + var errResp map[string]any + err = json.Unmarshal(reply.Data, &errResp) + require.NoError(t, err) + assert.Contains(t, errResp, "Code") +} + +func TestHandleEC2TerminateStoppedInstance_MissingInstance(t *testing.T) { + natsURL := sharedJSNATSURL + + daemon := createFullTestDaemonWithJetStream(t, natsURL) + + sub, err := daemon.natsConn.QueueSubscribe("ec2.terminate", "hive-workers", daemon.handleEC2TerminateStoppedInstance) + require.NoError(t, err) + defer sub.Unsubscribe() + + reqData, _ := json.Marshal(map[string]string{"instance_id": "i-nonexistent"}) + reply, err := daemon.natsConn.Request("ec2.terminate", reqData, 5*time.Second) + require.NoError(t, err) + + var errResp map[string]any + err = json.Unmarshal(reply.Data, &errResp) + require.NoError(t, err) + assert.Contains(t, errResp, "Code") +} + +func TestHandleEC2TerminateStoppedInstance_NotStoppedState(t *testing.T) { + natsURL := sharedJSNATSURL + + daemon := createFullTestDaemonWithJetStream(t, natsURL) + + // Write an instance in running state to shared KV + runningVM := &vm.VM{ + ID: "i-term-running", + Status: vm.StateRunning, + InstanceType: getTestInstanceType(), + } + err := daemon.jsManager.WriteStoppedInstance(runningVM.ID, runningVM) + require.NoError(t, err) + + sub, err := daemon.natsConn.QueueSubscribe("ec2.terminate", "hive-workers", daemon.handleEC2TerminateStoppedInstance) + require.NoError(t, err) + defer sub.Unsubscribe() + + reqData, _ := json.Marshal(map[string]string{"instance_id": "i-term-running"}) + reply, err := daemon.natsConn.Request("ec2.terminate", reqData, 5*time.Second) + require.NoError(t, err) + + var errResp map[string]any + err = json.Unmarshal(reply.Data, &errResp) + require.NoError(t, err) + assert.Contains(t, errResp, "Code") + + // Cleanup + _ = daemon.jsManager.DeleteStoppedInstance(runningVM.ID) +} + +func TestHandleEC2TerminateStoppedInstance_Success(t *testing.T) { + natsURL := sharedJSNATSURL + + daemon := createFullTestDaemonWithJetStream(t, natsURL) + + // Write a stopped instance to shared KV + stoppedVM := &vm.VM{ + ID: "i-term-stopped-001", + Status: vm.StateStopped, + InstanceType: getTestInstanceType(), + LastNode: "node-1", + Reservation: &ec2.Reservation{ + ReservationId: aws.String("r-term-001"), + OwnerId: aws.String("123456789012"), + }, + Instance: &ec2.Instance{ + InstanceId: aws.String("i-term-stopped-001"), + InstanceType: aws.String(getTestInstanceType()), + }, + } + err := daemon.jsManager.WriteStoppedInstance(stoppedVM.ID, stoppedVM) + require.NoError(t, err) + + sub, err := daemon.natsConn.QueueSubscribe("ec2.terminate", "hive-workers", daemon.handleEC2TerminateStoppedInstance) + require.NoError(t, err) + defer sub.Unsubscribe() + + reqData, _ := json.Marshal(map[string]string{"instance_id": "i-term-stopped-001"}) + reply, err := daemon.natsConn.Request("ec2.terminate", reqData, 5*time.Second) + require.NoError(t, err) + + var resp map[string]string + err = json.Unmarshal(reply.Data, &resp) + require.NoError(t, err) + assert.Equal(t, "terminated", resp["status"]) + assert.Equal(t, "i-term-stopped-001", resp["instanceId"]) + + // Verify instance was removed from shared KV + loaded, err := daemon.jsManager.LoadStoppedInstance("i-term-stopped-001") + require.NoError(t, err) + assert.Nil(t, loaded, "Instance should be removed from shared KV after termination") +} diff --git a/hive/daemon/jetstream.go b/hive/daemon/jetstream.go index 96a696619..a7ff3dbda 100644 --- a/hive/daemon/jetstream.go +++ b/hive/daemon/jetstream.go @@ -4,6 +4,7 @@ import ( "encoding/json" "errors" "log/slog" + "strings" "github.com/mulgadc/hive/hive/vm" "github.com/nats-io/nats.go" @@ -12,8 +13,10 @@ import ( const ( // InstanceStateBucket is the name of the KV bucket for storing instance state InstanceStateBucket = "hive-instance-state" - // InstanceStatePrefix is the key prefix for instance state entries + // InstanceStatePrefix is the key prefix for per-node instance state entries InstanceStatePrefix = "node." + // StoppedInstancePrefix is the key prefix for stopped instances in shared KV + StoppedInstancePrefix = "instance." ) // JetStreamManager manages JetStream KV store operations for instance state @@ -179,3 +182,105 @@ func (m *JetStreamManager) UpdateReplicas(newReplicas int) error { slog.Info("Updated JetStream KV bucket replicas", "bucket", InstanceStateBucket, "oldReplicas", currentReplicas, "newReplicas", newReplicas) return nil } + +// WriteStoppedInstance writes a stopped instance to the shared KV store. +func (m *JetStreamManager) WriteStoppedInstance(instanceID string, instance *vm.VM) error { + if m.kv == nil { + return errors.New("KV bucket not initialized") + } + + jsonData, err := json.Marshal(instance) + if err != nil { + return err + } + + key := StoppedInstancePrefix + instanceID + _, err = m.kv.Put(key, jsonData) + if err != nil { + return err + } + + slog.Debug("Wrote stopped instance to JetStream KV", "key", key, "instanceId", instanceID) + return nil +} + +// LoadStoppedInstance loads a stopped instance from the shared KV store. +// Returns nil, nil if the key does not exist. +func (m *JetStreamManager) LoadStoppedInstance(instanceID string) (*vm.VM, error) { + if m.kv == nil { + return nil, errors.New("KV bucket not initialized") + } + + key := StoppedInstancePrefix + instanceID + entry, err := m.kv.Get(key) + if err != nil { + if errors.Is(err, nats.ErrKeyNotFound) { + return nil, nil + } + return nil, err + } + + var instance vm.VM + if err := json.Unmarshal(entry.Value(), &instance); err != nil { + return nil, err + } + + return &instance, nil +} + +// DeleteStoppedInstance removes a stopped instance from the shared KV store. +// It is idempotent — deleting a non-existent key is not an error. +func (m *JetStreamManager) DeleteStoppedInstance(instanceID string) error { + if m.kv == nil { + return errors.New("KV bucket not initialized") + } + + key := StoppedInstancePrefix + instanceID + err := m.kv.Delete(key) + if err != nil && !errors.Is(err, nats.ErrKeyNotFound) { + return err + } + + slog.Debug("Deleted stopped instance from JetStream KV", "key", key) + return nil +} + +// ListStoppedInstances returns all stopped instances from the shared KV store. +func (m *JetStreamManager) ListStoppedInstances() ([]*vm.VM, error) { + if m.kv == nil { + return nil, errors.New("KV bucket not initialized") + } + + keys, err := m.kv.Keys() + if err != nil { + if errors.Is(err, nats.ErrNoKeysFound) { + return nil, nil + } + return nil, err + } + + var instances []*vm.VM + for _, key := range keys { + if !strings.HasPrefix(key, StoppedInstancePrefix) { + continue + } + + entry, err := m.kv.Get(key) + if err != nil { + if errors.Is(err, nats.ErrKeyNotFound) { + continue + } + return nil, err + } + + var instance vm.VM + if err := json.Unmarshal(entry.Value(), &instance); err != nil { + slog.Error("Failed to unmarshal stopped instance", "key", key, "err", err) + continue + } + + instances = append(instances, &instance) + } + + return instances, nil +} diff --git a/hive/daemon/jetstream_test.go b/hive/daemon/jetstream_test.go index 58d1d17c6..4bdbe57ef 100644 --- a/hive/daemon/jetstream_test.go +++ b/hive/daemon/jetstream_test.go @@ -378,3 +378,196 @@ func TestJetStreamManager_UpdateReplicas_NoInit(t *testing.T) { err := jsm.UpdateReplicas(3) assert.Error(t, err, "UpdateReplicas should error when JetStream not initialized") } + +// --- Stopped instance KV tests --- + +// TestJetStreamManager_WriteAndLoadStoppedInstance tests round-trip write and load of a stopped instance +func TestJetStreamManager_WriteAndLoadStoppedInstance(t *testing.T) { + nc, err := nats.Connect(sharedJSNATSURL) + require.NoError(t, err) + defer nc.Close() + + jsm, err := NewJetStreamManager(nc, 1) + require.NoError(t, err) + err = jsm.InitKVBucket() + require.NoError(t, err) + + testVM := &vm.VM{ + ID: "i-stopped-001", + Status: vm.StateStopped, + InstanceType: "t3.micro", + LastNode: "node-1", + } + + // Write stopped instance + err = jsm.WriteStoppedInstance(testVM.ID, testVM) + require.NoError(t, err) + + // Load stopped instance + loaded, err := jsm.LoadStoppedInstance(testVM.ID) + require.NoError(t, err) + require.NotNil(t, loaded) + assert.Equal(t, testVM.ID, loaded.ID) + assert.Equal(t, vm.StateStopped, loaded.Status) + assert.Equal(t, "t3.micro", loaded.InstanceType) + assert.Equal(t, "node-1", loaded.LastNode) + + // Cleanup + _ = jsm.DeleteStoppedInstance(testVM.ID) +} + +// TestJetStreamManager_LoadStoppedInstance_NotFound tests that LoadStoppedInstance returns nil for missing key +func TestJetStreamManager_LoadStoppedInstance_NotFound(t *testing.T) { + nc, err := nats.Connect(sharedJSNATSURL) + require.NoError(t, err) + defer nc.Close() + + jsm, err := NewJetStreamManager(nc, 1) + require.NoError(t, err) + err = jsm.InitKVBucket() + require.NoError(t, err) + + loaded, err := jsm.LoadStoppedInstance("i-nonexistent") + require.NoError(t, err) + assert.Nil(t, loaded) +} + +// TestJetStreamManager_DeleteStoppedInstance tests deleting a stopped instance (including non-existent) +func TestJetStreamManager_DeleteStoppedInstance(t *testing.T) { + nc, err := nats.Connect(sharedJSNATSURL) + require.NoError(t, err) + defer nc.Close() + + jsm, err := NewJetStreamManager(nc, 1) + require.NoError(t, err) + err = jsm.InitKVBucket() + require.NoError(t, err) + + testVM := &vm.VM{ + ID: "i-stopped-del", + Status: vm.StateStopped, + } + + // Write then delete + err = jsm.WriteStoppedInstance(testVM.ID, testVM) + require.NoError(t, err) + + err = jsm.DeleteStoppedInstance(testVM.ID) + require.NoError(t, err) + + // Verify it's gone + loaded, err := jsm.LoadStoppedInstance(testVM.ID) + require.NoError(t, err) + assert.Nil(t, loaded) + + // Delete non-existent should not error + err = jsm.DeleteStoppedInstance("i-does-not-exist") + require.NoError(t, err) +} + +// TestJetStreamManager_ListStoppedInstances tests listing multiple stopped instances +func TestJetStreamManager_ListStoppedInstances(t *testing.T) { + nc, err := nats.Connect(sharedJSNATSURL) + require.NoError(t, err) + defer nc.Close() + + jsm, err := NewJetStreamManager(nc, 1) + require.NoError(t, err) + err = jsm.InitKVBucket() + require.NoError(t, err) + + // Write multiple stopped instances + vms := []*vm.VM{ + {ID: "i-list-001", Status: vm.StateStopped, InstanceType: "t3.micro"}, + {ID: "i-list-002", Status: vm.StateStopped, InstanceType: "t3.small"}, + {ID: "i-list-003", Status: vm.StateStopped, InstanceType: "t3.medium"}, + } + + for _, v := range vms { + err = jsm.WriteStoppedInstance(v.ID, v) + require.NoError(t, err) + } + + // List stopped instances + instances, err := jsm.ListStoppedInstances() + require.NoError(t, err) + + // Should contain at least the 3 we wrote (may also contain instances from other tests) + instanceIDs := make(map[string]bool) + for _, inst := range instances { + instanceIDs[inst.ID] = true + } + assert.True(t, instanceIDs["i-list-001"]) + assert.True(t, instanceIDs["i-list-002"]) + assert.True(t, instanceIDs["i-list-003"]) + + // Cleanup + for _, v := range vms { + _ = jsm.DeleteStoppedInstance(v.ID) + } +} + +// TestJetStreamManager_StoppedInstances_NoInterference tests that stopped instances don't interfere with per-node state +func TestJetStreamManager_StoppedInstances_NoInterference(t *testing.T) { + nc, err := nats.Connect(sharedJSNATSURL) + require.NoError(t, err) + defer nc.Close() + + jsm, err := NewJetStreamManager(nc, 1) + require.NoError(t, err) + err = jsm.InitKVBucket() + 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}, + }, + } + err = jsm.WriteState("interference-test-node", nodeInstances) + require.NoError(t, err) + + // Write stopped instance + stoppedVM := &vm.VM{ID: "i-stopped-interf", Status: vm.StateStopped} + err = jsm.WriteStoppedInstance(stoppedVM.ID, stoppedVM) + require.NoError(t, err) + + // 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"]) + + // Verify stopped instance loads independently + loadedStopped, err := jsm.LoadStoppedInstance(stoppedVM.ID) + require.NoError(t, err) + require.NotNil(t, loadedStopped) + assert.Equal(t, "i-stopped-interf", loadedStopped.ID) + + // Cleanup + _ = jsm.DeleteStoppedInstance(stoppedVM.ID) + _ = jsm.DeleteState("interference-test-node") +} + +// TestJetStreamManager_StoppedInstance_KVNotInitialized tests error handling when KV is not initialized +func TestJetStreamManager_StoppedInstance_KVNotInitialized(t *testing.T) { + nc, err := nats.Connect(sharedJSNATSURL) + require.NoError(t, err) + defer nc.Close() + + jsm, err := NewJetStreamManager(nc, 1) + require.NoError(t, err) + // Don't call InitKVBucket + + err = jsm.WriteStoppedInstance("i-test", &vm.VM{}) + assert.Error(t, err) + + _, err = jsm.LoadStoppedInstance("i-test") + assert.Error(t, err) + + err = jsm.DeleteStoppedInstance("i-test") + assert.Error(t, err) + + _, err = jsm.ListStoppedInstances() + assert.Error(t, err) +} diff --git a/hive/daemon/state_test.go b/hive/daemon/state_test.go index de7521bf0..32ad7d3a3 100644 --- a/hive/daemon/state_test.go +++ b/hive/daemon/state_test.go @@ -1,7 +1,6 @@ package daemon import ( - "encoding/json" "fmt" "os" "sync" @@ -589,7 +588,8 @@ func TestTransitionState_MultipleInstancesIndependent(t *testing.T) { // exercises the recovery state resolution logic. // TestRestoreInstances_StoppingFinalizedToStopped verifies that an instance -// stuck in StateStopping when the daemon died gets finalized to StateStopped. +// stuck in StateStopping when the daemon died gets finalized to StateStopped +// and migrated to shared KV. func TestRestoreInstances_StoppingFinalizedToStopped(t *testing.T) { daemon := createDaemonWithJetStream(t) @@ -606,9 +606,14 @@ func TestRestoreInstances_StoppingFinalizedToStopped(t *testing.T) { // restoreInstances loads from JetStream and resolves transitional states. daemon.restoreInstances() - instance, ok := daemon.Instances.VMS["i-restore-stop"] - require.True(t, ok, "instance should be loaded from state") - assert.Equal(t, vm.StateStopped, instance.Status, + // Stopped instances are migrated to shared KV and removed from local map + _, ok := daemon.Instances.VMS["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") + require.NoError(t, err) + require.NotNil(t, stoppedInst, "stopping→stopped instance should exist in shared KV") + assert.Equal(t, vm.StateStopped, stoppedInst.Status, "stopping instance should be finalized to stopped on recovery") } @@ -652,9 +657,9 @@ func TestRestoreInstances_TerminatedSkipped(t *testing.T) { "terminated instance should remain terminated") } -// TestRestoreInstances_UserStoppedSkipped verifies that instances flagged as -// user-stopped are not relaunched. -func TestRestoreInstances_UserStoppedSkipped(t *testing.T) { +// TestRestoreInstances_UserStoppedMigratedToSharedKV verifies that instances +// flagged as user-stopped are migrated to shared KV and not relaunched. +func TestRestoreInstances_UserStoppedMigratedToSharedKV(t *testing.T) { daemon := createDaemonWithJetStream(t) daemon.Instances.VMS["i-restore-userstop"] = &vm.VM{ @@ -669,10 +674,15 @@ func TestRestoreInstances_UserStoppedSkipped(t *testing.T) { daemon.Instances.VMS = make(map[string]*vm.VM) daemon.restoreInstances() - instance, ok := daemon.Instances.VMS["i-restore-userstop"] - require.True(t, ok, "user-stopped instance should still be loaded") - assert.Equal(t, vm.StateStopped, instance.Status, - "user-stopped instance should remain stopped") + // Stopped instances should be migrated to shared KV + _, ok := daemon.Instances.VMS["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") + require.NoError(t, err) + require.NotNil(t, stoppedInst, "user-stopped instance should exist in shared KV") + assert.Equal(t, vm.StateStopped, stoppedInst.Status, + "user-stopped instance should remain stopped in shared KV") } // TestRestoreInstances_RunningResetToPending verifies that a running instance @@ -701,6 +711,7 @@ func TestRestoreInstances_RunningResetToPending(t *testing.T) { // TestRestoreInstances_MixedStates verifies recovery with multiple instances // in different states, ensuring each is handled correctly and independently. +// Stopped instances (including stopping→stopped transitions) are migrated to shared KV. func TestRestoreInstances_MixedStates(t *testing.T) { daemon := createDaemonWithJetStream(t) @@ -722,14 +733,25 @@ func TestRestoreInstances_MixedStates(t *testing.T) { daemon.Instances.VMS = make(map[string]*vm.VM) daemon.restoreInstances() - assert.Equal(t, vm.StateStopped, daemon.Instances.VMS["i-mix-stopping"].Status, - "stopping should finalize to stopped") + // Stopped instances are migrated to shared KV and removed from local map + assert.Nil(t, daemon.Instances.VMS["i-mix-stopping"], + "stopping→stopped should be migrated to shared KV") + stoppedFromStopping, err := daemon.jsManager.LoadStoppedInstance("i-mix-stopping") + require.NoError(t, err) + require.NotNil(t, stoppedFromStopping, "stopping→stopped should exist in shared KV") + assert.Equal(t, vm.StateStopped, stoppedFromStopping.Status) + assert.Equal(t, vm.StateTerminated, daemon.Instances.VMS["i-mix-shutting"].Status, "shutting-down should finalize to terminated") assert.Equal(t, vm.StateTerminated, daemon.Instances.VMS["i-mix-terminated"].Status, "terminated should remain terminated") - assert.Equal(t, vm.StateStopped, daemon.Instances.VMS["i-mix-stopped"].Status, - "user-stopped should remain stopped") + + assert.Nil(t, daemon.Instances.VMS["i-mix-stopped"], + "user-stopped should be migrated to shared KV") + stoppedFromUser, err := daemon.jsManager.LoadStoppedInstance("i-mix-stopped") + require.NoError(t, err) + require.NotNil(t, stoppedFromUser, "user-stopped should exist in shared KV") + assert.Equal(t, vm.StateStopped, stoppedFromUser.Status) } // TestRestoreInstances_StatePersistsAfterRecovery verifies that the finalized @@ -750,15 +772,27 @@ func TestRestoreInstances_StatePersistsAfterRecovery(t *testing.T) { daemon.Instances.VMS = make(map[string]*vm.VM) daemon.restoreInstances() - assert.Equal(t, vm.StateStopped, daemon.Instances.VMS["i-persist-stop"].Status) + // stopping→stopped migrates to shared KV + assert.Nil(t, daemon.Instances.VMS["i-persist-stop"], + "stopping→stopped should be migrated to shared KV") + stoppedInst, err := daemon.jsManager.LoadStoppedInstance("i-persist-stop") + require.NoError(t, err) + require.NotNil(t, stoppedInst) + assert.Equal(t, vm.StateStopped, stoppedInst.Status) + assert.Equal(t, vm.StateTerminated, daemon.Instances.VMS["i-persist-shut"].Status) - // Second restart: load again, verify states are already finalized. + // Second restart: stopped instance is in shared KV (not per-node state), + // so local map won't have it. Terminated instance persists in per-node state. daemon.Instances.VMS = make(map[string]*vm.VM) daemon.restoreInstances() - assert.Equal(t, vm.StateStopped, daemon.Instances.VMS["i-persist-stop"].Status, - "stopped state should persist through second restart") + // Stopped instance should still be in shared KV + stoppedInst2, err := daemon.jsManager.LoadStoppedInstance("i-persist-stop") + require.NoError(t, err) + require.NotNil(t, stoppedInst2, "stopped instance should persist in shared KV through second restart") + assert.Equal(t, vm.StateStopped, stoppedInst2.Status) + assert.Equal(t, vm.StateTerminated, daemon.Instances.VMS["i-persist-shut"].Status, "terminated state should persist through second restart") } @@ -792,10 +826,10 @@ func TestStatePersistence_RoundTrip(t *testing.T) { assert.Equal(t, original.Attributes.StopInstance, loaded.Attributes.StopInstance) } -// TestRestoreInstances_StoppedInstanceCanReceiveStartCommand verifies that after -// a daemon restart, a stopped instance has an active ec2.cmd. subscription -// so the gateway's StartInstances can reach it (no "no responders" error). -func TestRestoreInstances_StoppedInstanceCanReceiveStartCommand(t *testing.T) { +// TestRestoreInstances_StoppedInstanceMigratedToSharedKV verifies that after +// a daemon restart, a stopped instance is migrated to shared KV and can be +// retrieved by the ec2.start handler. +func TestRestoreInstances_StoppedInstanceMigratedToSharedKV(t *testing.T) { daemon := createDaemonWithJetStream(t) instanceID := "i-restore-start" @@ -811,31 +845,15 @@ func TestRestoreInstances_StoppedInstanceCanReceiveStartCommand(t *testing.T) { daemon.Instances.VMS = make(map[string]*vm.VM) daemon.restoreInstances() - // Instance should still be stopped. - instance, ok := daemon.Instances.VMS[instanceID] - require.True(t, ok, "stopped instance should be loaded after restore") - assert.Equal(t, vm.StateStopped, instance.Status) + // Instance should not be in local map (migrated to shared KV). + _, ok := daemon.Instances.VMS[instanceID] + assert.False(t, ok, "stopped instance should be migrated to shared KV, not in local map") - // Send a start command to the per-instance topic (ec2.cmd.). - // restoreInstances should have created this subscription for stopped instances. - command := qmp.Command{ - ID: instanceID, - QMPCommand: qmp.QMPCommand{ - Execute: "cont", - Arguments: map[string]any{}, - }, - Attributes: qmp.Attributes{ - StartInstance: true, - }, - } - cmdData, err := json.Marshal(command) + // Instance should be in shared KV + stoppedInst, err := daemon.jsManager.LoadStoppedInstance(instanceID) require.NoError(t, err) - - // The request should get a response (subscription exists), not "no responders". - // LaunchInstance will fail in tests (no QEMU), so we expect an error response, - // but the key assertion is that the subscription is active. - subject := fmt.Sprintf("ec2.cmd.%s", instanceID) - reply, err := daemon.natsConn.Request(subject, cmdData, 5*time.Second) - require.NoError(t, err, "ec2.cmd.%s should have a subscriber after restoreInstances", instanceID) - require.NotNil(t, reply) + require.NotNil(t, stoppedInst, "stopped instance should exist in shared KV") + assert.Equal(t, vm.StateStopped, stoppedInst.Status) + assert.Equal(t, "node-1", stoppedInst.LastNode, + "stopped instance should record last node") } diff --git a/hive/gateway/ec2/instance/DescribeInstances.go b/hive/gateway/ec2/instance/DescribeInstances.go index a48a658eb..2c01a892f 100644 --- a/hive/gateway/ec2/instance/DescribeInstances.go +++ b/hive/gateway/ec2/instance/DescribeInstances.go @@ -101,6 +101,25 @@ func DescribeInstances(input *ec2.DescribeInstancesInput, natsConn *nats.Conn, e } } + // Query stopped instances from shared KV (via queue group — only one daemon handles it) + stoppedMsg, err := natsConn.Request("ec2.DescribeStoppedInstances", jsonData, 3*time.Second) + if err != nil { + slog.Warn("DescribeInstances: Failed to query stopped instances (may not be available)", "err", err) + } else { + responseError, parseErr := utils.ValidateErrorPayload(stoppedMsg.Data) + if parseErr != nil { + slog.Warn("DescribeInstances: Stopped instances query returned error", "code", responseError.Code) + } else { + var stoppedOutput ec2.DescribeInstancesOutput + if err := json.Unmarshal(stoppedMsg.Data, &stoppedOutput); err != nil { + slog.Error("DescribeInstances: Failed to unmarshal stopped instances response", "err", err) + } else if stoppedOutput.Reservations != nil { + allReservations = append(allReservations, stoppedOutput.Reservations...) + slog.Info("DescribeInstances: Collected stopped instance reservations", "count", len(stoppedOutput.Reservations)) + } + } + } + // Build final aggregated response output := &ec2.DescribeInstancesOutput{ Reservations: allReservations, diff --git a/hive/gateway/ec2/instance/StartInstances.go b/hive/gateway/ec2/instance/StartInstances.go index c6a75a548..c6c4e410b 100644 --- a/hive/gateway/ec2/instance/StartInstances.go +++ b/hive/gateway/ec2/instance/StartInstances.go @@ -7,11 +7,17 @@ import ( "time" "github.com/aws/aws-sdk-go/service/ec2" - "github.com/mulgadc/hive/hive/qmp" + "github.com/mulgadc/hive/hive/utils" "github.com/nats-io/nats.go" ) -// StartInstances sends start commands to specified instances via NATS +// startStoppedInstanceRequest is the payload sent to the ec2.start topic +type startStoppedInstanceRequest struct { + InstanceID string `json:"instance_id"` +} + +// StartInstances sends start requests to the ec2.start queue group topic. +// Any available daemon can pick up the request and launch the stopped instance. func StartInstances(input *ec2.StartInstancesInput, natsConn *nats.Conn) (*ec2.StartInstancesOutput, error) { if len(input.InstanceIds) == 0 { return nil, fmt.Errorf("no instance IDs provided") @@ -21,48 +27,36 @@ func StartInstances(input *ec2.StartInstancesInput, natsConn *nats.Conn) (*ec2.S var stateChanges []*ec2.InstanceStateChange - // Process each instance for _, instanceIDPtr := range input.InstanceIds { if instanceIDPtr == nil { continue } instanceID := *instanceIDPtr - // Build the QMP command to resume the instance - command := qmp.Command{ - ID: instanceID, - QMPCommand: qmp.QMPCommand{ - Execute: "cont", - Arguments: map[string]any{}, - }, - Attributes: qmp.Attributes{ - StopInstance: false, - TerminateInstance: false, - StartInstance: true, - }, - } - - // Marshal the command - jsonData, err := json.Marshal(command) + req := startStoppedInstanceRequest{InstanceID: instanceID} + jsonData, err := json.Marshal(req) if err != nil { - slog.Error("StartInstances: Failed to marshal command", "instance_id", instanceID, "err", err) + slog.Error("StartInstances: Failed to marshal request", "instance_id", instanceID, "err", err) continue } - // Send NATS request to the specific instance topic - subject := fmt.Sprintf("ec2.cmd.%s", instanceID) - slog.Info("StartInstances: Sending NATS request", "subject", subject, "instance_id", instanceID) + slog.Info("StartInstances: Sending NATS request", "subject", "ec2.start", "instance_id", instanceID) - // TODO: Improve, stop/termination may take significant time - msg, err := natsConn.Request(subject, jsonData, 30*time.Second) + msg, err := natsConn.Request("ec2.start", jsonData, 30*time.Second) if err != nil { - slog.Error("StartInstances: Failed to send command", "instance_id", instanceID, "err", err) + slog.Error("StartInstances: Failed to send start request", "instance_id", instanceID, "err", err) + stateChanges = append(stateChanges, newStateChange(instanceID, 80, "stopped", 80, "stopped")) + continue + } + + // Check if the daemon returned an error response + if responseError, parseErr := utils.ValidateErrorPayload(msg.Data); parseErr != nil { + slog.Error("StartInstances: Daemon returned error", "instance_id", instanceID, "code", responseError.Code) stateChanges = append(stateChanges, newStateChange(instanceID, 80, "stopped", 80, "stopped")) continue } slog.Info("StartInstances: Command sent successfully", "instance_id", instanceID, "response", string(msg.Data)) - stateChanges = append(stateChanges, newStateChange(instanceID, 0, "pending", 80, "stopped")) } diff --git a/hive/gateway/ec2/instance/StartInstances_test.go b/hive/gateway/ec2/instance/StartInstances_test.go index 5920857bc..50d533d9b 100644 --- a/hive/gateway/ec2/instance/StartInstances_test.go +++ b/hive/gateway/ec2/instance/StartInstances_test.go @@ -6,7 +6,6 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/ec2" - "github.com/mulgadc/hive/hive/qmp" "github.com/nats-io/nats.go" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -17,19 +16,15 @@ func TestStartInstances_Success(t *testing.T) { instanceID := "i-0123456789abcdef0" - // Mock subscriber for the instance command topic - nc.Subscribe("ec2.cmd."+instanceID, func(msg *nats.Msg) { - var cmd qmp.Command - err := json.Unmarshal(msg.Data, &cmd) + // Mock subscriber for the ec2.start queue group topic + nc.QueueSubscribe("ec2.start", "hive-workers", func(msg *nats.Msg) { + var req startStoppedInstanceRequest + err := json.Unmarshal(msg.Data, &req) require.NoError(t, err) - assert.Equal(t, instanceID, cmd.ID) - assert.Equal(t, "cont", cmd.QMPCommand.Execute) - assert.True(t, cmd.Attributes.StartInstance) - assert.False(t, cmd.Attributes.StopInstance) - assert.False(t, cmd.Attributes.TerminateInstance) + assert.Equal(t, instanceID, req.InstanceID) - msg.Respond([]byte(`{"return":{}}`)) + msg.Respond([]byte(`{"status":"running","instanceId":"` + instanceID + `"}`)) }) input := &ec2.StartInstancesInput{ @@ -55,12 +50,12 @@ func TestStartInstances_MultipleInstances(t *testing.T) { ids := []string{"i-001", "i-002", "i-003"} - for _, id := range ids { - id := id - nc.Subscribe("ec2.cmd."+id, func(msg *nats.Msg) { - msg.Respond([]byte(`{"return":{}}`)) - }) - } + // Single subscriber handles all start requests via queue group + nc.QueueSubscribe("ec2.start", "hive-workers", func(msg *nats.Msg) { + var req startStoppedInstanceRequest + json.Unmarshal(msg.Data, &req) + msg.Respond([]byte(`{"status":"running","instanceId":"` + req.InstanceID + `"}`)) + }) input := &ec2.StartInstancesInput{ InstanceIds: []*string{aws.String(ids[0]), aws.String(ids[1]), aws.String(ids[2])}, @@ -94,8 +89,8 @@ func TestStartInstances_NilInstanceIdSkipped(t *testing.T) { _, nc := startTestNATSServer(t) instanceID := "i-valid" - nc.Subscribe("ec2.cmd."+instanceID, func(msg *nats.Msg) { - msg.Respond([]byte(`{"return":{}}`)) + nc.QueueSubscribe("ec2.start", "hive-workers", func(msg *nats.Msg) { + msg.Respond([]byte(`{"status":"running","instanceId":"` + instanceID + `"}`)) }) input := &ec2.StartInstancesInput{ @@ -114,7 +109,7 @@ func TestStartInstances_NilInstanceIdSkipped(t *testing.T) { func TestStartInstances_NATSRequestFails(t *testing.T) { _, nc := startTestNATSServer(t) - // No subscriber for this instance topic, so NATS request will fail + // No subscriber for ec2.start, so NATS request will fail instanceID := "i-nosubscriber" input := &ec2.StartInstancesInput{ @@ -140,10 +135,18 @@ func TestStartInstances_MixedSuccessAndFailure(t *testing.T) { _, nc := startTestNATSServer(t) goodID := "i-good" - badID := "i-bad" // no subscriber + badID := "i-bad" - nc.Subscribe("ec2.cmd."+goodID, func(msg *nats.Msg) { - msg.Respond([]byte(`{"return":{}}`)) + // Subscribe to ec2.start but only respond for goodID + nc.QueueSubscribe("ec2.start", "hive-workers", func(msg *nats.Msg) { + var req startStoppedInstanceRequest + json.Unmarshal(msg.Data, &req) + if req.InstanceID == goodID { + msg.Respond([]byte(`{"status":"running","instanceId":"` + goodID + `"}`)) + } else { + // Simulate error for bad instance + msg.Respond([]byte(`{"Code":"InvalidInstanceID.NotFound"}`)) + } }) input := &ec2.StartInstancesInput{ @@ -159,7 +162,8 @@ func TestStartInstances_MixedSuccessAndFailure(t *testing.T) { assert.Equal(t, goodID, *output.StartingInstances[0].InstanceId) assert.Equal(t, "pending", *output.StartingInstances[0].CurrentState.Name) - // Second instance failed: state → stopped (unchanged) + // Second instance: daemon returned an error, state remains stopped assert.Equal(t, badID, *output.StartingInstances[1].InstanceId) + assert.Equal(t, int64(80), *output.StartingInstances[1].CurrentState.Code) assert.Equal(t, "stopped", *output.StartingInstances[1].CurrentState.Name) } diff --git a/hive/gateway/ec2/instance/TerminateInstances.go b/hive/gateway/ec2/instance/TerminateInstances.go index 94ea14b1a..fb475f0c4 100644 --- a/hive/gateway/ec2/instance/TerminateInstances.go +++ b/hive/gateway/ec2/instance/TerminateInstances.go @@ -2,15 +2,22 @@ package gateway_ec2_instance import ( "encoding/json" + "errors" "fmt" "log/slog" "time" "github.com/aws/aws-sdk-go/service/ec2" "github.com/mulgadc/hive/hive/qmp" + "github.com/mulgadc/hive/hive/utils" "github.com/nats-io/nats.go" ) +// terminateStoppedInstanceRequest is the payload sent to the ec2.terminate topic +type terminateStoppedInstanceRequest struct { + InstanceID string `json:"instance_id"` +} + // TerminateInstances sends terminate commands to specified instances via NATS // Uses system_powerdown with stop_instance attribute to prevent restart func TerminateInstances(input *ec2.TerminateInstancesInput, natsConn *nats.Conn) (*ec2.TerminateInstancesOutput, error) { @@ -54,6 +61,21 @@ func TerminateInstances(input *ec2.TerminateInstancesInput, natsConn *nats.Conn) subject := fmt.Sprintf("ec2.cmd.%s", instanceID) msg, err := natsConn.Request(subject, jsonData, 5*time.Second) if err != nil { + // If no daemon owns this instance, try the ec2.terminate topic for stopped instances + if errors.Is(err, nats.ErrNoResponders) { + slog.Info("TerminateInstances: No responder on per-instance topic, trying ec2.terminate", "instance_id", instanceID) + + terminateReq, _ := json.Marshal(terminateStoppedInstanceRequest{InstanceID: instanceID}) + terminateMsg, terminateErr := natsConn.Request("ec2.terminate", terminateReq, 30*time.Second) + if terminateErr == nil { + if _, parseErr := utils.ValidateErrorPayload(terminateMsg.Data); parseErr == nil { + slog.Info("TerminateInstances: Stopped instance terminated via ec2.terminate", "instance_id", instanceID) + stateChanges = append(stateChanges, newStateChange(instanceID, 32, "shutting-down", 80, "stopped")) + continue + } + } + } + slog.Error("TerminateInstances: Failed to send command", "instance_id", instanceID, "err", err) stateChanges = append(stateChanges, newStateChange(instanceID, 16, "running", 16, "running")) continue diff --git a/hive/gateway/ec2/instance/TerminateInstances_test.go b/hive/gateway/ec2/instance/TerminateInstances_test.go index f22153d56..83b7f31be 100644 --- a/hive/gateway/ec2/instance/TerminateInstances_test.go +++ b/hive/gateway/ec2/instance/TerminateInstances_test.go @@ -181,3 +181,76 @@ func TestTerminateInstances_VerifiesQMPAttributes(t *testing.T) { assert.True(t, receivedCmd.Attributes.TerminateInstance, "TerminateInstance should be true") assert.False(t, receivedCmd.Attributes.StartInstance, "StartInstance should be false") } + +func TestTerminateInstances_StoppedInstanceFallback(t *testing.T) { + _, nc := startTestNATSServer(t) + + instanceID := "i-stopped-term" + + // No ec2.cmd. subscriber — simulate stopped instance with no daemon owning it. + // Subscribe to ec2.terminate to handle the fallback. + nc.QueueSubscribe("ec2.terminate", "hive-workers", func(msg *nats.Msg) { + var req terminateStoppedInstanceRequest + json.Unmarshal(msg.Data, &req) + assert.Equal(t, instanceID, req.InstanceID) + msg.Respond([]byte(`{"status":"terminated","instanceId":"` + instanceID + `"}`)) + }) + + input := &ec2.TerminateInstancesInput{ + InstanceIds: []*string{aws.String(instanceID)}, + } + + output, err := TerminateInstances(input, nc) + + require.NoError(t, err) + require.Len(t, output.TerminatingInstances, 1) + + sc := output.TerminatingInstances[0] + assert.Equal(t, instanceID, *sc.InstanceId) + assert.Equal(t, int64(32), *sc.CurrentState.Code) + assert.Equal(t, "shutting-down", *sc.CurrentState.Name) + assert.Equal(t, int64(80), *sc.PreviousState.Code) + assert.Equal(t, "stopped", *sc.PreviousState.Name) +} + +func TestTerminateInstances_MixedRunningAndStopped(t *testing.T) { + _, nc := startTestNATSServer(t) + + runningID := "i-running-mix" + stoppedID := "i-stopped-mix" + + // Running instance responds on ec2.cmd. + nc.Subscribe("ec2.cmd."+runningID, func(msg *nats.Msg) { + msg.Respond([]byte(`{"return":{}}`)) + }) + + // Stopped instance: no ec2.cmd subscriber, but ec2.terminate responds + nc.QueueSubscribe("ec2.terminate", "hive-workers", func(msg *nats.Msg) { + var req terminateStoppedInstanceRequest + json.Unmarshal(msg.Data, &req) + msg.Respond([]byte(`{"status":"terminated","instanceId":"` + req.InstanceID + `"}`)) + }) + + input := &ec2.TerminateInstancesInput{ + InstanceIds: []*string{aws.String(runningID), aws.String(stoppedID)}, + } + + output, err := TerminateInstances(input, nc) + + require.NoError(t, err) + require.Len(t, output.TerminatingInstances, 2) + + // Running instance: previous=running, current=shutting-down + assert.Equal(t, runningID, *output.TerminatingInstances[0].InstanceId) + assert.Equal(t, int64(32), *output.TerminatingInstances[0].CurrentState.Code) + assert.Equal(t, "shutting-down", *output.TerminatingInstances[0].CurrentState.Name) + assert.Equal(t, int64(16), *output.TerminatingInstances[0].PreviousState.Code) + assert.Equal(t, "running", *output.TerminatingInstances[0].PreviousState.Name) + + // Stopped instance: previous=stopped, current=shutting-down + assert.Equal(t, stoppedID, *output.TerminatingInstances[1].InstanceId) + assert.Equal(t, int64(32), *output.TerminatingInstances[1].CurrentState.Code) + assert.Equal(t, "shutting-down", *output.TerminatingInstances[1].CurrentState.Name) + assert.Equal(t, int64(80), *output.TerminatingInstances[1].PreviousState.Code) + assert.Equal(t, "stopped", *output.TerminatingInstances[1].PreviousState.Name) +} diff --git a/hive/gateway/ec2/volume/AttachVolume.go b/hive/gateway/ec2/volume/AttachVolume.go index 40672f8bc..7944221d2 100644 --- a/hive/gateway/ec2/volume/AttachVolume.go +++ b/hive/gateway/ec2/volume/AttachVolume.go @@ -7,6 +7,7 @@ import ( "log/slog" "time" + "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/ec2" "github.com/mulgadc/hive/hive/awserrors" "github.com/mulgadc/hive/hive/qmp" @@ -69,6 +70,10 @@ func AttachVolume(input *ec2.AttachVolumeInput, natsConn *nats.Conn) (ec2.Volume if err != nil { slog.Error("AttachVolume: NATS request failed", "instanceId", instanceID, "volumeId", volumeID, "err", err) if errors.Is(err, nats.ErrNoResponders) { + // Check shared KV for stopped instances before returning NotFound + if isStoppedInstance(instanceID, natsConn) { + return output, errors.New(awserrors.ErrorIncorrectInstanceState) + } return output, errors.New(awserrors.ErrorInvalidInstanceIDNotFound) } return output, errors.New(awserrors.ErrorServerInternal) @@ -87,3 +92,36 @@ func AttachVolume(input *ec2.AttachVolumeInput, natsConn *nats.Conn) (ec2.Volume slog.Info("AttachVolume completed", "instanceId", instanceID, "volumeId", volumeID) return output, nil } + +// isStoppedInstance checks the shared KV (via ec2.DescribeStoppedInstances) to +// determine whether instanceID exists as a stopped instance. +func isStoppedInstance(instanceID string, natsConn *nats.Conn) bool { + input := &ec2.DescribeInstancesInput{ + InstanceIds: []*string{aws.String(instanceID)}, + } + reqData, err := json.Marshal(input) + if err != nil { + return false + } + + msg, err := natsConn.Request("ec2.DescribeStoppedInstances", reqData, 3*time.Second) + if err != nil { + return false + } + + if _, err := utils.ValidateErrorPayload(msg.Data); err != nil { + return false + } + + var output ec2.DescribeInstancesOutput + if err := json.Unmarshal(msg.Data, &output); err != nil { + return false + } + + for _, res := range output.Reservations { + if len(res.Instances) > 0 { + return true + } + } + return false +} diff --git a/hive/vm/vm.go b/hive/vm/vm.go index 3614ef35e..9cf2e2b91 100644 --- a/hive/vm/vm.go +++ b/hive/vm/vm.go @@ -38,6 +38,10 @@ type VM struct { // Instance contains the current instance state and metadata Instance *ec2.Instance `json:"instance,omitempty"` + // LastNode records which daemon node last ran this instance. + // Set when ownership is released on stop for shared KV storage. + LastNode string `json:"last_node,omitempty"` + // User data for cloud-init (decoded from base64) UserData string `json:"user_data,omitempty"` @@ -45,6 +49,18 @@ type VM struct { MetadataServerAddress string `json:"metadata_server_address,omitempty"` } +// ResetNodeLocalState zeroes out fields that are specific to the daemon node +// that last ran this instance. Must be called after deserializing a VM from +// shared KV before launching it on a new node. +func (v *VM) ResetNodeLocalState() { + v.PID = 0 + v.PTS = 0 + v.Running = false + v.MetadataServerAddress = "" + v.QMPClient = &qmp.QMPClient{} + v.EBSRequests.Mu = sync.Mutex{} +} + type Instances struct { VMS map[string]*VM `json:"vms"` Mu sync.Mutex `json:"-"`