fix: guard RDS create rollback revisions (#728)

spinifex/handlers/rds/create.go

@@ -18,6 +18,9 @@ import (
 // so an agent report carrying an older generation is a superseded VM.
 const firstVMGeneration = 1
 
+// Bounds retries when same-owner record writes race rollback cleanup.
+const rollbackDeleteAttempts = 3
+
 // Assembles a live DB instance out of the launch primitives: validate, reserve
 // the identifier, place and launch the dual-NIC VM with its data volume and
 // customer ENI, seed the one-shot bootstrap config, and publish the endpoint
@@ -56,7 +59,8 @@ func (s *Service) CreateDBInstance(ctx context.Context, input *rds.CreateDBInsta
 	// uniqueness check atomic against a concurrent create on another node — a
 	// prior read-then-write would let both pass and both launch a VM.
 	rec := newDBInstanceRecord(accountID, req, placement, parameters)
-	if createErr := createJSON(ctx, kv, key, &rec); createErr != nil {
+	rollbackRev, createErr := createJSONRevision(ctx, kv, key, &rec)
+	if createErr != nil {
 		if errors.Is(createErr, jetstream.ErrKeyExists) {
 			return nil, awserrors.Errorf(awserrors.ErrorDBInstanceAlreadyExists,
 				"DB instance %s already exists", req.Identifier)
@@ -64,19 +68,13 @@ func (s *Service) CreateDBInstance(ctx context.Context, input *rds.CreateDBInsta
 		return nil, createErr
 	}
 
-	// The record is the reservation, so it is withdrawn on any failure below:
-	// leaving it behind would make the identifier permanently unusable while
-	// naming an instance that was never provisioned.
+	// The record is the reservation, so it is withdrawn on any failure below.
+	// The revision guard leaves a concurrent recreation of the identifier intact.
 	defer func() {
 		if err == nil {
 			return
 		}
-		rbCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), rollbackTimeout)
-		defer cancel()
-		if delErr := kv.Delete(rbCtx, key); delErr != nil {
-			slog.WarnContext(rbCtx, "rds: rollback delete of reserved DB instance record failed",
-				"dbInstance", req.Identifier, "err", delErr)
-		}
+		s.rollbackDBInstanceReservation(ctx, kv, key, req.Identifier, rec.DbiResourceID, rollbackRev)
 	}()
 
 	launched, err := LaunchDBInstanceVM(ctx, s.deps.Launch, LaunchInput{
@@ -106,11 +104,12 @@ func (s *Service) CreateDBInstance(ctx context.Context, input *rds.CreateDBInsta
 
 	// The launch already unwound everything on failure, so from here the resources
 	// exist and the record has to catch up with them.
-	stored, err := s.recordLaunch(ctx, kv, key, accountID, launched)
+	stored, launchRev, err := s.recordLaunch(ctx, kv, key, accountID, rec.DbiResourceID, launched)
 	if err != nil {
 		s.unwindLaunched(ctx, launched)
 		return nil, err
 	}
+	rollbackRev = launchRev
 
 	// Written before the DNS publish because an agent that boots fast enough to
 	// call the gateway before this exists is denied and has to retry.
@@ -178,14 +177,18 @@ func newDBInstanceRecord(accountID string, req *validatedCreate, placement *endp
 // Folds the launch results into the reserved record and returns it as written.
 // The endpoint is settled here because the ENI IP — and so the vanity name — is
 // only known now.
-func (s *Service) recordLaunch(ctx context.Context, kv jetstream.KeyValue, key, accountID string, launched *LaunchOutput) (*DBInstanceRecord, error) {
+func (s *Service) recordLaunch(ctx context.Context, kv jetstream.KeyValue, key, accountID,
+	expectedResourceID string, launched *LaunchOutput) (*DBInstanceRecord, uint64, error) {
 	var rec DBInstanceRecord
 	rev, found, err := getJSONRevision(ctx, kv, key, &rec)
 	if err != nil {
-		return nil, err
+		return nil, 0, err
 	}
 	if !found {
-		return nil, errors.New(awserrors.ErrorDBInstanceNotFound)
+		return nil, 0, errors.New(awserrors.ErrorDBInstanceNotFound)
+	}
+	if rec.DbiResourceID != expectedResourceID {
+		return nil, 0, fmt.Errorf("rds: DB instance reservation %s changed ownership during launch", key)
 	}
 
 	rec.InstanceID = launched.InstanceID
@@ -208,10 +211,44 @@ func (s *Service) recordLaunch(ctx context.Context, kv jetstream.KeyValue, key,
 	}
 	rec.UpdatedAt = time.Now().UTC()
 
-	if err := updateJSON(ctx, kv, key, rev, &rec); err != nil {
-		return nil, err
+	updatedRev, err := updateJSONRevision(ctx, kv, key, rev, &rec)
+	if err != nil {
+		return nil, 0, err
 	}
-	return &rec, nil
+	return &rec, updatedRev, nil
+}
+
+func (s *Service) rollbackDBInstanceReservation(ctx context.Context, kv jetstream.KeyValue,
+	key, identifier, resourceID string, rev uint64) {
+	rbCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), rollbackTimeout)
+	defer cancel()
+
+	var rollbackErr error
+	for range rollbackDeleteAttempts {
+		rollbackErr = kv.Delete(rbCtx, key, jetstream.LastRevision(rev))
+		if rollbackErr == nil || errors.Is(rollbackErr, jetstream.ErrKeyNotFound) {
+			return
+		}
+		if !errors.Is(rollbackErr, jetstream.ErrKeyExists) {
+			break
+		}
+
+		// A same-owner update remains part of this failed create and must be
+		// withdrawn. A different resource ID is a replacement to preserve.
+		var current DBInstanceRecord
+		currentRev, found, getErr := getJSONRevision(rbCtx, kv, key, &current)
+		if getErr != nil {
+			rollbackErr = getErr
+			break
+		}
+		if !found || current.DbiResourceID != resourceID {
+			return
+		}
+		rev = currentRev
+	}
+
+	slog.WarnContext(rbCtx, "rds: rollback delete of reserved DB instance record failed",
+		"dbInstance", identifier, "err", rollbackErr)
 }
 
 // The launch's own rollback only covers failures inside it. A failure after it

spinifex/handlers/rds/create_test.go

@@ -193,6 +193,22 @@ func (h *createHarness) recordExists(t *testing.T, id string) bool {
 	return found
 }
 
+func replaceInstanceRecord(t *testing.T, svc *Service, id string) DBInstanceRecord {
+	t.Helper()
+	kv, err := svc.bucket(t.Context(), testAccountID)
+	require.NoError(t, err)
+	key := DBInstanceKey(id)
+	require.NoError(t, kv.Delete(t.Context(), key))
+	replacement := DBInstanceRecord{
+		DBInstanceIdentifier: id,
+		DbiResourceID:        "db-replacement-owner",
+		InstanceID:           "i-replacement",
+		Status:               StatusCreating,
+	}
+	require.NoError(t, createJSON(t.Context(), kv, key, &replacement))
+	return replacement
+}
+
 func validCreateInput() *rds.CreateDBInstanceInput {
 	return &rds.CreateDBInstanceInput{
 		DBInstanceIdentifier: aws.String(testDBInstanceID),
@@ -372,6 +388,71 @@ func TestCreateDBInstance_FailureAfterLaunchUnwindsEveryResource(t *testing.T) {
 	assert.Len(t, h.launch.enis.deleted, 2, "both the customer and system ENI are deleted")
 }
 
+func TestCreateDBInstance_IndexFailureWithdrawsTheRecordedLaunch(t *testing.T) {
+	h := newCreateHarness(t, "")
+	// This invalid KV key makes the instance-index write fail after recordLaunch
+	// has advanced the reservation revision.
+	h.launch.launcher.instanceID = "invalid instance id"
+
+	_, err := h.svc.CreateDBInstance(t.Context(), validCreateInput(), testAccountID)
+	require.Error(t, err)
+	assert.False(t, h.recordExists(t, testDBInstanceID))
+
+	h.launch.launcher.instanceID = ""
+	_, err = h.svc.CreateDBInstance(t.Context(), validCreateInput(), testAccountID)
+	require.NoError(t, err, "the failed create must release the identifier for reuse")
+}
+
+func TestCreateDBInstance_RecordLaunchDoesNotOverwriteAConcurrentReplacement(t *testing.T) {
+	h := newCreateHarness(t, "")
+	var replacement DBInstanceRecord
+	h.launch.launcher.onLaunch = func() {
+		replacement = replaceInstanceRecord(t, h.svc, testDBInstanceID)
+	}
+
+	_, err := h.svc.CreateDBInstance(t.Context(), validCreateInput(), testAccountID)
+	require.Error(t, err)
+
+	assert.Equal(t, replacement, h.record(t, testDBInstanceID))
+}
+
+func TestCreateDBInstance_RollbackDoesNotDeleteAConcurrentReplacement(t *testing.T) {
+	h := newCreateHarness(t, "")
+	h.launch.launcher.instanceID = "invalid instance id"
+	var replacement DBInstanceRecord
+	h.launch.launcher.onTerminate = func() {
+		assert.Equal(t, "invalid instance id", h.record(t, testDBInstanceID).InstanceID,
+			"recordLaunch must finish before the replacement race")
+		replacement = replaceInstanceRecord(t, h.svc, testDBInstanceID)
+	}
+
+	_, err := h.svc.CreateDBInstance(t.Context(), validCreateInput(), testAccountID)
+	require.Error(t, err)
+
+	assert.Equal(t, replacement, h.record(t, testDBInstanceID))
+}
+
+func TestCreateDBInstance_RollbackFollowsSameOwnerUpdates(t *testing.T) {
+	h := newCreateHarness(t, "")
+	h.launch.launcher.instanceID = "invalid instance id"
+	h.launch.launcher.onTerminate = func() {
+		kv, err := h.svc.bucket(t.Context(), testAccountID)
+		require.NoError(t, err)
+		key := DBInstanceKey(testDBInstanceID)
+		var current DBInstanceRecord
+		rev, found, err := getJSONRevision(t.Context(), kv, key, &current)
+		require.NoError(t, err)
+		require.True(t, found)
+		current.Tags = map[string]string{"updated": "during-rollback"}
+		require.NoError(t, updateJSON(t.Context(), kv, key, rev, &current))
+	}
+
+	_, err := h.svc.CreateDBInstance(t.Context(), validCreateInput(), testAccountID)
+	require.Error(t, err)
+
+	assert.False(t, h.recordExists(t, testDBInstanceID))
+}
+
 // The request may not be answered with an instance whose storage is
 // not actually encrypted, so an unencrypted volume fails the whole create.
 func TestCreateDBInstance_UnencryptedDataVolumeFailsTheCreate(t *testing.T) {

spinifex/handlers/rds/launch_test.go

@@ -252,6 +252,7 @@ func (f *fakeENIs) ModifyNetworkInterfaceAttribute(_ context.Context, in *ec2.Mo
 // fakeLauncher stands in for the system-instance launcher.
 type fakeLauncher struct {
 	input      *sysinstance.SystemInstanceInput
+	instanceID string
 	err        error
 	terminated []string
 	unwind     *[]string
@@ -262,7 +263,8 @@ type fakeLauncher struct {
 
 	// onLaunch runs once the VM exists, for tests that need to disturb state
 	// after the launch has committed resources but before the caller records it.
-	onLaunch func()
+	onLaunch    func()
+	onTerminate func()
 }
 
 var _ launchInstanceLauncher = (*fakeLauncher)(nil)
@@ -275,7 +277,11 @@ func (f *fakeLauncher) LaunchSystemInstance(in *sysinstance.SystemInstanceInput)
 	if f.onLaunch != nil {
 		f.onLaunch()
 	}
-	return &sysinstance.SystemInstanceOutput{InstanceID: "i-rds0001"}, nil
+	instanceID := f.instanceID
+	if instanceID == "" {
+		instanceID = "i-rds0001"
+	}
+	return &sysinstance.SystemInstanceOutput{InstanceID: instanceID}, nil
 }
 
 func (f *fakeLauncher) TerminateSystemInstance(instanceID string) error {
@@ -283,6 +289,9 @@ func (f *fakeLauncher) TerminateSystemInstance(instanceID string) error {
 	if f.unwind != nil {
 		*f.unwind = append(*f.unwind, "terminate")
 	}
+	if f.onTerminate != nil {
+		f.onTerminate()
+	}
 	return f.terminateErr
 }
 

spinifex/handlers/rds/restore.go

@@ -74,7 +74,8 @@ func (s *Service) RestoreDBInstanceFromDBSnapshot(ctx context.Context, input *rd
 	// withdrawn on any failure below.
 	key := DBInstanceKey(req.Identifier)
 	rec := newRestoredDBInstanceRecord(accountID, req, placement, parameters, snapshot)
-	if createErr := createJSON(ctx, kv, key, &rec); createErr != nil {
+	rollbackRev, createErr := createJSONRevision(ctx, kv, key, &rec)
+	if createErr != nil {
 		if errors.Is(createErr, jetstream.ErrKeyExists) {
 			return nil, awserrors.Errorf(awserrors.ErrorDBInstanceAlreadyExists,
 				"DB instance %s already exists", req.Identifier)
@@ -85,12 +86,7 @@ func (s *Service) RestoreDBInstanceFromDBSnapshot(ctx context.Context, input *rd
 		if err == nil {
 			return
 		}
-		rbCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), rollbackTimeout)
-		defer cancel()
-		if delErr := kv.Delete(rbCtx, key); delErr != nil {
-			slog.WarnContext(rbCtx, "rds: rollback delete of reserved restored DB instance record failed",
-				"dbInstance", req.Identifier, "err", delErr)
-		}
+		s.rollbackDBInstanceReservation(ctx, kv, key, req.Identifier, rec.DbiResourceID, rollbackRev)
 	}()
 
 	// The datadir the restored engine comes up on. Created before the VM because
@@ -136,11 +132,12 @@ func (s *Service) RestoreDBInstanceFromDBSnapshot(ctx context.Context, input *rd
 		return nil, err
 	}
 
-	stored, err := s.recordLaunch(ctx, kv, key, accountID, launched)
+	stored, launchRev, err := s.recordLaunch(ctx, kv, key, accountID, rec.DbiResourceID, launched)
 	if err != nil {
 		s.unwindLaunched(ctx, launched)
 		return nil, err
 	}
+	rollbackRev = launchRev
 	if err = s.PutInstanceIndex(ctx, launched.InstanceID, InstanceIndexEntry{
 		AccountID:            accountID,
 		DBInstanceIdentifier: req.Identifier,

spinifex/handlers/rds/restore_test.go

@@ -206,6 +206,53 @@ func TestRestoreDBInstanceFromDBSnapshot_UnwindsTheVolumeAndTheRecordWhenTheLaun
 	assert.NotEmpty(t, h.launch.unwind)
 }
 
+func TestRestoreDBInstanceFromDBSnapshot_IndexFailureWithdrawsTheRecordedLaunch(t *testing.T) {
+	h := newSnapshotHarness(t, false)
+	h.seedSnapshot(t)
+	// This invalid KV key makes the instance-index write fail after recordLaunch
+	// has advanced the reservation revision.
+	h.launch.launcher.instanceID = "invalid instance id"
+
+	_, err := h.svc.RestoreDBInstanceFromDBSnapshot(t.Context(), restoreInput(), testAccountID)
+	require.Error(t, err)
+	assert.False(t, h.instanceExists(t, testRestoredID))
+
+	h.launch.launcher.instanceID = ""
+	_, err = h.svc.RestoreDBInstanceFromDBSnapshot(t.Context(), restoreInput(), testAccountID)
+	require.NoError(t, err, "the failed restore must release the identifier for reuse")
+}
+
+func TestRestoreDBInstanceFromDBSnapshot_RecordLaunchDoesNotOverwriteAConcurrentReplacement(t *testing.T) {
+	h := newSnapshotHarness(t, false)
+	h.seedSnapshot(t)
+	var replacement DBInstanceRecord
+	h.launch.launcher.onLaunch = func() {
+		replacement = replaceInstanceRecord(t, h.svc, testRestoredID)
+	}
+
+	_, err := h.svc.RestoreDBInstanceFromDBSnapshot(t.Context(), restoreInput(), testAccountID)
+	require.Error(t, err)
+
+	assert.Equal(t, replacement, h.instance(t, testRestoredID))
+}
+
+func TestRestoreDBInstanceFromDBSnapshot_RollbackDoesNotDeleteAConcurrentReplacement(t *testing.T) {
+	h := newSnapshotHarness(t, false)
+	h.seedSnapshot(t)
+	h.launch.launcher.instanceID = "invalid instance id"
+	var replacement DBInstanceRecord
+	h.launch.launcher.onTerminate = func() {
+		assert.Equal(t, "invalid instance id", h.instance(t, testRestoredID).InstanceID,
+			"recordLaunch must finish before the replacement race")
+		replacement = replaceInstanceRecord(t, h.svc, testRestoredID)
+	}
+
+	_, err := h.svc.RestoreDBInstanceFromDBSnapshot(t.Context(), restoreInput(), testAccountID)
+	require.Error(t, err)
+
+	assert.Equal(t, replacement, h.instance(t, testRestoredID))
+}
+
 func TestRestoreDBInstanceFromDBSnapshot_RejectsATakenIdentifier(t *testing.T) {
 	h := newSnapshotHarness(t, false)
 	h.seedSnapshot(t)

spinifex/handlers/rds/service.go

@@ -236,12 +236,19 @@ func createJSONRevision(ctx context.Context, kv jetstream.KeyValue, key string,
 
 // Writes v at key only if the stored entry is still at rev.
 func updateJSON(ctx context.Context, kv jetstream.KeyValue, key string, rev uint64, v any) error {
+	_, err := updateJSONRevision(ctx, kv, key, rev, v)
+	return err
+}
+
+// updateJSON plus the revision written by the successful CAS update.
+func updateJSONRevision(ctx context.Context, kv jetstream.KeyValue, key string, rev uint64, v any) (uint64, error) {
 	data, err := json.Marshal(v)
 	if err != nil {
-		return err
+		return 0, err
 	}
-	if _, err := kv.Update(ctx, key, data, rev); err != nil {
-		return fmt.Errorf("update %s: %w", key, err)
+	updatedRev, err := kv.Update(ctx, key, data, rev)
+	if err != nil {
+		return 0, fmt.Errorf("update %s: %w", key, err)
 	}
-	return nil
+	return updatedRev, nil
 }