fix(ec2): serialise tag writes across the cluster (#1070)

CreateTags is a read-modify-write of one resource's whole tag map. It was
serialised by a sync.Mutex on the service struct, which exists once per
process, while the handler answers on the spinifex-workers queue group and so
runs on whichever node NATS picks. Two nodes read the same map, each added its
own key, and the later write won. Every call returned success.

Six parallel writers against one security group on a three-node cluster left
four tags. On a single node the mutex hides it entirely, which is why no test
caught it.

The tag index moves from one S3 object per resource to a JetStream KV bucket,
where the merge runs inside kvutil.Update and commits at the revision it read.
A writer that loses the race re-reads the winner's keys and adds its own on
top. S3 could not host the fix: predastore's PutObject has no If-Match, so
there is nothing to compare and swap against.

Tags written before the move are not migrated in a batch. A batch that
half-failed would leave a resource looking untagged, and the next write would
commit that as the truth. Instead the CAS seeds a missing entry from the old
object inside the attempt, so a resource migrates on its first write and a
failure to read the old object fails that one call. DescribeTags reads the
union of both key sets, preferring the bucket, and DeleteAllTags clears both so
old tags cannot reappear through the fallback.

The process-local mutex is removed rather than left beside the CAS.

The regression test drives six separate service values over one bucket,
because that is what makes it able to fail: with the old code every caller
inside one process was already serialised, and only a second node could lose an
update.

spinifex/daemon/daemon.go

@@ -1594,7 +1594,17 @@ func (d *Daemon) startCluster() error {
 	if err := d.configureEBSProvider(); err != nil {
 		return fmt.Errorf("configure EBS provider: %w", err)
 	}
-	d.tagsService = handlers_ec2_tags.NewTagsServiceImpl(d.config)
+	tagsKV, err := initServiceWithRetry("tags KV bucket", func() (jetstream.KeyValue, error) {
+		tagsJS, jsErr := jetstream.New(d.natsConn)
+		if jsErr != nil {
+			return nil, fmt.Errorf("jetstream handle: %w", jsErr)
+		}
+		return handlers_ec2_tags.GetOrCreateTagsBucket(d.ctx, tagsJS)
+	})
+	if err != nil {
+		return fmt.Errorf("failed to get tags KV bucket: %w", err)
+	}
+	d.tagsService = handlers_ec2_tags.NewTagsServiceImpl(d.config, tagsKV)
 
 	d.eigwService, err = initServiceWithRetry("EIGW service", func() (*handlers_ec2_eigw.EgressOnlyIGWServiceImpl, error) {
 		return handlers_ec2_eigw.NewEgressOnlyIGWServiceImplWithNATS(d.ctx, d.config, d.natsConn)

spinifex/daemon/daemon_handlers_instance_tags_test.go

@@ -30,7 +30,7 @@ func tagTestDaemon(t *testing.T, instanceID string, initial map[string]string) *
 func tagTestDaemonWithStopped(t *testing.T, instanceID string, initial map[string]string) (*Daemon, *vmmock.StateStore) {
 	t.Helper()
 	d := createTestDaemon(t, sharedNATSURL)
-	d.tagsService = handlers_ec2_tags.NewTagsServiceImplWithStore(d.config, objectstore.NewMemoryObjectStore())
+	d.tagsService = handlers_ec2_tags.NewTagsServiceImplWithStore(d.config, objectstore.NewMemoryObjectStore(), testTagsKV(t))
 
 	stopped := vmmock.New()
 	d.instanceService = handlers_ec2_instance.NewInstanceServiceImpl(

spinifex/daemon/daemon_handlers_test.go

@@ -42,6 +42,7 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/vm"
 	vmmock "github.com/mulgadc/spinifex/spinifex/vm/mock"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -49,6 +50,19 @@ import (
 // testAccountID is the default account ID used in daemon tests.
 const testAccountID = "123456789012"
 
+// testTagsKV returns a tag bucket on a JetStream of its own. The daemon tests'
+// shared NATS server runs without JetStream, and the tag store needs a real one
+// because its writes are compare-and-swap rather than plain puts. A server per
+// test also keeps one test's tags out of the next one's listing.
+func testTagsKV(t *testing.T) jetstream.KeyValue {
+	t.Helper()
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	kv, err := handlers_ec2_tags.GetOrCreateTagsBucket(t.Context(), js)
+	require.NoError(t, err)
+	return kv
+}
+
 // natsRequest sends a NATS request with the X-Account-ID header set.
 func natsRequest(nc *nats.Conn, subject string, data []byte, timeout time.Duration) (*nats.Msg, error) {
 	msg := nats.NewMsg(subject)
@@ -69,7 +83,7 @@ func createFullTestDaemonWithStore(t *testing.T, natsURL string) (*Daemon, *obje
 	daemon.imageService = handlers_ec2_image.NewImageServiceImplWithStore(memStore, cfg.Predastore.Bucket)
 	daemon.volumeService = handlers_ec2_volume.NewVolumeServiceImplWithStore(cfg, memStore, daemon.natsConn)
 	daemon.snapshotService = handlers_ec2_snapshot.NewSnapshotServiceImplWithStore(cfg, memStore, daemon.natsConn)
-	daemon.tagsService = handlers_ec2_tags.NewTagsServiceImplWithStore(cfg, memStore)
+	daemon.tagsService = handlers_ec2_tags.NewTagsServiceImplWithStore(cfg, memStore, testTagsKV(t))
 	wireTestEBSProvider(daemon, memStore)
 	initAccountServiceForTest(t, daemon)
 

spinifex/handlers/ec2/tags/service_impl.go

@@ -1,13 +1,15 @@
 package handlers_ec2_tags
 
 import (
-	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
+	"fmt"
+	"io"
 	"log/slog"
+	"maps"
+	"slices"
 	"strings"
-	"sync"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ec2"
@@ -17,8 +19,10 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/config"
 	"github.com/mulgadc/spinifex/spinifex/filterutil"
 	handlers_ec2_instance "github.com/mulgadc/spinifex/spinifex/handlers/ec2/instance"
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/objectstore"
 	"github.com/mulgadc/spinifex/spinifex/utils"
+	"github.com/nats-io/nats.go/jetstream"
 )
 
 // Ensure TagsServiceImpl implements TagsService.
@@ -27,17 +31,22 @@ var _ TagsService = (*TagsServiceImpl)(nil)
 // Ensure TagsServiceImpl can project instance record tags into the store.
 var _ handlers_ec2_instance.InstanceTagWriter = (*TagsServiceImpl)(nil)
 
-// TagsServiceImpl implements TagsService with S3-backed storage.
-// Tags are stored per-account in S3 (tags/{accountID}/{resourceID}.json),
-// so account scoping is enforced at the storage layer.
+// TagsServiceImpl implements TagsService over a JetStream KV bucket, one entry
+// per resource, scoped by account in the key.
+//
+// The bucket replaced a per-resource S3 object because a tag write is a
+// read-modify-write of the whole set and the handler answers on a cluster-wide
+// queue group: only a compare-and-swap keeps two nodes from losing each other's
+// keys. The object store is still held to read resources written before the
+// move; nothing writes to it.
 type TagsServiceImpl struct {
 	config *config.Config
 	store  objectstore.ObjectStore
-	mutex  sync.RWMutex
+	kv     jetstream.KeyValue
 }
 
 // NewTagsServiceImpl creates a new tags service implementation.
-func NewTagsServiceImpl(cfg *config.Config) *TagsServiceImpl {
+func NewTagsServiceImpl(cfg *config.Config, kv jetstream.KeyValue) *TagsServiceImpl {
 	store := objectstore.NewS3ObjectStoreFromConfig(
 		cfg.Predastore.Host,
 		cfg.Predastore.Region,
@@ -48,14 +57,16 @@ func NewTagsServiceImpl(cfg *config.Config) *TagsServiceImpl {
 	return &TagsServiceImpl{
 		config: cfg,
 		store:  store,
+		kv:     kv,
 	}
 }
 
 // NewTagsServiceImplWithStore creates a tags service with a custom ObjectStore (for testing).
-func NewTagsServiceImplWithStore(cfg *config.Config, store objectstore.ObjectStore) *TagsServiceImpl {
+func NewTagsServiceImplWithStore(cfg *config.Config, store objectstore.ObjectStore, kv jetstream.KeyValue) *TagsServiceImpl {
 	return &TagsServiceImpl{
 		config: cfg,
 		store:  store,
+		kv:     kv,
 	}
 }
 
@@ -80,13 +91,41 @@ func getTagsPrefix(accountID string) string {
 	return "tags/" + accountID + "/"
 }
 
-// getResourceTags retrieves tags for a specific resource from S3.
-func (s *TagsServiceImpl) getResourceTags(ctx context.Context, accountID, resourceID string) (map[string]string, error) {
-	key := getTagsKey(accountID, resourceID)
+// decodeTags decodes a stored tag map into a non-nil map, so a stored JSON
+// "null" (which Unmarshal leaves untouched without error) yields an empty map
+// rather than a nil one a caller would panic on when indexing.
+func decodeTags(data []byte) (map[string]string, error) {
+	tags := make(map[string]string)
+	if len(data) == 0 {
+		return tags, nil
+	}
+	if err := json.Unmarshal(data, &tags); err != nil {
+		return nil, err
+	}
+	if tags == nil {
+		tags = make(map[string]string)
+	}
+	return tags, nil
+}
+
+func encodeTags(tags map[string]string) ([]byte, error) {
+	if tags == nil {
+		tags = make(map[string]string)
+	}
+	return json.Marshal(tags)
+}
+
+// legacyTags reads a resource's tags from the object store the tag index used
+// before it moved to KV. An absent object means the resource was never written
+// there, which is the ordinary case for anything tagged since.
+func (s *TagsServiceImpl) legacyTags(ctx context.Context, accountID, resourceID string) (map[string]string, error) {
+	if s.store == nil {
+		return make(map[string]string), nil
+	}
 
 	result, err := s.store.GetObject(ctx, &s3.GetObjectInput{
 		Bucket: aws.String(s.config.Predastore.Bucket),
-		Key:    aws.String(key),
+		Key:    aws.String(getTagsKey(accountID, resourceID)),
 	})
 	if err != nil {
 		if objectstore.IsNoSuchKeyError(err) {
@@ -96,53 +135,20 @@ func (s *TagsServiceImpl) getResourceTags(ctx context.Context, accountID, resour
 	}
 	defer result.Body.Close()
 
-	// Decode into a non-nil map so a stored JSON "null" body (which Decode leaves
-	// untouched without error) yields an empty map, not a nil map that callers
-	// would panic on when indexing.
-	tags := make(map[string]string)
-	if err := json.NewDecoder(result.Body).Decode(&tags); err != nil {
+	data, err := io.ReadAll(result.Body)
+	if err != nil {
 		return nil, err
 	}
-
-	return tags, nil
+	return decodeTags(data)
 }
 
-// putResourceTags stores tags for a specific resource in S3.
-func (s *TagsServiceImpl) putResourceTags(ctx context.Context, accountID, resourceID string, tags map[string]string) error {
-	key := getTagsKey(accountID, resourceID)
-
-	data, err := json.Marshal(tags)
-	if err != nil {
-		return err
+// deleteLegacyTags removes a resource's pre-move object. Without it a resource
+// whose KV entry is deleted would have its old tags reappear from the fallback
+// read on the next describe.
+func (s *TagsServiceImpl) deleteLegacyTags(ctx context.Context, accountID, resourceID string) error {
+	if s.store == nil {
+		return nil
 	}
-
-	_, err = s.store.PutObject(ctx, &s3.PutObjectInput{
-		Bucket:      aws.String(s.config.Predastore.Bucket),
-		Key:         aws.String(key),
-		Body:        bytes.NewReader(data),
-		ContentType: aws.String("application/json"),
-	})
-
-	return err
-}
-
-// PutResourceTags overwrites the stored tag set for a resource. Used to
-// project an instance record's tags (the source of truth) into the central
-// store so describe-tags agrees with describe-instances.
-func (s *TagsServiceImpl) PutResourceTags(ctx context.Context, accountID, resourceID string, tags map[string]string) error {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-	return s.putResourceTags(ctx, accountID, resourceID, tags)
-}
-
-// DeleteAllTags removes the stored tag object for a resource. Used on
-// instance terminate so describe-tags stops reporting the instance while the
-// terminated record keeps its tags until TTL. Idempotent: a missing object
-// is not an error.
-func (s *TagsServiceImpl) DeleteAllTags(ctx context.Context, accountID, resourceID string) error {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
 	_, err := s.store.DeleteObject(ctx, &s3.DeleteObjectInput{
 		Bucket: aws.String(s.config.Predastore.Bucket),
 		Key:    aws.String(getTagsKey(accountID, resourceID)),
@@ -153,6 +159,36 @@ func (s *TagsServiceImpl) DeleteAllTags(ctx context.Context, accountID, resource
 	return nil
 }
 
+// getResourceTags returns a resource's tags, preferring the KV entry and falling
+// back to the pre-move object for a resource not yet written since.
+func (s *TagsServiceImpl) getResourceTags(ctx context.Context, accountID, resourceID string) (map[string]string, error) {
+	tags, found, err := s.readKVTags(ctx, accountID, resourceID)
+	if err != nil {
+		return nil, err
+	}
+	if found {
+		return tags, nil
+	}
+	return s.legacyTags(ctx, accountID, resourceID)
+}
+
+// PutResourceTags overwrites the stored tag set for a resource. Used to
+// project an instance record's tags (the source of truth) into the central
+// store so describe-tags agrees with describe-instances.
+func (s *TagsServiceImpl) PutResourceTags(ctx context.Context, accountID, resourceID string, tags map[string]string) error {
+	return s.putTags(ctx, accountID, resourceID, tags)
+}
+
+// DeleteAllTags removes the stored tags for a resource. Used on instance
+// terminate so describe-tags stops reporting the instance while the terminated
+// record keeps its tags until TTL. Idempotent: an absent entry is not an error.
+func (s *TagsServiceImpl) DeleteAllTags(ctx context.Context, accountID, resourceID string) error {
+	if err := s.deleteTagsEntry(ctx, accountID, resourceID); err != nil {
+		return err
+	}
+	return s.deleteLegacyTags(ctx, accountID, resourceID)
+}
+
 // CreateTags adds or overwrites tags for the specified resources.
 func (s *TagsServiceImpl) CreateTags(ctx context.Context, input *ec2.CreateTagsInput, accountID string) (*ec2.CreateTagsOutput, error) {
 	if input == nil {
@@ -167,9 +203,6 @@ func (s *TagsServiceImpl) CreateTags(ctx context.Context, input *ec2.CreateTagsI
 		return nil, errors.New(awserrors.ErrorMissingParameter)
 	}
 
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
 	slog.InfoContext(ctx, "CreateTags request", "resources", len(input.Resources), "tags", len(input.Tags))
 
 	for _, resourceID := range input.Resources {
@@ -177,32 +210,71 @@ func (s *TagsServiceImpl) CreateTags(ctx context.Context, input *ec2.CreateTagsI
 			continue
 		}
 
-		// Get existing tags
-		existingTags, err := s.getResourceTags(ctx, accountID, *resourceID)
-		if err != nil {
-			slog.ErrorContext(ctx, "CreateTags failed to get existing tags", "resourceId", *resourceID, "err", err)
-			return nil, errors.New(awserrors.ErrorServerInternal)
-		}
-
-		// Add/update new tags
-		for _, tag := range input.Tags {
-			if tag.Key != nil && tag.Value != nil {
-				existingTags[*tag.Key] = *tag.Value
+		// The merge runs inside the CAS attempt, so a writer that loses the race
+		// re-reads the other writer's keys and adds its own on top of them.
+		err := s.mutateTags(ctx, accountID, *resourceID, func(tags map[string]string) {
+			for _, tag := range input.Tags {
+				if tag.Key != nil && tag.Value != nil {
+					tags[*tag.Key] = *tag.Value
+				}
 			}
-		}
-
-		// Save tags
-		if err := s.putResourceTags(ctx, accountID, *resourceID, existingTags); err != nil {
+		})
+		if err != nil {
 			slog.ErrorContext(ctx, "CreateTags failed to save tags", "resourceId", *resourceID, "err", err)
 			return nil, errors.New(awserrors.ErrorServerInternal)
 		}
 
-		slog.InfoContext(ctx, "CreateTags applied", "resourceId", *resourceID, "tagCount", len(existingTags))
+		slog.InfoContext(ctx, "CreateTags applied", "resourceId", *resourceID, "tagCount", len(input.Tags))
 	}
 
 	return &ec2.CreateTagsOutput{}, nil
 }
 
+// taggedResourceIDs lists every resource the account has tags for, across both
+// the KV bucket and the objects written before the store moved.
+//
+// The two are unioned rather than read in preference order because a resource
+// migrated by a write since the move exists in both, and a resource never
+// touched since exists only in the object store. Sorting keeps a describe's
+// order stable between calls, which the KV listing does not guarantee on its
+// own.
+func (s *TagsServiceImpl) taggedResourceIDs(ctx context.Context, accountID string) ([]string, error) {
+	seen := make(map[string]struct{})
+
+	keys, err := kvutil.Keys(ctx, s.kv)
+	if err != nil && !errors.Is(err, jetstream.ErrNoKeysFound) {
+		return nil, fmt.Errorf("list tag keys: %w", err)
+	}
+	for _, key := range keys {
+		if resourceID, ok := resourceIDFromKVKey(key, accountID); ok {
+			seen[resourceID] = struct{}{}
+		}
+	}
+
+	if s.store != nil {
+		objects, _, lerr := objectstore.ListAll(ctx, s.store, &s3.ListObjectsV2Input{
+			Bucket: aws.String(s.config.Predastore.Bucket),
+			Prefix: aws.String(getTagsPrefix(accountID)),
+		})
+		if lerr != nil {
+			return nil, fmt.Errorf("list legacy tag objects: %w", lerr)
+		}
+		for _, obj := range objects {
+			if obj == nil || obj.Key == nil {
+				continue
+			}
+			resourceID := strings.TrimPrefix(*obj.Key, getTagsPrefix(accountID))
+			resourceID = strings.TrimSuffix(resourceID, ".json")
+			if resourceID == "" {
+				continue
+			}
+			seen[resourceID] = struct{}{}
+		}
+	}
+
+	return slices.Sorted(maps.Keys(seen)), nil
+}
+
 var describeTagsValidFilters = map[string]bool{
 	"resource-id":   true,
 	"resource-type": true,
@@ -222,32 +294,17 @@ func (s *TagsServiceImpl) DescribeTags(ctx context.Context, input *ec2.DescribeT
 		}
 	}
 
-	s.mutex.RLock()
-	defer s.mutex.RUnlock()
-
 	slog.InfoContext(ctx, "DescribeTags request")
 
 	var tags []*ec2.TagDescription
 
-	// List all tag files from S3 scoped to this account
-	objects, _, err := objectstore.ListAll(ctx, s.store, &s3.ListObjectsV2Input{
-		Bucket: aws.String(s.config.Predastore.Bucket),
-		Prefix: aws.String(getTagsPrefix(accountID)),
-	})
+	resourceIDs, err := s.taggedResourceIDs(ctx, accountID)
 	if err != nil {
-		slog.ErrorContext(ctx, "DescribeTags failed to list objects", "err", err)
+		slog.ErrorContext(ctx, "DescribeTags failed to list tagged resources", "err", err)
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
-	// Process each tag file
-	for _, obj := range objects {
-		if obj == nil || obj.Key == nil {
-			continue
-		}
-
-		// Extract resource ID from key (tags/{accountID}/i-xxx.json -> i-xxx)
-		resourceID := strings.TrimPrefix(*obj.Key, getTagsPrefix(accountID))
-		resourceID = strings.TrimSuffix(resourceID, ".json")
+	for _, resourceID := range resourceIDs {
 		resourceType := describeTagsResourceType(resourceID)
 
 		if !filterutil.MatchesAny(filters["resource-id"], resourceID) {
@@ -298,32 +355,26 @@ func (s *TagsServiceImpl) DeleteTags(ctx context.Context, input *ec2.DeleteTagsI
 		return nil, errors.New(awserrors.ErrorMissingParameter)
 	}
 
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
 	slog.InfoContext(ctx, "DeleteTags request", "resources", len(input.Resources), "tags", len(input.Tags))
 
+	remove := utils.RemoveTagsMut(input)
+
 	for _, resourceID := range input.Resources {
 		if resourceID == nil {
 			continue
 		}
 
-		// Get existing tags
-		existingTags, err := s.getResourceTags(ctx, accountID, *resourceID)
+		var remaining int
+		err := s.mutateTags(ctx, accountID, *resourceID, func(tags map[string]string) {
+			remove(tags)
+			remaining = len(tags)
+		})
 		if err != nil {
-			slog.ErrorContext(ctx, "DeleteTags failed to get existing tags", "resourceId", *resourceID, "err", err)
-			return nil, errors.New(awserrors.ErrorServerInternal)
-		}
-
-		utils.RemoveTagsMut(input)(existingTags)
-
-		// Save updated tags
-		if err := s.putResourceTags(ctx, accountID, *resourceID, existingTags); err != nil {
 			slog.ErrorContext(ctx, "DeleteTags failed to save tags", "resourceId", *resourceID, "err", err)
 			return nil, errors.New(awserrors.ErrorServerInternal)
 		}
 
-		slog.InfoContext(ctx, "DeleteTags applied", "resourceId", *resourceID, "remainingTags", len(existingTags))
+		slog.InfoContext(ctx, "DeleteTags applied", "resourceId", *resourceID, "remainingTags", remaining)
 	}
 
 	return &ec2.DeleteTagsOutput{}, nil

spinifex/handlers/ec2/tags/service_impl_test.go

@@ -10,12 +10,26 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	"github.com/mulgadc/spinifex/spinifex/config"
 	"github.com/mulgadc/spinifex/spinifex/objectstore"
+	"github.com/mulgadc/spinifex/spinifex/testutil"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 const testAccountID = "111111111111"
 
+// testTagsBucket returns a bucket backed by an embedded JetStream, so these
+// tests exercise the same compare-and-swap the cluster does rather than a stub
+// that cannot lose a race.
+func testTagsBucket(t *testing.T) jetstream.KeyValue {
+	t.Helper()
+	_, nc, _ := testutil.StartTestJetStream(t)
+	js := testutil.NewJetStream(t, nc)
+	kv, err := GetOrCreateTagsBucket(t.Context(), js)
+	require.NoError(t, err)
+	return kv
+}
+
 // setupTestTagsService creates a tags service with in-memory storage for testing.
 func setupTestTagsService(t *testing.T) (*TagsServiceImpl, *objectstore.MemoryObjectStore) {
 	store := objectstore.NewMemoryObjectStore()
@@ -25,7 +39,7 @@ func setupTestTagsService(t *testing.T) (*TagsServiceImpl, *objectstore.MemoryOb
 		},
 	}
 
-	svc := NewTagsServiceImplWithStore(cfg, store)
+	svc := NewTagsServiceImplWithStore(cfg, store, testTagsBucket(t))
 	return svc, store
 }
 

spinifex/handlers/ec2/tags/store.go

@@ -0,0 +1,121 @@
+package handlers_ec2_tags
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+
+	"github.com/mulgadc/spinifex/spinifex/kvutil"
+	"github.com/nats-io/nats.go/jetstream"
+)
+
+// KVBucketTags holds one entry per tagged resource, keyed {accountID}.{resourceID}.
+// The entry is the resource's whole tag map, so every mutation is a read-modify-
+// write and has to be committed at the revision it was read at.
+const KVBucketTags = "ec2-tags"
+
+// kvTagsHistory is the bucket's history depth. One is enough: nothing reads a
+// superseded tag map, and a deeper history only slows the CAS retries down.
+const kvTagsHistory = 1
+
+// tagCASAttempts is the retry budget for one tag mutation, set well above the
+// shared default because a Terraform apply tags one resource from several
+// parallel calls and a writer can lose the race repeatedly. Spending the budget
+// returns an error the caller retries, which is the right failure for this
+// store: a slow tag write is recoverable and a lost one is not.
+const tagCASAttempts = 12
+
+// tagsKVKey is the bucket key for a resource. The dot is JetStream KV's own
+// subject separator, and neither half can contain one, so the two parts stay
+// recoverable from the key.
+func tagsKVKey(accountID, resourceID string) string {
+	return accountID + "." + resourceID
+}
+
+// resourceIDFromKVKey recovers the resource ID from a bucket key, reporting
+// false for a key belonging to another account or carrying no resource at all.
+func resourceIDFromKVKey(key, accountID string) (string, bool) {
+	rest, found := strings.CutPrefix(key, accountID+".")
+	if !found || rest == "" {
+		return "", false
+	}
+	return rest, true
+}
+
+// GetOrCreateTagsBucket opens the tag bucket, creating it at the cluster's
+// replica count if this is the first node to reach it.
+func GetOrCreateTagsBucket(ctx context.Context, js jetstream.KeyValueManager) (jetstream.KeyValue, error) {
+	return kvutil.GetOrCreateBucket(ctx, js, KVBucketTags, kvTagsHistory)
+}
+
+// readKVTags returns a resource's tags from the bucket, and whether the key
+// exists at all. An absent key is not an error: it means the resource has never
+// been written here, and the caller falls back to the legacy object.
+func (s *TagsServiceImpl) readKVTags(ctx context.Context, accountID, resourceID string) (map[string]string, bool, error) {
+	entry, err := s.kv.Get(ctx, tagsKVKey(accountID, resourceID))
+	if err != nil {
+		if errors.Is(err, jetstream.ErrKeyNotFound) {
+			return nil, false, nil
+		}
+		return nil, false, err
+	}
+	tags, err := decodeTags(entry.Value())
+	if err != nil {
+		return nil, false, err
+	}
+	return tags, true, nil
+}
+
+// mutateTags applies mutate to a resource's tag map under optimistic
+// concurrency, so writers on different nodes cannot lose each other's keys.
+//
+// A resource with no KV entry yet is seeded from its legacy S3 object inside the
+// attempt rather than before it, so a migration failure fails this one call
+// instead of quietly committing a map that has dropped everything written
+// before the store moved.
+func (s *TagsServiceImpl) mutateTags(ctx context.Context, accountID, resourceID string, mutate func(map[string]string)) error {
+	key := tagsKVKey(accountID, resourceID)
+	_, err := kvutil.Update(ctx, s.kv, key, kvutil.CASConfig{
+		Attempts:       tagCASAttempts,
+		CreateIfAbsent: true,
+		Exhausted: func(key string, attempts int) error {
+			return fmt.Errorf("tag store contended for %s after %d attempts", key, attempts)
+		},
+	}, func(tags *map[string]string) (bool, error) {
+		if *tags == nil {
+			seeded, serr := s.legacyTags(ctx, accountID, resourceID)
+			if serr != nil {
+				return false, serr
+			}
+			*tags = seeded
+		}
+		mutate(*tags)
+		return true, nil
+	})
+	return err
+}
+
+// putTags replaces a resource's whole tag set. It still goes through the CAS
+// loop so a concurrent merge is overwritten as a whole rather than interleaved
+// with it, which is what a caller projecting an authoritative record wants.
+func (s *TagsServiceImpl) putTags(ctx context.Context, accountID, resourceID string, tags map[string]string) error {
+	data, err := encodeTags(tags)
+	if err != nil {
+		return err
+	}
+	return kvutil.Put(ctx, s.kv, tagsKVKey(accountID, resourceID), kvutil.CASConfig{
+		Attempts:       tagCASAttempts,
+		CreateIfAbsent: true,
+	}, data)
+}
+
+// deleteTagsEntry removes a resource's bucket entry, treating an absent key as
+// success so teardown stays idempotent.
+func (s *TagsServiceImpl) deleteTagsEntry(ctx context.Context, accountID, resourceID string) error {
+	err := s.kv.Delete(ctx, tagsKVKey(accountID, resourceID))
+	if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
+		return err
+	}
+	return nil
+}

spinifex/handlers/ec2/tags/store_test.go

@@ -0,0 +1,276 @@
+// unexported, and a test that could only reach them through DescribeTags could
+// not tell a migrated resource from one that was never written.
+//
+//test:in-package — the read-through fallback and the key layout are
+package handlers_ec2_tags
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"sync"
+	"testing"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/ec2"
+	"github.com/aws/aws-sdk-go/service/s3"
+	"github.com/mulgadc/spinifex/spinifex/config"
+	"github.com/mulgadc/spinifex/spinifex/objectstore"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// newTestTagsCluster returns n independent service values sharing one bucket and
+// one object store — the shape a multi-node cluster presents, where each node
+// has its own service struct and only the store is common.
+func newTestTagsCluster(t *testing.T, n int) []*TagsServiceImpl {
+	t.Helper()
+	kv := testTagsBucket(t)
+	store := objectstore.NewMemoryObjectStore()
+	cfg := &config.Config{Predastore: config.PredastoreConfig{Bucket: "test-bucket"}}
+
+	nodes := make([]*TagsServiceImpl, n)
+	for i := range nodes {
+		nodes[i] = NewTagsServiceImplWithStore(cfg, store, kv)
+	}
+	return nodes
+}
+
+// seedLegacyTags writes a resource's tags where the store kept them before the
+// move to KV, so a test can stand in for a cluster upgraded with tags already in
+// place.
+func seedLegacyTags(t *testing.T, store *objectstore.MemoryObjectStore, bucket, accountID, resourceID string, tags map[string]string) {
+	t.Helper()
+	data, err := json.Marshal(tags)
+	require.NoError(t, err)
+	_, err = store.PutObject(t.Context(), &s3.PutObjectInput{
+		Bucket: aws.String(bucket),
+		Key:    aws.String(getTagsKey(accountID, resourceID)),
+		Body:   bytes.NewReader(data),
+	})
+	require.NoError(t, err)
+}
+
+// tagMap reads a describe back into a map, so an assertion names the keys it
+// wants rather than an index into a slice whose order is not promised.
+func tagMap(t *testing.T, svc *TagsServiceImpl, resourceID string) map[string]string {
+	t.Helper()
+	out, err := svc.DescribeTags(t.Context(), &ec2.DescribeTagsInput{
+		Filters: []*ec2.Filter{
+			{Name: aws.String("resource-id"), Values: []*string{aws.String(resourceID)}},
+		},
+	}, testAccountID)
+	require.NoError(t, err)
+
+	got := make(map[string]string, len(out.Tags))
+	for _, tag := range out.Tags {
+		got[aws.StringValue(tag.Key)] = aws.StringValue(tag.Value)
+	}
+	return got
+}
+
+// This is the gap itself. Six writers each adding one key used to leave four.
+//
+// Each writer gets its own service value, because that is what makes the test
+// able to fail: the old code serialised the merge on a mutex held by the service
+// struct, so every caller inside one process was already safe and only a second
+// node could lose an update. Driving one service from six goroutines would pass
+// against the very code this test exists to catch.
+func TestCreateTags_ConcurrentWritersOnSeparateNodesAllLand(t *testing.T) {
+	const resourceID = "sg-race"
+	const writers = 6
+
+	nodes := newTestTagsCluster(t, writers)
+
+	var wg sync.WaitGroup
+	errs := make([]error, writers)
+	for i, node := range nodes {
+		wg.Go(func() {
+			_, errs[i] = node.CreateTags(context.Background(), &ec2.CreateTagsInput{
+				Resources: []*string{aws.String(resourceID)},
+				Tags: []*ec2.Tag{
+					{Key: aws.String(fmt.Sprintf("k%d", i)), Value: aws.String(fmt.Sprintf("v%d", i))},
+				},
+			}, testAccountID)
+		})
+	}
+	wg.Wait()
+
+	for i, err := range errs {
+		require.NoError(t, err, "writer %d", i)
+	}
+
+	got := tagMap(t, nodes[0], resourceID)
+	require.Len(t, got, writers)
+	for i := range writers {
+		assert.Equal(t, fmt.Sprintf("v%d", i), got[fmt.Sprintf("k%d", i)])
+	}
+}
+
+// A concurrent delete and add must not resurrect the deleted key or drop the
+// added one, which is the same lost update seen from the other direction.
+func TestDeleteTags_ConcurrentWithCreateKeepsBothDecisions(t *testing.T) {
+	nodes := newTestTagsCluster(t, 2)
+	svc, other := nodes[0], nodes[1]
+	const resourceID = "sg-mixed"
+
+	_, err := svc.CreateTags(t.Context(), &ec2.CreateTagsInput{
+		Resources: []*string{aws.String(resourceID)},
+		Tags:      []*ec2.Tag{{Key: aws.String("doomed"), Value: aws.String("1")}},
+	}, testAccountID)
+	require.NoError(t, err)
+
+	var wg sync.WaitGroup
+	var addErr, delErr error
+	wg.Go(func() {
+		_, addErr = svc.CreateTags(context.Background(), &ec2.CreateTagsInput{
+			Resources: []*string{aws.String(resourceID)},
+			Tags:      []*ec2.Tag{{Key: aws.String("kept"), Value: aws.String("2")}},
+		}, testAccountID)
+	})
+	wg.Go(func() {
+		_, delErr = other.DeleteTags(context.Background(), &ec2.DeleteTagsInput{
+			Resources: []*string{aws.String(resourceID)},
+			Tags:      []*ec2.Tag{{Key: aws.String("doomed")}},
+		}, testAccountID)
+	})
+	wg.Wait()
+	require.NoError(t, addErr)
+	require.NoError(t, delErr)
+
+	got := tagMap(t, svc, resourceID)
+	assert.Equal(t, map[string]string{"kept": "2"}, got)
+}
+
+// A cluster upgraded with tags already written must not lose them the first time
+// something writes a new one.
+func TestCreateTags_ReadsThroughToPreMoveTags(t *testing.T) {
+	svc, store := setupTestTagsService(t)
+	const resourceID = "vpc-legacy"
+	seedLegacyTags(t, store, svc.config.Predastore.Bucket, testAccountID, resourceID, map[string]string{
+		"Name": "old", "owner": "platform",
+	})
+
+	assert.Equal(t, map[string]string{"Name": "old", "owner": "platform"}, tagMap(t, svc, resourceID))
+
+	_, err := svc.CreateTags(t.Context(), &ec2.CreateTagsInput{
+		Resources: []*string{aws.String(resourceID)},
+		Tags:      []*ec2.Tag{{Key: aws.String("env"), Value: aws.String("prod")}},
+	}, testAccountID)
+	require.NoError(t, err)
+
+	assert.Equal(t, map[string]string{"Name": "old", "owner": "platform", "env": "prod"}, tagMap(t, svc, resourceID))
+
+	// The resource is migrated: its tags now answer from the bucket, so a second
+	// write no longer depends on the object store being reachable.
+	_, found, err := svc.readKVTags(t.Context(), testAccountID, resourceID)
+	require.NoError(t, err)
+	assert.True(t, found)
+}
+
+// DeleteTags on a resource that has never been written since the move has to see
+// the pre-move keys too, or a delete silently does nothing.
+func TestDeleteTags_RemovesAPreMoveTag(t *testing.T) {
+	svc, store := setupTestTagsService(t)
+	const resourceID = "subnet-legacy"
+	seedLegacyTags(t, store, svc.config.Predastore.Bucket, testAccountID, resourceID, map[string]string{
+		"Name": "old", "temp": "yes",
+	})
+
+	_, err := svc.DeleteTags(t.Context(), &ec2.DeleteTagsInput{
+		Resources: []*string{aws.String(resourceID)},
+		Tags:      []*ec2.Tag{{Key: aws.String("temp")}},
+	}, testAccountID)
+	require.NoError(t, err)
+
+	assert.Equal(t, map[string]string{"Name": "old"}, tagMap(t, svc, resourceID))
+}
+
+// A resource in one store, the other, or both is listed once.
+func TestDescribeTags_UnionsBothStoresWithoutDuplicating(t *testing.T) {
+	svc, store := setupTestTagsService(t)
+	bucket := svc.config.Predastore.Bucket
+
+	seedLegacyTags(t, store, bucket, testAccountID, "vpc-only-legacy", map[string]string{"a": "1"})
+	seedLegacyTags(t, store, bucket, testAccountID, "vpc-both", map[string]string{"a": "stale"})
+
+	_, err := svc.CreateTags(t.Context(), &ec2.CreateTagsInput{
+		Resources: []*string{aws.String("vpc-only-kv")},
+		Tags:      []*ec2.Tag{{Key: aws.String("b"), Value: aws.String("2")}},
+	}, testAccountID)
+	require.NoError(t, err)
+	_, err = svc.CreateTags(t.Context(), &ec2.CreateTagsInput{
+		Resources: []*string{aws.String("vpc-both")},
+		Tags:      []*ec2.Tag{{Key: aws.String("a"), Value: aws.String("fresh")}},
+	}, testAccountID)
+	require.NoError(t, err)
+
+	ids, err := svc.taggedResourceIDs(t.Context(), testAccountID)
+	require.NoError(t, err)
+	assert.Equal(t, []string{"vpc-both", "vpc-only-kv", "vpc-only-legacy"}, ids)
+
+	// The migrated resource answers from the bucket, so the object it was
+	// seeded from cannot shadow the value written since.
+	assert.Equal(t, map[string]string{"a": "fresh"}, tagMap(t, svc, "vpc-both"))
+}
+
+// Deleting a migrated resource's tags has to clear the pre-move object too,
+// or the fallback read hands the old set straight back.
+func TestDeleteAllTags_ClearsBothStores(t *testing.T) {
+	svc, store := setupTestTagsService(t)
+	const resourceID = "i-terminated"
+	seedLegacyTags(t, store, svc.config.Predastore.Bucket, testAccountID, resourceID, map[string]string{"Name": "gone"})
+
+	_, err := svc.CreateTags(t.Context(), &ec2.CreateTagsInput{
+		Resources: []*string{aws.String(resourceID)},
+		Tags:      []*ec2.Tag{{Key: aws.String("env"), Value: aws.String("test")}},
+	}, testAccountID)
+	require.NoError(t, err)
+
+	require.NoError(t, svc.DeleteAllTags(t.Context(), testAccountID, resourceID))
+
+	assert.Empty(t, tagMap(t, svc, resourceID))
+
+	ids, err := svc.taggedResourceIDs(t.Context(), testAccountID)
+	require.NoError(t, err)
+	assert.NotContains(t, ids, resourceID)
+
+	// Idempotent: teardown runs more than once.
+	require.NoError(t, svc.DeleteAllTags(t.Context(), testAccountID, resourceID))
+}
+
+// One account's keys must not be readable through another's listing, which the
+// key layout is the only thing enforcing now the account is a key prefix rather
+// than a directory.
+func TestTaggedResourceIDs_ScopesToTheAccount(t *testing.T) {
+	svc, _ := setupTestTagsService(t)
+	const otherAccount = "222222222222"
+
+	_, err := svc.CreateTags(t.Context(), &ec2.CreateTagsInput{
+		Resources: []*string{aws.String("vpc-mine")},
+		Tags:      []*ec2.Tag{{Key: aws.String("a"), Value: aws.String("1")}},
+	}, testAccountID)
+	require.NoError(t, err)
+	_, err = svc.CreateTags(t.Context(), &ec2.CreateTagsInput{
+		Resources: []*string{aws.String("vpc-theirs")},
+		Tags:      []*ec2.Tag{{Key: aws.String("a"), Value: aws.String("1")}},
+	}, otherAccount)
+	require.NoError(t, err)
+
+	ids, err := svc.taggedResourceIDs(t.Context(), testAccountID)
+	require.NoError(t, err)
+	assert.Equal(t, []string{"vpc-mine"}, ids)
+}
+
+// PutResourceTags replaces the whole set rather than merging, because its caller
+// is projecting a record that is already the source of truth.
+func TestPutResourceTags_ReplacesTheWholeSet(t *testing.T) {
+	svc, _ := setupTestTagsService(t)
+	const resourceID = "i-projected"
+
+	require.NoError(t, svc.PutResourceTags(t.Context(), testAccountID, resourceID, map[string]string{"a": "1", "b": "2"}))
+	require.NoError(t, svc.PutResourceTags(t.Context(), testAccountID, resourceID, map[string]string{"a": "3"}))
+
+	assert.Equal(t, map[string]string{"a": "3"}, tagMap(t, svc, resourceID))
+}

tests/integration/daemon_lite.go

@@ -21,6 +21,7 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/objectstore"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/nats-io/nats.go"
+	"github.com/nats-io/nats.go/jetstream"
 	"github.com/stretchr/testify/require"
 )
 
@@ -108,7 +109,12 @@ func StartDaemonLite(t *testing.T, gw *Gateway, opts ...DaemonLiteOption) *Daemo
 	}
 
 	keySvc := handlers_ec2_key.NewKeyServiceImplWithStore(memStore, cfg.Predastore.Bucket)
-	tagsSvc := handlers_ec2_tags.NewTagsServiceImplWithStore(cfg, memStore)
+
+	tagsJS, err := jetstream.New(nc)
+	require.NoError(t, err, "jetstream handle for the tag store")
+	tagsKV, err := handlers_ec2_tags.GetOrCreateTagsBucket(t.Context(), tagsJS)
+	require.NoError(t, err, "tag store bucket")
+	tagsSvc := handlers_ec2_tags.NewTagsServiceImplWithStore(cfg, memStore, tagsKV)
 
 	vpcSvc, err := handlers_ec2_vpc.NewVPCServiceImplWithNATS(t.Context(), cfg, nc)
 	require.NoError(t, err, "construct VPC service")