diff --git a/spinifex/handlers/ec2/placementgroup/service_impl.go b/spinifex/handlers/ec2/placementgroup/service_impl.go index 6f3003150..1d4e42b4d 100644 --- a/spinifex/handlers/ec2/placementgroup/service_impl.go +++ b/spinifex/handlers/ec2/placementgroup/service_impl.go @@ -313,43 +313,51 @@ func pgMatchesAnyTag(tags map[string]string, values []string, field func(k, v st return false } -// GetPlacementGroupRecord reads a placement group record from KV with its revision for CAS operations. -// Returns the record and the KV entry (for revision). -func (s *PlacementGroupServiceImpl) GetPlacementGroupRecord(ctx context.Context, accountID, groupName string) (*PlacementGroupRecord, jetstream.KeyValueEntry, error) { - key := utils.AccountKey(accountID, groupName) - entry, err := s.kv.Get(ctx, key) - if err != nil { - if errors.Is(err, jetstream.ErrKeyNotFound) { - return nil, nil, errors.New(awserrors.ErrorInvalidPlacementGroupUnknown) - } +const maxCASRetries = 5 + +// errPGAbsent and errPGContended mark the two kvutil outcomes that map to +// specific AWS errors rather than a generic internal failure. Neither +// escapes this file. +var ( + errPGAbsent = errors.New("placementgroup: record absent") + errPGContended = errors.New("placementgroup: CAS retries exhausted") +) + +// mapCASError turns a kvutil failure into the AWS error the caller returns. +// A non-kvutil error came from mutate and is already an AWS code. +func mapCASError(ctx context.Context, op, groupName, accountID string, err error) error { + switch { + case errors.Is(err, errPGAbsent): + return errors.New(awserrors.ErrorInvalidPlacementGroupUnknown) + case errors.Is(err, errPGContended): + slog.ErrorContext(ctx, "Placement group CAS retries exhausted under contention", + "op", op, "groupName", groupName, "accountID", accountID) + case errors.Is(err, kvutil.ErrRead): slog.ErrorContext(ctx, "Failed to read placement group from KV", - "groupName", groupName, "accountID", accountID, "err", err) - return nil, nil, errors.New(awserrors.ErrorServerInternal) - } - - 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(ctx context.Context, 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(ctx, key, data, revision); err != nil { + "op", op, "groupName", groupName, "accountID", accountID, "err", err) + case errors.Is(err, kvutil.ErrDecode): + slog.ErrorContext(ctx, "Corrupt placement group record in KV", + "op", op, "groupName", groupName, "accountID", accountID, "err", err) + case errors.Is(err, kvutil.ErrEncode): + slog.ErrorContext(ctx, "Failed to marshal placement group record", + "op", op, "groupName", groupName, "accountID", accountID, "err", err) + case errors.Is(err, kvutil.ErrWrite): + slog.ErrorContext(ctx, "Failed to write placement group to KV", + "op", op, "groupName", groupName, "accountID", accountID, "err", err) + default: return err } - return nil + return errors.New(awserrors.ErrorServerInternal) } -const maxCASRetries = 5 +// casConfig is the shared CAS policy for every placement group mutation. +func casConfig() kvutil.CASConfig { + return kvutil.CASConfig{ + Attempts: maxCASRetries, + NotFound: errPGAbsent, + Exhausted: func(string, int) error { return errPGContended }, + } +} // ReserveSpreadNodes atomically reserves node slots for a spread placement group launch. // Filters occupied nodes, selects up to MaxCount, writes placeholders via CAS with retries. @@ -358,58 +366,41 @@ func (s *PlacementGroupServiceImpl) ReserveSpreadNodes(ctx context.Context, inpu return nil, errors.New(awserrors.ErrorMissingParameter) } - for attempt := range maxCASRetries { - record, entry, err := s.GetPlacementGroupRecord(ctx, 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) + var selected []string + _, err := kvutil.Update(ctx, s.kv, utils.AccountKey(accountID, input.GroupName), casConfig(), + func(record *PlacementGroupRecord) (bool, error) { + if record.State != ec2.PlacementGroupStateAvailable { + return false, errors.New(awserrors.ErrorInvalidPlacementGroupUnknown) + } + if record.Strategy != ec2.PlacementStrategySpread { + return false, errors.New(awserrors.ErrorInvalidParameterValue) } - } - if len(available) < input.MinCount { - return nil, errors.New(awserrors.ErrorInsufficientInstanceCapacity) - } + // Eligible means it has capacity and hosts nothing from this group yet. + var available []string + for _, node := range input.EligibleNodes { + if _, occupied := record.NodeInstances[node]; !occupied { + available = append(available, node) + } + } + if len(available) < input.MinCount { + return false, errors.New(awserrors.ErrorInsufficientInstanceCapacity) + } - // Select nodes: up to MaxCount, at least MinCount - launchCount := min(input.MaxCount, len(available)) - selected := available[:launchCount] + // An empty instance list reserves the node without claiming a launch. + selected = available[:min(input.MaxCount, len(available))] + for _, node := range selected { + record.NodeInstances[node] = []string{} + } + return true, nil + }) - // 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(ctx, accountID, input.GroupName, record, entry.Revision()); err != nil { - slog.DebugContext(ctx, "ReserveSpreadNodes: CAS conflict, retrying", "attempt", attempt, "err", err) - continue - } - - slog.InfoContext(ctx, "ReserveSpreadNodes completed", "groupName", input.GroupName, "nodes", selected, "accountID", accountID) - return &ReserveSpreadNodesOutput{ReservedNodes: selected}, nil + if err != nil { + return nil, mapCASError(ctx, "ReserveSpreadNodes", input.GroupName, accountID, err) } - slog.ErrorContext(ctx, "ReserveSpreadNodes: CAS retries exhausted", "groupName", input.GroupName, "accountID", accountID) - return nil, errors.New(awserrors.ErrorServerInternal) + slog.InfoContext(ctx, "ReserveSpreadNodes completed", "groupName", input.GroupName, "nodes", selected, "accountID", accountID) + return &ReserveSpreadNodesOutput{ReservedNodes: selected}, nil } // FinalizeSpreadInstances replaces placeholder entries with actual instance IDs. @@ -419,26 +410,17 @@ func (s *PlacementGroupServiceImpl) FinalizeSpreadInstances(ctx context.Context, return nil, errors.New(awserrors.ErrorMissingParameter) } - for attempt := range maxCASRetries { - record, entry, err := s.GetPlacementGroupRecord(ctx, 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(ctx, accountID, input.GroupName, record, entry.Revision()); err != nil { - slog.DebugContext(ctx, "FinalizeSpreadInstances: CAS conflict, retrying", "attempt", attempt, "err", err) - continue - } - - slog.InfoContext(ctx, "FinalizeSpreadInstances completed", "groupName", input.GroupName, "accountID", accountID) - return &FinalizeSpreadInstancesOutput{}, nil + _, err := kvutil.Update(ctx, s.kv, utils.AccountKey(accountID, input.GroupName), casConfig(), + func(record *PlacementGroupRecord) (bool, error) { + maps.Copy(record.NodeInstances, input.NodeInstances) + return true, nil + }) + if err != nil { + return nil, mapCASError(ctx, "FinalizeSpreadInstances", input.GroupName, accountID, err) } - slog.ErrorContext(ctx, "FinalizeSpreadInstances: CAS retries exhausted", "groupName", input.GroupName, "accountID", accountID) - return nil, errors.New(awserrors.ErrorServerInternal) + slog.InfoContext(ctx, "FinalizeSpreadInstances completed", "groupName", input.GroupName, "accountID", accountID) + return &FinalizeSpreadInstancesOutput{}, nil } // ReleaseSpreadNodes removes placeholder entries for nodes that failed to launch. @@ -453,31 +435,24 @@ func (s *PlacementGroupServiceImpl) ReleaseSpreadNodes(ctx context.Context, inpu releaseSet[n] = true } - for attempt := range maxCASRetries { - record, entry, err := s.GetPlacementGroupRecord(ctx, 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) + _, err := kvutil.Update(ctx, s.kv, utils.AccountKey(accountID, input.GroupName), casConfig(), + func(record *PlacementGroupRecord) (bool, error) { + changed := false + for node := range releaseSet { + // Only placeholders are releasable. A node holding instance IDs is in use and must not be droppped. + if ids, ok := record.NodeInstances[node]; ok && len(ids) == 0 { + delete(record.NodeInstances, node) + changed = true + } } - } - - if err := s.UpdatePlacementGroupRecord(ctx, accountID, input.GroupName, record, entry.Revision()); err != nil { - slog.DebugContext(ctx, "ReleaseSpreadNodes: CAS conflict, retrying", "attempt", attempt, "err", err) - continue - } - - slog.InfoContext(ctx, "ReleaseSpreadNodes completed", "groupName", input.GroupName, "nodes", input.Nodes, "accountID", accountID) - return &ReleaseSpreadNodesOutput{}, nil + return changed, nil + }) + if err != nil { + return nil, mapCASError(ctx, "ReleaseSpreadNodes", input.GroupName, accountID, err) } - slog.ErrorContext(ctx, "ReleaseSpreadNodes: CAS retries exhausted", "groupName", input.GroupName, "accountID", accountID) - return nil, errors.New(awserrors.ErrorServerInternal) + slog.InfoContext(ctx, "ReleaseSpreadNodes completed", "groupName", input.GroupName, "nodes", input.Nodes, "accountID", accountID) + return &ReleaseSpreadNodesOutput{}, nil } // RemoveInstance removes a specific instance from a placement group's NodeInstances. @@ -488,48 +463,38 @@ func (s *PlacementGroupServiceImpl) RemoveInstance(ctx context.Context, input *R return nil, errors.New(awserrors.ErrorMissingParameter) } - for attempt := range maxCASRetries { - record, entry, lookupErr := s.GetPlacementGroupRecord(ctx, accountID, input.GroupName) - if lookupErr != nil { - if !awserrors.IsErrorCode(lookupErr, awserrors.ErrorInvalidPlacementGroupUnknown) { - return nil, lookupErr + _, err := kvutil.Update(ctx, s.kv, utils.AccountKey(accountID, input.GroupName), casConfig(), + func(record *PlacementGroupRecord) (bool, error) { + instances, exists := record.NodeInstances[input.NodeName] + if !exists { + return false, nil // node not tracked, nothing to remove } - // Group may have been deleted already — treat as success - slog.DebugContext(ctx, "RemoveInstance: group not found, treating as success", "groupName", input.GroupName) - return &RemoveInstanceOutput{}, nil - } - 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) + 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(ctx, accountID, input.GroupName, record, entry.Revision()); err != nil { - slog.DebugContext(ctx, "RemoveInstance: CAS conflict, retrying", "attempt", attempt, "err", err) - continue - } - - slog.InfoContext(ctx, "RemoveInstance completed", "groupName", input.GroupName, "instanceId", input.InstanceID, "node", input.NodeName, "accountID", accountID) + // A node with no instances left stops occupying a spread slot. + if len(filtered) == 0 { + delete(record.NodeInstances, input.NodeName) + } else { + record.NodeInstances[input.NodeName] = filtered + } + return true, nil + }) + if errors.Is(err, errPGAbsent) { + // Group already deleted, so there is nothing left to remove. + slog.DebugContext(ctx, "RemoveInstance: group not found, treating as success", "groupName", input.GroupName) return &RemoveInstanceOutput{}, nil } + if err != nil { + return nil, mapCASError(ctx, "RemoveInstance", input.GroupName, accountID, err) + } - slog.ErrorContext(ctx, "RemoveInstance: CAS retries exhausted", "groupName", input.GroupName, "instanceId", input.InstanceID, "accountID", accountID) - return nil, errors.New(awserrors.ErrorServerInternal) + slog.InfoContext(ctx, "RemoveInstance completed", "groupName", input.GroupName, "instanceId", input.InstanceID, "node", input.NodeName, "accountID", accountID) + return &RemoveInstanceOutput{}, nil } // ReserveClusterNode determines the target node for a cluster placement group launch. @@ -540,42 +505,35 @@ func (s *PlacementGroupServiceImpl) ReserveClusterNode(ctx context.Context, inpu return nil, errors.New(awserrors.ErrorMissingParameter) } - for attempt := range maxCASRetries { - record, entry, err := s.GetPlacementGroupRecord(ctx, 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) - } + var targetNode string + _, err := kvutil.Update(ctx, s.kv, utils.AccountKey(accountID, input.GroupName), casConfig(), + func(record *PlacementGroupRecord) (bool, error) { + if record.State != ec2.PlacementGroupStateAvailable { + return false, errors.New(awserrors.ErrorInvalidPlacementGroupUnknown) + } + if record.Strategy != ec2.PlacementStrategyCluster { + return false, errors.New(awserrors.ErrorInvalidParameterValue) + } - // If existing node, return it immediately (no CAS write needed) - for node := range record.NodeInstances { - return &ReserveClusterNodeOutput{TargetNode: node}, nil - } + // A cluster group keeps every instance on one node, so an existing node is the answer and nothing needs writing. + for node := range record.NodeInstances { + targetNode = node + return false, 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(ctx, accountID, input.GroupName, record, entry.Revision()); err != nil { - slog.DebugContext(ctx, "ReserveClusterNode: CAS conflict, retrying", "attempt", attempt, "err", err) - continue - } - - slog.InfoContext(ctx, "ReserveClusterNode completed", "groupName", input.GroupName, "targetNode", targetNode, "accountID", accountID) - return &ReserveClusterNodeOutput{TargetNode: targetNode}, nil + if len(input.EligibleNodes) == 0 { + return false, errors.New(awserrors.ErrorInsufficientInstanceCapacity) + } + targetNode = input.EligibleNodes[0] // sorted desc by capacity by the caller + record.NodeInstances[targetNode] = []string{} + return true, nil + }) + if err != nil { + return nil, mapCASError(ctx, "ReserveClusterNode", input.GroupName, accountID, err) } - slog.ErrorContext(ctx, "ReserveClusterNode: CAS retries exhausted", "groupName", input.GroupName, "accountID", accountID) - return nil, errors.New(awserrors.ErrorServerInternal) + slog.InfoContext(ctx, "ReserveClusterNode completed", "groupName", input.GroupName, "targetNode", targetNode, "accountID", accountID) + return &ReserveClusterNodeOutput{TargetNode: targetNode}, nil } // FinalizeClusterInstances appends launched instance IDs to the cluster placement group record. @@ -585,28 +543,20 @@ func (s *PlacementGroupServiceImpl) FinalizeClusterInstances(ctx context.Context return nil, errors.New(awserrors.ErrorMissingParameter) } - for attempt := range maxCASRetries { - record, entry, err := s.GetPlacementGroupRecord(ctx, 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(ctx, accountID, input.GroupName, record, entry.Revision()); err != nil { - slog.DebugContext(ctx, "FinalizeClusterInstances: CAS conflict, retrying", "attempt", attempt, "err", err) - continue - } - - slog.InfoContext(ctx, "FinalizeClusterInstances completed", "groupName", input.GroupName, "accountID", accountID) - return &FinalizeClusterInstancesOutput{}, nil + _, err := kvutil.Update(ctx, s.kv, utils.AccountKey(accountID, input.GroupName), casConfig(), + func(record *PlacementGroupRecord) (bool, error) { + // Append rather than replace: a cluster group can have several launches in flight at once. + for node, ids := range input.NodeInstances { + record.NodeInstances[node] = append(record.NodeInstances[node], ids...) + } + return true, nil + }) + if err != nil { + return nil, mapCASError(ctx, "FinalizeClusterInstances", input.GroupName, accountID, err) } - slog.ErrorContext(ctx, "FinalizeClusterInstances: CAS retries exhausted", "groupName", input.GroupName, "accountID", accountID) - return nil, errors.New(awserrors.ErrorServerInternal) + slog.InfoContext(ctx, "FinalizeClusterInstances completed", "groupName", input.GroupName, "accountID", accountID) + return &FinalizeClusterInstancesOutput{}, nil } // recordToEC2 converts an internal record to the AWS SDK PlacementGroup type. diff --git a/spinifex/handlers/ec2/placementgroup/service_impl_test.go b/spinifex/handlers/ec2/placementgroup/service_impl_test.go index d95cac24d..c77734265 100644 --- a/spinifex/handlers/ec2/placementgroup/service_impl_test.go +++ b/spinifex/handlers/ec2/placementgroup/service_impl_test.go @@ -2,12 +2,15 @@ package handlers_ec2_placementgroup import ( "context" + "encoding/json" "testing" "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/kvutil" "github.com/mulgadc/spinifex/spinifex/testutil" + "github.com/mulgadc/spinifex/spinifex/utils" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -33,6 +36,28 @@ func createTestGroup(t *testing.T, svc *PlacementGroupServiceImpl, name, strateg return out.PlacementGroup } +// readRecord returns the stored record for assertions. +func readRecord(t *testing.T, svc *PlacementGroupServiceImpl, groupName string) *PlacementGroupRecord { + t.Helper() + entry, err := svc.kv.Get(t.Context(), utils.AccountKey(testAccountID, groupName)) + require.NoError(t, err) + var record PlacementGroupRecord + require.NoError(t, json.Unmarshal(entry.Value(), &record)) + return &record +} + +// seedRecord applies mutate to the stored record, for tests that need a +// starting state the public API cannot produce. +func seedRecord(t *testing.T, svc *PlacementGroupServiceImpl, groupName string, mutate func(*PlacementGroupRecord)) { + t.Helper() + _, err := kvutil.Update(t.Context(), svc.kv, utils.AccountKey(testAccountID, groupName), kvutil.CASConfig{}, + func(record *PlacementGroupRecord) (bool, error) { + mutate(record) + return true, nil + }) + require.NoError(t, err) +} + // --- CreatePlacementGroup Tests --- func TestCreatePlacementGroup_Spread(t *testing.T) { @@ -162,13 +187,9 @@ func TestDeletePlacementGroup_InUse(t *testing.T) { createTestGroup(t, svc, "in-use-group", "spread") // Simulate instances by updating the record directly - record, entry, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "in-use-group") - require.NoError(t, err) - record.NodeInstances["node1"] = []string{"i-123"} - err = svc.UpdatePlacementGroupRecord(t.Context(), testAccountID, "in-use-group", record, entry.Revision()) - require.NoError(t, err) + seedRecord(t, svc, "in-use-group", func(r *PlacementGroupRecord) { r.NodeInstances["node1"] = []string{"i-123"} }) - _, err = svc.DeletePlacementGroup(context.Background(), &ec2.DeletePlacementGroupInput{ + _, err := svc.DeletePlacementGroup(context.Background(), &ec2.DeletePlacementGroupInput{ GroupName: aws.String("in-use-group"), }, testAccountID) require.Error(t, err) @@ -308,57 +329,6 @@ func TestDescribePlacementGroups_AccountScoped(t *testing.T) { 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(t.Context(), 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(t.Context(), testAccountID, "cas-group", record, entry.Revision()) - require.NoError(t, err) - - // Verify update - record2, _, err := svc.GetPlacementGroupRecord(t.Context(), 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(t.Context(), testAccountID, "conflict-group") - require.NoError(t, err) - record2, _, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "conflict-group") - require.NoError(t, err) - - // First update succeeds - record1.NodeInstances["node1"] = []string{"i-111"} - err = svc.UpdatePlacementGroupRecord(t.Context(), 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(t.Context(), testAccountID, "conflict-group", record2, entry1.Revision()) - require.Error(t, err) -} - -func TestGetPlacementGroupRecord_NotFound(t *testing.T) { - svc := setupTestService(t) - _, _, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "nonexistent") - require.Error(t, err) - assert.Equal(t, awserrors.ErrorInvalidPlacementGroupUnknown, err.Error()) -} - // --- ReserveSpreadNodes Tests --- func TestReserveSpreadNodes_Success(t *testing.T) { @@ -375,7 +345,7 @@ func TestReserveSpreadNodes_Success(t *testing.T) { assert.Len(t, out.ReservedNodes, 3) // Verify placeholders in record - record, _, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "reserve-group") + record := readRecord(t, svc, "reserve-group") require.NoError(t, err) assert.Len(t, record.NodeInstances, 3) for _, node := range out.ReservedNodes { @@ -389,10 +359,7 @@ func TestReserveSpreadNodes_ExcludesOccupiedNodes(t *testing.T) { createTestGroup(t, svc, "occupied-group", "spread") // Pre-occupy node-a - record, entry, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "occupied-group") - require.NoError(t, err) - record.NodeInstances["node-a"] = []string{"i-existing"} - require.NoError(t, svc.UpdatePlacementGroupRecord(t.Context(), testAccountID, "occupied-group", record, entry.Revision())) + seedRecord(t, svc, "occupied-group", func(r *PlacementGroupRecord) { r.NodeInstances["node-a"] = []string{"i-existing"} }) out, err := svc.ReserveSpreadNodes(context.Background(), &ReserveSpreadNodesInput{ GroupName: "occupied-group", @@ -413,13 +380,12 @@ func TestReserveSpreadNodes_InsufficientNodes(t *testing.T) { createTestGroup(t, svc, "insufficient-group", "spread") // Pre-occupy node-a and node-b - record, entry, err := svc.GetPlacementGroupRecord(t.Context(), 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(t.Context(), testAccountID, "insufficient-group", record, entry.Revision())) + seedRecord(t, svc, "insufficient-group", func(r *PlacementGroupRecord) { + r.NodeInstances["node-a"] = []string{"i-1"} + r.NodeInstances["node-b"] = []string{"i-2"} + }) - _, err = svc.ReserveSpreadNodes(context.Background(), &ReserveSpreadNodesInput{ + _, err := svc.ReserveSpreadNodes(context.Background(), &ReserveSpreadNodesInput{ GroupName: "insufficient-group", EligibleNodes: []string{"node-a", "node-b", "node-c"}, MinCount: 2, @@ -481,7 +447,7 @@ func TestFinalizeSpreadInstances_Success(t *testing.T) { require.NoError(t, err) // Verify - record, _, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "finalize-group") + record := readRecord(t, svc, "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"]) @@ -510,7 +476,7 @@ func TestReleaseSpreadNodes_Success(t *testing.T) { require.NoError(t, err) // Verify only node-a remains - record, _, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "release-group") + record := readRecord(t, svc, "release-group") require.NoError(t, err) assert.Len(t, record.NodeInstances, 1) _, ok := record.NodeInstances["node-a"] @@ -538,7 +504,7 @@ func TestReleaseSpreadNodes_AllNodes(t *testing.T) { require.NoError(t, err) // Verify empty - record, _, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "release-all-group") + record := readRecord(t, svc, "release-all-group") require.NoError(t, err) assert.Empty(t, record.NodeInstances) } @@ -550,14 +516,13 @@ func TestRemoveInstance_Success(t *testing.T) { createTestGroup(t, svc, "remove-inst-group", "spread") // Add instances to two nodes - record, entry, err := svc.GetPlacementGroupRecord(t.Context(), 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(t.Context(), testAccountID, "remove-inst-group", record, entry.Revision())) + seedRecord(t, svc, "remove-inst-group", func(r *PlacementGroupRecord) { + r.NodeInstances["node-a"] = []string{"i-aaa"} + r.NodeInstances["node-b"] = []string{"i-bbb"} + }) // Remove i-aaa from node-a - _, err = svc.RemoveInstance(context.Background(), &RemoveInstanceInput{ + _, err := svc.RemoveInstance(context.Background(), &RemoveInstanceInput{ GroupName: "remove-inst-group", NodeName: "node-a", InstanceID: "i-aaa", @@ -565,7 +530,7 @@ func TestRemoveInstance_Success(t *testing.T) { require.NoError(t, err) // Verify node-a is removed (was the only instance), node-b remains - record, _, err = svc.GetPlacementGroupRecord(t.Context(), testAccountID, "remove-inst-group") + record := readRecord(t, svc, "remove-inst-group") require.NoError(t, err) assert.Len(t, record.NodeInstances, 1) _, hasA := record.NodeInstances["node-a"] @@ -605,13 +570,10 @@ func TestRemoveInstance_MultipleInstancesOnNode(t *testing.T) { createTestGroup(t, svc, "multi-inst-group", "cluster") // Add multiple instances to same node (cluster scenario) - record, entry, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "multi-inst-group") - require.NoError(t, err) - record.NodeInstances["node-a"] = []string{"i-111", "i-222", "i-333"} - require.NoError(t, svc.UpdatePlacementGroupRecord(t.Context(), testAccountID, "multi-inst-group", record, entry.Revision())) + seedRecord(t, svc, "multi-inst-group", func(r *PlacementGroupRecord) { r.NodeInstances["node-a"] = []string{"i-111", "i-222", "i-333"} }) // Remove i-222 - _, err = svc.RemoveInstance(context.Background(), &RemoveInstanceInput{ + _, err := svc.RemoveInstance(context.Background(), &RemoveInstanceInput{ GroupName: "multi-inst-group", NodeName: "node-a", InstanceID: "i-222", @@ -619,7 +581,7 @@ func TestRemoveInstance_MultipleInstancesOnNode(t *testing.T) { require.NoError(t, err) // Verify i-111 and i-333 remain - record, _, err = svc.GetPlacementGroupRecord(t.Context(), testAccountID, "multi-inst-group") + record := readRecord(t, svc, "multi-inst-group") require.NoError(t, err) assert.Equal(t, []string{"i-111", "i-333"}, record.NodeInstances["node-a"]) } @@ -639,7 +601,7 @@ func TestReserveClusterNode_EmptyGroup(t *testing.T) { assert.Equal(t, "node-a", out.TargetNode) // Verify placeholder written - record, _, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "cluster-empty") + record := readRecord(t, svc, "cluster-empty") require.NoError(t, err) assert.Len(t, record.NodeInstances, 1) _, ok := record.NodeInstances["node-a"] @@ -651,10 +613,7 @@ func TestReserveClusterNode_ExistingNode(t *testing.T) { createTestGroup(t, svc, "cluster-existing", "cluster") // Pre-populate with instances on node-b - record, entry, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "cluster-existing") - require.NoError(t, err) - record.NodeInstances["node-b"] = []string{"i-existing"} - require.NoError(t, svc.UpdatePlacementGroupRecord(t.Context(), testAccountID, "cluster-existing", record, entry.Revision())) + seedRecord(t, svc, "cluster-existing", func(r *PlacementGroupRecord) { r.NodeInstances["node-b"] = []string{"i-existing"} }) // Even though node-a is in eligible list, should return existing node-b out, err := svc.ReserveClusterNode(context.Background(), &ReserveClusterNodeInput{ @@ -731,7 +690,7 @@ func TestFinalizeClusterInstances_Success(t *testing.T) { require.NoError(t, err) // Verify instances recorded - record, _, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "cluster-finalize") + record := readRecord(t, svc, "cluster-finalize") require.NoError(t, err) assert.Equal(t, []string{"i-c1", "i-c2", "i-c3"}, record.NodeInstances["node-a"]) } @@ -741,13 +700,10 @@ func TestFinalizeClusterInstances_AppendsToExisting(t *testing.T) { createTestGroup(t, svc, "cluster-append", "cluster") // Pre-populate with existing instances - record, entry, err := svc.GetPlacementGroupRecord(t.Context(), testAccountID, "cluster-append") - require.NoError(t, err) - record.NodeInstances["node-a"] = []string{"i-existing1", "i-existing2"} - require.NoError(t, svc.UpdatePlacementGroupRecord(t.Context(), testAccountID, "cluster-append", record, entry.Revision())) + seedRecord(t, svc, "cluster-append", func(r *PlacementGroupRecord) { r.NodeInstances["node-a"] = []string{"i-existing1", "i-existing2"} }) // Finalize with new instances (should append, not overwrite) - _, err = svc.FinalizeClusterInstances(context.Background(), &FinalizeClusterInstancesInput{ + _, err := svc.FinalizeClusterInstances(context.Background(), &FinalizeClusterInstancesInput{ GroupName: "cluster-append", NodeInstances: map[string][]string{ "node-a": {"i-new1", "i-new2"}, @@ -756,7 +712,7 @@ func TestFinalizeClusterInstances_AppendsToExisting(t *testing.T) { require.NoError(t, err) // Verify all instances present - record, _, err = svc.GetPlacementGroupRecord(t.Context(), testAccountID, "cluster-append") + record := readRecord(t, svc, "cluster-append") require.NoError(t, err) assert.Equal(t, []string{"i-existing1", "i-existing2", "i-new1", "i-new2"}, record.NodeInstances["node-a"]) }