From 35d0df5b5681acbbe8b467fa7792c4167f72c04f Mon Sep 17 00:00:00 2001 From: Julian Sommer <115516918+juliansommer@users.noreply.github.com> Date: Tue, 24 Mar 2026 09:52:00 +1000 Subject: [PATCH] feat: placement groups (#85) --- docs/COMMANDS.md | 12 +- scripts/reset-dev-env.sh | 10 + spinifex/daemon/daemon.go | 113 +- spinifex/daemon/daemon_handlers_instance.go | 27 +- .../daemon/daemon_handlers_placementgroup.go | 39 + spinifex/daemon/daemon_test.go | 19 +- spinifex/gateway/ec2.go | 10 + spinifex/gateway/ec2/instance/RunInstances.go | 76 +- spinifex/gateway/ec2/instance/placement.go | 555 ++++++++ .../gateway/ec2/instance/placement_test.go | 1134 +++++++++++++++++ .../placementgroup/CreatePlacementGroup.go | 33 + .../placementgroup/DeletePlacementGroup.go | 30 + .../placementgroup/DescribePlacementGroups.go | 20 + spinifex/gateway/gateway_test.go | 1 + .../handlers/ec2/placementgroup/service.go | 77 ++ .../ec2/placementgroup/service_impl.go | 572 +++++++++ .../ec2/placementgroup/service_impl_test.go | 840 ++++++++++++ .../ec2/placementgroup/service_nats.go | 55 + spinifex/instancetypes/generate.go | 6 + spinifex/instancetypes/generate_test.go | 19 + spinifex/utils/pidfile.go | 4 + spinifex/vm/vm.go | 4 + tests/e2e/lib/multinode-helpers.sh | 58 +- tests/e2e/run-pseudo-multinode-e2e.sh | 694 +++++++++- 24 files changed, 4343 insertions(+), 65 deletions(-) create mode 100644 spinifex/daemon/daemon_handlers_placementgroup.go create mode 100644 spinifex/gateway/ec2/instance/placement.go create mode 100644 spinifex/gateway/ec2/instance/placement_test.go create mode 100644 spinifex/gateway/ec2/placementgroup/CreatePlacementGroup.go create mode 100644 spinifex/gateway/ec2/placementgroup/DeletePlacementGroup.go create mode 100644 spinifex/gateway/ec2/placementgroup/DescribePlacementGroups.go create mode 100644 spinifex/handlers/ec2/placementgroup/service.go create mode 100644 spinifex/handlers/ec2/placementgroup/service_impl.go create mode 100644 spinifex/handlers/ec2/placementgroup/service_impl_test.go create mode 100644 spinifex/handlers/ec2/placementgroup/service_nats.go diff --git a/docs/COMMANDS.md b/docs/COMMANDS.md index 4265b66ea..1c6750835 100644 --- a/docs/COMMANDS.md +++ b/docs/COMMANDS.md @@ -80,12 +80,12 @@ Operational commands for inspecting cluster state. These fan out NATS requests t | Command | Implemented Flags | Missing Flags | Prerequisites | Basic Logic | Test Cases | Status | |---------|-------------------|---------------|---------------|-------------|------------|--------| -| `run-instances` | `--image-id`, `--instance-type`, `--count` (Min/MaxCount), `--key-name`, `--user-data`, `--subnet-id` (auto-creates ENI, assigns private IP), `--block-device-mappings` (DeviceName, VolumeSize, VolumeType, Iops, DeleteOnTermination) | `--security-group-ids`, `--placement`, `--tag-specifications`, `--dry-run`, `--client-token`, `--disable-api-termination`, `--ebs-optimized`, `--iam-instance-profile`, `--network-interfaces`, `--private-ip-address`, `--monitoring`, `--credit-specification`, `--cpu-options`, `--metadata-options`, `--launch-template`, `--hibernate-options` | `describe-images` (AMI must exist), `create-key-pair` (optional), VPC/SG (optional) | Gateway parses AWS query → NATS `ec2.runinstances` → daemon creates QEMU/KVM VM with viperblock-backed root volume via NBD → if SubnetId provided, auto-creates ENI with private IP → cloud-init injects user-data/keys → returns reservation with instance ID | 1. Launch with valid AMI and key pair
2. Launch with invalid AMI ID (error)
3. Launch with block device mappings (custom volume size)
4. Launch multiple instances (MinCount/MaxCount)
5. Launch with subnet-id (auto-creates ENI)
6. Invalid instance type returns error | **DONE** | +| `run-instances` | `--image-id`, `--instance-type`, `--count` (Min/MaxCount), `--key-name`, `--user-data`, `--subnet-id` (auto-creates ENI, assigns private IP), `--block-device-mappings` (DeviceName, VolumeSize, VolumeType, Iops, DeleteOnTermination), `--placement` (GroupName only — routes via spread or cluster strategy) | `--security-group-ids`, `--tag-specifications`, `--dry-run`, `--client-token`, `--disable-api-termination`, `--ebs-optimized`, `--iam-instance-profile`, `--network-interfaces`, `--private-ip-address`, `--monitoring`, `--credit-specification`, `--cpu-options`, `--metadata-options`, `--launch-template`, `--hibernate-options` | `describe-images` (AMI must exist), `create-key-pair` (optional), VPC/SG (optional) | Gateway parses AWS query → if Placement.GroupName set, looks up strategy: spread → `distributeInstancesSpread()` (1 instance per node, atomic CAS reservation), cluster → `distributeInstancesCluster()` (pin all to single node); otherwise NATS `ec2.runinstances` → daemon creates QEMU/KVM VM with viperblock-backed root volume via NBD → if SubnetId provided, auto-creates ENI with private IP → cloud-init injects user-data/keys → on termination, removes instance from placement group → returns reservation with instance ID | 1. Launch with valid AMI and key pair
2. Launch with invalid AMI ID (error)
3. Launch with block device mappings (custom volume size)
4. Launch multiple instances (MinCount/MaxCount)
5. Launch with subnet-id (auto-creates ENI)
6. Invalid instance type returns error
7. Launch with spread placement group (1 per node)
8. Launch with cluster placement group (all on one node)
9. Insufficient capacity for placement group (error) | **DONE** | | `describe-instances` | `--instance-ids` | `--filters`, `--max-results`, `--next-token`, `--dry-run` | None | Gateway fans out NATS `ec2.DescribeInstances` to all nodes (no queue group) → each daemon returns local instances → gateway aggregates and returns combined list | 1. Describe all instances (no filter)
2. Describe by instance ID
3. Describe with filters (e.g. instance-state-name)
4. Instance not found returns empty set
5. Multi-node aggregation returns instances from all nodes | **DONE** | | `start-instances` | `--instance-ids` | `--dry-run`, `--force` | `run-instances` (instance must exist in stopped state) | Gateway sends NATS `ec2.cmd.{instance-id}` → daemon restarts stopped QEMU process with same config → state transitions stopped→pending→running | 1. Start a stopped instance
2. Start already-running instance (error: IncorrectInstanceState)
3. Start with invalid instance ID
4. Verify volumes re-mount on start | **DONE** | | `stop-instances` | `--instance-ids` | `--force`, `--hibernate`, `--dry-run` | `run-instances` (instance must be running) | Gateway sends NATS to target node → daemon issues QMP `system_powerdown` for graceful shutdown → monitors heartbeat until QEMU exits → state transitions running→stopping→stopped | 1. Graceful stop of running instance
2. Force stop (kills QEMU process)
3. Stop already-stopped instance (error)
4. Verify ~30s heartbeat detection | **DONE** | | `terminate-instances` | `--instance-ids`, `DeleteOnTermination` (per-volume flag, default true) | `--dry-run` | `run-instances` (instance must exist) | Gateway sends NATS to target node → daemon kills QEMU process → cleans up NBD mounts → deletes volumes with `DeleteOnTermination=true` via `volumeService.DeleteVolume()` (S3 cleanup of vol/, vol-efi/, vol-cloudinit/) → internal volumes (EFI, cloud-init) always cleaned up via `ebs.delete` NATS → volumes with `DeleteOnTermination=false` left in available state → state→terminated | 1. Terminate running instance
2. Terminate stopped instance
3. Terminate with DeleteOnTermination=true deletes volumes
4. Terminate with DeleteOnTermination=false preserves volumes
5. Terminate already-terminated (idempotent)
6. Internal volumes (EFI, cloud-init) always cleaned up
7. Invalid instance ID | **DONE** | -| `reboot-instances` | — | `--instance-ids`, `--dry-run` | `run-instances` (instance must be running) | Gateway sends NATS to target node → daemon issues QMP `system_reset` → instance reboots without stopping | 1. Reboot running instance
2. Reboot stopped instance (error)
3. Verify instance stays in running state after reboot | **NOT STARTED** (parser test exists) | +| `reboot-instances` | `--instance-ids` | `--dry-run` | `run-instances` (instance must be running) | Gateway validates instance IDs → sends EC2InstanceCommand with `RebootInstance=true` via NATS `ec2.cmd.{instanceId}` → daemon validates instance is in StateRunning (returns IncorrectInstanceState if stopped) → issues QMP `system_reset` → instance reboots without stopping, stays in running state | 1. Reboot running instance
2. Reboot multiple instances
3. Reboot stopped instance (error: IncorrectInstanceState)
4. Instance not found (error: InvalidInstanceID.NotFound)
5. Verify instance stays in running state after reboot | **DONE** | | `describe-instance-types` | `--filters` (capacity filter only) | `--instance-types`, `--max-results`, `--next-token`, `--dry-run`, all other filters | None | Gateway fans out NATS `ec2.DescribeInstanceTypes` to all nodes → each daemon reports supported types (t3.micro/small/medium/large) with vCPU/memory specs → gateway deduplicates and returns | 1. List all instance types
2. Filter by specific type
3. Filter with `capacity=true` shows available slots
4. Verify vCPU/memory specs match hardware | **DONE** | | `modify-instance-attribute` | `--instance-id`, `--instance-type`, `--user-data` | `--disable-api-termination`, `--ebs-optimized`, `--source-dest-check`, `--instance-initiated-shutdown-behavior`, `--block-device-mappings`, `--groups`, `--ena-support`, `--sriov-net-support` | Instance must be stopped (in NATS KV) | Gateway validates input (exactly one attribute per call, instance ID format) → NATS `ec2.ModifyInstanceAttribute` with `spinifex-workers` queue group → daemon loads stopped instance from JetStream KV → applies attribute change → writes back to KV → returns `{}` on success. **InstanceType**: updates vm.InstanceType, Config, and Instance fields; clears StateReason (enables recovery from instance-type-missing bug). **UserData**: stores decoded content in vm.UserData and re-encodes to base64 for RunInstancesInput (cloud-init on next start). No instance type pre-validation (matches AWS — invalid types accepted, fail at StartInstances time). | 1. Change instance type while stopped
2. Change user data while stopped
3. Modify running instance (error: NotFound — running instances not in KV)
4. Instance not found (error: InvalidInstanceID.NotFound)
5. Instance not stopped (error: IncorrectInstanceState)
6. Invalid instance type accepted (fails on start with InsufficientInstanceCapacity)
7. StateReason cleared on type change (recovery from capacity-unavailable)
8. Missing/malformed instance ID (error: InvalidInstanceID.Malformed)
9. No attribute set (error: InvalidParameterValue)
10. Multiple attributes in one call (error: InvalidParameterValue) | **DONE** | | `get-console-output` | `--instance-id` | `--latest` (always returns latest), `--dry-run` | Instance must be running on a node | Gateway sends NATS `ec2.{instanceId}.GetConsoleOutput` (per-instance topic, routed to owning node) → daemon reads console log file from disk → returns last 64KB base64-encoded with timestamp. Always available regardless of serial console access setting (matches AWS behavior). | 1. Get output from running instance
2. Empty log file returns empty output
3. Instance not found (error: InvalidInstanceID.NotFound) | **DONE** | @@ -192,11 +192,13 @@ Gateway routing and KV CRUD are wired up but no OVN/OVS integration exists. EIGW ### EC2 - Placement Groups +Placement groups are stored in `spinifex-placement-groups` JetStream KV bucket, keyed by `{accountID}.{groupName}`. All mutations use CAS (compare-and-swap) with 5-retry limit for optimistic concurrency. Supported strategies: **spread** (strict 1 instance per node) and **cluster** (all instances pinned to single node). Partition strategy is rejected. Internal NATS operations (ReserveSpreadNodes, FinalizeSpreadInstances, ReleaseSpreadNodes, RemoveInstanceFromPlacementGroup, ReserveClusterNode, FinalizeClusterInstances) support atomic multi-node coordination for RunInstances and TerminateInstances. + | Command | Implemented Flags | Missing Flags | Prerequisites | Basic Logic | Test Cases | Status | |---------|-------------------|---------------|---------------|-------------|------------|--------| -| `create-placement-group` | — | `--group-name`, `--strategy` (cluster/spread/partition), `--partition-count`, `--spread-level`, `--tag-specifications` | None | NATS `ec2.CreatePlacementGroup` → daemon stores placement group in JetStream KV with strategy, state=available → return group metadata | 1. Create with cluster strategy
2. Create with spread strategy
3. Create with partition strategy + partition count
4. Duplicate name (error)
5. Invalid strategy (error) | **NOT STARTED** | -| `delete-placement-group` | — | `--group-name`, `--dry-run` | Group must exist, no instances in group | NATS `ec2.DeletePlacementGroup` → daemon verifies group exists and is empty → delete from KV → return success | 1. Delete empty group
2. Delete group with instances (error)
3. Delete non-existent group (error) | **NOT STARTED** | -| `describe-placement-groups` | — | `--group-names`, `--group-ids`, `--filters` | None | NATS `ec2.DescribePlacementGroups` → daemon lists groups from KV → return list with strategy and state | 1. List all groups
2. Filter by name
3. Filter by strategy | **NOT STARTED** | +| `create-placement-group` | `--group-name`, `--strategy` (spread/cluster) | `--partition-count`, `--spread-level` (hardcoded to "host"), `--tag-specifications`, `--dry-run` | None | Gateway validates GroupName + Strategy → NATS `ec2.CreatePlacementGroup` with `spinifex-workers` queue group → daemon generates pg-ID (pg-{16hex}) → stores PlacementGroupRecord in `spinifex-placement-groups` KV with state=available, SpreadLevel="host" for spread → rejects partition strategy (InvalidParameterValue) → rejects duplicate names (InvalidPlacementGroup.Duplicate) → returns PlacementGroup | 1. Create with spread strategy
2. Create with cluster strategy
3. Create with partition strategy (error: InvalidParameterValue)
4. Duplicate name (error: InvalidPlacementGroup.Duplicate)
5. Missing GroupName (error)
6. Missing Strategy (error) | **DONE** | +| `delete-placement-group` | `--group-name` | `--dry-run` | Group must exist, no instances in group | NATS `ec2.DeletePlacementGroup` → daemon verifies group exists (InvalidPlacementGroup.Unknown) → counts instances in NodeInstances map → rejects if non-empty (InvalidPlacementGroup.InUse) → deletes from KV → returns success | 1. Delete empty group
2. Delete group with instances (error: InvalidPlacementGroup.InUse)
3. Delete non-existent group (error: InvalidPlacementGroup.Unknown) | **DONE** | +| `describe-placement-groups` | `--group-names`, `--group-ids`, `--filters` (strategy, state, spread-level, group-name) | `--dry-run` | None | NATS `ec2.DescribePlacementGroups` → daemon lists all keys from `spinifex-placement-groups` KV → filters by GroupNames, GroupIds, or Filters → returns list with Strategy, State, SpreadLevel, GroupId | 1. List all groups
2. Filter by name
3. Filter by group ID
4. Filter by strategy
5. Filter by state
6. Empty result when no groups exist | **DONE** | ### EC2 - Dedicated Hosts diff --git a/scripts/reset-dev-env.sh b/scripts/reset-dev-env.sh index 483052bfc..762de0a87 100755 --- a/scripts/reset-dev-env.sh +++ b/scripts/reset-dev-env.sh @@ -154,6 +154,16 @@ if [ -n "$WAN_IFACE" ] && [ -n "$WAN_GW" ]; then fi ./bin/spx admin init $ADMIN_INIT_ARGS +# Re-initialize platform (generates fresh credentials, certs, config, and updates ~/.aws/credentials) +echo "Re-initializing platform..." +./bin/spx admin init --force + +# Generate SSH key if it doesn't exist +if [ ! -f ~/.ssh/spinifex-key.pub ]; then + echo "Generating SSH key pair..." + ssh-keygen -t ed25519 -f ~/.ssh/spinifex-key -N "" +fi + # Enable pprof for development PPROF_ENABLED=1 PPROF_OUTPUT=/tmp/spinifex-vm.prof ./scripts/start-dev.sh --build diff --git a/spinifex/daemon/daemon.go b/spinifex/daemon/daemon.go index 23d695671..8c1a69f8c 100644 --- a/spinifex/daemon/daemon.go +++ b/spinifex/daemon/daemon.go @@ -38,6 +38,7 @@ import ( handlers_ec2_image "github.com/mulgadc/spinifex/spinifex/handlers/ec2/image" handlers_ec2_instance "github.com/mulgadc/spinifex/spinifex/handlers/ec2/instance" handlers_ec2_key "github.com/mulgadc/spinifex/spinifex/handlers/ec2/key" + handlers_ec2_placementgroup "github.com/mulgadc/spinifex/spinifex/handlers/ec2/placementgroup" handlers_ec2_snapshot "github.com/mulgadc/spinifex/spinifex/handlers/ec2/snapshot" handlers_ec2_tags "github.com/mulgadc/spinifex/spinifex/handlers/ec2/tags" handlers_ec2_volume "github.com/mulgadc/spinifex/spinifex/handlers/ec2/volume" @@ -88,30 +89,32 @@ type ResourceManager struct { natsConn *nats.Conn instanceSubs map[string]*nats.Subscription handler nats.MsgHandler + nodeID string // node identifier for node-specific topic subscriptions } // Daemon represents the main daemon service type Daemon struct { - node string - clusterConfig *config.ClusterConfig - config *config.Config - natsConn *nats.Conn - resourceMgr *ResourceManager - instanceService *handlers_ec2_instance.InstanceServiceImpl - keyService *handlers_ec2_key.KeyServiceImpl - imageService *handlers_ec2_image.ImageServiceImpl - volumeService *handlers_ec2_volume.VolumeServiceImpl - accountService *handlers_ec2_account.AccountSettingsServiceImpl - snapshotService *handlers_ec2_snapshot.SnapshotServiceImpl - tagsService *handlers_ec2_tags.TagsServiceImpl - eigwService *handlers_ec2_eigw.EgressOnlyIGWServiceImpl - igwService *handlers_ec2_igw.IGWServiceImpl - vpcService *handlers_ec2_vpc.VPCServiceImpl - eipService *handlers_ec2_eip.EIPServiceImpl - externalIPAM *handlers_ec2_vpc.ExternalIPAM - ctx context.Context - cancel context.CancelFunc - shutdownWg sync.WaitGroup + node string + clusterConfig *config.ClusterConfig + config *config.Config + natsConn *nats.Conn + resourceMgr *ResourceManager + instanceService *handlers_ec2_instance.InstanceServiceImpl + keyService *handlers_ec2_key.KeyServiceImpl + imageService *handlers_ec2_image.ImageServiceImpl + volumeService *handlers_ec2_volume.VolumeServiceImpl + accountService *handlers_ec2_account.AccountSettingsServiceImpl + snapshotService *handlers_ec2_snapshot.SnapshotServiceImpl + tagsService *handlers_ec2_tags.TagsServiceImpl + eigwService *handlers_ec2_eigw.EgressOnlyIGWServiceImpl + igwService *handlers_ec2_igw.IGWServiceImpl + placementGroupService *handlers_ec2_placementgroup.PlacementGroupServiceImpl + vpcService *handlers_ec2_vpc.VPCServiceImpl + eipService *handlers_ec2_eip.EIPServiceImpl + externalIPAM *handlers_ec2_vpc.ExternalIPAM + ctx context.Context + cancel context.CancelFunc + shutdownWg sync.WaitGroup // Local VM Instances Instances vm.Instances @@ -390,6 +393,15 @@ func (d *Daemon) subscribeAll() error { {"ec2.DescribeInternetGateways", d.handleEC2DescribeInternetGateways, "spinifex-workers"}, {"ec2.AttachInternetGateway", d.handleEC2AttachInternetGateway, "spinifex-workers"}, {"ec2.DetachInternetGateway", d.handleEC2DetachInternetGateway, "spinifex-workers"}, + {"ec2.CreatePlacementGroup", d.handleEC2CreatePlacementGroup, "spinifex-workers"}, + {"ec2.DeletePlacementGroup", d.handleEC2DeletePlacementGroup, "spinifex-workers"}, + {"ec2.DescribePlacementGroups", d.handleEC2DescribePlacementGroups, "spinifex-workers"}, + {"ec2.ReserveSpreadNodes", d.handleEC2ReserveSpreadNodes, "spinifex-workers"}, + {"ec2.FinalizeSpreadInstances", d.handleEC2FinalizeSpreadInstances, "spinifex-workers"}, + {"ec2.ReleaseSpreadNodes", d.handleEC2ReleaseSpreadNodes, "spinifex-workers"}, + {"ec2.RemoveInstanceFromPlacementGroup", d.handleEC2RemoveInstanceFromPlacementGroup, "spinifex-workers"}, + {"ec2.ReserveClusterNode", d.handleEC2ReserveClusterNode, "spinifex-workers"}, + {"ec2.FinalizeClusterInstances", d.handleEC2FinalizeClusterInstances, "spinifex-workers"}, {"ec2.CreateVpc", d.handleEC2CreateVpc, "spinifex-workers"}, {"ec2.DeleteVpc", d.handleEC2DeleteVpc, "spinifex-workers"}, {"ec2.DescribeVpcs", d.handleEC2DescribeVpcs, "spinifex-workers"}, @@ -521,6 +533,13 @@ func (d *Daemon) Start() error { return fmt.Errorf("failed to initialize IGW service: %w", err) } + d.placementGroupService, err = initServiceWithRetry("placement group service", func() (*handlers_ec2_placementgroup.PlacementGroupServiceImpl, error) { + return handlers_ec2_placementgroup.NewPlacementGroupServiceImplWithNATS(d.config, d.natsConn) + }) + if err != nil { + return fmt.Errorf("failed to initialize placement group service: %w", err) + } + d.vpcService, err = initServiceWithRetry("VPC service", func() (*handlers_ec2_vpc.VPCServiceImpl, error) { return handlers_ec2_vpc.NewVPCServiceImplWithNATS(d.config, d.natsConn) }) @@ -616,7 +635,7 @@ func (d *Daemon) Start() error { // Initialize dynamic per-instance-type subscriptions for capacity-aware routing. // Each instance type gets its own NATS topic (ec2.RunInstances.{type}) so requests // are only routed to nodes with available capacity. - d.resourceMgr.initSubscriptions(d.natsConn, d.handleEC2RunInstances) + d.resourceMgr.initSubscriptions(d.natsConn, d.handleEC2RunInstances, d.node) d.startHeartbeat() d.startPendingWatchdog() @@ -2552,17 +2571,21 @@ func (rm *ResourceManager) deallocate(instanceType *ec2.InstanceTypeInfo) { // initSubscriptions sets up dynamic per-instance-type NATS subscriptions. // Called once during daemon startup after NATS is connected. -func (rm *ResourceManager) initSubscriptions(nc *nats.Conn, handler nats.MsgHandler) { +func (rm *ResourceManager) initSubscriptions(nc *nats.Conn, handler nats.MsgHandler, nodeID string) { rm.natsConn = nc rm.handler = handler + rm.nodeID = nodeID rm.instanceSubs = make(map[string]*nats.Subscription) rm.updateInstanceSubscriptions() } // updateInstanceSubscriptions recalculates which instance types can fit on this // node and subscribes/unsubscribes from the corresponding NATS topics. Each type -// gets its own topic (ec2.RunInstances.{type}) with the spinifex-workers queue group, -// so NATS only routes requests to nodes that can actually serve them. +// gets two topics: +// - ec2.RunInstances.{type} with spinifex-workers queue group (load-balanced, for single-instance launches) +// - ec2.RunInstances.{type}.{nodeId} without queue group (targeted, for multi-node distribution) +// +// Both use the same handler. NATS only routes requests to nodes with available capacity. func (rm *ResourceManager) updateInstanceSubscriptions() { if rm.natsConn == nil { return @@ -2572,24 +2595,46 @@ func (rm *ResourceManager) updateInstanceSubscriptions() { defer rm.subsMu.Unlock() for typeName, typeInfo := range rm.instanceTypes { - topic := fmt.Sprintf("ec2.RunInstances.%s", typeName) + queueTopic := fmt.Sprintf("ec2.RunInstances.%s", typeName) canFit := rm.canAllocate(typeInfo, 1) >= 1 - _, subscribed := rm.instanceSubs[topic] + // Queue group subscription (load-balanced across nodes) + _, subscribed := rm.instanceSubs[queueTopic] if canFit && !subscribed { - sub, err := rm.natsConn.QueueSubscribe(topic, "spinifex-workers", rm.handler) + sub, err := rm.natsConn.QueueSubscribe(queueTopic, "spinifex-workers", rm.handler) if err != nil { - slog.Error("Failed to subscribe to instance type topic", "topic", topic, "err", err) + slog.Error("Failed to subscribe to instance type topic", "topic", queueTopic, "err", err) continue } - rm.instanceSubs[topic] = sub - slog.Debug("Subscribed to instance type", "topic", topic) + rm.instanceSubs[queueTopic] = sub + slog.Debug("Subscribed to instance type", "topic", queueTopic) } else if !canFit && subscribed { - if err := rm.instanceSubs[topic].Unsubscribe(); err != nil { - slog.Error("Failed to unsubscribe from instance type topic", "topic", topic, "err", err) + if err := rm.instanceSubs[queueTopic].Unsubscribe(); err != nil { + slog.Error("Failed to unsubscribe from instance type topic", "topic", queueTopic, "err", err) + } + delete(rm.instanceSubs, queueTopic) + slog.Info("Unsubscribed from instance type (capacity full)", "topic", queueTopic) + } + + // Node-specific subscription (targeted routing for multi-node distribution) + if rm.nodeID != "" { + nodeTopic := fmt.Sprintf("ec2.RunInstances.%s.%s", typeName, rm.nodeID) + _, nodeSubscribed := rm.instanceSubs[nodeTopic] + if canFit && !nodeSubscribed { + sub, err := rm.natsConn.Subscribe(nodeTopic, rm.handler) + if err != nil { + slog.Error("Failed to subscribe to node-specific topic", "topic", nodeTopic, "err", err) + continue + } + rm.instanceSubs[nodeTopic] = sub + slog.Debug("Subscribed to node-specific instance type", "topic", nodeTopic) + } else if !canFit && nodeSubscribed { + if err := rm.instanceSubs[nodeTopic].Unsubscribe(); err != nil { + slog.Error("Failed to unsubscribe from node-specific topic", "topic", nodeTopic, "err", err) + } + delete(rm.instanceSubs, nodeTopic) + slog.Info("Unsubscribed from node-specific instance type (capacity full)", "topic", nodeTopic) } - delete(rm.instanceSubs, topic) - slog.Info("Unsubscribed from instance type (capacity full)", "topic", topic) } } } diff --git a/spinifex/daemon/daemon_handlers_instance.go b/spinifex/daemon/daemon_handlers_instance.go index d688eb22a..ff66e5c9a 100644 --- a/spinifex/daemon/daemon_handlers_instance.go +++ b/spinifex/daemon/daemon_handlers_instance.go @@ -11,6 +11,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/ec2" "github.com/mulgadc/spinifex/spinifex/awserrors" + handlers_ec2_placementgroup "github.com/mulgadc/spinifex/spinifex/handlers/ec2/placementgroup" "github.com/mulgadc/spinifex/spinifex/qmp" "github.com/mulgadc/spinifex/spinifex/types" "github.com/mulgadc/spinifex/spinifex/utils" @@ -271,10 +272,14 @@ func (d *Daemon) handleEC2RunInstances(msg *nats.Msg) { reservation.SetOwnerId(accountID) reservation.Instances = allEC2Instances - // Store reservation reference and account ID in all VMs + // Store reservation reference, account ID, and placement group in all VMs for _, instance := range instances { instance.Reservation = &reservation instance.AccountID = accountID + if runInstancesInput.Placement != nil && runInstancesInput.Placement.GroupName != nil && *runInstancesInput.Placement.GroupName != "" { + instance.PlacementGroupName = *runInstancesInput.Placement.GroupName + instance.PlacementGroupNode = d.node + } } // Respond to NATS immediately with reservation (instances are provisioning) @@ -528,6 +533,18 @@ func (d *Daemon) handleStopOrTerminateInstance(msg *nats.Msg, command types.EC2I } slog.Info("Instance "+string(finalState), "id", inst.ID) + // Remove instance from placement group on terminate + if isTerminate && inst.PlacementGroupName != "" && d.placementGroupService != nil { + if _, pgErr := d.placementGroupService.RemoveInstance(&handlers_ec2_placementgroup.RemoveInstanceInput{ + GroupName: inst.PlacementGroupName, + NodeName: inst.PlacementGroupNode, + InstanceID: inst.ID, + }, inst.AccountID); pgErr != nil { + slog.Error("Failed to remove instance from placement group", + "instanceId", inst.ID, "groupName", inst.PlacementGroupName, "err", pgErr) + } + } + if d.jsManager != nil { if isTerminate { // Write to terminated KV bucket (auto-expires after 1 hour via TTL). @@ -688,6 +705,14 @@ func (d *Daemon) handleEC2DescribeInstances(msg *nats.Msg) { 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), + } + } + // Add instance to its reservation reservationMap[resID].Instances = append(reservationMap[resID].Instances, &instanceCopy) } diff --git a/spinifex/daemon/daemon_handlers_placementgroup.go b/spinifex/daemon/daemon_handlers_placementgroup.go new file mode 100644 index 000000000..b57bc4804 --- /dev/null +++ b/spinifex/daemon/daemon_handlers_placementgroup.go @@ -0,0 +1,39 @@ +package daemon + +import "github.com/nats-io/nats.go" + +func (d *Daemon) handleEC2CreatePlacementGroup(msg *nats.Msg) { + handleNATSRequest(msg, d.placementGroupService.CreatePlacementGroup) +} + +func (d *Daemon) handleEC2DeletePlacementGroup(msg *nats.Msg) { + handleNATSRequest(msg, d.placementGroupService.DeletePlacementGroup) +} + +func (d *Daemon) handleEC2DescribePlacementGroups(msg *nats.Msg) { + handleNATSRequest(msg, d.placementGroupService.DescribePlacementGroups) +} + +func (d *Daemon) handleEC2ReserveSpreadNodes(msg *nats.Msg) { + handleNATSRequest(msg, d.placementGroupService.ReserveSpreadNodes) +} + +func (d *Daemon) handleEC2FinalizeSpreadInstances(msg *nats.Msg) { + handleNATSRequest(msg, d.placementGroupService.FinalizeSpreadInstances) +} + +func (d *Daemon) handleEC2ReleaseSpreadNodes(msg *nats.Msg) { + handleNATSRequest(msg, d.placementGroupService.ReleaseSpreadNodes) +} + +func (d *Daemon) handleEC2RemoveInstanceFromPlacementGroup(msg *nats.Msg) { + handleNATSRequest(msg, d.placementGroupService.RemoveInstance) +} + +func (d *Daemon) handleEC2ReserveClusterNode(msg *nats.Msg) { + handleNATSRequest(msg, d.placementGroupService.ReserveClusterNode) +} + +func (d *Daemon) handleEC2FinalizeClusterInstances(msg *nats.Msg) { + handleNATSRequest(msg, d.placementGroupService.FinalizeClusterInstances) +} diff --git a/spinifex/daemon/daemon_test.go b/spinifex/daemon/daemon_test.go index edb0da23d..b5f7f32ba 100644 --- a/spinifex/daemon/daemon_test.go +++ b/spinifex/daemon/daemon_test.go @@ -1394,7 +1394,7 @@ func TestInstanceTypeSubscriptions(t *testing.T) { defer nc.Close() handler := func(msg *nats.Msg) {} - rm.initSubscriptions(nc, handler) + rm.initSubscriptions(nc, handler, "test-node") // Count how many types actually fit on this machine fittableTypes := 0 @@ -1404,8 +1404,9 @@ func TestInstanceTypeSubscriptions(t *testing.T) { } } - assert.Equal(t, fittableTypes, len(rm.instanceSubs), - "should subscribe to all instance types that fit") + // Each fittable type gets 2 subscriptions: queue group + node-specific + assert.Equal(t, fittableTypes*2, len(rm.instanceSubs), + "should subscribe to all instance types that fit (queue + node-specific)") assert.Greater(t, len(rm.instanceSubs), 0, "should subscribe to at least some instance types") @@ -1424,7 +1425,7 @@ func TestInstanceTypeSubscriptions(t *testing.T) { defer nc.Close() handler := func(msg *nats.Msg) {} - rm.initSubscriptions(nc, handler) + rm.initSubscriptions(nc, handler, "test-node") initialCount := len(rm.instanceSubs) require.Greater(t, initialCount, 0) @@ -1449,7 +1450,7 @@ func TestInstanceTypeSubscriptions(t *testing.T) { defer nc.Close() handler := func(msg *nats.Msg) {} - rm.initSubscriptions(nc, handler) + rm.initSubscriptions(nc, handler, "test-node") expectedCount := len(rm.instanceSubs) @@ -1480,7 +1481,7 @@ func TestInstanceTypeSubscriptions(t *testing.T) { defer nc.Close() handler := func(msg *nats.Msg) {} - rm.initSubscriptions(nc, handler) + rm.initSubscriptions(nc, handler, "test-node") // Leave only 2 vCPUs and 1 GB free — enough for nano/micro but not larger types rm.mu.Lock() @@ -1492,7 +1493,7 @@ func TestInstanceTypeSubscriptions(t *testing.T) { // Count subscribed types — should be less than total but more than zero assert.Greater(t, len(rm.instanceSubs), 0, "should still be subscribed to small instance types") - assert.Less(t, len(rm.instanceSubs), len(rm.instanceTypes), + assert.Less(t, len(rm.instanceSubs), len(rm.instanceTypes)*2, "should not be subscribed to large instance types") // Verify nano (0.5 GB) and micro (1 GB) are subscribed @@ -1509,7 +1510,7 @@ func TestInstanceTypeSubscriptions(t *testing.T) { defer nc.Close() handler := func(msg *nats.Msg) {} - rm.initSubscriptions(nc, handler) + rm.initSubscriptions(nc, handler, "test-node") initialCount := len(rm.instanceSubs) require.Greater(t, initialCount, 0) @@ -1553,7 +1554,7 @@ func TestInstanceTypeSubscriptions(t *testing.T) { defer nc.Close() handler := func(msg *nats.Msg) {} - rm.initSubscriptions(nc, handler) + rm.initSubscriptions(nc, handler, "test-node") // Fill the node completely rm.mu.Lock() diff --git a/spinifex/gateway/ec2.go b/spinifex/gateway/ec2.go index 904de7a06..c779ae49d 100644 --- a/spinifex/gateway/ec2.go +++ b/spinifex/gateway/ec2.go @@ -18,6 +18,7 @@ import ( gateway_ec2_image "github.com/mulgadc/spinifex/spinifex/gateway/ec2/image" gateway_ec2_instance "github.com/mulgadc/spinifex/spinifex/gateway/ec2/instance" gateway_ec2_key "github.com/mulgadc/spinifex/spinifex/gateway/ec2/key" + gateway_ec2_placementgroup "github.com/mulgadc/spinifex/spinifex/gateway/ec2/placementgroup" gateway_ec2_snapshot "github.com/mulgadc/spinifex/spinifex/gateway/ec2/snapshot" gateway_ec2_tags "github.com/mulgadc/spinifex/spinifex/gateway/ec2/tags" gateway_ec2_volume "github.com/mulgadc/spinifex/spinifex/gateway/ec2/volume" @@ -198,6 +199,15 @@ var ec2Actions = map[string]EC2Handler{ "DescribeEgressOnlyInternetGateways": ec2Handler(func(input *ec2.DescribeEgressOnlyInternetGatewaysInput, gw *GatewayConfig, accountID string) (any, error) { return gateway_ec2_eigw.DescribeEgressOnlyInternetGateways(input, gw.NATSConn, accountID) }), + "CreatePlacementGroup": ec2Handler(func(input *ec2.CreatePlacementGroupInput, gw *GatewayConfig, accountID string) (any, error) { + return gateway_ec2_placementgroup.CreatePlacementGroup(input, gw.NATSConn, accountID) + }), + "DeletePlacementGroup": ec2Handler(func(input *ec2.DeletePlacementGroupInput, gw *GatewayConfig, accountID string) (any, error) { + return gateway_ec2_placementgroup.DeletePlacementGroup(input, gw.NATSConn, accountID) + }), + "DescribePlacementGroups": ec2Handler(func(input *ec2.DescribePlacementGroupsInput, gw *GatewayConfig, accountID string) (any, error) { + return gateway_ec2_placementgroup.DescribePlacementGroups(input, gw.NATSConn, accountID) + }), "CreateVpc": ec2Handler(func(input *ec2.CreateVpcInput, gw *GatewayConfig, accountID string) (any, error) { return gateway_ec2_vpc.CreateVpc(input, gw.NATSConn, accountID) }), diff --git a/spinifex/gateway/ec2/instance/RunInstances.go b/spinifex/gateway/ec2/instance/RunInstances.go index d14833fcf..8ab401669 100644 --- a/spinifex/gateway/ec2/instance/RunInstances.go +++ b/spinifex/gateway/ec2/instance/RunInstances.go @@ -5,11 +5,11 @@ import ( "strings" "time" + "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/ec2" "github.com/mulgadc/spinifex/spinifex/awserrors" + handlers_ec2_placementgroup "github.com/mulgadc/spinifex/spinifex/handlers/ec2/placementgroup" "github.com/mulgadc/spinifex/spinifex/utils" - - handlers_ec2_instance "github.com/mulgadc/spinifex/spinifex/handlers/ec2/instance" "github.com/nats-io/nats.go" ) @@ -66,28 +66,78 @@ func RunInstances(input *ec2.RunInstancesInput, natsConn *nats.Conn, accountID s return reservation, err } - // Create NATS-based instance service - service := handlers_ec2_instance.NewNATSInstanceService(natsConn) + // Placement group routing: when a placement group is specified, validate it + // and route based on its strategy (spread or cluster). + groupName := placementGroupName(input) + if groupName != "" { + strategy, err := lookupPlacementGroupStrategy(natsConn, accountID, groupName) + if err != nil { + return reservation, err + } - // Call the service directly (no need for JSON marshaling/unmarshaling in same process) - reservationPtr, err := service.RunInstances(input, accountID) + switch strategy { + case ec2.PlacementStrategySpread: + reservationPtr, err := distributeInstancesSpread(input, natsConn, accountID, groupName) + if err != nil { + return reservation, err + } + return *reservationPtr, nil + case ec2.PlacementStrategyCluster: + reservationPtr, err := distributeInstancesCluster(input, natsConn, accountID, groupName) + if err != nil { + return reservation, err + } + return *reservationPtr, nil + default: + return reservation, errors.New(awserrors.ErrorInvalidParameterValue) + } + } + + // Capacity-aware routing: query all nodes for capacity and distribute + // instances across nodes with best-effort spread. This applies to both + // single-instance (count=1) and batch (count>1) launches, ensuring fair + // distribution across the cluster. + reservationPtr, err := distributeInstances(input, natsConn, accountID) if err != nil { - if errors.Is(err, nats.ErrNoResponders) { - // ErrNoResponders means no daemon subscribes to ec2.RunInstances.{type}. - // This happens when either: (a) the type is unknown, or (b) all nodes - // are at capacity. Query DescribeInstanceTypes to differentiate. + // When no nodes have capacity, distinguish between "unknown instance type" + // and "all nodes full" by checking DescribeInstanceTypes. + if err.Error() == awserrors.ErrorInsufficientInstanceCapacity { if !isKnownInstanceType(natsConn, *input.InstanceType) { return reservation, errors.New(awserrors.ErrorInvalidInstanceType) } - return reservation, errors.New(awserrors.ErrorInsufficientInstanceCapacity) } return reservation, err } - - // Dereference pointer to return value return *reservationPtr, nil } +// placementGroupName extracts the placement group name from RunInstancesInput. +func placementGroupName(input *ec2.RunInstancesInput) string { + if input.Placement != nil && input.Placement.GroupName != nil { + return aws.StringValue(input.Placement.GroupName) + } + return "" +} + +// lookupPlacementGroupStrategy validates that a placement group exists and returns its strategy. +func lookupPlacementGroupStrategy(natsConn *nats.Conn, accountID, groupName string) (string, error) { + pgSvc := handlers_ec2_placementgroup.NewNATSPlacementGroupService(natsConn) + out, err := pgSvc.DescribePlacementGroups(&ec2.DescribePlacementGroupsInput{ + GroupNames: []*string{aws.String(groupName)}, + }, accountID) + if err != nil { + return "", err + } + if len(out.PlacementGroups) == 0 { + return "", errors.New(awserrors.ErrorInvalidPlacementGroupUnknown) + } + pg := out.PlacementGroups[0] + if pg.State == nil || *pg.State != ec2.PlacementGroupStateAvailable { + return "", errors.New(awserrors.ErrorInvalidPlacementGroupUnknown) + } + return aws.StringValue(pg.Strategy), nil +} + // isKnownInstanceType checks whether any daemon recognizes the given instance type. func isKnownInstanceType(natsConn *nats.Conn, instanceType string) bool { result, err := utils.NATSRequest[ec2.DescribeInstanceTypesOutput]( diff --git a/spinifex/gateway/ec2/instance/placement.go b/spinifex/gateway/ec2/instance/placement.go new file mode 100644 index 000000000..62a4e35d7 --- /dev/null +++ b/spinifex/gateway/ec2/instance/placement.go @@ -0,0 +1,555 @@ +package gateway_ec2_instance + +import ( + "encoding/json" + "errors" + "fmt" + "log/slog" + "math/rand" + "sort" + "sync" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/ec2" + "github.com/mulgadc/spinifex/spinifex/awserrors" + handlers_ec2_placementgroup "github.com/mulgadc/spinifex/spinifex/handlers/ec2/placementgroup" + "github.com/mulgadc/spinifex/spinifex/types" + "github.com/mulgadc/spinifex/spinifex/utils" + + "github.com/nats-io/nats.go" +) + +// nodeAllocation tracks how many instances to launch on a specific node. +type nodeAllocation struct { + NodeID string + Available int // capacity for the requested instance type + Assigned int // instances assigned to this node +} + +// distributeInstances implements the best-effort spread algorithm for multi-node +// instance distribution. It queries cluster capacity, filters eligible nodes, +// distributes instances across nodes (1 per node first, then pack extras onto +// nodes with most remaining capacity), launches in parallel, and handles +// partial failures with rollback. +// +// Returns the merged reservation on success or an error. +func distributeInstances(input *ec2.RunInstancesInput, natsConn *nats.Conn, accountID string) (*ec2.Reservation, error) { + instanceType := aws.StringValue(input.InstanceType) + minCount := int(aws.Int64Value(input.MinCount)) + maxCount := int(aws.Int64Value(input.MaxCount)) + + // Step 1: Query capacity from all nodes via fan-out + nodes, err := queryNodeCapacity(natsConn, instanceType) + if err != nil { + return nil, err + } + + if len(nodes) == 0 { + return nil, errors.New(awserrors.ErrorInsufficientInstanceCapacity) + } + + // Step 2: Calculate total capacity and check feasibility + totalCapacity := 0 + for _, n := range nodes { + totalCapacity += n.Available + } + if totalCapacity < minCount { + return nil, errors.New(awserrors.ErrorInsufficientInstanceCapacity) + } + + // Step 3: Determine launch count (capped to MaxCount and available capacity) + launchCount := min(maxCount, totalCapacity) + + // Step 4: Distribute instances across nodes using best-effort spread + allocations := spreadAllocate(nodes, launchCount) + + // Step 5: Launch instances on each node in parallel + results := launchOnNodes(allocations, input, natsConn, accountID) + + // Step 6: Aggregate results and handle partial failure + return aggregateResults(results, minCount, natsConn, accountID) +} + +// queryNodeCapacity fans out spinifex.node.status to all daemons and returns +// eligible nodes (those with Available >= 1 for the requested instance type), +// sorted by available capacity descending with random tiebreaking for fair +// distribution among equal-capacity nodes. +// +// Uses a collection window: after the first response arrives, only waits an +// additional 200ms for remaining responses (instead of the full 3s timeout). +func queryNodeCapacity(natsConn *nats.Conn, instanceType string) ([]nodeAllocation, error) { + inbox := nats.NewInbox() + sub, err := natsConn.SubscribeSync(inbox) + if err != nil { + return nil, fmt.Errorf("failed to create inbox: %w", err) + } + defer sub.Unsubscribe() + + pubMsg := nats.NewMsg("spinifex.node.status") + pubMsg.Reply = inbox + pubMsg.Data = []byte("{}") + if err := natsConn.PublishMsg(pubMsg); err != nil { + return nil, fmt.Errorf("failed to publish node status request: %w", err) + } + + const ( + initialTimeout = 3 * time.Second + collectWindow = 200 * time.Millisecond + ) + + deadline := time.Now().Add(initialTimeout) + gotFirst := false + var nodes []nodeAllocation + + for time.Now().Before(deadline) { + remaining := time.Until(deadline) + if remaining <= 0 { + break + } + msg, err := sub.NextMsg(remaining) + if err != nil { + if err == nats.ErrTimeout { + break + } + slog.Debug("queryNodeCapacity: error receiving message", "err", err) + break + } + + var status types.NodeStatusResponse + if err := json.Unmarshal(msg.Data, &status); err != nil { + slog.Debug("queryNodeCapacity: failed to unmarshal response", "err", err) + continue + } + if status.Node == "" { + slog.Debug("queryNodeCapacity: skipping response with empty node ID") + continue + } + + // Find capacity for the requested instance type on this node + for _, cap := range status.InstanceTypes { + if cap.Name == instanceType && cap.Available >= 1 { + nodes = append(nodes, nodeAllocation{ + NodeID: status.Node, + Available: cap.Available, + }) + break + } + } + + // After the first valid response, tighten the deadline so we don't + // wait the full 3s for stragglers. + if !gotFirst { + gotFirst = true + collectDeadline := time.Now().Add(collectWindow) + if collectDeadline.Before(deadline) { + deadline = collectDeadline + } + } + } + + // Shuffle first for random tiebreaking, then stable-sort by capacity + // descending. This ensures fair distribution among equal-capacity nodes. + rand.Shuffle(len(nodes), func(i, j int) { + nodes[i], nodes[j] = nodes[j], nodes[i] + }) + sort.SliceStable(nodes, func(i, j int) bool { + return nodes[i].Available > nodes[j].Available + }) + + return nodes, nil +} + +// spreadAllocate distributes count instances across nodes using best-effort spread: +// - Round 1: assign 1 instance to each node (up to count) +// - Round 2+: assign remaining to nodes with most remaining capacity +func spreadAllocate(nodes []nodeAllocation, count int) []nodeAllocation { + // Make a working copy + allocs := make([]nodeAllocation, len(nodes)) + copy(allocs, nodes) + + remaining := count + + // Round 1: 1 instance per node + for i := range allocs { + if remaining <= 0 { + break + } + allocs[i].Assigned = 1 + remaining-- + } + + // Round 2+: pack remaining onto nodes with most remaining capacity. + // Ties are broken by preferring the node with fewer assigned instances, + // which produces a more balanced spread. + for remaining > 0 { + bestIdx := -1 + bestRemaining := 0 + for i, a := range allocs { + nodeRemaining := a.Available - a.Assigned + if nodeRemaining > bestRemaining { + bestRemaining = nodeRemaining + bestIdx = i + } else if nodeRemaining == bestRemaining && nodeRemaining > 0 && + (bestIdx < 0 || a.Assigned < allocs[bestIdx].Assigned) { + bestIdx = i + } + } + if bestIdx < 0 { + break // no more capacity anywhere + } + allocs[bestIdx].Assigned++ + remaining-- + } + + // Filter out nodes with no assignments + result := make([]nodeAllocation, 0, len(allocs)) + for _, a := range allocs { + if a.Assigned > 0 { + result = append(result, a) + } + } + return result +} + +// nodeLaunchResult holds the outcome of launching instances on a single node. +type nodeLaunchResult struct { + NodeID string + Reservation *ec2.Reservation + Err error +} + +// launchOnNodes sends targeted RunInstances requests to specific nodes in parallel. +// Each node gets MinCount=MaxCount=assignedCount so the daemon treats it as all-or-nothing. +func launchOnNodes(allocations []nodeAllocation, input *ec2.RunInstancesInput, natsConn *nats.Conn, accountID string) []nodeLaunchResult { + instanceType := aws.StringValue(input.InstanceType) + + results := make([]nodeLaunchResult, len(allocations)) + var wg sync.WaitGroup + + for i, alloc := range allocations { + wg.Add(1) + go func(idx int, a nodeAllocation) { + defer wg.Done() + + // Build per-node input with MinCount=MaxCount=assigned + nodeInput := *input + nodeInput.MinCount = aws.Int64(int64(a.Assigned)) + nodeInput.MaxCount = aws.Int64(int64(a.Assigned)) + + topic := fmt.Sprintf("ec2.RunInstances.%s.%s", instanceType, a.NodeID) + reservation, err := utils.NATSRequest[ec2.Reservation](natsConn, topic, &nodeInput, 5*time.Minute, accountID) + if err != nil { + results[idx] = nodeLaunchResult{NodeID: a.NodeID, Err: fmt.Errorf("launch on %s: %w", a.NodeID, err)} + return + } + results[idx] = nodeLaunchResult{NodeID: a.NodeID, Reservation: reservation} + }(i, alloc) + } + + wg.Wait() + return results +} + +// aggregateResults merges successful node launches into a single reservation. +// If total launched instances < minCount, all successfully launched instances +// are terminated (rollback) and InsufficientInstanceCapacity is returned. +func aggregateResults(results []nodeLaunchResult, minCount int, natsConn *nats.Conn, accountID string) (*ec2.Reservation, error) { + var allInstances []*ec2.Instance + var reservationID *string + + for _, r := range results { + if r.Err != nil { + slog.Warn("distributeInstances: node launch failed", "node", r.NodeID, "err", r.Err) + continue + } + if r.Reservation != nil { + allInstances = append(allInstances, r.Reservation.Instances...) + if reservationID == nil { + reservationID = r.Reservation.ReservationId + } + } + } + + totalLaunched := len(allInstances) + + if totalLaunched < minCount { + // Rollback: terminate all successfully launched instances + if totalLaunched > 0 { + rollbackInstances(allInstances, natsConn, accountID) + } + // Propagate specific client errors (e.g. InvalidAMIID.NotFound) instead + // of the generic InsufficientInstanceCapacity. + if clientErr := extractClientError(results); clientErr != nil { + return nil, clientErr + } + return nil, errors.New(awserrors.ErrorInsufficientInstanceCapacity) + } + + return &ec2.Reservation{ + ReservationId: reservationID, + Instances: allInstances, + }, nil +} + +// distributeInstancesSpread implements strict 1-per-node spread for placement groups. +// It queries capacity, reserves unused nodes via CAS, launches 1 instance per node, +// and finalizes or rolls back the placement group record. +func distributeInstancesSpread(input *ec2.RunInstancesInput, natsConn *nats.Conn, accountID string, groupName string) (*ec2.Reservation, error) { + instanceType := aws.StringValue(input.InstanceType) + minCount := int(aws.Int64Value(input.MinCount)) + maxCount := int(aws.Int64Value(input.MaxCount)) + + pgSvc := handlers_ec2_placementgroup.NewNATSPlacementGroupService(natsConn) + + // Step 1: Query capacity from all nodes + nodes, err := queryNodeCapacity(natsConn, instanceType) + if err != nil { + return nil, err + } + if len(nodes) == 0 { + return nil, errors.New(awserrors.ErrorInsufficientInstanceCapacity) + } + + // Build list of eligible node IDs + eligibleNodeIDs := make([]string, len(nodes)) + for i, n := range nodes { + eligibleNodeIDs[i] = n.NodeID + } + + // Step 2: Reserve nodes atomically (CAS-based, handles concurrent launches) + reserveOut, err := pgSvc.ReserveSpreadNodes(&handlers_ec2_placementgroup.ReserveSpreadNodesInput{ + GroupName: groupName, + EligibleNodes: eligibleNodeIDs, + MinCount: minCount, + MaxCount: maxCount, + }, accountID) + if err != nil { + return nil, err + } + + reservedNodes := reserveOut.ReservedNodes + + // Step 3: Build allocations (1 instance per reserved node) + allocations := make([]nodeAllocation, len(reservedNodes)) + for i, nodeID := range reservedNodes { + allocations[i] = nodeAllocation{NodeID: nodeID, Assigned: 1} + } + + // Step 4: Launch instances on reserved nodes in parallel + results := launchOnNodes(allocations, input, natsConn, accountID) + + // Step 5: Collect results + var allInstances []*ec2.Instance + var reservationID *string + nodeInstances := make(map[string][]string) + var failedNodes []string + + for _, r := range results { + if r.Err != nil { + slog.Warn("distributeInstancesSpread: node launch failed", "node", r.NodeID, "err", r.Err) + failedNodes = append(failedNodes, r.NodeID) + continue + } + if r.Reservation != nil { + for _, inst := range r.Reservation.Instances { + allInstances = append(allInstances, inst) + if inst.InstanceId != nil { + nodeInstances[r.NodeID] = append(nodeInstances[r.NodeID], *inst.InstanceId) + } + } + if reservationID == nil { + reservationID = r.Reservation.ReservationId + } + } + } + + totalLaunched := len(allInstances) + + if totalLaunched < minCount { + // Rollback: terminate launched instances and release all reserved nodes + if totalLaunched > 0 { + rollbackInstances(allInstances, natsConn, accountID) + } + // Release all reserved nodes from the placement group + if _, err := pgSvc.ReleaseSpreadNodes(&handlers_ec2_placementgroup.ReleaseSpreadNodesInput{ + GroupName: groupName, + Nodes: reservedNodes, + }, accountID); err != nil { + slog.Error("distributeInstancesSpread: failed to release nodes after rollback", "err", err) + } + if clientErr := extractClientError(results); clientErr != nil { + return nil, clientErr + } + return nil, errors.New(awserrors.ErrorInsufficientInstanceCapacity) + } + + // Step 6: Finalize — replace placeholders with actual instance IDs. + // If finalization fails, the placement group record has stale placeholders; + // roll back the launched instances to avoid untracked instances. + if _, err := pgSvc.FinalizeSpreadInstances(&handlers_ec2_placementgroup.FinalizeSpreadInstancesInput{ + GroupName: groupName, + NodeInstances: nodeInstances, + }, accountID); err != nil { + slog.Error("distributeInstancesSpread: finalize failed, rolling back instances", "err", err) + rollbackInstances(allInstances, natsConn, accountID) + // Release all reserved nodes (both succeeded and failed) + allReleaseNodes := append(reservedNodes[:0:0], reservedNodes...) + if _, releaseErr := pgSvc.ReleaseSpreadNodes(&handlers_ec2_placementgroup.ReleaseSpreadNodesInput{ + GroupName: groupName, + Nodes: allReleaseNodes, + }, accountID); releaseErr != nil { + slog.Error("distributeInstancesSpread: failed to release nodes after finalize rollback", "err", releaseErr) + } + return nil, fmt.Errorf("failed to finalize placement group record: %w", err) + } + + // Release any failed nodes that didn't launch + if len(failedNodes) > 0 { + if _, err := pgSvc.ReleaseSpreadNodes(&handlers_ec2_placementgroup.ReleaseSpreadNodesInput{ + GroupName: groupName, + Nodes: failedNodes, + }, accountID); err != nil { + slog.Error("distributeInstancesSpread: failed to release failed nodes", "err", err) + } + } + + return &ec2.Reservation{ + ReservationId: reservationID, + Instances: allInstances, + }, nil +} + +// distributeInstancesCluster implements cluster placement group routing. +// All instances are pinned to a single node. If the group already has instances, +// subsequent launches go to the same node. If empty, picks the node with most capacity. +func distributeInstancesCluster(input *ec2.RunInstancesInput, natsConn *nats.Conn, accountID string, groupName string) (*ec2.Reservation, error) { + instanceType := aws.StringValue(input.InstanceType) + minCount := int(aws.Int64Value(input.MinCount)) + maxCount := int(aws.Int64Value(input.MaxCount)) + + pgSvc := handlers_ec2_placementgroup.NewNATSPlacementGroupService(natsConn) + + // Step 1: Query capacity from all nodes + nodes, err := queryNodeCapacity(natsConn, instanceType) + if err != nil { + return nil, err + } + + // Build eligible node IDs (already sorted by capacity desc) + eligibleNodeIDs := make([]string, len(nodes)) + for i, n := range nodes { + eligibleNodeIDs[i] = n.NodeID + } + + // Step 2: Reserve the cluster target node (CAS-based for empty groups) + reserveOut, err := pgSvc.ReserveClusterNode(&handlers_ec2_placementgroup.ReserveClusterNodeInput{ + GroupName: groupName, + EligibleNodes: eligibleNodeIDs, + }, accountID) + if err != nil { + return nil, err + } + + targetNode := reserveOut.TargetNode + + // Step 3: Check capacity on the target node + var targetCapacity int + for _, n := range nodes { + if n.NodeID == targetNode { + targetCapacity = n.Available + break + } + } + + if targetCapacity < minCount { + return nil, errors.New(awserrors.ErrorInsufficientInstanceCapacity) + } + + // Cap launch count to available capacity and MaxCount + launchCount := min(maxCount, targetCapacity) + + // Step 4: Targeted launch — all instances to the pinned node + allocations := []nodeAllocation{{ + NodeID: targetNode, + Assigned: launchCount, + }} + results := launchOnNodes(allocations, input, natsConn, accountID) + + // Step 5: Handle result (single node, no partial failure logic needed) + if results[0].Err != nil { + if clientErr := extractClientError(results); clientErr != nil { + return nil, clientErr + } + return nil, results[0].Err + } + + reservation := results[0].Reservation + if reservation == nil { + return nil, errors.New(awserrors.ErrorServerInternal) + } + + // Step 6: CAS-update record with launched instance IDs + nodeInstances := make(map[string][]string) + for _, inst := range reservation.Instances { + if inst.InstanceId != nil { + nodeInstances[targetNode] = append(nodeInstances[targetNode], *inst.InstanceId) + } + } + + if _, err := pgSvc.FinalizeClusterInstances(&handlers_ec2_placementgroup.FinalizeClusterInstancesInput{ + GroupName: groupName, + NodeInstances: nodeInstances, + }, accountID); err != nil { + slog.Error("distributeInstancesCluster: finalize failed, rolling back instances", "err", err) + rollbackInstances(reservation.Instances, natsConn, accountID) + return nil, fmt.Errorf("failed to finalize cluster placement group record: %w", err) + } + + return reservation, nil +} + +// extractClientError scans node launch results for specific client validation +// errors (e.g. InvalidAMIID.NotFound) that should be propagated instead of the +// generic InsufficientInstanceCapacity. Returns the first match, or nil. +func extractClientError(results []nodeLaunchResult) error { + for _, r := range results { + if r.Err == nil { + continue + } + inner := errors.Unwrap(r.Err) + if inner == nil { + continue + } + switch inner.Error() { + case awserrors.ErrorInvalidAMIIDNotFound, + awserrors.ErrorInvalidAMIIDMalformed, + awserrors.ErrorInvalidAMIIDUnavailable, + awserrors.ErrorInvalidKeyPairNotFound, + awserrors.ErrorMissingParameter: + return inner + } + } + return nil +} + +// rollbackInstances terminates all instances from a failed multi-node launch. +func rollbackInstances(instances []*ec2.Instance, natsConn *nats.Conn, accountID string) { + var ids []*string + for _, inst := range instances { + if inst.InstanceId != nil { + ids = append(ids, inst.InstanceId) + } + } + if len(ids) == 0 { + return + } + + slog.Warn("distributeInstances: rolling back launched instances", "count", len(ids)) + _, err := TerminateInstances(&ec2.TerminateInstancesInput{ + InstanceIds: ids, + }, natsConn, accountID) + if err != nil { + slog.Error("distributeInstances: rollback failed", "err", err) + } +} diff --git a/spinifex/gateway/ec2/instance/placement_test.go b/spinifex/gateway/ec2/instance/placement_test.go new file mode 100644 index 000000000..a3c00b352 --- /dev/null +++ b/spinifex/gateway/ec2/instance/placement_test.go @@ -0,0 +1,1134 @@ +package gateway_ec2_instance + +import ( + "encoding/json" + "errors" + "fmt" + "testing" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/ec2" + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/types" + "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/nats-io/nats.go" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// --- spreadAllocate tests (pure algorithm) --- + +func TestSpreadAllocate_EqualDistribution(t *testing.T) { + // 3 instances across 3 nodes → 1 per node + nodes := []nodeAllocation{ + {NodeID: "A", Available: 4}, + {NodeID: "B", Available: 3}, + {NodeID: "C", Available: 2}, + } + result := spreadAllocate(nodes, 3) + + assert.Len(t, result, 3) + for _, a := range result { + assert.Equal(t, 1, a.Assigned, "node %s should get exactly 1", a.NodeID) + } +} + +func TestSpreadAllocate_SpreadThenPack(t *testing.T) { + // 5 instances across 3 nodes (capacities: A=4, B=3, C=2) + // Round 1: A=1, B=1, C=1 (3 assigned, 2 remaining) + // Round 2: A gets 1 (remaining cap 3), B gets 1 (remaining cap 2) + nodes := []nodeAllocation{ + {NodeID: "A", Available: 4}, + {NodeID: "B", Available: 3}, + {NodeID: "C", Available: 2}, + } + result := spreadAllocate(nodes, 5) + + assert.Len(t, result, 3) + byNode := make(map[string]int) + for _, a := range result { + byNode[a.NodeID] = a.Assigned + } + assert.Equal(t, 2, byNode["A"], "node A (cap 4) should get 2") + assert.Equal(t, 2, byNode["B"], "node B (cap 3) should get 2") + assert.Equal(t, 1, byNode["C"], "node C (cap 2) should get 1") +} + +func TestSpreadAllocate_SingleNode(t *testing.T) { + // All 3 instances on 1 node + nodes := []nodeAllocation{ + {NodeID: "A", Available: 5}, + } + result := spreadAllocate(nodes, 3) + + assert.Len(t, result, 1) + assert.Equal(t, "A", result[0].NodeID) + assert.Equal(t, 3, result[0].Assigned) +} + +func TestSpreadAllocate_MoreNodesThanInstances(t *testing.T) { + // 2 instances across 5 nodes → only 2 get assigned + nodes := []nodeAllocation{ + {NodeID: "A", Available: 4}, + {NodeID: "B", Available: 3}, + {NodeID: "C", Available: 2}, + {NodeID: "D", Available: 2}, + {NodeID: "E", Available: 1}, + } + result := spreadAllocate(nodes, 2) + + assert.Len(t, result, 2) + totalAssigned := 0 + for _, a := range result { + assert.Equal(t, 1, a.Assigned) + totalAssigned += a.Assigned + } + assert.Equal(t, 2, totalAssigned) +} + +func TestSpreadAllocate_HeavyPacking(t *testing.T) { + // 10 instances across 2 nodes (A=8, B=6) + // Round 1: A=1, B=1 + // Packing: each round picks node with most remaining + nodes := []nodeAllocation{ + {NodeID: "A", Available: 8}, + {NodeID: "B", Available: 6}, + } + result := spreadAllocate(nodes, 10) + + assert.Len(t, result, 2) + byNode := make(map[string]int) + for _, a := range result { + byNode[a.NodeID] = a.Assigned + } + total := byNode["A"] + byNode["B"] + assert.Equal(t, 10, total) + // A has more capacity so should get more + assert.GreaterOrEqual(t, byNode["A"], byNode["B"]) +} + +func TestSpreadAllocate_ExactCapacity(t *testing.T) { + // Request exactly matches total capacity + nodes := []nodeAllocation{ + {NodeID: "A", Available: 2}, + {NodeID: "B", Available: 1}, + } + result := spreadAllocate(nodes, 3) + + assert.Len(t, result, 2) + byNode := make(map[string]int) + for _, a := range result { + byNode[a.NodeID] = a.Assigned + } + assert.Equal(t, 2, byNode["A"]) + assert.Equal(t, 1, byNode["B"]) +} + +func TestSpreadAllocate_ZeroCount(t *testing.T) { + nodes := []nodeAllocation{ + {NodeID: "A", Available: 4}, + } + result := spreadAllocate(nodes, 0) + assert.Len(t, result, 0) +} + +func TestSpreadAllocate_EmptyNodes(t *testing.T) { + result := spreadAllocate(nil, 3) + assert.Len(t, result, 0) +} + +// --- queryNodeCapacity tests (NATS-based) --- + +func TestQueryNodeCapacity_FiltersEligibleNodes(t *testing.T) { + _, nc := startTestNATSServer(t) + + // Simulate 3 daemons responding to spinifex.node.status + sub, err := nc.Subscribe("spinifex.node.status", func(msg *nats.Msg) { + // Respond 3 times with different node statuses + responses := []types.NodeStatusResponse{ + { + Node: "node-1", + InstanceTypes: []types.InstanceTypeCap{ + {Name: "t3.micro", Available: 4}, + {Name: "t3.small", Available: 2}, + }, + }, + { + Node: "node-2", + InstanceTypes: []types.InstanceTypeCap{ + {Name: "t3.micro", Available: 0}, // no capacity + {Name: "t3.small", Available: 3}, + }, + }, + { + Node: "node-3", + InstanceTypes: []types.InstanceTypeCap{ + {Name: "t3.micro", Available: 2}, + }, + }, + } + for _, resp := range responses { + data, _ := json.Marshal(resp) + _ = nc.Publish(msg.Reply, data) + } + }) + require.NoError(t, err) + defer sub.Unsubscribe() + + // Query for t3.micro — should get node-1 (cap 4) and node-3 (cap 2), not node-2 (cap 0) + nodes, err := queryNodeCapacity(nc, "t3.micro") + require.NoError(t, err) + + assert.Len(t, nodes, 2) + // Should be sorted by capacity descending + assert.Equal(t, "node-1", nodes[0].NodeID) + assert.Equal(t, 4, nodes[0].Available) + assert.Equal(t, "node-3", nodes[1].NodeID) + assert.Equal(t, 2, nodes[1].Available) +} + +func TestQueryNodeCapacity_NoNodes(t *testing.T) { + _, nc := startTestNATSServer(t) + + // No subscribers → timeout, empty result + nodes, err := queryNodeCapacity(nc, "t3.micro") + require.NoError(t, err) + assert.Len(t, nodes, 0) +} + +// --- aggregateResults tests --- + +func TestAggregateResults_AllSucceed(t *testing.T) { + results := []nodeLaunchResult{ + { + NodeID: "node-1", + Reservation: &ec2.Reservation{ + ReservationId: aws.String("r-abc"), + Instances: []*ec2.Instance{ + {InstanceId: aws.String("i-001")}, + }, + }, + }, + { + NodeID: "node-2", + Reservation: &ec2.Reservation{ + ReservationId: aws.String("r-def"), + Instances: []*ec2.Instance{ + {InstanceId: aws.String("i-002")}, + {InstanceId: aws.String("i-003")}, + }, + }, + }, + } + + // No NATS needed — all succeed, no rollback + reservation, err := aggregateResults(results, 2, nil, "") + require.NoError(t, err) + assert.Len(t, reservation.Instances, 3) + assert.Equal(t, "r-abc", aws.StringValue(reservation.ReservationId)) +} + +func TestAggregateResults_PartialSuccessMeetsMinCount(t *testing.T) { + results := []nodeLaunchResult{ + { + NodeID: "node-1", + Reservation: &ec2.Reservation{ + ReservationId: aws.String("r-abc"), + Instances: []*ec2.Instance{ + {InstanceId: aws.String("i-001")}, + {InstanceId: aws.String("i-002")}, + }, + }, + }, + { + NodeID: "node-2", + Err: assert.AnError, + }, + } + + // MinCount=2, got 2 from node-1 → success + reservation, err := aggregateResults(results, 2, nil, "") + require.NoError(t, err) + assert.Len(t, reservation.Instances, 2) +} + +func TestAggregateResults_PartialFailureBelowMinCount(t *testing.T) { + _, nc := startTestNATSServer(t) + + results := []nodeLaunchResult{ + { + NodeID: "node-1", + Reservation: &ec2.Reservation{ + Instances: []*ec2.Instance{ + {InstanceId: aws.String("i-001")}, + }, + }, + }, + { + NodeID: "node-2", + Err: assert.AnError, + }, + } + + // MinCount=3, only got 1 → should fail with InsufficientInstanceCapacity + // Note: rollback will attempt to terminate i-001 but we don't have a + // daemon responding, so it will fail silently — that's OK for this test + _, err := aggregateResults(results, 3, nc, "test-account") + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInsufficientInstanceCapacity, err.Error()) +} + +func TestAggregateResults_AllFail(t *testing.T) { + results := []nodeLaunchResult{ + {NodeID: "node-1", Err: assert.AnError}, + {NodeID: "node-2", Err: assert.AnError}, + } + + _, err := aggregateResults(results, 1, nil, "") + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInsufficientInstanceCapacity, err.Error()) +} + +// --- extractClientError tests --- + +func TestExtractClientError_NoErrors(t *testing.T) { + results := []nodeLaunchResult{ + {NodeID: "node-1", Reservation: &ec2.Reservation{}}, + } + assert.Nil(t, extractClientError(results)) +} + +func TestExtractClientError_GenericError(t *testing.T) { + results := []nodeLaunchResult{ + {NodeID: "node-1", Err: assert.AnError}, + } + assert.Nil(t, extractClientError(results)) +} + +func TestExtractClientError_AMINotFound(t *testing.T) { + // Simulate the error wrapping that launchOnNodes does + inner := errors.New(awserrors.ErrorInvalidAMIIDNotFound) + wrapped := fmt.Errorf("launch on node-1: %w", inner) + results := []nodeLaunchResult{ + {NodeID: "node-1", Err: wrapped}, + } + err := extractClientError(results) + require.NotNil(t, err) + assert.Equal(t, awserrors.ErrorInvalidAMIIDNotFound, err.Error()) +} + +func TestExtractClientError_KeyPairNotFound(t *testing.T) { + inner := errors.New(awserrors.ErrorInvalidKeyPairNotFound) + wrapped := fmt.Errorf("launch on node-1: %w", inner) + results := []nodeLaunchResult{ + {NodeID: "node-1", Err: wrapped}, + } + err := extractClientError(results) + require.NotNil(t, err) + assert.Equal(t, awserrors.ErrorInvalidKeyPairNotFound, err.Error()) +} + +func TestAggregateResults_PropagatesClientError(t *testing.T) { + inner := errors.New(awserrors.ErrorInvalidAMIIDNotFound) + wrapped := fmt.Errorf("launch on node-1: %w", inner) + results := []nodeLaunchResult{ + {NodeID: "node-1", Err: wrapped}, + {NodeID: "node-2", Err: wrapped}, + } + + _, err := aggregateResults(results, 1, nil, "") + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInvalidAMIIDNotFound, err.Error()) +} + +// --- distributeInstances integration tests (end-to-end with mock daemons) --- + +func TestDistributeInstances_SuccessfulSpread(t *testing.T) { + _, nc := startTestNATSServer(t) + + // Mock node.status responder + statusSub, err := nc.Subscribe("spinifex.node.status", func(msg *nats.Msg) { + for _, resp := range []types.NodeStatusResponse{ + {Node: "node-1", InstanceTypes: []types.InstanceTypeCap{{Name: "t3.micro", Available: 2}}}, + {Node: "node-2", InstanceTypes: []types.InstanceTypeCap{{Name: "t3.micro", Available: 2}}}, + } { + data, _ := json.Marshal(resp) + _ = nc.Publish(msg.Reply, data) + } + }) + require.NoError(t, err) + defer statusSub.Unsubscribe() + + // Mock daemon on node-1 + sub1, err := nc.Subscribe("ec2.RunInstances.t3.micro.node-1", func(msg *nats.Msg) { + reservation := ec2.Reservation{ + ReservationId: aws.String("r-test1"), + Instances: []*ec2.Instance{{InstanceId: aws.String("i-n1")}}, + } + data, _ := json.Marshal(reservation) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer sub1.Unsubscribe() + + // Mock daemon on node-2 + sub2, err := nc.Subscribe("ec2.RunInstances.t3.micro.node-2", func(msg *nats.Msg) { + reservation := ec2.Reservation{ + ReservationId: aws.String("r-test2"), + Instances: []*ec2.Instance{{InstanceId: aws.String("i-n2")}}, + } + data, _ := json.Marshal(reservation) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer sub2.Unsubscribe() + + time.Sleep(50 * time.Millisecond) // let subscriptions propagate + + input := &ec2.RunInstancesInput{ + ImageId: aws.String("ami-test"), + InstanceType: aws.String("t3.micro"), + MinCount: aws.Int64(2), + MaxCount: aws.Int64(2), + } + + reservation, err := distributeInstances(input, nc, "test-account") + require.NoError(t, err) + assert.Len(t, reservation.Instances, 2) + + // Verify instances came from different nodes + ids := make(map[string]bool) + for _, inst := range reservation.Instances { + ids[aws.StringValue(inst.InstanceId)] = true + } + assert.True(t, ids["i-n1"], "should have instance from node-1") + assert.True(t, ids["i-n2"], "should have instance from node-2") +} + +func TestDistributeInstances_InsufficientCapacity(t *testing.T) { + _, nc := startTestNATSServer(t) + + // Mock node.status with only 1 available slot total + statusSub, err := nc.Subscribe("spinifex.node.status", func(msg *nats.Msg) { + resp := types.NodeStatusResponse{ + Node: "node-1", + InstanceTypes: []types.InstanceTypeCap{{Name: "t3.micro", Available: 1}}, + } + data, _ := json.Marshal(resp) + _ = nc.Publish(msg.Reply, data) + }) + require.NoError(t, err) + defer statusSub.Unsubscribe() + + time.Sleep(50 * time.Millisecond) + + input := &ec2.RunInstancesInput{ + ImageId: aws.String("ami-test"), + InstanceType: aws.String("t3.micro"), + MinCount: aws.Int64(3), + MaxCount: aws.Int64(3), + } + + _, err = distributeInstances(input, nc, "test-account") + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInsufficientInstanceCapacity, err.Error()) +} + +func TestDistributeInstances_PropagatesAMINotFound(t *testing.T) { + _, nc := startTestNATSServer(t) + + // Mock node.status with capacity available + statusSub, err := nc.Subscribe("spinifex.node.status", func(msg *nats.Msg) { + resp := types.NodeStatusResponse{ + Node: "node-1", + InstanceTypes: []types.InstanceTypeCap{{Name: "t3.micro", Available: 2}}, + } + data, _ := json.Marshal(resp) + _ = nc.Publish(msg.Reply, data) + }) + require.NoError(t, err) + defer statusSub.Unsubscribe() + + // Mock daemon responds with InvalidAMIID.NotFound (AMI doesn't exist) + sub1, err := nc.Subscribe("ec2.RunInstances.t3.micro.node-1", func(msg *nats.Msg) { + _ = msg.Respond(utils.GenerateErrorPayload(awserrors.ErrorInvalidAMIIDNotFound)) + }) + require.NoError(t, err) + defer sub1.Unsubscribe() + + time.Sleep(50 * time.Millisecond) + + input := &ec2.RunInstancesInput{ + ImageId: aws.String("ami-0000000000000dead"), + InstanceType: aws.String("t3.micro"), + MinCount: aws.Int64(1), + MaxCount: aws.Int64(1), + } + + _, err = distributeInstances(input, nc, "test-account") + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInvalidAMIIDNotFound, err.Error(), + "should propagate InvalidAMIID.NotFound, not InsufficientInstanceCapacity") +} + +func TestDistributeInstances_LaunchCountCappedToMaxCount(t *testing.T) { + _, nc := startTestNATSServer(t) + + // 3 nodes with capacity, but MaxCount=2 + statusSub, err := nc.Subscribe("spinifex.node.status", func(msg *nats.Msg) { + for _, resp := range []types.NodeStatusResponse{ + {Node: "node-1", InstanceTypes: []types.InstanceTypeCap{{Name: "t3.micro", Available: 4}}}, + {Node: "node-2", InstanceTypes: []types.InstanceTypeCap{{Name: "t3.micro", Available: 3}}}, + {Node: "node-3", InstanceTypes: []types.InstanceTypeCap{{Name: "t3.micro", Available: 2}}}, + } { + data, _ := json.Marshal(resp) + _ = nc.Publish(msg.Reply, data) + } + }) + require.NoError(t, err) + defer statusSub.Unsubscribe() + + // Mock daemons — each returns 1 instance + for _, nodeID := range []string{"node-1", "node-2"} { + sub, err := nc.Subscribe("ec2.RunInstances.t3.micro."+nodeID, func(msg *nats.Msg) { + reservation := ec2.Reservation{ + ReservationId: aws.String("r-" + nodeID), + Instances: []*ec2.Instance{{InstanceId: aws.String("i-" + nodeID)}}, + } + data, _ := json.Marshal(reservation) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer sub.Unsubscribe() + } + + time.Sleep(50 * time.Millisecond) + + input := &ec2.RunInstancesInput{ + ImageId: aws.String("ami-test"), + InstanceType: aws.String("t3.micro"), + MinCount: aws.Int64(1), + MaxCount: aws.Int64(2), + } + + reservation, err := distributeInstances(input, nc, "test-account") + require.NoError(t, err) + // Should launch exactly 2 (MaxCount), not 3 (total capacity) + assert.Len(t, reservation.Instances, 2) +} + +func TestDistributeInstances_NoNodesAvailable(t *testing.T) { + _, nc := startTestNATSServer(t) + + // No responders to node.status + time.Sleep(50 * time.Millisecond) + + input := &ec2.RunInstancesInput{ + ImageId: aws.String("ami-test"), + InstanceType: aws.String("t3.micro"), + MinCount: aws.Int64(2), + MaxCount: aws.Int64(2), + } + + _, err := distributeInstances(input, nc, "test-account") + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInsufficientInstanceCapacity, err.Error()) +} + +// --- RunInstances routing tests --- + +func TestRunInstances_SingleInstanceDistributes(t *testing.T) { + // For MinCount=MaxCount=1, RunInstances should still query node capacity + // and route to a specific node via targeted topic (not queue group). + _, nc := startTestNATSServer(t) + + // Mock node status response (1 node with capacity) + statusSub, err := nc.Subscribe("spinifex.node.status", func(msg *nats.Msg) { + resp := types.NodeStatusResponse{ + Node: "node-1", + InstanceTypes: []types.InstanceTypeCap{{Name: "t3.micro", Available: 4}}, + } + data, _ := json.Marshal(resp) + _ = nc.Publish(msg.Reply, data) + }) + require.NoError(t, err) + defer statusSub.Unsubscribe() + + // Mock the node-specific handler (targeted topic) + nodeSub, err := nc.Subscribe("ec2.RunInstances.t3.micro.node-1", func(msg *nats.Msg) { + reservation := ec2.Reservation{ + ReservationId: aws.String("r-single"), + Instances: []*ec2.Instance{{InstanceId: aws.String("i-single")}}, + } + data, _ := json.Marshal(reservation) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer nodeSub.Unsubscribe() + + time.Sleep(50 * time.Millisecond) + + input := &ec2.RunInstancesInput{ + ImageId: aws.String("ami-test"), + InstanceType: aws.String("t3.micro"), + MinCount: aws.Int64(1), + MaxCount: aws.Int64(1), + } + + reservation, err := RunInstances(input, nc, "test-account") + require.NoError(t, err) + assert.Len(t, reservation.Instances, 1) + assert.Equal(t, "i-single", aws.StringValue(reservation.Instances[0].InstanceId)) +} + +// --- placementGroupName tests --- + +func TestPlacementGroupName_WithGroupName(t *testing.T) { + input := &ec2.RunInstancesInput{ + Placement: &ec2.Placement{ + GroupName: aws.String("my-group"), + }, + } + assert.Equal(t, "my-group", placementGroupName(input)) +} + +func TestPlacementGroupName_NilPlacement(t *testing.T) { + input := &ec2.RunInstancesInput{} + assert.Equal(t, "", placementGroupName(input)) +} + +func TestPlacementGroupName_NilGroupName(t *testing.T) { + input := &ec2.RunInstancesInput{ + Placement: &ec2.Placement{}, + } + assert.Equal(t, "", placementGroupName(input)) +} + +func TestPlacementGroupName_EmptyGroupName(t *testing.T) { + input := &ec2.RunInstancesInput{ + Placement: &ec2.Placement{ + GroupName: aws.String(""), + }, + } + assert.Equal(t, "", placementGroupName(input)) +} + +// --- lookupPlacementGroupStrategy tests --- + +func TestLookupPlacementGroupStrategy_Success(t *testing.T) { + _, nc := startTestNATSServer(t) + + // Mock DescribePlacementGroups responder + sub, err := nc.QueueSubscribe("ec2.DescribePlacementGroups", "spinifex-workers", func(msg *nats.Msg) { + out := ec2.DescribePlacementGroupsOutput{ + PlacementGroups: []*ec2.PlacementGroup{ + { + GroupName: aws.String("my-group"), + Strategy: aws.String("spread"), + State: aws.String("available"), + }, + }, + } + data, _ := json.Marshal(out) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer sub.Unsubscribe() + + time.Sleep(50 * time.Millisecond) + + strategy, err := lookupPlacementGroupStrategy(nc, "test-account", "my-group") + require.NoError(t, err) + assert.Equal(t, "spread", strategy) +} + +func TestLookupPlacementGroupStrategy_NotFound(t *testing.T) { + _, nc := startTestNATSServer(t) + + // Mock responder returns error + sub, err := nc.QueueSubscribe("ec2.DescribePlacementGroups", "spinifex-workers", func(msg *nats.Msg) { + errPayload := `{"Code":"InvalidPlacementGroup.Unknown","Message":"not found"}` + _ = msg.Respond([]byte(errPayload)) + }) + require.NoError(t, err) + defer sub.Unsubscribe() + + time.Sleep(50 * time.Millisecond) + + _, err = lookupPlacementGroupStrategy(nc, "test-account", "ghost-group") + require.Error(t, err) +} + +func TestLookupPlacementGroupStrategy_NotAvailable(t *testing.T) { + _, nc := startTestNATSServer(t) + + sub, err := nc.QueueSubscribe("ec2.DescribePlacementGroups", "spinifex-workers", func(msg *nats.Msg) { + out := ec2.DescribePlacementGroupsOutput{ + PlacementGroups: []*ec2.PlacementGroup{ + { + GroupName: aws.String("my-group"), + Strategy: aws.String("spread"), + State: aws.String("deleting"), + }, + }, + } + data, _ := json.Marshal(out) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer sub.Unsubscribe() + + time.Sleep(50 * time.Millisecond) + + _, err = lookupPlacementGroupStrategy(nc, "test-account", "my-group") + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInvalidPlacementGroupUnknown, err.Error()) +} + +func TestLookupPlacementGroupStrategy_EmptyResult(t *testing.T) { + _, nc := startTestNATSServer(t) + + sub, err := nc.QueueSubscribe("ec2.DescribePlacementGroups", "spinifex-workers", func(msg *nats.Msg) { + out := ec2.DescribePlacementGroupsOutput{ + PlacementGroups: []*ec2.PlacementGroup{}, + } + data, _ := json.Marshal(out) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer sub.Unsubscribe() + + time.Sleep(50 * time.Millisecond) + + _, err = lookupPlacementGroupStrategy(nc, "test-account", "my-group") + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInvalidPlacementGroupUnknown, err.Error()) +} + +// --- distributeInstancesCluster tests --- + +func TestDistributeInstancesCluster_FirstLaunchPicksBestNode(t *testing.T) { + // First launch on empty cluster group should pick node with most capacity + _, nc := startTestNATSServer(t) + + // Mock node.status: node-1 has 4, node-2 has 6 (node-2 should be picked) + statusSub, err := nc.Subscribe("spinifex.node.status", func(msg *nats.Msg) { + for _, resp := range []types.NodeStatusResponse{ + {Node: "node-1", InstanceTypes: []types.InstanceTypeCap{{Name: "t3.micro", Available: 4}}}, + {Node: "node-2", InstanceTypes: []types.InstanceTypeCap{{Name: "t3.micro", Available: 6}}}, + } { + data, _ := json.Marshal(resp) + _ = nc.Publish(msg.Reply, data) + } + }) + require.NoError(t, err) + defer statusSub.Unsubscribe() + + // Mock ReserveClusterNode — returns node-2 (highest capacity) + reserveSub, err := nc.QueueSubscribe("ec2.ReserveClusterNode", "spinifex-workers", func(msg *nats.Msg) { + out := struct { + TargetNode string `json:"target_node"` + }{TargetNode: "node-2"} + data, _ := json.Marshal(out) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer reserveSub.Unsubscribe() + + // Mock daemon on node-2 + daemonSub, err := nc.Subscribe("ec2.RunInstances.t3.micro.node-2", func(msg *nats.Msg) { + reservation := ec2.Reservation{ + ReservationId: aws.String("r-cluster"), + Instances: []*ec2.Instance{ + {InstanceId: aws.String("i-c1")}, + {InstanceId: aws.String("i-c2")}, + {InstanceId: aws.String("i-c3")}, + }, + } + data, _ := json.Marshal(reservation) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer daemonSub.Unsubscribe() + + // Mock FinalizeClusterInstances + finalizeSub, err := nc.QueueSubscribe("ec2.FinalizeClusterInstances", "spinifex-workers", func(msg *nats.Msg) { + data, _ := json.Marshal(struct{}{}) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer finalizeSub.Unsubscribe() + + time.Sleep(50 * time.Millisecond) + + input := &ec2.RunInstancesInput{ + ImageId: aws.String("ami-test"), + InstanceType: aws.String("t3.micro"), + MinCount: aws.Int64(3), + MaxCount: aws.Int64(3), + } + + reservation, err := distributeInstancesCluster(input, nc, "test-account", "my-cluster-group") + require.NoError(t, err) + assert.Len(t, reservation.Instances, 3) + + // All instances should be from node-2 + for _, inst := range reservation.Instances { + assert.NotNil(t, inst.InstanceId) + } +} + +func TestDistributeInstancesCluster_SubsequentLaunchPinsToExistingNode(t *testing.T) { + // Subsequent launch should go to the same node the group already uses + _, nc := startTestNATSServer(t) + + // Mock node.status: both nodes have capacity + statusSub, err := nc.Subscribe("spinifex.node.status", func(msg *nats.Msg) { + for _, resp := range []types.NodeStatusResponse{ + {Node: "node-1", InstanceTypes: []types.InstanceTypeCap{{Name: "t3.micro", Available: 5}}}, + {Node: "node-2", InstanceTypes: []types.InstanceTypeCap{{Name: "t3.micro", Available: 3}}}, + } { + data, _ := json.Marshal(resp) + _ = nc.Publish(msg.Reply, data) + } + }) + require.NoError(t, err) + defer statusSub.Unsubscribe() + + // Mock ReserveClusterNode — returns node-2 (existing pinned node, even though node-1 has more capacity) + reserveSub, err := nc.QueueSubscribe("ec2.ReserveClusterNode", "spinifex-workers", func(msg *nats.Msg) { + out := struct { + TargetNode string `json:"target_node"` + }{TargetNode: "node-2"} + data, _ := json.Marshal(out) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer reserveSub.Unsubscribe() + + // Mock daemon on node-2 only — node-1 should NOT be contacted + node1Contacted := false + node1Sub, err := nc.Subscribe("ec2.RunInstances.t3.micro.node-1", func(msg *nats.Msg) { + node1Contacted = true + }) + require.NoError(t, err) + defer node1Sub.Unsubscribe() + + daemonSub, err := nc.Subscribe("ec2.RunInstances.t3.micro.node-2", func(msg *nats.Msg) { + reservation := ec2.Reservation{ + ReservationId: aws.String("r-cluster2"), + Instances: []*ec2.Instance{ + {InstanceId: aws.String("i-c4")}, + {InstanceId: aws.String("i-c5")}, + }, + } + data, _ := json.Marshal(reservation) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer daemonSub.Unsubscribe() + + // Mock FinalizeClusterInstances + finalizeSub, err := nc.QueueSubscribe("ec2.FinalizeClusterInstances", "spinifex-workers", func(msg *nats.Msg) { + data, _ := json.Marshal(struct{}{}) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer finalizeSub.Unsubscribe() + + time.Sleep(50 * time.Millisecond) + + input := &ec2.RunInstancesInput{ + ImageId: aws.String("ami-test"), + InstanceType: aws.String("t3.micro"), + MinCount: aws.Int64(2), + MaxCount: aws.Int64(2), + } + + reservation, err := distributeInstancesCluster(input, nc, "test-account", "my-cluster-group") + require.NoError(t, err) + assert.Len(t, reservation.Instances, 2) + assert.False(t, node1Contacted, "cluster should only contact the pinned node") +} + +func TestDistributeInstancesCluster_InsufficientCapacityOnPinnedNode(t *testing.T) { + // Pinned node doesn't have enough capacity → InsufficientInstanceCapacity + _, nc := startTestNATSServer(t) + + // Mock node.status: pinned node-2 has only 1 slot, node-1 has plenty + statusSub, err := nc.Subscribe("spinifex.node.status", func(msg *nats.Msg) { + for _, resp := range []types.NodeStatusResponse{ + {Node: "node-1", InstanceTypes: []types.InstanceTypeCap{{Name: "t3.micro", Available: 10}}}, + {Node: "node-2", InstanceTypes: []types.InstanceTypeCap{{Name: "t3.micro", Available: 1}}}, + } { + data, _ := json.Marshal(resp) + _ = nc.Publish(msg.Reply, data) + } + }) + require.NoError(t, err) + defer statusSub.Unsubscribe() + + // Mock ReserveClusterNode — returns node-2 (existing pinned node) + reserveSub, err := nc.QueueSubscribe("ec2.ReserveClusterNode", "spinifex-workers", func(msg *nats.Msg) { + out := struct { + TargetNode string `json:"target_node"` + }{TargetNode: "node-2"} + data, _ := json.Marshal(out) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer reserveSub.Unsubscribe() + + time.Sleep(50 * time.Millisecond) + + input := &ec2.RunInstancesInput{ + ImageId: aws.String("ami-test"), + InstanceType: aws.String("t3.micro"), + MinCount: aws.Int64(3), + MaxCount: aws.Int64(3), + } + + _, err = distributeInstancesCluster(input, nc, "test-account", "my-cluster-group") + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInsufficientInstanceCapacity, err.Error()) +} + +func TestDistributeInstancesCluster_PinnedNodeNotInCapacityResults(t *testing.T) { + // Pinned node has no capacity at all (not in fan-out results) → InsufficientInstanceCapacity + _, nc := startTestNATSServer(t) + + // Mock node.status: only node-1 has capacity (pinned node-2 is at 0) + statusSub, err := nc.Subscribe("spinifex.node.status", func(msg *nats.Msg) { + resp := types.NodeStatusResponse{ + Node: "node-1", + InstanceTypes: []types.InstanceTypeCap{{Name: "t3.micro", Available: 5}}, + } + data, _ := json.Marshal(resp) + _ = nc.Publish(msg.Reply, data) + }) + require.NoError(t, err) + defer statusSub.Unsubscribe() + + // Mock ReserveClusterNode — returns node-2 (pinned but no capacity) + reserveSub, err := nc.QueueSubscribe("ec2.ReserveClusterNode", "spinifex-workers", func(msg *nats.Msg) { + out := struct { + TargetNode string `json:"target_node"` + }{TargetNode: "node-2"} + data, _ := json.Marshal(out) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer reserveSub.Unsubscribe() + + time.Sleep(50 * time.Millisecond) + + input := &ec2.RunInstancesInput{ + ImageId: aws.String("ami-test"), + InstanceType: aws.String("t3.micro"), + MinCount: aws.Int64(1), + MaxCount: aws.Int64(1), + } + + _, err = distributeInstancesCluster(input, nc, "test-account", "my-cluster-group") + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInsufficientInstanceCapacity, err.Error()) +} + +func TestDistributeInstancesCluster_LaunchCountCappedByCapacityAndMaxCount(t *testing.T) { + // Target node has 3 available but MaxCount=2 → should launch 2 + _, nc := startTestNATSServer(t) + + statusSub, err := nc.Subscribe("spinifex.node.status", func(msg *nats.Msg) { + resp := types.NodeStatusResponse{ + Node: "node-1", + InstanceTypes: []types.InstanceTypeCap{{Name: "t3.micro", Available: 3}}, + } + data, _ := json.Marshal(resp) + _ = nc.Publish(msg.Reply, data) + }) + require.NoError(t, err) + defer statusSub.Unsubscribe() + + reserveSub, err := nc.QueueSubscribe("ec2.ReserveClusterNode", "spinifex-workers", func(msg *nats.Msg) { + out := struct { + TargetNode string `json:"target_node"` + }{TargetNode: "node-1"} + data, _ := json.Marshal(out) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer reserveSub.Unsubscribe() + + // Daemon returns 2 instances (matching assigned count) + daemonSub, err := nc.Subscribe("ec2.RunInstances.t3.micro.node-1", func(msg *nats.Msg) { + var reqInput ec2.RunInstancesInput + _ = json.Unmarshal(msg.Data, &reqInput) + count := int(aws.Int64Value(reqInput.MaxCount)) + instances := make([]*ec2.Instance, count) + for i := range instances { + instances[i] = &ec2.Instance{InstanceId: aws.String(fmt.Sprintf("i-%d", i))} + } + reservation := ec2.Reservation{ + ReservationId: aws.String("r-capped"), + Instances: instances, + } + data, _ := json.Marshal(reservation) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer daemonSub.Unsubscribe() + + finalizeSub, err := nc.QueueSubscribe("ec2.FinalizeClusterInstances", "spinifex-workers", func(msg *nats.Msg) { + data, _ := json.Marshal(struct{}{}) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer finalizeSub.Unsubscribe() + + time.Sleep(50 * time.Millisecond) + + input := &ec2.RunInstancesInput{ + ImageId: aws.String("ami-test"), + InstanceType: aws.String("t3.micro"), + MinCount: aws.Int64(1), + MaxCount: aws.Int64(2), + } + + reservation, err := distributeInstancesCluster(input, nc, "test-account", "my-cluster-group") + require.NoError(t, err) + assert.Len(t, reservation.Instances, 2, "should launch min(MaxCount=2, capacity=3) = 2") +} + +func TestRunInstances_ClusterPlacementGroupRouting(t *testing.T) { + // RunInstances with a cluster placement group should route through + // the cluster path (lookupPlacementGroupStrategy → distributeInstancesCluster). + _, nc := startTestNATSServer(t) + + // Mock DescribePlacementGroups (for lookupPlacementGroupStrategy) + pgSub, err := nc.QueueSubscribe("ec2.DescribePlacementGroups", "spinifex-workers", func(msg *nats.Msg) { + out := ec2.DescribePlacementGroupsOutput{ + PlacementGroups: []*ec2.PlacementGroup{ + { + GroupName: aws.String("my-cluster"), + Strategy: aws.String("cluster"), + State: aws.String("available"), + }, + }, + } + data, _ := json.Marshal(out) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer pgSub.Unsubscribe() + + // Mock node.status + statusSub, err := nc.Subscribe("spinifex.node.status", func(msg *nats.Msg) { + resp := types.NodeStatusResponse{ + Node: "node-1", + InstanceTypes: []types.InstanceTypeCap{{Name: "t3.micro", Available: 5}}, + } + data, _ := json.Marshal(resp) + _ = nc.Publish(msg.Reply, data) + }) + require.NoError(t, err) + defer statusSub.Unsubscribe() + + // Mock ReserveClusterNode + reserveSub, err := nc.QueueSubscribe("ec2.ReserveClusterNode", "spinifex-workers", func(msg *nats.Msg) { + out := struct { + TargetNode string `json:"target_node"` + }{TargetNode: "node-1"} + data, _ := json.Marshal(out) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer reserveSub.Unsubscribe() + + // Mock daemon + daemonSub, err := nc.Subscribe("ec2.RunInstances.t3.micro.node-1", func(msg *nats.Msg) { + reservation := ec2.Reservation{ + ReservationId: aws.String("r-cluster"), + Instances: []*ec2.Instance{ + {InstanceId: aws.String("i-c1")}, + {InstanceId: aws.String("i-c2")}, + }, + } + data, _ := json.Marshal(reservation) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer daemonSub.Unsubscribe() + + // Mock FinalizeClusterInstances + finalizeSub, err := nc.QueueSubscribe("ec2.FinalizeClusterInstances", "spinifex-workers", func(msg *nats.Msg) { + data, _ := json.Marshal(struct{}{}) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer finalizeSub.Unsubscribe() + + time.Sleep(50 * time.Millisecond) + + input := &ec2.RunInstancesInput{ + ImageId: aws.String("ami-test"), + InstanceType: aws.String("t3.micro"), + MinCount: aws.Int64(2), + MaxCount: aws.Int64(2), + Placement: &ec2.Placement{ + GroupName: aws.String("my-cluster"), + }, + } + + reservation, err := RunInstances(input, nc, "test-account") + require.NoError(t, err) + assert.Len(t, reservation.Instances, 2) +} + +func TestRunInstances_MultiInstanceUsesDistribution(t *testing.T) { + // For MaxCount > 1, RunInstances should use the distribution path, + // which queries spinifex.node.status. + _, nc := startTestNATSServer(t) + + statusQueried := false + statusSub, err := nc.Subscribe("spinifex.node.status", func(msg *nats.Msg) { + statusQueried = true + resp := types.NodeStatusResponse{ + Node: "node-1", + InstanceTypes: []types.InstanceTypeCap{{Name: "t3.micro", Available: 3}}, + } + data, _ := json.Marshal(resp) + _ = nc.Publish(msg.Reply, data) + }) + require.NoError(t, err) + defer statusSub.Unsubscribe() + + // Mock node-specific handler + nodeSub, err := nc.Subscribe("ec2.RunInstances.t3.micro.node-1", func(msg *nats.Msg) { + reservation := ec2.Reservation{ + ReservationId: aws.String("r-multi"), + Instances: []*ec2.Instance{ + {InstanceId: aws.String("i-001")}, + {InstanceId: aws.String("i-002")}, + }, + } + data, _ := json.Marshal(reservation) + _ = msg.Respond(data) + }) + require.NoError(t, err) + defer nodeSub.Unsubscribe() + + time.Sleep(50 * time.Millisecond) + + input := &ec2.RunInstancesInput{ + ImageId: aws.String("ami-test"), + InstanceType: aws.String("t3.micro"), + MinCount: aws.Int64(2), + MaxCount: aws.Int64(2), + } + + reservation, err := RunInstances(input, nc, "test-account") + require.NoError(t, err) + assert.Len(t, reservation.Instances, 2) + assert.True(t, statusQueried, "multi-instance launch should query node status") +} diff --git a/spinifex/gateway/ec2/placementgroup/CreatePlacementGroup.go b/spinifex/gateway/ec2/placementgroup/CreatePlacementGroup.go new file mode 100644 index 000000000..b3e8be0a4 --- /dev/null +++ b/spinifex/gateway/ec2/placementgroup/CreatePlacementGroup.go @@ -0,0 +1,33 @@ +package gateway_ec2_placementgroup + +import ( + "errors" + + "github.com/aws/aws-sdk-go/service/ec2" + "github.com/mulgadc/spinifex/spinifex/awserrors" + handlers_ec2_placementgroup "github.com/mulgadc/spinifex/spinifex/handlers/ec2/placementgroup" + "github.com/nats-io/nats.go" +) + +// CreatePlacementGroup handles the EC2 CreatePlacementGroup API call. +func CreatePlacementGroup(input *ec2.CreatePlacementGroupInput, natsConn *nats.Conn, accountID string) (ec2.CreatePlacementGroupOutput, error) { + var output ec2.CreatePlacementGroupOutput + + if input == nil { + return output, errors.New(awserrors.ErrorInvalidParameterValue) + } + if input.GroupName == nil || *input.GroupName == "" { + return output, errors.New(awserrors.ErrorMissingParameter) + } + if input.Strategy == nil || *input.Strategy == "" { + return output, errors.New(awserrors.ErrorMissingParameter) + } + + svc := handlers_ec2_placementgroup.NewNATSPlacementGroupService(natsConn) + result, err := svc.CreatePlacementGroup(input, accountID) + if err != nil { + return output, err + } + + return *result, nil +} diff --git a/spinifex/gateway/ec2/placementgroup/DeletePlacementGroup.go b/spinifex/gateway/ec2/placementgroup/DeletePlacementGroup.go new file mode 100644 index 000000000..33ca644be --- /dev/null +++ b/spinifex/gateway/ec2/placementgroup/DeletePlacementGroup.go @@ -0,0 +1,30 @@ +package gateway_ec2_placementgroup + +import ( + "errors" + + "github.com/aws/aws-sdk-go/service/ec2" + "github.com/mulgadc/spinifex/spinifex/awserrors" + handlers_ec2_placementgroup "github.com/mulgadc/spinifex/spinifex/handlers/ec2/placementgroup" + "github.com/nats-io/nats.go" +) + +// DeletePlacementGroup handles the EC2 DeletePlacementGroup API call. +func DeletePlacementGroup(input *ec2.DeletePlacementGroupInput, natsConn *nats.Conn, accountID string) (ec2.DeletePlacementGroupOutput, error) { + var output ec2.DeletePlacementGroupOutput + + if input == nil { + return output, errors.New(awserrors.ErrorInvalidParameterValue) + } + if input.GroupName == nil || *input.GroupName == "" { + return output, errors.New(awserrors.ErrorMissingParameter) + } + + svc := handlers_ec2_placementgroup.NewNATSPlacementGroupService(natsConn) + result, err := svc.DeletePlacementGroup(input, accountID) + if err != nil { + return output, err + } + + return *result, nil +} diff --git a/spinifex/gateway/ec2/placementgroup/DescribePlacementGroups.go b/spinifex/gateway/ec2/placementgroup/DescribePlacementGroups.go new file mode 100644 index 000000000..95b5fea18 --- /dev/null +++ b/spinifex/gateway/ec2/placementgroup/DescribePlacementGroups.go @@ -0,0 +1,20 @@ +package gateway_ec2_placementgroup + +import ( + "github.com/aws/aws-sdk-go/service/ec2" + handlers_ec2_placementgroup "github.com/mulgadc/spinifex/spinifex/handlers/ec2/placementgroup" + "github.com/nats-io/nats.go" +) + +// DescribePlacementGroups handles the EC2 DescribePlacementGroups API call. +func DescribePlacementGroups(input *ec2.DescribePlacementGroupsInput, natsConn *nats.Conn, accountID string) (ec2.DescribePlacementGroupsOutput, error) { + var output ec2.DescribePlacementGroupsOutput + + svc := handlers_ec2_placementgroup.NewNATSPlacementGroupService(natsConn) + result, err := svc.DescribePlacementGroups(input, accountID) + if err != nil { + return output, err + } + + return *result, nil +} diff --git a/spinifex/gateway/gateway_test.go b/spinifex/gateway/gateway_test.go index 4f0203416..82f8ab74f 100644 --- a/spinifex/gateway/gateway_test.go +++ b/spinifex/gateway/gateway_test.go @@ -1020,6 +1020,7 @@ func TestEC2ActionMapCompleteness(t *testing.T) { "DescribeInternetGateways", "AttachInternetGateway", "DetachInternetGateway", "CreateEgressOnlyInternetGateway", "DeleteEgressOnlyInternetGateway", "DescribeEgressOnlyInternetGateways", + "CreatePlacementGroup", "DeletePlacementGroup", "DescribePlacementGroups", "CreateVpc", "DeleteVpc", "DescribeVpcs", "CreateSubnet", "DeleteSubnet", "DescribeSubnets", "ModifySubnetAttribute", "CreateNetworkInterface", "DeleteNetworkInterface", "DescribeNetworkInterfaces", diff --git a/spinifex/handlers/ec2/placementgroup/service.go b/spinifex/handlers/ec2/placementgroup/service.go new file mode 100644 index 000000000..c8ecabc2c --- /dev/null +++ b/spinifex/handlers/ec2/placementgroup/service.go @@ -0,0 +1,77 @@ +package handlers_ec2_placementgroup + +import "github.com/aws/aws-sdk-go/service/ec2" + +// PlacementGroupService defines the interface for placement group operations. +type PlacementGroupService interface { + CreatePlacementGroup(input *ec2.CreatePlacementGroupInput, accountID string) (*ec2.CreatePlacementGroupOutput, error) + DeletePlacementGroup(input *ec2.DeletePlacementGroupInput, accountID string) (*ec2.DeletePlacementGroupOutput, error) + DescribePlacementGroups(input *ec2.DescribePlacementGroupsInput, accountID string) (*ec2.DescribePlacementGroupsOutput, error) + ReserveSpreadNodes(input *ReserveSpreadNodesInput, accountID string) (*ReserveSpreadNodesOutput, error) + FinalizeSpreadInstances(input *FinalizeSpreadInstancesInput, accountID string) (*FinalizeSpreadInstancesOutput, error) + ReleaseSpreadNodes(input *ReleaseSpreadNodesInput, accountID string) (*ReleaseSpreadNodesOutput, error) + RemoveInstance(input *RemoveInstanceInput, accountID string) (*RemoveInstanceOutput, error) + ReserveClusterNode(input *ReserveClusterNodeInput, accountID string) (*ReserveClusterNodeOutput, error) + FinalizeClusterInstances(input *FinalizeClusterInstancesInput, accountID string) (*FinalizeClusterInstancesOutput, error) +} + +// ReserveSpreadNodesInput requests atomic node reservation for a spread placement group. +type ReserveSpreadNodesInput struct { + GroupName string `json:"group_name"` + EligibleNodes []string `json:"eligible_nodes"` // nodes with capacity (from fan-out) + MinCount int `json:"min_count"` + MaxCount int `json:"max_count"` +} + +// ReserveSpreadNodesOutput contains the nodes selected for launch. +type ReserveSpreadNodesOutput struct { + ReservedNodes []string `json:"reserved_nodes"` +} + +// FinalizeSpreadInstancesInput records actual instance IDs on reserved nodes. +type FinalizeSpreadInstancesInput struct { + GroupName string `json:"group_name"` + NodeInstances map[string][]string `json:"node_instances"` // node -> instance IDs +} + +// FinalizeSpreadInstancesOutput is empty on success. +type FinalizeSpreadInstancesOutput struct{} + +// ReleaseSpreadNodesInput releases previously reserved node slots (rollback). +type ReleaseSpreadNodesInput struct { + GroupName string `json:"group_name"` + Nodes []string `json:"nodes"` +} + +// ReleaseSpreadNodesOutput is empty on success. +type ReleaseSpreadNodesOutput struct{} + +// RemoveInstanceInput removes a specific instance from its placement group's NodeInstances. +type RemoveInstanceInput struct { + GroupName string `json:"group_name"` + NodeName string `json:"node_name"` + InstanceID string `json:"instance_id"` +} + +// RemoveInstanceOutput is empty on success. +type RemoveInstanceOutput struct{} + +// ReserveClusterNodeInput requests target node determination for a cluster placement group. +type ReserveClusterNodeInput struct { + GroupName string `json:"group_name"` + EligibleNodes []string `json:"eligible_nodes"` // nodes with capacity, sorted by capacity desc +} + +// ReserveClusterNodeOutput contains the target node for cluster placement. +type ReserveClusterNodeOutput struct { + TargetNode string `json:"target_node"` +} + +// FinalizeClusterInstancesInput records instance IDs launched in a cluster placement group. +type FinalizeClusterInstancesInput struct { + GroupName string `json:"group_name"` + NodeInstances map[string][]string `json:"node_instances"` // node -> instance IDs to append +} + +// FinalizeClusterInstancesOutput is empty on success. +type FinalizeClusterInstancesOutput struct{} diff --git a/spinifex/handlers/ec2/placementgroup/service_impl.go b/spinifex/handlers/ec2/placementgroup/service_impl.go new file mode 100644 index 000000000..b5d181b7e --- /dev/null +++ b/spinifex/handlers/ec2/placementgroup/service_impl.go @@ -0,0 +1,572 @@ +package handlers_ec2_placementgroup + +import ( + "encoding/json" + "errors" + "fmt" + "log/slog" + "maps" + "strings" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/ec2" + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/config" + "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/nats-io/nats.go" +) + +// Ensure PlacementGroupServiceImpl implements PlacementGroupService +var _ PlacementGroupService = (*PlacementGroupServiceImpl)(nil) + +const ( + KVBucketPlacementGroups = "spinifex-placement-groups" + KVBucketPlacementGroupsVersion = 1 +) + +// PlacementGroupRecord represents a stored placement group. +type PlacementGroupRecord struct { + GroupId string `json:"group_id"` + GroupName string `json:"group_name"` + Strategy string `json:"strategy"` + State string `json:"state"` + // SpreadLevel is always "host" for bare-metal Spinifex clusters. + SpreadLevel string `json:"spread_level"` + AccountID string `json:"account_id"` + // NodeInstances tracks which node hosts which instances in this group. + // Key = node name, Value = list of instance IDs on that node. + NodeInstances map[string][]string `json:"node_instances"` +} + +// PlacementGroupServiceImpl implements placement group operations with NATS JetStream persistence. +type PlacementGroupServiceImpl struct { + config *config.Config + natsConn *nats.Conn + kv nats.KeyValue +} + +// NewPlacementGroupServiceImplWithNATS creates a placement group service with NATS JetStream. +func NewPlacementGroupServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn) (*PlacementGroupServiceImpl, error) { + js, err := natsConn.JetStream() + if err != nil { + return nil, fmt.Errorf("failed to get JetStream context: %w", err) + } + + kv, err := utils.GetOrCreateKVBucket(js, KVBucketPlacementGroups, 10) + if err != nil { + return nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketPlacementGroups, err) + } + if err := utils.WriteVersion(kv, KVBucketPlacementGroupsVersion); err != nil { + return nil, fmt.Errorf("write version to %s: %w", KVBucketPlacementGroups, err) + } + + slog.Info("Placement group service initialized with JetStream KV", "bucket", KVBucketPlacementGroups) + + return &PlacementGroupServiceImpl{ + config: cfg, + natsConn: natsConn, + kv: kv, + }, nil +} + +// CreatePlacementGroup creates a new placement group. +func (s *PlacementGroupServiceImpl) CreatePlacementGroup(input *ec2.CreatePlacementGroupInput, accountID string) (*ec2.CreatePlacementGroupOutput, error) { + if input.GroupName == nil || *input.GroupName == "" { + return nil, errors.New(awserrors.ErrorMissingParameter) + } + + strategy := aws.StringValue(input.Strategy) + if strategy == "" { + return nil, errors.New(awserrors.ErrorMissingParameter) + } + + // Only spread and cluster are supported; partition is rejected. + if strategy == ec2.PlacementStrategyPartition { + return nil, errors.New(awserrors.ErrorInvalidParameterValue) + } + if strategy != ec2.PlacementStrategySpread && strategy != ec2.PlacementStrategyCluster { + return nil, errors.New(awserrors.ErrorInvalidParameterValue) + } + + groupName := *input.GroupName + key := utils.AccountKey(accountID, groupName) + groupID := utils.GenerateResourceID("pg") + + record := PlacementGroupRecord{ + GroupId: groupID, + GroupName: groupName, + Strategy: strategy, + State: ec2.PlacementGroupStateAvailable, + SpreadLevel: ec2.SpreadLevelHost, + AccountID: accountID, + NodeInstances: make(map[string][]string), + } + + data, err := json.Marshal(record) + if err != nil { + return nil, errors.New(awserrors.ErrorServerInternal) + } + // Atomic create-if-not-exists to prevent TOCTOU race on duplicate names + if _, err := s.kv.Create(key, data); err != nil { + // Create fails if key already exists + return nil, errors.New(awserrors.ErrorInvalidPlacementGroupDuplicate) + } + + slog.Info("CreatePlacementGroup completed", "groupId", groupID, "groupName", groupName, "strategy", strategy, "accountID", accountID) + + return &ec2.CreatePlacementGroupOutput{ + PlacementGroup: s.recordToEC2(&record), + }, nil +} + +// DeletePlacementGroup deletes a placement group. +func (s *PlacementGroupServiceImpl) DeletePlacementGroup(input *ec2.DeletePlacementGroupInput, accountID string) (*ec2.DeletePlacementGroupOutput, error) { + if input.GroupName == nil || *input.GroupName == "" { + return nil, errors.New(awserrors.ErrorMissingParameter) + } + + groupName := *input.GroupName + key := utils.AccountKey(accountID, groupName) + + entry, err := s.kv.Get(key) + if err != nil { + return nil, errors.New(awserrors.ErrorInvalidPlacementGroupUnknown) + } + + var record PlacementGroupRecord + if err := json.Unmarshal(entry.Value(), &record); err != nil { + return nil, errors.New(awserrors.ErrorServerInternal) + } + + // Check for running instances + instanceCount := 0 + for _, ids := range record.NodeInstances { + instanceCount += len(ids) + } + if instanceCount > 0 { + return nil, errors.New(awserrors.ErrorInvalidPlacementGroupInUse) + } + + if err := s.kv.Delete(key); err != nil { + return nil, errors.New(awserrors.ErrorServerInternal) + } + + slog.Info("DeletePlacementGroup completed", "groupName", groupName, "accountID", accountID) + + return &ec2.DeletePlacementGroupOutput{}, nil +} + +// DescribePlacementGroups lists placement groups with optional filters. +func (s *PlacementGroupServiceImpl) DescribePlacementGroups(input *ec2.DescribePlacementGroupsInput, accountID string) (*ec2.DescribePlacementGroupsOutput, error) { + // Build filter maps + nameSet := make(map[string]bool) + for _, name := range input.GroupNames { + if name != nil { + nameSet[*name] = true + } + } + idSet := make(map[string]bool) + for _, id := range input.GroupIds { + if id != nil { + idSet[*id] = true + } + } + + // Extract filter values + filterStrategy := "" + filterState := "" + filterSpreadLevel := "" + filterGroupName := "" + for _, f := range input.Filters { + if f.Name == nil || len(f.Values) == 0 || f.Values[0] == nil { + continue + } + switch *f.Name { + case "strategy": + filterStrategy = *f.Values[0] + case "state": + filterState = *f.Values[0] + case "spread-level": + filterSpreadLevel = *f.Values[0] + case "group-name": + filterGroupName = *f.Values[0] + } + } + + prefix := accountID + "." + keys, err := s.kv.Keys() + if err != nil && !errors.Is(err, nats.ErrNoKeysFound) { + return nil, errors.New(awserrors.ErrorServerInternal) + } + + var groups []*ec2.PlacementGroup + for _, k := range keys { + if k == utils.VersionKey { + continue + } + if !strings.HasPrefix(k, prefix) { + continue + } + + entry, err := s.kv.Get(k) + if err != nil { + slog.Warn("Failed to get placement group record", "key", k, "error", err) + continue + } + + var record PlacementGroupRecord + if err := json.Unmarshal(entry.Value(), &record); err != nil { + slog.Warn("Failed to unmarshal placement group record", "key", k, "error", err) + continue + } + + // Apply name filter (from GroupNames parameter) + if len(nameSet) > 0 && !nameSet[record.GroupName] { + continue + } + // Apply ID filter (from GroupIds parameter) + if len(idSet) > 0 && !idSet[record.GroupId] { + continue + } + // Apply Filters + if filterStrategy != "" && record.Strategy != filterStrategy { + continue + } + if filterState != "" && record.State != filterState { + continue + } + if filterSpreadLevel != "" && record.SpreadLevel != filterSpreadLevel { + continue + } + if filterGroupName != "" && record.GroupName != filterGroupName { + continue + } + + groups = append(groups, s.recordToEC2(&record)) + } + + // If specific names were requested but not found, return error + if len(nameSet) > 0 { + found := make(map[string]bool) + for _, g := range groups { + if g.GroupName != nil { + found[*g.GroupName] = true + } + } + for name := range nameSet { + if !found[name] { + return nil, errors.New(awserrors.ErrorInvalidPlacementGroupUnknown) + } + } + } + + slog.Info("DescribePlacementGroups completed", "count", len(groups), "accountID", accountID) + + return &ec2.DescribePlacementGroupsOutput{ + PlacementGroups: groups, + }, nil +} + +// GetPlacementGroupRecord reads a placement group record from KV with its revision for CAS operations. +// Returns the record and the KV entry (for revision). Exported for use by gateway spread routing. +func (s *PlacementGroupServiceImpl) GetPlacementGroupRecord(accountID, groupName string) (*PlacementGroupRecord, nats.KeyValueEntry, error) { + key := utils.AccountKey(accountID, groupName) + entry, err := s.kv.Get(key) + if err != nil { + return nil, nil, errors.New(awserrors.ErrorInvalidPlacementGroupUnknown) + } + + var record PlacementGroupRecord + if err := json.Unmarshal(entry.Value(), &record); err != nil { + return nil, nil, errors.New(awserrors.ErrorServerInternal) + } + + return &record, entry, nil +} + +// UpdatePlacementGroupRecord writes a placement group record using CAS (optimistic concurrency). +// Returns nil on success or the error on CAS conflict. +func (s *PlacementGroupServiceImpl) UpdatePlacementGroupRecord(accountID, groupName string, record *PlacementGroupRecord, revision uint64) error { + key := utils.AccountKey(accountID, groupName) + data, err := json.Marshal(record) + if err != nil { + return errors.New(awserrors.ErrorServerInternal) + } + if _, err := s.kv.Update(key, data, revision); err != nil { + return err + } + return nil +} + +const maxCASRetries = 5 + +// ReserveSpreadNodes atomically reserves node slots for a spread placement group launch. +// It reads the group record, filters eligible nodes (excluding already-occupied ones), +// selects nodes, writes placeholder entries, and returns the selected nodes. +// Uses CAS with retries following the IPAM pattern. +func (s *PlacementGroupServiceImpl) ReserveSpreadNodes(input *ReserveSpreadNodesInput, accountID string) (*ReserveSpreadNodesOutput, error) { + if input.GroupName == "" { + return nil, errors.New(awserrors.ErrorMissingParameter) + } + + for attempt := range maxCASRetries { + record, entry, err := s.GetPlacementGroupRecord(accountID, input.GroupName) + if err != nil { + return nil, err + } + + if record.State != ec2.PlacementGroupStateAvailable { + return nil, errors.New(awserrors.ErrorInvalidPlacementGroupUnknown) + } + if record.Strategy != ec2.PlacementStrategySpread { + return nil, errors.New(awserrors.ErrorInvalidParameterValue) + } + + // Build set of nodes already hosting instances in this group + occupiedNodes := make(map[string]bool) + for node := range record.NodeInstances { + occupiedNodes[node] = true + } + + // Filter eligible nodes: must have capacity AND not already occupied + var available []string + for _, node := range input.EligibleNodes { + if !occupiedNodes[node] { + available = append(available, node) + } + } + + if len(available) < input.MinCount { + return nil, errors.New(awserrors.ErrorInsufficientInstanceCapacity) + } + + // Select nodes: up to MaxCount, at least MinCount + launchCount := min(input.MaxCount, len(available)) + selected := available[:launchCount] + + // Add placeholder entries (empty instance list = reserved but not yet launched) + for _, node := range selected { + record.NodeInstances[node] = []string{} + } + + // CAS write — retry on conflict + if err := s.UpdatePlacementGroupRecord(accountID, input.GroupName, record, entry.Revision()); err != nil { + slog.Debug("ReserveSpreadNodes: CAS conflict, retrying", "attempt", attempt, "err", err) + continue + } + + slog.Info("ReserveSpreadNodes completed", "groupName", input.GroupName, "nodes", selected, "accountID", accountID) + return &ReserveSpreadNodesOutput{ReservedNodes: selected}, nil + } + + slog.Error("ReserveSpreadNodes: CAS retries exhausted", "groupName", input.GroupName, "accountID", accountID) + return nil, errors.New(awserrors.ErrorServerInternal) +} + +// FinalizeSpreadInstances replaces placeholder entries with actual instance IDs. +// Uses CAS with retries following the IPAM pattern. +func (s *PlacementGroupServiceImpl) FinalizeSpreadInstances(input *FinalizeSpreadInstancesInput, accountID string) (*FinalizeSpreadInstancesOutput, error) { + if input.GroupName == "" { + return nil, errors.New(awserrors.ErrorMissingParameter) + } + + for attempt := range maxCASRetries { + record, entry, err := s.GetPlacementGroupRecord(accountID, input.GroupName) + if err != nil { + return nil, err + } + + // Replace placeholder entries with actual instance IDs per node + maps.Copy(record.NodeInstances, input.NodeInstances) + + if err := s.UpdatePlacementGroupRecord(accountID, input.GroupName, record, entry.Revision()); err != nil { + slog.Debug("FinalizeSpreadInstances: CAS conflict, retrying", "attempt", attempt, "err", err) + continue + } + + slog.Info("FinalizeSpreadInstances completed", "groupName", input.GroupName, "accountID", accountID) + return &FinalizeSpreadInstancesOutput{}, nil + } + + slog.Error("FinalizeSpreadInstances: CAS retries exhausted", "groupName", input.GroupName, "accountID", accountID) + return nil, errors.New(awserrors.ErrorServerInternal) +} + +// ReleaseSpreadNodes removes placeholder entries for nodes that failed to launch. +// Uses CAS with retries following the IPAM pattern. +func (s *PlacementGroupServiceImpl) ReleaseSpreadNodes(input *ReleaseSpreadNodesInput, accountID string) (*ReleaseSpreadNodesOutput, error) { + if input.GroupName == "" { + return nil, errors.New(awserrors.ErrorMissingParameter) + } + + releaseSet := make(map[string]bool, len(input.Nodes)) + for _, n := range input.Nodes { + releaseSet[n] = true + } + + for attempt := range maxCASRetries { + record, entry, err := s.GetPlacementGroupRecord(accountID, input.GroupName) + if err != nil { + return nil, err + } + + for node := range releaseSet { + // Only release placeholder entries (empty instance list). + // Entries with real instance IDs must not be deleted. + if ids, ok := record.NodeInstances[node]; ok && len(ids) == 0 { + delete(record.NodeInstances, node) + } + } + + if err := s.UpdatePlacementGroupRecord(accountID, input.GroupName, record, entry.Revision()); err != nil { + slog.Debug("ReleaseSpreadNodes: CAS conflict, retrying", "attempt", attempt, "err", err) + continue + } + + slog.Info("ReleaseSpreadNodes completed", "groupName", input.GroupName, "nodes", input.Nodes, "accountID", accountID) + return &ReleaseSpreadNodesOutput{}, nil + } + + slog.Error("ReleaseSpreadNodes: CAS retries exhausted", "groupName", input.GroupName, "accountID", accountID) + return nil, errors.New(awserrors.ErrorServerInternal) +} + +// RemoveInstance removes a specific instance from a placement group's NodeInstances. +// If the node's instance list becomes empty after removal, the node key is deleted. +// Uses CAS with retries following the IPAM pattern. +func (s *PlacementGroupServiceImpl) RemoveInstance(input *RemoveInstanceInput, accountID string) (*RemoveInstanceOutput, error) { + if input.GroupName == "" || input.InstanceID == "" || input.NodeName == "" { + return nil, errors.New(awserrors.ErrorMissingParameter) + } + + for attempt := range maxCASRetries { + record, entry, lookupErr := s.GetPlacementGroupRecord(accountID, input.GroupName) + if lookupErr != nil { + // Group may have been deleted already — treat as success + slog.Debug("RemoveInstance: group not found, treating as success", "groupName", input.GroupName) + return &RemoveInstanceOutput{}, nil //nolint:nilerr // intentional: deleted group = success + } + + instances, exists := record.NodeInstances[input.NodeName] + if !exists { + // Node not tracked — nothing to remove + return &RemoveInstanceOutput{}, nil + } + + // Remove the specific instance ID + filtered := make([]string, 0, len(instances)) + for _, id := range instances { + if id != input.InstanceID { + filtered = append(filtered, id) + } + } + + if len(filtered) == 0 { + delete(record.NodeInstances, input.NodeName) + } else { + record.NodeInstances[input.NodeName] = filtered + } + + if err := s.UpdatePlacementGroupRecord(accountID, input.GroupName, record, entry.Revision()); err != nil { + slog.Debug("RemoveInstance: CAS conflict, retrying", "attempt", attempt, "err", err) + continue + } + + slog.Info("RemoveInstance completed", "groupName", input.GroupName, "instanceId", input.InstanceID, "node", input.NodeName, "accountID", accountID) + return &RemoveInstanceOutput{}, nil + } + + slog.Error("RemoveInstance: CAS retries exhausted", "groupName", input.GroupName, "instanceId", input.InstanceID, "accountID", accountID) + return nil, errors.New(awserrors.ErrorServerInternal) +} + +// ReserveClusterNode determines the target node for a cluster placement group launch. +// If the group already has instances, it returns the existing node (cluster = all on one node). +// If empty, it picks the first eligible node (highest capacity) and writes a placeholder via CAS. +func (s *PlacementGroupServiceImpl) ReserveClusterNode(input *ReserveClusterNodeInput, accountID string) (*ReserveClusterNodeOutput, error) { + if input.GroupName == "" { + return nil, errors.New(awserrors.ErrorMissingParameter) + } + + for attempt := range maxCASRetries { + record, entry, err := s.GetPlacementGroupRecord(accountID, input.GroupName) + if err != nil { + return nil, err + } + if record.State != ec2.PlacementGroupStateAvailable { + return nil, errors.New(awserrors.ErrorInvalidPlacementGroupUnknown) + } + if record.Strategy != ec2.PlacementStrategyCluster { + return nil, errors.New(awserrors.ErrorInvalidParameterValue) + } + + // If existing node, return it immediately (no CAS write needed) + for node := range record.NodeInstances { + return &ReserveClusterNodeOutput{TargetNode: node}, nil + } + + // Empty group: need to pick and claim a node via CAS + if len(input.EligibleNodes) == 0 { + return nil, errors.New(awserrors.ErrorInsufficientInstanceCapacity) + } + + targetNode := input.EligibleNodes[0] // first = highest capacity (sorted desc by caller) + record.NodeInstances[targetNode] = []string{} + + if err := s.UpdatePlacementGroupRecord(accountID, input.GroupName, record, entry.Revision()); err != nil { + slog.Debug("ReserveClusterNode: CAS conflict, retrying", "attempt", attempt, "err", err) + continue + } + + slog.Info("ReserveClusterNode completed", "groupName", input.GroupName, "targetNode", targetNode, "accountID", accountID) + return &ReserveClusterNodeOutput{TargetNode: targetNode}, nil + } + + slog.Error("ReserveClusterNode: CAS retries exhausted", "groupName", input.GroupName, "accountID", accountID) + return nil, errors.New(awserrors.ErrorServerInternal) +} + +// FinalizeClusterInstances appends launched instance IDs to the cluster placement group record. +// Uses CAS with retries following the IPAM pattern. +func (s *PlacementGroupServiceImpl) FinalizeClusterInstances(input *FinalizeClusterInstancesInput, accountID string) (*FinalizeClusterInstancesOutput, error) { + if input.GroupName == "" { + return nil, errors.New(awserrors.ErrorMissingParameter) + } + + for attempt := range maxCASRetries { + record, entry, err := s.GetPlacementGroupRecord(accountID, input.GroupName) + if err != nil { + return nil, err + } + + // Append new instance IDs to existing entries (cluster may have concurrent launches) + for node, ids := range input.NodeInstances { + record.NodeInstances[node] = append(record.NodeInstances[node], ids...) + } + + if err := s.UpdatePlacementGroupRecord(accountID, input.GroupName, record, entry.Revision()); err != nil { + slog.Debug("FinalizeClusterInstances: CAS conflict, retrying", "attempt", attempt, "err", err) + continue + } + + slog.Info("FinalizeClusterInstances completed", "groupName", input.GroupName, "accountID", accountID) + return &FinalizeClusterInstancesOutput{}, nil + } + + slog.Error("FinalizeClusterInstances: CAS retries exhausted", "groupName", input.GroupName, "accountID", accountID) + return nil, errors.New(awserrors.ErrorServerInternal) +} + +// recordToEC2 converts an internal record to the AWS SDK PlacementGroup type. +func (s *PlacementGroupServiceImpl) recordToEC2(record *PlacementGroupRecord) *ec2.PlacementGroup { + pg := &ec2.PlacementGroup{ + GroupId: aws.String(record.GroupId), + GroupName: aws.String(record.GroupName), + Strategy: aws.String(record.Strategy), + State: aws.String(record.State), + } + if record.Strategy == ec2.PlacementStrategySpread { + pg.SpreadLevel = aws.String(record.SpreadLevel) + } + return pg +} diff --git a/spinifex/handlers/ec2/placementgroup/service_impl_test.go b/spinifex/handlers/ec2/placementgroup/service_impl_test.go new file mode 100644 index 000000000..3ea8d8718 --- /dev/null +++ b/spinifex/handlers/ec2/placementgroup/service_impl_test.go @@ -0,0 +1,840 @@ +package handlers_ec2_placementgroup + +import ( + "testing" + "time" + + "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/ec2" + "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/nats-io/nats-server/v2/server" + "github.com/nats-io/nats.go" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const testAccountID = "123456789012" + +func setupTestService(t *testing.T) *PlacementGroupServiceImpl { + t.Helper() + opts := &server.Options{ + Host: "127.0.0.1", + Port: -1, + JetStream: true, + StoreDir: t.TempDir(), + NoLog: true, + NoSigs: true, + } + ns, err := server.NewServer(opts) + require.NoError(t, err) + go ns.Start() + require.True(t, ns.ReadyForConnections(5*time.Second)) + t.Cleanup(func() { ns.Shutdown() }) + + nc, err := nats.Connect(ns.ClientURL()) + require.NoError(t, err) + t.Cleanup(func() { nc.Close() }) + + svc, err := NewPlacementGroupServiceImplWithNATS(nil, nc) + require.NoError(t, err) + return svc +} + +func createTestGroup(t *testing.T, svc *PlacementGroupServiceImpl, name, strategy string) *ec2.PlacementGroup { + t.Helper() + out, err := svc.CreatePlacementGroup(&ec2.CreatePlacementGroupInput{ + GroupName: aws.String(name), + Strategy: aws.String(strategy), + }, testAccountID) + require.NoError(t, err) + return out.PlacementGroup +} + +// --- CreatePlacementGroup Tests --- + +func TestCreatePlacementGroup_Spread(t *testing.T) { + svc := setupTestService(t) + pg := createTestGroup(t, svc, "my-spread-group", "spread") + + assert.Equal(t, "pg-", (*pg.GroupId)[:3]) + assert.Equal(t, "my-spread-group", *pg.GroupName) + assert.Equal(t, "spread", *pg.Strategy) + assert.Equal(t, "available", *pg.State) + assert.Equal(t, "host", *pg.SpreadLevel) +} + +func TestCreatePlacementGroup_Cluster(t *testing.T) { + svc := setupTestService(t) + pg := createTestGroup(t, svc, "my-cluster-group", "cluster") + + assert.Equal(t, "pg-", (*pg.GroupId)[:3]) + assert.Equal(t, "my-cluster-group", *pg.GroupName) + assert.Equal(t, "cluster", *pg.Strategy) + assert.Equal(t, "available", *pg.State) + // SpreadLevel should be nil for cluster strategy + assert.Nil(t, pg.SpreadLevel) +} + +func TestCreatePlacementGroup_DuplicateName(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "dup-group", "spread") + + _, err := svc.CreatePlacementGroup(&ec2.CreatePlacementGroupInput{ + GroupName: aws.String("dup-group"), + Strategy: aws.String("spread"), + }, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInvalidPlacementGroupDuplicate, err.Error()) +} + +func TestCreatePlacementGroup_PartitionRejected(t *testing.T) { + svc := setupTestService(t) + _, err := svc.CreatePlacementGroup(&ec2.CreatePlacementGroupInput{ + GroupName: aws.String("part-group"), + Strategy: aws.String("partition"), + }, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInvalidParameterValue, err.Error()) +} + +func TestCreatePlacementGroup_InvalidStrategy(t *testing.T) { + svc := setupTestService(t) + _, err := svc.CreatePlacementGroup(&ec2.CreatePlacementGroupInput{ + GroupName: aws.String("bad-group"), + Strategy: aws.String("invalid"), + }, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInvalidParameterValue, err.Error()) +} + +func TestCreatePlacementGroup_MissingName(t *testing.T) { + svc := setupTestService(t) + _, err := svc.CreatePlacementGroup(&ec2.CreatePlacementGroupInput{ + Strategy: aws.String("spread"), + }, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorMissingParameter, err.Error()) +} + +func TestCreatePlacementGroup_MissingStrategy(t *testing.T) { + svc := setupTestService(t) + _, err := svc.CreatePlacementGroup(&ec2.CreatePlacementGroupInput{ + GroupName: aws.String("no-strat"), + }, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorMissingParameter, err.Error()) +} + +func TestCreatePlacementGroup_SameNameDifferentAccounts(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "shared-name", "spread") + + // Different account should succeed + out, err := svc.CreatePlacementGroup(&ec2.CreatePlacementGroupInput{ + GroupName: aws.String("shared-name"), + Strategy: aws.String("cluster"), + }, "999999999999") + require.NoError(t, err) + assert.Equal(t, "shared-name", *out.PlacementGroup.GroupName) +} + +// --- DeletePlacementGroup Tests --- + +func TestDeletePlacementGroup_Success(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "del-group", "spread") + + _, err := svc.DeletePlacementGroup(&ec2.DeletePlacementGroupInput{ + GroupName: aws.String("del-group"), + }, testAccountID) + require.NoError(t, err) + + // Verify it's gone + out, err := svc.DescribePlacementGroups(&ec2.DescribePlacementGroupsInput{}, testAccountID) + require.NoError(t, err) + assert.Empty(t, out.PlacementGroups) +} + +func TestDeletePlacementGroup_NotFound(t *testing.T) { + svc := setupTestService(t) + _, err := svc.DeletePlacementGroup(&ec2.DeletePlacementGroupInput{ + GroupName: aws.String("nonexistent"), + }, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInvalidPlacementGroupUnknown, err.Error()) +} + +func TestDeletePlacementGroup_InUse(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "in-use-group", "spread") + + // Simulate instances by updating the record directly + record, entry, err := svc.GetPlacementGroupRecord(testAccountID, "in-use-group") + require.NoError(t, err) + record.NodeInstances["node1"] = []string{"i-123"} + err = svc.UpdatePlacementGroupRecord(testAccountID, "in-use-group", record, entry.Revision()) + require.NoError(t, err) + + _, err = svc.DeletePlacementGroup(&ec2.DeletePlacementGroupInput{ + GroupName: aws.String("in-use-group"), + }, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInvalidPlacementGroupInUse, err.Error()) +} + +func TestDeletePlacementGroup_MissingName(t *testing.T) { + svc := setupTestService(t) + _, err := svc.DeletePlacementGroup(&ec2.DeletePlacementGroupInput{}, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorMissingParameter, err.Error()) +} + +// --- DescribePlacementGroups Tests --- + +func TestDescribePlacementGroups_All(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "group-a", "spread") + createTestGroup(t, svc, "group-b", "cluster") + + out, err := svc.DescribePlacementGroups(&ec2.DescribePlacementGroupsInput{}, testAccountID) + require.NoError(t, err) + assert.Len(t, out.PlacementGroups, 2) +} + +func TestDescribePlacementGroups_FilterByName(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "alpha", "spread") + createTestGroup(t, svc, "beta", "cluster") + + out, err := svc.DescribePlacementGroups(&ec2.DescribePlacementGroupsInput{ + GroupNames: []*string{aws.String("alpha")}, + }, testAccountID) + require.NoError(t, err) + require.Len(t, out.PlacementGroups, 1) + assert.Equal(t, "alpha", *out.PlacementGroups[0].GroupName) +} + +func TestDescribePlacementGroups_FilterByGroupId(t *testing.T) { + svc := setupTestService(t) + pg := createTestGroup(t, svc, "id-filter", "spread") + + out, err := svc.DescribePlacementGroups(&ec2.DescribePlacementGroupsInput{ + GroupIds: []*string{pg.GroupId}, + }, testAccountID) + require.NoError(t, err) + require.Len(t, out.PlacementGroups, 1) + assert.Equal(t, *pg.GroupId, *out.PlacementGroups[0].GroupId) +} + +func TestDescribePlacementGroups_FilterByStrategy(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "spread1", "spread") + createTestGroup(t, svc, "cluster1", "cluster") + + out, err := svc.DescribePlacementGroups(&ec2.DescribePlacementGroupsInput{ + Filters: []*ec2.Filter{ + {Name: aws.String("strategy"), Values: []*string{aws.String("spread")}}, + }, + }, testAccountID) + require.NoError(t, err) + require.Len(t, out.PlacementGroups, 1) + assert.Equal(t, "spread1", *out.PlacementGroups[0].GroupName) +} + +func TestDescribePlacementGroups_FilterByState(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "avail-group", "spread") + + out, err := svc.DescribePlacementGroups(&ec2.DescribePlacementGroupsInput{ + Filters: []*ec2.Filter{ + {Name: aws.String("state"), Values: []*string{aws.String("available")}}, + }, + }, testAccountID) + require.NoError(t, err) + assert.Len(t, out.PlacementGroups, 1) + + out, err = svc.DescribePlacementGroups(&ec2.DescribePlacementGroupsInput{ + Filters: []*ec2.Filter{ + {Name: aws.String("state"), Values: []*string{aws.String("deleting")}}, + }, + }, testAccountID) + require.NoError(t, err) + assert.Empty(t, out.PlacementGroups) +} + +func TestDescribePlacementGroups_NameNotFound(t *testing.T) { + svc := setupTestService(t) + _, err := svc.DescribePlacementGroups(&ec2.DescribePlacementGroupsInput{ + GroupNames: []*string{aws.String("ghost")}, + }, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInvalidPlacementGroupUnknown, err.Error()) +} + +func TestDescribePlacementGroups_AccountScoped(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "acct-group", "spread") + + // Different account should see nothing + out, err := svc.DescribePlacementGroups(&ec2.DescribePlacementGroupsInput{}, "999999999999") + require.NoError(t, err) + assert.Empty(t, out.PlacementGroups) +} + +// --- NodeInstances CAS Tests --- + +func TestGetAndUpdatePlacementGroupRecord(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "cas-group", "spread") + + // Get with revision + record, entry, err := svc.GetPlacementGroupRecord(testAccountID, "cas-group") + require.NoError(t, err) + assert.Equal(t, "spread", record.Strategy) + assert.Empty(t, record.NodeInstances) + + // Update with CAS + record.NodeInstances["node1"] = []string{"i-abc"} + err = svc.UpdatePlacementGroupRecord(testAccountID, "cas-group", record, entry.Revision()) + require.NoError(t, err) + + // Verify update + record2, _, err := svc.GetPlacementGroupRecord(testAccountID, "cas-group") + require.NoError(t, err) + assert.Equal(t, []string{"i-abc"}, record2.NodeInstances["node1"]) +} + +func TestUpdatePlacementGroupRecord_CASConflict(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "conflict-group", "spread") + + // Get the record twice (same revision) + record1, entry1, err := svc.GetPlacementGroupRecord(testAccountID, "conflict-group") + require.NoError(t, err) + record2, _, err := svc.GetPlacementGroupRecord(testAccountID, "conflict-group") + require.NoError(t, err) + + // First update succeeds + record1.NodeInstances["node1"] = []string{"i-111"} + err = svc.UpdatePlacementGroupRecord(testAccountID, "conflict-group", record1, entry1.Revision()) + require.NoError(t, err) + + // Second update with stale revision fails + record2.NodeInstances["node2"] = []string{"i-222"} + err = svc.UpdatePlacementGroupRecord(testAccountID, "conflict-group", record2, entry1.Revision()) + require.Error(t, err) +} + +func TestGetPlacementGroupRecord_NotFound(t *testing.T) { + svc := setupTestService(t) + _, _, err := svc.GetPlacementGroupRecord(testAccountID, "nonexistent") + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInvalidPlacementGroupUnknown, err.Error()) +} + +// --- ReserveSpreadNodes Tests --- + +func TestReserveSpreadNodes_Success(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "reserve-group", "spread") + + out, err := svc.ReserveSpreadNodes(&ReserveSpreadNodesInput{ + GroupName: "reserve-group", + EligibleNodes: []string{"node-a", "node-b", "node-c"}, + MinCount: 2, + MaxCount: 3, + }, testAccountID) + require.NoError(t, err) + assert.Len(t, out.ReservedNodes, 3) + + // Verify placeholders in record + record, _, err := svc.GetPlacementGroupRecord(testAccountID, "reserve-group") + require.NoError(t, err) + assert.Len(t, record.NodeInstances, 3) + for _, node := range out.ReservedNodes { + _, ok := record.NodeInstances[node] + assert.True(t, ok, "node %s should be in NodeInstances", node) + } +} + +func TestReserveSpreadNodes_ExcludesOccupiedNodes(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "occupied-group", "spread") + + // Pre-occupy node-a + record, entry, err := svc.GetPlacementGroupRecord(testAccountID, "occupied-group") + require.NoError(t, err) + record.NodeInstances["node-a"] = []string{"i-existing"} + require.NoError(t, svc.UpdatePlacementGroupRecord(testAccountID, "occupied-group", record, entry.Revision())) + + out, err := svc.ReserveSpreadNodes(&ReserveSpreadNodesInput{ + GroupName: "occupied-group", + EligibleNodes: []string{"node-a", "node-b", "node-c"}, + MinCount: 1, + MaxCount: 2, + }, testAccountID) + require.NoError(t, err) + assert.Len(t, out.ReservedNodes, 2) + // node-a should NOT be in reserved nodes + for _, n := range out.ReservedNodes { + assert.NotEqual(t, "node-a", n) + } +} + +func TestReserveSpreadNodes_InsufficientNodes(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "insufficient-group", "spread") + + // Pre-occupy node-a and node-b + record, entry, err := svc.GetPlacementGroupRecord(testAccountID, "insufficient-group") + require.NoError(t, err) + record.NodeInstances["node-a"] = []string{"i-1"} + record.NodeInstances["node-b"] = []string{"i-2"} + require.NoError(t, svc.UpdatePlacementGroupRecord(testAccountID, "insufficient-group", record, entry.Revision())) + + _, err = svc.ReserveSpreadNodes(&ReserveSpreadNodesInput{ + GroupName: "insufficient-group", + EligibleNodes: []string{"node-a", "node-b", "node-c"}, + MinCount: 2, + MaxCount: 2, + }, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInsufficientInstanceCapacity, err.Error()) +} + +func TestReserveSpreadNodes_WrongStrategy(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "cluster-group", "cluster") + + _, err := svc.ReserveSpreadNodes(&ReserveSpreadNodesInput{ + GroupName: "cluster-group", + EligibleNodes: []string{"node-a"}, + MinCount: 1, + MaxCount: 1, + }, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInvalidParameterValue, err.Error()) +} + +func TestReserveSpreadNodes_GroupNotFound(t *testing.T) { + svc := setupTestService(t) + _, err := svc.ReserveSpreadNodes(&ReserveSpreadNodesInput{ + GroupName: "ghost-group", + EligibleNodes: []string{"node-a"}, + MinCount: 1, + MaxCount: 1, + }, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInvalidPlacementGroupUnknown, err.Error()) +} + +// --- FinalizeSpreadInstances Tests --- + +func TestFinalizeSpreadInstances_Success(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "finalize-group", "spread") + + // Reserve nodes first + _, err := svc.ReserveSpreadNodes(&ReserveSpreadNodesInput{ + GroupName: "finalize-group", + EligibleNodes: []string{"node-a", "node-b"}, + MinCount: 2, + MaxCount: 2, + }, testAccountID) + require.NoError(t, err) + + // Finalize with instance IDs + _, err = svc.FinalizeSpreadInstances(&FinalizeSpreadInstancesInput{ + GroupName: "finalize-group", + NodeInstances: map[string][]string{ + "node-a": {"i-aaa"}, + "node-b": {"i-bbb"}, + }, + }, testAccountID) + require.NoError(t, err) + + // Verify + record, _, err := svc.GetPlacementGroupRecord(testAccountID, "finalize-group") + require.NoError(t, err) + assert.Equal(t, []string{"i-aaa"}, record.NodeInstances["node-a"]) + assert.Equal(t, []string{"i-bbb"}, record.NodeInstances["node-b"]) +} + +// --- ReleaseSpreadNodes Tests --- + +func TestReleaseSpreadNodes_Success(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "release-group", "spread") + + // Reserve nodes + _, err := svc.ReserveSpreadNodes(&ReserveSpreadNodesInput{ + GroupName: "release-group", + EligibleNodes: []string{"node-a", "node-b"}, + MinCount: 2, + MaxCount: 2, + }, testAccountID) + require.NoError(t, err) + + // Release node-b + _, err = svc.ReleaseSpreadNodes(&ReleaseSpreadNodesInput{ + GroupName: "release-group", + Nodes: []string{"node-b"}, + }, testAccountID) + require.NoError(t, err) + + // Verify only node-a remains + record, _, err := svc.GetPlacementGroupRecord(testAccountID, "release-group") + require.NoError(t, err) + assert.Len(t, record.NodeInstances, 1) + _, ok := record.NodeInstances["node-a"] + assert.True(t, ok) +} + +func TestReleaseSpreadNodes_AllNodes(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "release-all-group", "spread") + + // Reserve nodes + _, err := svc.ReserveSpreadNodes(&ReserveSpreadNodesInput{ + GroupName: "release-all-group", + EligibleNodes: []string{"node-a", "node-b"}, + MinCount: 2, + MaxCount: 2, + }, testAccountID) + require.NoError(t, err) + + // Release all + _, err = svc.ReleaseSpreadNodes(&ReleaseSpreadNodesInput{ + GroupName: "release-all-group", + Nodes: []string{"node-a", "node-b"}, + }, testAccountID) + require.NoError(t, err) + + // Verify empty + record, _, err := svc.GetPlacementGroupRecord(testAccountID, "release-all-group") + require.NoError(t, err) + assert.Empty(t, record.NodeInstances) +} + +// --- RemoveInstance Tests --- + +func TestRemoveInstance_Success(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "remove-inst-group", "spread") + + // Add instances to two nodes + record, entry, err := svc.GetPlacementGroupRecord(testAccountID, "remove-inst-group") + require.NoError(t, err) + record.NodeInstances["node-a"] = []string{"i-aaa"} + record.NodeInstances["node-b"] = []string{"i-bbb"} + require.NoError(t, svc.UpdatePlacementGroupRecord(testAccountID, "remove-inst-group", record, entry.Revision())) + + // Remove i-aaa from node-a + _, err = svc.RemoveInstance(&RemoveInstanceInput{ + GroupName: "remove-inst-group", + NodeName: "node-a", + InstanceID: "i-aaa", + }, testAccountID) + require.NoError(t, err) + + // Verify node-a is removed (was the only instance), node-b remains + record, _, err = svc.GetPlacementGroupRecord(testAccountID, "remove-inst-group") + require.NoError(t, err) + assert.Len(t, record.NodeInstances, 1) + _, hasA := record.NodeInstances["node-a"] + assert.False(t, hasA) + assert.Equal(t, []string{"i-bbb"}, record.NodeInstances["node-b"]) +} + +func TestRemoveInstance_NodeNotFound(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "remove-nonode-group", "spread") + + // Removing from a non-existent node is a no-op + _, err := svc.RemoveInstance(&RemoveInstanceInput{ + GroupName: "remove-nonode-group", + NodeName: "ghost-node", + InstanceID: "i-xxx", + }, testAccountID) + require.NoError(t, err) +} + +func TestRemoveInstance_GroupNotFound(t *testing.T) { + svc := setupTestService(t) + + // Removing from a non-existent group is a no-op (group may have been deleted) + _, err := svc.RemoveInstance(&RemoveInstanceInput{ + GroupName: "deleted-group", + NodeName: "node-a", + InstanceID: "i-xxx", + }, testAccountID) + require.NoError(t, err) +} + +func TestRemoveInstance_MultipleInstancesOnNode(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "multi-inst-group", "cluster") + + // Add multiple instances to same node (cluster scenario) + record, entry, err := svc.GetPlacementGroupRecord(testAccountID, "multi-inst-group") + require.NoError(t, err) + record.NodeInstances["node-a"] = []string{"i-111", "i-222", "i-333"} + require.NoError(t, svc.UpdatePlacementGroupRecord(testAccountID, "multi-inst-group", record, entry.Revision())) + + // Remove i-222 + _, err = svc.RemoveInstance(&RemoveInstanceInput{ + GroupName: "multi-inst-group", + NodeName: "node-a", + InstanceID: "i-222", + }, testAccountID) + require.NoError(t, err) + + // Verify i-111 and i-333 remain + record, _, err = svc.GetPlacementGroupRecord(testAccountID, "multi-inst-group") + require.NoError(t, err) + assert.Equal(t, []string{"i-111", "i-333"}, record.NodeInstances["node-a"]) +} + +// --- ReserveClusterNode Tests --- + +func TestReserveClusterNode_EmptyGroup(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "cluster-empty", "cluster") + + out, err := svc.ReserveClusterNode(&ReserveClusterNodeInput{ + GroupName: "cluster-empty", + EligibleNodes: []string{"node-a", "node-b", "node-c"}, + }, testAccountID) + require.NoError(t, err) + // Should pick first eligible node (highest capacity, sorted by caller) + assert.Equal(t, "node-a", out.TargetNode) + + // Verify placeholder written + record, _, err := svc.GetPlacementGroupRecord(testAccountID, "cluster-empty") + require.NoError(t, err) + assert.Len(t, record.NodeInstances, 1) + _, ok := record.NodeInstances["node-a"] + assert.True(t, ok) +} + +func TestReserveClusterNode_ExistingNode(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "cluster-existing", "cluster") + + // Pre-populate with instances on node-b + record, entry, err := svc.GetPlacementGroupRecord(testAccountID, "cluster-existing") + require.NoError(t, err) + record.NodeInstances["node-b"] = []string{"i-existing"} + require.NoError(t, svc.UpdatePlacementGroupRecord(testAccountID, "cluster-existing", record, entry.Revision())) + + // Even though node-a is in eligible list, should return existing node-b + out, err := svc.ReserveClusterNode(&ReserveClusterNodeInput{ + GroupName: "cluster-existing", + EligibleNodes: []string{"node-a", "node-c"}, + }, testAccountID) + require.NoError(t, err) + assert.Equal(t, "node-b", out.TargetNode) +} + +func TestReserveClusterNode_NoEligibleNodes(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "cluster-noelig", "cluster") + + _, err := svc.ReserveClusterNode(&ReserveClusterNodeInput{ + GroupName: "cluster-noelig", + EligibleNodes: []string{}, + }, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInsufficientInstanceCapacity, err.Error()) +} + +func TestReserveClusterNode_WrongStrategy(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "spread-for-cluster", "spread") + + _, err := svc.ReserveClusterNode(&ReserveClusterNodeInput{ + GroupName: "spread-for-cluster", + EligibleNodes: []string{"node-a"}, + }, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInvalidParameterValue, err.Error()) +} + +func TestReserveClusterNode_GroupNotFound(t *testing.T) { + svc := setupTestService(t) + _, err := svc.ReserveClusterNode(&ReserveClusterNodeInput{ + GroupName: "ghost-cluster", + EligibleNodes: []string{"node-a"}, + }, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInvalidPlacementGroupUnknown, err.Error()) +} + +func TestReserveClusterNode_MissingGroupName(t *testing.T) { + svc := setupTestService(t) + _, err := svc.ReserveClusterNode(&ReserveClusterNodeInput{ + EligibleNodes: []string{"node-a"}, + }, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorMissingParameter, err.Error()) +} + +// --- FinalizeClusterInstances Tests --- + +func TestFinalizeClusterInstances_Success(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "cluster-finalize", "cluster") + + // Reserve a node first + _, err := svc.ReserveClusterNode(&ReserveClusterNodeInput{ + GroupName: "cluster-finalize", + EligibleNodes: []string{"node-a"}, + }, testAccountID) + require.NoError(t, err) + + // Finalize with instance IDs + _, err = svc.FinalizeClusterInstances(&FinalizeClusterInstancesInput{ + GroupName: "cluster-finalize", + NodeInstances: map[string][]string{ + "node-a": {"i-c1", "i-c2", "i-c3"}, + }, + }, testAccountID) + require.NoError(t, err) + + // Verify instances recorded + record, _, err := svc.GetPlacementGroupRecord(testAccountID, "cluster-finalize") + require.NoError(t, err) + assert.Equal(t, []string{"i-c1", "i-c2", "i-c3"}, record.NodeInstances["node-a"]) +} + +func TestFinalizeClusterInstances_AppendsToExisting(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "cluster-append", "cluster") + + // Pre-populate with existing instances + record, entry, err := svc.GetPlacementGroupRecord(testAccountID, "cluster-append") + require.NoError(t, err) + record.NodeInstances["node-a"] = []string{"i-existing1", "i-existing2"} + require.NoError(t, svc.UpdatePlacementGroupRecord(testAccountID, "cluster-append", record, entry.Revision())) + + // Finalize with new instances (should append, not overwrite) + _, err = svc.FinalizeClusterInstances(&FinalizeClusterInstancesInput{ + GroupName: "cluster-append", + NodeInstances: map[string][]string{ + "node-a": {"i-new1", "i-new2"}, + }, + }, testAccountID) + require.NoError(t, err) + + // Verify all instances present + record, _, err = svc.GetPlacementGroupRecord(testAccountID, "cluster-append") + require.NoError(t, err) + assert.Equal(t, []string{"i-existing1", "i-existing2", "i-new1", "i-new2"}, record.NodeInstances["node-a"]) +} + +func TestFinalizeClusterInstances_MissingGroupName(t *testing.T) { + svc := setupTestService(t) + _, err := svc.FinalizeClusterInstances(&FinalizeClusterInstancesInput{ + NodeInstances: map[string][]string{"node-a": {"i-1"}}, + }, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorMissingParameter, err.Error()) +} + +// --- Cluster Lifecycle Test --- + +func TestClusterLifecycle_ReserveFinalizeRemoveDelete(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "cluster-lifecycle", "cluster") + + // Reserve node + reserveOut, err := svc.ReserveClusterNode(&ReserveClusterNodeInput{ + GroupName: "cluster-lifecycle", + EligibleNodes: []string{"node-1"}, + }, testAccountID) + require.NoError(t, err) + assert.Equal(t, "node-1", reserveOut.TargetNode) + + // Finalize with instances + _, err = svc.FinalizeClusterInstances(&FinalizeClusterInstancesInput{ + GroupName: "cluster-lifecycle", + NodeInstances: map[string][]string{ + "node-1": {"i-c1", "i-c2"}, + }, + }, testAccountID) + require.NoError(t, err) + + // Can't delete — instances present + _, err = svc.DeletePlacementGroup(&ec2.DeletePlacementGroupInput{ + GroupName: aws.String("cluster-lifecycle"), + }, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInvalidPlacementGroupInUse, err.Error()) + + // Remove instances via RemoveInstance (simulating terminate) + _, err = svc.RemoveInstance(&RemoveInstanceInput{ + GroupName: "cluster-lifecycle", NodeName: "node-1", InstanceID: "i-c1", + }, testAccountID) + require.NoError(t, err) + _, err = svc.RemoveInstance(&RemoveInstanceInput{ + GroupName: "cluster-lifecycle", NodeName: "node-1", InstanceID: "i-c2", + }, testAccountID) + require.NoError(t, err) + + // Now delete succeeds + _, err = svc.DeletePlacementGroup(&ec2.DeletePlacementGroupInput{ + GroupName: aws.String("cluster-lifecycle"), + }, testAccountID) + require.NoError(t, err) +} + +// --- End-to-End Spread Lifecycle Test --- + +func TestSpreadLifecycle_ReserveFinalizeDelete(t *testing.T) { + svc := setupTestService(t) + createTestGroup(t, svc, "lifecycle-group", "spread") + + // Reserve 2 nodes + reserveOut, err := svc.ReserveSpreadNodes(&ReserveSpreadNodesInput{ + GroupName: "lifecycle-group", + EligibleNodes: []string{"node-1", "node-2", "node-3"}, + MinCount: 2, + MaxCount: 2, + }, testAccountID) + require.NoError(t, err) + assert.Len(t, reserveOut.ReservedNodes, 2) + + // Finalize with instance IDs + nodeInstances := make(map[string][]string) + for _, n := range reserveOut.ReservedNodes { + nodeInstances[n] = []string{"i-" + n} + } + _, err = svc.FinalizeSpreadInstances(&FinalizeSpreadInstancesInput{ + GroupName: "lifecycle-group", + NodeInstances: nodeInstances, + }, testAccountID) + require.NoError(t, err) + + // Can't delete — instances present + _, err = svc.DeletePlacementGroup(&ec2.DeletePlacementGroupInput{ + GroupName: aws.String("lifecycle-group"), + }, testAccountID) + require.Error(t, err) + assert.Equal(t, awserrors.ErrorInvalidPlacementGroupInUse, err.Error()) + + // Remove instances via RemoveInstance (simulating terminate) + for _, n := range reserveOut.ReservedNodes { + _, err = svc.RemoveInstance(&RemoveInstanceInput{ + GroupName: "lifecycle-group", + NodeName: n, + InstanceID: "i-" + n, + }, testAccountID) + require.NoError(t, err) + } + + // Now delete succeeds + _, err = svc.DeletePlacementGroup(&ec2.DeletePlacementGroupInput{ + GroupName: aws.String("lifecycle-group"), + }, testAccountID) + require.NoError(t, err) +} diff --git a/spinifex/handlers/ec2/placementgroup/service_nats.go b/spinifex/handlers/ec2/placementgroup/service_nats.go new file mode 100644 index 000000000..ef23185c5 --- /dev/null +++ b/spinifex/handlers/ec2/placementgroup/service_nats.go @@ -0,0 +1,55 @@ +package handlers_ec2_placementgroup + +import ( + "time" + + "github.com/aws/aws-sdk-go/service/ec2" + "github.com/mulgadc/spinifex/spinifex/utils" + "github.com/nats-io/nats.go" +) + +// NATSPlacementGroupService handles placement group operations via NATS messaging. +type NATSPlacementGroupService struct { + natsConn *nats.Conn +} + +// NewNATSPlacementGroupService creates a new NATS-based placement group service. +func NewNATSPlacementGroupService(conn *nats.Conn) PlacementGroupService { + return &NATSPlacementGroupService{natsConn: conn} +} + +func (s *NATSPlacementGroupService) CreatePlacementGroup(input *ec2.CreatePlacementGroupInput, accountID string) (*ec2.CreatePlacementGroupOutput, error) { + return utils.NATSRequest[ec2.CreatePlacementGroupOutput](s.natsConn, "ec2.CreatePlacementGroup", input, 30*time.Second, accountID) +} + +func (s *NATSPlacementGroupService) DeletePlacementGroup(input *ec2.DeletePlacementGroupInput, accountID string) (*ec2.DeletePlacementGroupOutput, error) { + return utils.NATSRequest[ec2.DeletePlacementGroupOutput](s.natsConn, "ec2.DeletePlacementGroup", input, 30*time.Second, accountID) +} + +func (s *NATSPlacementGroupService) DescribePlacementGroups(input *ec2.DescribePlacementGroupsInput, accountID string) (*ec2.DescribePlacementGroupsOutput, error) { + return utils.NATSRequest[ec2.DescribePlacementGroupsOutput](s.natsConn, "ec2.DescribePlacementGroups", input, 30*time.Second, accountID) +} + +func (s *NATSPlacementGroupService) ReserveSpreadNodes(input *ReserveSpreadNodesInput, accountID string) (*ReserveSpreadNodesOutput, error) { + return utils.NATSRequest[ReserveSpreadNodesOutput](s.natsConn, "ec2.ReserveSpreadNodes", input, 30*time.Second, accountID) +} + +func (s *NATSPlacementGroupService) FinalizeSpreadInstances(input *FinalizeSpreadInstancesInput, accountID string) (*FinalizeSpreadInstancesOutput, error) { + return utils.NATSRequest[FinalizeSpreadInstancesOutput](s.natsConn, "ec2.FinalizeSpreadInstances", input, 30*time.Second, accountID) +} + +func (s *NATSPlacementGroupService) ReleaseSpreadNodes(input *ReleaseSpreadNodesInput, accountID string) (*ReleaseSpreadNodesOutput, error) { + return utils.NATSRequest[ReleaseSpreadNodesOutput](s.natsConn, "ec2.ReleaseSpreadNodes", input, 30*time.Second, accountID) +} + +func (s *NATSPlacementGroupService) RemoveInstance(input *RemoveInstanceInput, accountID string) (*RemoveInstanceOutput, error) { + return utils.NATSRequest[RemoveInstanceOutput](s.natsConn, "ec2.RemoveInstanceFromPlacementGroup", input, 30*time.Second, accountID) +} + +func (s *NATSPlacementGroupService) ReserveClusterNode(input *ReserveClusterNodeInput, accountID string) (*ReserveClusterNodeOutput, error) { + return utils.NATSRequest[ReserveClusterNodeOutput](s.natsConn, "ec2.ReserveClusterNode", input, 30*time.Second, accountID) +} + +func (s *NATSPlacementGroupService) FinalizeClusterInstances(input *FinalizeClusterInstancesInput, accountID string) (*FinalizeClusterInstancesOutput, error) { + return utils.NATSRequest[FinalizeClusterInstancesOutput](s.natsConn, "ec2.FinalizeClusterInstances", input, 30*time.Second, accountID) +} diff --git a/spinifex/instancetypes/generate.go b/spinifex/instancetypes/generate.go index 8d5ff39c5..a1a0d626e 100644 --- a/spinifex/instancetypes/generate.go +++ b/spinifex/instancetypes/generate.go @@ -44,6 +44,12 @@ func generateForGeneration(gen cpuGeneration, arch string) map[string]*ec2.Insta Hypervisor: aws.String("kvm"), SupportedVirtualizationTypes: []*string{aws.String("hvm")}, SupportedRootDeviceTypes: []*string{aws.String("ebs")}, + PlacementGroupInfo: &ec2.PlacementGroupInfo{ + SupportedStrategies: []*string{ + aws.String("cluster"), + aws.String("spread"), + }, + }, } } } diff --git a/spinifex/instancetypes/generate_test.go b/spinifex/instancetypes/generate_test.go index 9e9e22396..a7b2a21d9 100644 --- a/spinifex/instancetypes/generate_test.go +++ b/spinifex/instancetypes/generate_test.go @@ -313,3 +313,22 @@ func TestGenerateInstanceTypes_BurstableFlag(t *testing.T) { "%s should be current generation", name) } } + +func TestGenerateInstanceTypes_PlacementGroupInfo(t *testing.T) { + types := generateForGeneration(genIntelIceLake, "x86_64") + require.NotEmpty(t, types) + + for name, info := range types { + require.NotNil(t, info.PlacementGroupInfo, + "%s should have PlacementGroupInfo", name) + assert.Len(t, info.PlacementGroupInfo.SupportedStrategies, 2, + "%s should support 2 strategies", name) + + strategies := make(map[string]bool) + for _, s := range info.PlacementGroupInfo.SupportedStrategies { + strategies[*s] = true + } + assert.True(t, strategies["cluster"], "%s should support cluster", name) + assert.True(t, strategies["spread"], "%s should support spread", name) + } +} diff --git a/spinifex/utils/pidfile.go b/spinifex/utils/pidfile.go index e2471601a..68ce1a7a2 100644 --- a/spinifex/utils/pidfile.go +++ b/spinifex/utils/pidfile.go @@ -72,6 +72,10 @@ func WritePidFileTo(dir string, name string, pid int) error { return WritePidFile(name, pid) } + if err := os.MkdirAll(dir, 0750); err != nil { + return fmt.Errorf("create pid directory %s: %w", dir, err) + } + pidFilename := filepath.Join(dir, fmt.Sprintf("%s.pid", name)) pidFile, err := os.Create(pidFilename) diff --git a/spinifex/vm/vm.go b/spinifex/vm/vm.go index ae1aa5a21..ba5a9a61b 100644 --- a/spinifex/vm/vm.go +++ b/spinifex/vm/vm.go @@ -75,6 +75,10 @@ type VM struct { // DevMAC is the MAC for the dev/hostfwd NIC (DEV_NETWORKING mode). // Set before cloud-init ISO generation so netplan can suppress its default route. DevMAC string `json:"dev_mac,omitempty"` + + // Placement group tracking (set during RunInstances when a placement group is specified) + PlacementGroupName string `json:"placement_group_name,omitempty"` + PlacementGroupNode string `json:"placement_group_node,omitempty"` } // ResetNodeLocalState zeroes out fields that are specific to the daemon node diff --git a/tests/e2e/lib/multinode-helpers.sh b/tests/e2e/lib/multinode-helpers.sh index 8394079e2..046900f9e 100755 --- a/tests/e2e/lib/multinode-helpers.sh +++ b/tests/e2e/lib/multinode-helpers.sh @@ -267,10 +267,15 @@ wait_for_daemon_ready() { # Check instance distribution across nodes # Counts QEMU processes per node IP from hostfwd bindings +# Usage: check_instance_distribution [expected_nodes] +# If expected_nodes is provided, fails when instances don't span that many nodes check_instance_distribution() { + local expected_nodes="${1:-0}" + echo "Checking instance distribution across nodes..." local total=0 + local nodes_used=0 for i in 1 2 3; do local node_ip="${SIMULATED_NETWORK}.$i" # Count QEMU processes with hostfwd bound to this node's IP @@ -278,9 +283,17 @@ check_instance_distribution() { count=$(ps auxw | grep qemu-system | grep -v grep | grep -c "hostfwd=tcp:${node_ip}:" 2>/dev/null) || count=0 echo " Node$i ($node_ip): $count instances" total=$((total + count)) + if [ "$count" -gt 0 ]; then + nodes_used=$((nodes_used + 1)) + fi done - echo " Total instances: $total" + echo " Total instances: $total (across $nodes_used nodes)" + + if [ "$expected_nodes" -gt 0 ] && [ "$nodes_used" -lt "$expected_nodes" ]; then + echo " FAIL: Expected instances on $expected_nodes nodes, only used $nodes_used" + return 1 + fi } # Find the SSH port for an instance from the QEMU process @@ -807,3 +820,46 @@ verify_predastore_cluster() { return 1 fi } + +# Count instances per node for given instance IDs +# Usage: count_instances_per_node INSTANCE_ID [INSTANCE_ID...] +# Sets NODE1_COUNT, NODE2_COUNT, NODE3_COUNT +count_instances_per_node() { + local ids=("$@") + NODE1_COUNT=0; NODE2_COUNT=0; NODE3_COUNT=0 + for i in 1 2 3; do + local node_ip="${SIMULATED_NETWORK}.$i" + local count=0 + for id in "${ids[@]}"; do + if ps auxw | grep "$id" | grep qemu-system | grep -v grep | grep -q "hostfwd=tcp:${node_ip}:"; then + count=$((count + 1)) + fi + done + eval "NODE${i}_COUNT=$count" + echo " Node$i ($node_ip): $count instances" + done +} + +# Verify Placement field in DescribeInstances response +# Usage: verify_placement GROUP_NAME INSTANCE_ID [INSTANCE_ID...] +verify_placement() { + local group_name="$1"; shift + local ids=("$@") + for id in "${ids[@]}"; do + local placement_group + placement_group=$($AWS_EC2 describe-instances --instance-ids "$id" \ + --query 'Reservations[0].Instances[0].Placement.GroupName' --output text) + if [ "$placement_group" != "$group_name" ]; then + echo " FAIL: Instance $id Placement.GroupName='$placement_group', expected '$group_name'" + return 1 + fi + local az + az=$($AWS_EC2 describe-instances --instance-ids "$id" \ + --query 'Reservations[0].Instances[0].Placement.AvailabilityZone' --output text) + if [ -z "$az" ] || [ "$az" = "None" ]; then + echo " FAIL: Instance $id missing Placement.AvailabilityZone" + return 1 + fi + done + echo " All instances show Placement.GroupName=$group_name" +} diff --git a/tests/e2e/run-pseudo-multinode-e2e.sh b/tests/e2e/run-pseudo-multinode-e2e.sh index a72f79e94..98722b3e3 100755 --- a/tests/e2e/run-pseudo-multinode-e2e.sh +++ b/tests/e2e/run-pseudo-multinode-e2e.sh @@ -319,9 +319,13 @@ for instance_id in "${INSTANCE_IDS[@]}"; do } done -# Check distribution +# Check distribution — enforce all 3 nodes received an instance echo "" -check_instance_distribution +check_instance_distribution 3 || { + echo "ERROR: Default spread did not distribute instances across all 3 nodes" + exit 1 +} +echo " Default spread distribution verified" # Verify spx get vms shows all running instances echo "" @@ -476,6 +480,692 @@ verify_nats_cluster 3 || { echo "WARNING: NATS cluster verification failed after operations" } +# Phase 5b: Multi-Node Batch Distribution Tests +# Tests batch launches with --count N (placement groups phase 1 routing) +echo "" +echo "Phase 5b: Multi-Node Batch Distribution Tests" +echo "========================================" + +# Clean up instances from Phase 5 before batch tests +echo "Cleaning up Phase 5 instances..." +if ! terminate_and_wait "${INSTANCE_IDS[@]}"; then + echo "WARNING: Some Phase 5 instances failed to terminate cleanly" +fi +unset INSTANCE_IDS + +# Wait for clean state +sleep 2 + +# Test 5b-1: Batch launch spreads across nodes +echo "" +echo "Test 5b-1: Batch Launch Spread (3 instances on 3 nodes)" +echo "----------------------------------------" +echo "Launching 3 instances in a single API call..." + +RUN_OUTPUT=$($AWS_EC2 run-instances \ + --image-id "$AMI_ID" \ + --instance-type "$INSTANCE_TYPE" \ + --key-name multinode-test-key \ + --count 3) + +BATCH_COUNT=$(echo "$RUN_OUTPUT" | jq '.Instances | length') +if [ "$BATCH_COUNT" -ne 3 ]; then + echo "FAIL: Expected 3 instances, got $BATCH_COUNT" + echo "Output: $RUN_OUTPUT" + exit 1 +fi + +# Verify single reservation in response +RESERVATION_ID=$(echo "$RUN_OUTPUT" | jq -r '.ReservationId // empty') +if [ -z "$RESERVATION_ID" ]; then + echo "FAIL: No ReservationId in response" + exit 1 +fi +echo " ReservationId: $RESERVATION_ID" +echo " Instances: $BATCH_COUNT" + +BATCH_IDS=($(echo "$RUN_OUTPUT" | jq -r '.Instances[].InstanceId')) +echo " Launched: ${BATCH_IDS[*]}" + +# Wait for all to reach running +for id in "${BATCH_IDS[@]}"; do + wait_for_instance_state "$id" "running" 60 || { + echo "ERROR: Instance $id failed to reach running state" + exit 1 + } +done + +# Verify distribution: count QEMU processes per node IP +echo " Checking distribution..." +NODES_USED=0 +for i in 1 2 3; do + NODE_IP="${SIMULATED_NETWORK}.$i" + NODE_COUNT=0 + for id in "${BATCH_IDS[@]}"; do + if ps auxw | grep "$id" | grep qemu-system | grep -v grep | grep -q "hostfwd=tcp:${NODE_IP}:"; then + NODE_COUNT=$((NODE_COUNT + 1)) + fi + done + echo " Node$i ($NODE_IP): $NODE_COUNT instances" + if [ "$NODE_COUNT" -gt 0 ]; then + NODES_USED=$((NODES_USED + 1)) + fi +done + +if [ "$NODES_USED" -lt 3 ]; then + echo "FAIL: Expected instances on 3 nodes, only used $NODES_USED" + exit 1 +fi +echo " PASS: Instances distributed across $NODES_USED nodes" + +# Verify DescribeInstances from each node sees all batch-launched instances +echo " Verifying DescribeInstances aggregation from all gateways..." +for i in 1 2 3; do + NODE_GW_IP="${SIMULATED_NETWORK}.$i" + DESCRIBE=$(aws --endpoint-url "https://${NODE_GW_IP}:${AWSGW_PORT}" \ + ec2 describe-instances \ + --query 'Reservations[*].Instances[?State.Name!=`terminated`].InstanceId' \ + --output text 2>/dev/null) || { + echo " WARNING: Could not query gateway on node$i" + continue + } + + MISSING=0 + for id in "${BATCH_IDS[@]}"; do + if ! echo "$DESCRIBE" | grep -q "$id"; then + echo " Node$i missing instance $id" + MISSING=$((MISSING + 1)) + fi + done + + if [ "$MISSING" -gt 0 ]; then + echo "FAIL: Node$i gateway missing $MISSING instances" + exit 1 + fi + echo " Node$i gateway: all ${#BATCH_IDS[@]} instances visible" +done +echo " PASS: All gateways see all batch-launched instances" + +# Cleanup +echo " Cleaning up..." +if ! terminate_and_wait "${BATCH_IDS[@]}"; then + echo "WARNING: Some batch instances failed to terminate" +fi +sleep 2 + +# Test 5b-2: Batch launch with overflow packing (5 on 3 nodes) +echo "" +echo "Test 5b-2: Batch Launch Overflow (5 instances on 3 nodes)" +echo "----------------------------------------" +echo "Launching 5 instances in a single API call..." + +RUN_OUTPUT=$($AWS_EC2 run-instances \ + --image-id "$AMI_ID" \ + --instance-type "$INSTANCE_TYPE" \ + --key-name multinode-test-key \ + --count 5) + +BATCH_COUNT=$(echo "$RUN_OUTPUT" | jq '.Instances | length') +if [ "$BATCH_COUNT" -ne 5 ]; then + echo "FAIL: Expected 5 instances, got $BATCH_COUNT" + exit 1 +fi + +BATCH_IDS=($(echo "$RUN_OUTPUT" | jq -r '.Instances[].InstanceId')) +echo " Launched: ${BATCH_IDS[*]}" + +for id in "${BATCH_IDS[@]}"; do + wait_for_instance_state "$id" "running" 60 || { + echo "ERROR: Instance $id failed to reach running state" + exit 1 + } +done + +# All 3 nodes should have at least 1 instance (spread), some get 2 (pack) +echo " Checking distribution..." +ALL_NODES_USED=true +for i in 1 2 3; do + NODE_IP="${SIMULATED_NETWORK}.$i" + NODE_COUNT=0 + for id in "${BATCH_IDS[@]}"; do + if ps auxw | grep "$id" | grep qemu-system | grep -v grep | grep -q "hostfwd=tcp:${NODE_IP}:"; then + NODE_COUNT=$((NODE_COUNT + 1)) + fi + done + echo " Node$i ($NODE_IP): $NODE_COUNT instances" + if [ "$NODE_COUNT" -eq 0 ]; then + ALL_NODES_USED=false + fi +done + +if [ "$ALL_NODES_USED" = false ]; then + echo "FAIL: Not all nodes received instances" + exit 1 +fi +echo " PASS: All nodes received at least 1 instance, 5 total distributed" + +# Cleanup +echo " Cleaning up..." +if ! terminate_and_wait "${BATCH_IDS[@]}"; then + echo "WARNING: Some batch instances failed to terminate" +fi +sleep 2 + +# Test 5b-3: Insufficient capacity error +echo "" +echo "Test 5b-3: Insufficient Capacity Error" +echo "----------------------------------------" +echo "Requesting more instances than cluster can provide..." + +# Request an absurdly high count that definitely exceeds capacity +if RUN_OUTPUT=$($AWS_EC2 run-instances \ + --image-id "$AMI_ID" \ + --instance-type "$INSTANCE_TYPE" \ + --key-name multinode-test-key \ + --count 100 2>&1); then + echo "FAIL: Expected error but got success" + echo "Output: $RUN_OUTPUT" + exit 1 +fi + +if echo "$RUN_OUTPUT" | grep -q "InsufficientInstanceCapacity"; then + echo " PASS: Got InsufficientInstanceCapacity as expected" +else + echo "FAIL: Wrong error: $RUN_OUTPUT" + exit 1 +fi + +# Verify no orphaned instances were left behind +ORPHAN_CHECK=$($AWS_EC2 describe-instances \ + --query 'Reservations[*].Instances[?State.Name!=`terminated`].InstanceId' --output text) +if [ -n "$ORPHAN_CHECK" ] && [ "$ORPHAN_CHECK" != "None" ]; then + echo " WARNING: Found non-terminated instances after capacity error: $ORPHAN_CHECK" +else + echo " PASS: No orphaned instances after capacity error" +fi + +# Test 5b-4: MinCount/MaxCount capping +echo "" +echo "Test 5b-4: MinCount/MaxCount Capping" +echo "----------------------------------------" +echo "Launching with --count 2..." + +RUN_OUTPUT=$($AWS_EC2 run-instances \ + --image-id "$AMI_ID" \ + --instance-type "$INSTANCE_TYPE" \ + --key-name multinode-test-key \ + --count 2) + +BATCH_COUNT=$(echo "$RUN_OUTPUT" | jq '.Instances | length') +if [ "$BATCH_COUNT" -ne 2 ]; then + echo "FAIL: Expected exactly 2 instances, got $BATCH_COUNT" + exit 1 +fi +echo " PASS: Launched exactly 2 instances (capped to MaxCount)" + +# Cleanup +BATCH_IDS=($(echo "$RUN_OUTPUT" | jq -r '.Instances[].InstanceId')) +for id in "${BATCH_IDS[@]}"; do + wait_for_instance_state "$id" "running" 60 || true +done +if ! terminate_and_wait "${BATCH_IDS[@]}"; then + echo "WARNING: Some batch instances failed to terminate" +fi +sleep 2 + +# Test 5b-5: MinCount failure threshold +echo "" +echo "Test 5b-5: MinCount Failure Threshold" +echo "----------------------------------------" +echo "Requesting --count 100 (exceeds cluster capacity)..." + +if RUN_OUTPUT=$($AWS_EC2 run-instances \ + --image-id "$AMI_ID" \ + --instance-type "$INSTANCE_TYPE" \ + --key-name multinode-test-key \ + --count 100 2>&1); then + echo "FAIL: Expected InsufficientInstanceCapacity" + exit 1 +fi + +if echo "$RUN_OUTPUT" | grep -q "InsufficientInstanceCapacity"; then + echo " PASS: Got InsufficientInstanceCapacity for MinCount > capacity" +else + echo "FAIL: Wrong error: $RUN_OUTPUT" + exit 1 +fi + +echo "" +echo "Phase 5b: All batch distribution tests passed" + +# Phase 5c: Placement Group E2E Tests +echo "" +echo "Phase 5c: Placement Group E2E Tests (CRUD + Spread + Cluster)" +echo "==============================================================" + +# Test 5c-1: Create spread group +echo "" +echo "Test 5c-1: CreatePlacementGroup (spread)" +echo "----------------------------------------" +CREATE_OUTPUT=$($AWS_EC2 create-placement-group \ + --group-name test-spread --strategy spread) +SPREAD_GROUP_ID=$(echo "$CREATE_OUTPUT" | jq -r '.PlacementGroup.GroupId') +SPREAD_STATE=$(echo "$CREATE_OUTPUT" | jq -r '.PlacementGroup.State') +if [ -z "$SPREAD_GROUP_ID" ] || [ "$SPREAD_GROUP_ID" = "null" ]; then + echo "FAIL: No GroupId returned"; exit 1 +fi +if [ "$SPREAD_STATE" != "available" ]; then + echo "FAIL: Expected state=available, got $SPREAD_STATE"; exit 1 +fi +echo " PASS: Created spread group $SPREAD_GROUP_ID (state=$SPREAD_STATE)" + +# Test 5c-2: Create cluster group +echo "" +echo "Test 5c-2: CreatePlacementGroup (cluster)" +echo "----------------------------------------" +CREATE_OUTPUT=$($AWS_EC2 create-placement-group \ + --group-name test-cluster --strategy cluster) +CLUSTER_GROUP_ID=$(echo "$CREATE_OUTPUT" | jq -r '.PlacementGroup.GroupId') +CLUSTER_STATE=$(echo "$CREATE_OUTPUT" | jq -r '.PlacementGroup.State') +if [ -z "$CLUSTER_GROUP_ID" ] || [ "$CLUSTER_GROUP_ID" = "null" ]; then + echo "FAIL: No GroupId returned"; exit 1 +fi +if [ "$CLUSTER_STATE" != "available" ]; then + echo "FAIL: Expected state=available, got $CLUSTER_STATE"; exit 1 +fi +echo " PASS: Created cluster group $CLUSTER_GROUP_ID (state=$CLUSTER_STATE)" + +# Test 5c-3: Partition strategy rejected +echo "" +echo "Test 5c-3: CreatePlacementGroup (partition rejected)" +echo "----------------------------------------" +if PARTITION_OUTPUT=$($AWS_EC2 create-placement-group \ + --group-name test-partition --strategy partition 2>&1); then + echo "FAIL: Expected error for partition strategy"; exit 1 +fi +if echo "$PARTITION_OUTPUT" | grep -q "InvalidParameterValue"; then + echo " PASS: partition strategy correctly rejected" +else + echo "FAIL: Wrong error: $PARTITION_OUTPUT"; exit 1 +fi + +# Test 5c-4: Duplicate name rejected +echo "" +echo "Test 5c-4: CreatePlacementGroup (duplicate name)" +echo "----------------------------------------" +if DUP_OUTPUT=$($AWS_EC2 create-placement-group \ + --group-name test-spread --strategy spread 2>&1); then + echo "FAIL: Expected error for duplicate name"; exit 1 +fi +if echo "$DUP_OUTPUT" | grep -q "InvalidPlacementGroup.Duplicate"; then + echo " PASS: Duplicate name correctly rejected" +else + echo "FAIL: Wrong error: $DUP_OUTPUT"; exit 1 +fi + +# Test 5c-5: DescribePlacementGroups +echo "" +echo "Test 5c-5: DescribePlacementGroups" +echo "----------------------------------------" +DESCRIBE_OUTPUT=$($AWS_EC2 describe-placement-groups) +PG_COUNT=$(echo "$DESCRIBE_OUTPUT" | jq '.PlacementGroups | length') +if [ "$PG_COUNT" -lt 2 ]; then + echo "FAIL: Expected at least 2 placement groups, got $PG_COUNT"; exit 1 +fi + +SPREAD_STRATEGY=$(echo "$DESCRIBE_OUTPUT" | jq -r '.PlacementGroups[] | select(.GroupName=="test-spread") | .Strategy') +SPREAD_LEVEL=$(echo "$DESCRIBE_OUTPUT" | jq -r '.PlacementGroups[] | select(.GroupName=="test-spread") | .SpreadLevel') +if [ "$SPREAD_STRATEGY" != "spread" ]; then + echo "FAIL: Spread group strategy=$SPREAD_STRATEGY"; exit 1 +fi +if [ "$SPREAD_LEVEL" != "host" ]; then + echo "FAIL: Spread group SpreadLevel=$SPREAD_LEVEL, expected host"; exit 1 +fi + +CLUSTER_STRATEGY=$(echo "$DESCRIBE_OUTPUT" | jq -r '.PlacementGroups[] | select(.GroupName=="test-cluster") | .Strategy') +if [ "$CLUSTER_STRATEGY" != "cluster" ]; then + echo "FAIL: Cluster group strategy=$CLUSTER_STRATEGY"; exit 1 +fi +echo " PASS: Both groups described with correct fields" + +FILTER_OUTPUT=$($AWS_EC2 describe-placement-groups --group-names test-spread) +FILTER_COUNT=$(echo "$FILTER_OUTPUT" | jq '.PlacementGroups | length') +if [ "$FILTER_COUNT" -ne 1 ]; then + echo "FAIL: Name filter returned $FILTER_COUNT groups, expected 1"; exit 1 +fi +echo " PASS: Name filter works correctly" + +# Test 5c-6: DescribeInstanceTypes PlacementGroupInfo +echo "" +echo "Test 5c-6: DescribeInstanceTypes PlacementGroupInfo" +echo "----------------------------------------" +DIT_OUTPUT=$($AWS_EC2 describe-instance-types --instance-types "$INSTANCE_TYPE") +STRATEGIES=$(echo "$DIT_OUTPUT" | jq -r '.InstanceTypes[0].PlacementGroupInfo.SupportedStrategies[]' 2>/dev/null | sort | tr '\n' ',') +if ! echo "$STRATEGIES" | grep -q "cluster" || ! echo "$STRATEGIES" | grep -q "spread"; then + echo "FAIL: PlacementGroupInfo.SupportedStrategies missing spread/cluster: $STRATEGIES"; exit 1 +fi +echo " PASS: $INSTANCE_TYPE supports strategies: $STRATEGIES" + +# Test 5c-7: Spread — 3 instances on 3 nodes (strict 1-per-node) +echo "" +echo "Test 5c-7: Spread — Launch 3 Instances (1 per node)" +echo "----------------------------------------" +RUN_OUTPUT=$($AWS_EC2 run-instances \ + --image-id "$AMI_ID" \ + --instance-type "$INSTANCE_TYPE" \ + --key-name multinode-test-key \ + --count 3 \ + --placement GroupName=test-spread) + +SPREAD_COUNT=$(echo "$RUN_OUTPUT" | jq '.Instances | length') +if [ "$SPREAD_COUNT" -ne 3 ]; then + echo "FAIL: Expected 3 instances, got $SPREAD_COUNT"; exit 1 +fi + +SPREAD_IDS=($(echo "$RUN_OUTPUT" | jq -r '.Instances[].InstanceId')) +echo " Launched: ${SPREAD_IDS[*]}" + +for id in "${SPREAD_IDS[@]}"; do + wait_for_instance_state "$id" "running" 60 || { + echo "ERROR: Instance $id failed to reach running state"; exit 1 + } +done + +echo " Checking strict 1-per-node distribution..." +count_instances_per_node "${SPREAD_IDS[@]}" +if [ "$NODE1_COUNT" -ne 1 ] || [ "$NODE2_COUNT" -ne 1 ] || [ "$NODE3_COUNT" -ne 1 ]; then + echo "FAIL: Spread group did not enforce strict 1-per-node"; exit 1 +fi +echo " PASS: Strict 1-per-node spread enforced" + +# Test 5c-8: DescribeInstances shows Placement field +echo "" +echo "Test 5c-8: Spread — DescribeInstances Placement Field" +echo "----------------------------------------" +verify_placement "test-spread" "${SPREAD_IDS[@]}" || exit 1 +echo " PASS: Placement field correctly populated" + +# Test 5c-9: 4th instance fails (all nodes occupied) +echo "" +echo "Test 5c-9: Spread — 4th Instance Fails (all nodes used)" +echo "----------------------------------------" +if FOURTH_OUTPUT=$($AWS_EC2 run-instances \ + --image-id "$AMI_ID" \ + --instance-type "$INSTANCE_TYPE" \ + --key-name multinode-test-key \ + --count 1 \ + --placement GroupName=test-spread 2>&1); then + echo "FAIL: Expected InsufficientInstanceCapacity for 4th spread instance" + ACCIDENTAL_ID=$(echo "$FOURTH_OUTPUT" | jq -r '.Instances[0].InstanceId // empty') + [ -n "$ACCIDENTAL_ID" ] && $AWS_EC2 terminate-instances --instance-ids "$ACCIDENTAL_ID" > /dev/null + exit 1 +fi +if echo "$FOURTH_OUTPUT" | grep -q "InsufficientInstanceCapacity"; then + echo " PASS: 4th instance correctly rejected (all nodes occupied)" +else + echo "FAIL: Wrong error: $FOURTH_OUTPUT"; exit 1 +fi + +# Test 5c-10: Terminate frees node slot, replacement succeeds +echo "" +echo "Test 5c-10: Spread — Terminate Frees Node, Replacement Succeeds" +echo "----------------------------------------" +TERMINATED_ID="${SPREAD_IDS[0]}" +echo " Terminating $TERMINATED_ID to free a node slot..." +$AWS_EC2 terminate-instances --instance-ids "$TERMINATED_ID" > /dev/null +wait_for_instance_state "$TERMINATED_ID" "terminated" 30 || true +sleep 2 + +echo " Launching replacement instance..." +REPLACE_OUTPUT=$($AWS_EC2 run-instances \ + --image-id "$AMI_ID" \ + --instance-type "$INSTANCE_TYPE" \ + --key-name multinode-test-key \ + --count 1 \ + --placement GroupName=test-spread) +REPLACE_ID=$(echo "$REPLACE_OUTPUT" | jq -r '.Instances[0].InstanceId') +if [ -z "$REPLACE_ID" ] || [ "$REPLACE_ID" = "null" ]; then + echo "FAIL: Replacement launch failed"; exit 1 +fi +wait_for_instance_state "$REPLACE_ID" "running" 60 || { + echo "ERROR: Replacement instance failed to reach running state"; exit 1 +} +echo " PASS: Replacement launched successfully after terminate freed slot" + +# Cleanup spread instances +echo " Cleaning up spread instances..." +SPREAD_CLEANUP_IDS=("${SPREAD_IDS[@]:1}" "$REPLACE_ID") +if ! terminate_and_wait "${SPREAD_CLEANUP_IDS[@]}"; then + echo "WARNING: Some spread instances failed to terminate" +fi +sleep 2 + +# Test 5c-11: Cluster — all instances on same node +echo "" +echo "Test 5c-11: Cluster — Launch 3 Instances (all on one node)" +echo "----------------------------------------" +RUN_OUTPUT=$($AWS_EC2 run-instances \ + --image-id "$AMI_ID" \ + --instance-type "$INSTANCE_TYPE" \ + --key-name multinode-test-key \ + --count 3 \ + --placement GroupName=test-cluster) + +CLUSTER_COUNT=$(echo "$RUN_OUTPUT" | jq '.Instances | length') +if [ "$CLUSTER_COUNT" -ne 3 ]; then + echo "FAIL: Expected 3 instances, got $CLUSTER_COUNT"; exit 1 +fi + +CLUSTER_IDS=($(echo "$RUN_OUTPUT" | jq -r '.Instances[].InstanceId')) +echo " Launched: ${CLUSTER_IDS[*]}" + +for id in "${CLUSTER_IDS[@]}"; do + wait_for_instance_state "$id" "running" 60 || { + echo "ERROR: Instance $id failed to reach running state"; exit 1 + } +done + +echo " Checking cluster pinning..." +count_instances_per_node "${CLUSTER_IDS[@]}" +NODES_WITH_INSTANCES=0 +PINNED_NODE="" +for i in 1 2 3; do + count_var="NODE${i}_COUNT" + if [ "${!count_var}" -gt 0 ]; then + NODES_WITH_INSTANCES=$((NODES_WITH_INSTANCES + 1)) + PINNED_NODE=$i + fi +done +if [ "$NODES_WITH_INSTANCES" -ne 1 ]; then + echo "FAIL: Cluster instances spread across $NODES_WITH_INSTANCES nodes (expected 1)"; exit 1 +fi +echo " PASS: All 3 instances pinned to Node$PINNED_NODE" + +# Test 5c-12: DescribeInstances shows Placement field +echo "" +echo "Test 5c-12: Cluster — DescribeInstances Placement Field" +echo "----------------------------------------" +verify_placement "test-cluster" "${CLUSTER_IDS[@]}" || exit 1 +echo " PASS: Placement field correctly populated" + +# Test 5c-13: Subsequent launch pins to same node +echo "" +echo "Test 5c-13: Cluster — Subsequent Launch Pins to Same Node" +echo "----------------------------------------" +echo " Launching 1 more instance into cluster group..." +EXTRA_OUTPUT=$($AWS_EC2 run-instances \ + --image-id "$AMI_ID" \ + --instance-type "$INSTANCE_TYPE" \ + --key-name multinode-test-key \ + --count 1 \ + --placement GroupName=test-cluster) +EXTRA_ID=$(echo "$EXTRA_OUTPUT" | jq -r '.Instances[0].InstanceId') +if [ -z "$EXTRA_ID" ] || [ "$EXTRA_ID" = "null" ]; then + echo "FAIL: Subsequent launch failed"; exit 1 +fi +wait_for_instance_state "$EXTRA_ID" "running" 60 || { + echo "ERROR: Extra instance failed to reach running state"; exit 1 +} + +echo " Checking pinning..." +ALL_CLUSTER_IDS=("${CLUSTER_IDS[@]}" "$EXTRA_ID") +count_instances_per_node "${ALL_CLUSTER_IDS[@]}" +NODES_WITH_INSTANCES=0 +for i in 1 2 3; do + count_var="NODE${i}_COUNT" + if [ "${!count_var}" -gt 0 ]; then + NODES_WITH_INSTANCES=$((NODES_WITH_INSTANCES + 1)) + fi +done +if [ "$NODES_WITH_INSTANCES" -ne 1 ]; then + echo "FAIL: Extra instance not pinned to same node"; exit 1 +fi +echo " PASS: Subsequent launch pinned to same Node$PINNED_NODE" + +# Cleanup cluster instances +echo " Cleaning up cluster instances..." +if ! terminate_and_wait "${ALL_CLUSTER_IDS[@]}"; then + echo "WARNING: Some cluster instances failed to terminate" +fi +sleep 2 + +# Test 5c-14: Cluster capacity exhausted on pinned node +echo "" +echo "Test 5c-14: Cluster — Capacity Exhausted on Pinned Node" +echo "----------------------------------------" +FILL_OUTPUT=$($AWS_EC2 run-instances \ + --image-id "$AMI_ID" \ + --instance-type "$INSTANCE_TYPE" \ + --key-name multinode-test-key \ + --count 1 \ + --placement GroupName=test-cluster) +FILL_ID=$(echo "$FILL_OUTPUT" | jq -r '.Instances[0].InstanceId') +wait_for_instance_state "$FILL_ID" "running" 60 || true + +if CAP_OUTPUT=$($AWS_EC2 run-instances \ + --image-id "$AMI_ID" \ + --instance-type "$INSTANCE_TYPE" \ + --key-name multinode-test-key \ + --count 100 \ + --placement GroupName=test-cluster 2>&1); then + echo "FAIL: Expected InsufficientInstanceCapacity"; exit 1 +fi +if echo "$CAP_OUTPUT" | grep -q "InsufficientInstanceCapacity"; then + echo " PASS: Capacity exhaustion correctly detected on pinned node" +else + echo "FAIL: Wrong error: $CAP_OUTPUT"; exit 1 +fi + +if [ -n "$FILL_ID" ] && [ "$FILL_ID" != "null" ]; then + terminate_and_wait "$FILL_ID" || true +fi +sleep 2 + +# Test 5c-15: DeletePlacementGroup with running instances fails +echo "" +echo "Test 5c-15: DeletePlacementGroup with Instances (InUse)" +echo "----------------------------------------" +INUSE_OUTPUT=$($AWS_EC2 run-instances \ + --image-id "$AMI_ID" \ + --instance-type "$INSTANCE_TYPE" \ + --key-name multinode-test-key \ + --count 1 \ + --placement GroupName=test-spread) +INUSE_ID=$(echo "$INUSE_OUTPUT" | jq -r '.Instances[0].InstanceId') +wait_for_instance_state "$INUSE_ID" "running" 60 || true + +if DEL_OUTPUT=$($AWS_EC2 delete-placement-group \ + --group-name test-spread 2>&1); then + echo "FAIL: Expected InvalidPlacementGroup.InUse"; exit 1 +fi +if echo "$DEL_OUTPUT" | grep -q "InvalidPlacementGroup.InUse"; then + echo " PASS: Delete correctly blocked while instances running" +else + echo "FAIL: Wrong error: $DEL_OUTPUT"; exit 1 +fi + +# Test 5c-16: Terminate all + delete succeeds +echo "" +echo "Test 5c-16: Terminate All + DeletePlacementGroup" +echo "----------------------------------------" +echo " Terminating instance $INUSE_ID..." +terminate_and_wait "$INUSE_ID" || true +sleep 2 + +$AWS_EC2 delete-placement-group --group-name test-spread +echo " PASS: Spread group deleted after terminating instances" + +$AWS_EC2 delete-placement-group --group-name test-cluster +echo " PASS: Cluster group deleted" + +REMAINING=$($AWS_EC2 describe-placement-groups \ + --query 'PlacementGroups[?GroupName==`test-spread` || GroupName==`test-cluster`] | length(@)' \ + --output text 2>/dev/null || echo "0") +if [ "$REMAINING" != "0" ] && [ "$REMAINING" != "None" ]; then + echo "FAIL: Groups still exist after deletion"; exit 1 +fi +echo " PASS: Both groups deleted and verified gone" + +# Test 5c-17: Default spread unchanged (regression) +echo "" +echo "Test 5c-17: Default Spread (No Placement Group) Unchanged" +echo "----------------------------------------" +RUN_OUTPUT=$($AWS_EC2 run-instances \ + --image-id "$AMI_ID" \ + --instance-type "$INSTANCE_TYPE" \ + --key-name multinode-test-key \ + --count 3) +REG_COUNT=$(echo "$RUN_OUTPUT" | jq '.Instances | length') +if [ "$REG_COUNT" -ne 3 ]; then + echo "FAIL: Default spread launch failed, got $REG_COUNT instances"; exit 1 +fi +REG_IDS=($(echo "$RUN_OUTPUT" | jq -r '.Instances[].InstanceId')) +for id in "${REG_IDS[@]}"; do + wait_for_instance_state "$id" "running" 60 || true +done + +PLACEMENT=$($AWS_EC2 describe-instances --instance-ids "${REG_IDS[0]}" \ + --query 'Reservations[0].Instances[0].Placement.GroupName' --output text) +if [ -n "$PLACEMENT" ] && [ "$PLACEMENT" != "None" ] && [ "$PLACEMENT" != "null" ]; then + echo "FAIL: Default spread instance has Placement.GroupName='$PLACEMENT' (should be empty)"; exit 1 +fi +echo " PASS: Default spread works, no Placement.GroupName set" + +if ! terminate_and_wait "${REG_IDS[@]}"; then + echo "WARNING: Some regression test instances failed to terminate" +fi +sleep 2 + +echo "" +echo "Phase 5c: All placement group E2E tests passed" + +# Re-launch 3 instances for Phase 6 shutdown/restart tests +echo "" +echo "Launching 3 instances for Phase 6 shutdown/restart tests..." +INSTANCE_IDS=() +for i in 1 2 3; do + RUN_OUTPUT=$($AWS_EC2 run-instances \ + --image-id "$AMI_ID" \ + --instance-type "$INSTANCE_TYPE" \ + --key-name multinode-test-key) + INSTANCE_ID=$(echo "$RUN_OUTPUT" | jq -r '.Instances[0].InstanceId') + if [ -z "$INSTANCE_ID" ] || [ "$INSTANCE_ID" == "null" ]; then + echo " ERROR: Failed to launch instance $i for Phase 6 setup" + exit 1 + fi + echo " Launched: $INSTANCE_ID" + INSTANCE_IDS+=("$INSTANCE_ID") + sleep 1 +done +for id in "${INSTANCE_IDS[@]}"; do + wait_for_instance_state "$id" "running" 30 || { + echo "ERROR: Instance $id failed to start for Phase 6 setup" + exit 1 + } +done +echo " Phase 6 setup: ${#INSTANCE_IDS[@]} instances running" +echo "========================================" + # Phase 6: Cluster Shutdown + Restart echo "" echo "Phase 6: Cluster Shutdown + Restart"