diff --git a/spinifex/daemon/jetstream.go b/spinifex/daemon/jetstream.go index 9ae004d47..210808f75 100644 --- a/spinifex/daemon/jetstream.go +++ b/spinifex/daemon/jetstream.go @@ -182,34 +182,31 @@ func (m *JetStreamManager) InitKVBucket() error { return err } -// InitClusterStateBucket initializes the cluster-state KV bucket, creating it if it doesn't exist. -func (m *JetStreamManager) InitClusterStateBucket() error { - ctx := context.Background() - kv, err := m.js.KeyValue(ctx, ClusterStateBucket) - if err != nil { - if errors.Is(err, jetstream.ErrBucketNotFound) { - slog.Debug("Creating JetStream KV bucket", "bucket", ClusterStateBucket, "replicas", m.replicas) - kv, err = m.js.CreateKeyValue(ctx, jetstream.KeyValueConfig{ - Bucket: ClusterStateBucket, - Description: "Spinifex cluster state (heartbeats, shutdown markers, service maps)", - History: 1, - Replicas: m.replicas, - TTL: 1 * time.Hour, - }) - if err != nil { - return err - } - } else { - return err - } - } else { - slog.Debug("Connected to existing JetStream KV bucket", "bucket", ClusterStateBucket) +// clusterStateConfig describes the cluster-state bucket. +func clusterStateConfig(replicas int) kvstore.Config { + return kvstore.Config{ + Name: ClusterStateBucket, + Description: "Spinifex cluster state (heartbeats, shutdown markers, service maps)", + History: 1, + Replicas: replicas, + TTL: 1 * time.Hour, + OnOpen: func(ctx context.Context, kv jetstream.KeyValue) error { + return migrate.DefaultRegistry.RunKV(ctx, ClusterStateBucket, kv, ClusterStateBucketVersion) + }, + Missing: "cluster state KV bucket not initialized", } +} - m.clusterKV = kv - if err := migrate.DefaultRegistry.RunKV(ctx, ClusterStateBucket, kv, ClusterStateBucketVersion); err != nil { - return fmt.Errorf("migrate %s: %w", ClusterStateBucket, err) +// InitClusterStateBucket initializes the cluster-state KV bucket, creating it if it doesn't exist. +// +// The handle is kept rather than the Bucket: this one holds several unrelated +// record types under one namespace, and its callers still read it as raw keys. +func (m *JetStreamManager) InitClusterStateBucket() error { + kv, err := kvstore.NewBucket(m.js, clusterStateConfig(m.replicas)).KV(context.Background()) + if err != nil { + return err } + m.clusterKV = kv return nil } @@ -679,7 +676,7 @@ func (m *JetStreamManager) ClaimStoppedInstance(instanceID string) (*vm.VM, erro claimed := *record claimed.Spec.DesiredState = vm.DesiredRunning - if err := m.records.CompareAndSet(context.Background(), key, &claimed, rev); err != nil { + if _, err := m.records.CompareAndSet(context.Background(), key, &claimed, rev); err != nil { if errors.Is(err, kvstore.ErrConflict) { return nil, vm.ErrStoppedInstanceClaimed } diff --git a/spinifex/gateway/bedrock/guardrail.go b/spinifex/gateway/bedrock/guardrail.go index 070a4b888..a24aa3dd9 100644 --- a/spinifex/gateway/bedrock/guardrail.go +++ b/spinifex/gateway/bedrock/guardrail.go @@ -148,7 +148,7 @@ func (s *GuardrailStore) put(ctx context.Context, rec GuardrailRecord) error { // concurrent mutations of the same guardrail never silently clobber one // another. A lost race is reported as a retryable ConflictException. func (s *GuardrailStore) update(ctx context.Context, key string, rec GuardrailRecord, rev uint64) error { - err := s.store.CompareAndSet(ctx, key, &rec, rev) + _, err := s.store.CompareAndSet(ctx, key, &rec, rev) if errors.Is(err, kvstore.ErrConflict) { return awserrors.Errorf(awserrors.ErrorConflictException, "bedrock: guardrail %s was modified concurrently; retry the request", rec.GuardrailID) diff --git a/spinifex/gateway/bedrock/provisioned.go b/spinifex/gateway/bedrock/provisioned.go index 01c0262ff..ed4832024 100644 --- a/spinifex/gateway/bedrock/provisioned.go +++ b/spinifex/gateway/bedrock/provisioned.go @@ -128,7 +128,7 @@ func (s *ProvisionedStore) getRevision(ctx context.Context, key string) (Provisi // race is reported as a retryable ConflictException rather than silently // clobbering the winner's write. func (s *ProvisionedStore) update(ctx context.Context, key string, rec ProvisionedModelRecord, rev uint64) error { - err := s.store.CompareAndSet(ctx, key, &rec, rev) + _, err := s.store.CompareAndSet(ctx, key, &rec, rev) if errors.Is(err, kvstore.ErrConflict) { return awserrors.Errorf(awserrors.ErrorConflictException, "bedrock: provisioned throughput %s was modified concurrently; retry the request", rec.ARN) diff --git a/spinifex/handlers/acm/store.go b/spinifex/handlers/acm/store.go index 5f2f4c8fb..f4a5043c1 100644 --- a/spinifex/handlers/acm/store.go +++ b/spinifex/handlers/acm/store.go @@ -520,7 +520,7 @@ func (s *Store) AcquireLease(ctx context.Context, certArn, holderID string, ttl } rec.LeaseHolder = holderID rec.LeaseExpiresAt = now.Add(ttl) - if err := s.store.CompareAndSet(ctx, key, rec, rev); err != nil { + if _, err := s.store.CompareAndSet(ctx, key, rec, rev); err != nil { if errors.Is(err, kvstore.ErrConflict) { // Lost the race to a concurrent acquirer between Get and Update; // the caller skips this tick rather than retrying immediately. @@ -552,7 +552,7 @@ func (s *Store) ReleaseLease(ctx context.Context, certArn, holderID string) erro } rec.LeaseHolder = "" rec.LeaseExpiresAt = time.Time{} - if err := s.store.CompareAndSet(ctx, key, rec, rev); err != nil { + if _, err := s.store.CompareAndSet(ctx, key, rec, rev); err != nil { if errors.Is(err, kvstore.ErrConflict) { return nil // record changed concurrently; nothing to clean up } diff --git a/spinifex/handlers/bedrock/eviction_test.go b/spinifex/handlers/bedrock/eviction_test.go index c7e4757ee..bcb926fb3 100644 --- a/spinifex/handlers/bedrock/eviction_test.go +++ b/spinifex/handlers/bedrock/eviction_test.go @@ -234,5 +234,6 @@ func pinEndpoint(t *testing.T, s *Service, modelID string) { require.NoError(t, err) require.True(t, found) rec.Pinned = true - require.NoError(t, s.store.CompareAndSet(t.Context(), key, &rec, rev)) + _, casErr := s.store.CompareAndSet(t.Context(), key, &rec, rev) + require.NoError(t, casErr) } diff --git a/spinifex/handlers/bedrock/reaper.go b/spinifex/handlers/bedrock/reaper.go index 3e79c2de9..8acfa111d 100644 --- a/spinifex/handlers/bedrock/reaper.go +++ b/spinifex/handlers/bedrock/reaper.go @@ -286,7 +286,7 @@ func (r *Reaper) persistObservation(ctx context.Context, prev, next EndpointReco return nil } next.Generation = current.Generation + 1 - if err := r.svc.store.CompareAndSet(ctx, key, &next, rev); err != nil { + if _, err := r.svc.store.CompareAndSet(ctx, key, &next, rev); err != nil { if errors.Is(err, kvstore.ErrConflict) { return nil } diff --git a/spinifex/handlers/bedrock/reaper_test.go b/spinifex/handlers/bedrock/reaper_test.go index 7ad9519cd..7d39f73a1 100644 --- a/spinifex/handlers/bedrock/reaper_test.go +++ b/spinifex/handlers/bedrock/reaper_test.go @@ -130,7 +130,8 @@ func (f *reaperFixture) age(t *testing.T, readyAgo, activeAgo time.Duration) End if activeAgo > 0 { rec.LastActiveAt = now.Add(-activeAgo) } - require.NoError(t, f.store.CompareAndSet(t.Context(), key, &rec, rev)) + _, casErr := f.store.CompareAndSet(t.Context(), key, &rec, rev) + require.NoError(t, casErr) return rec } @@ -178,7 +179,8 @@ func TestReaper_NeverReapsPinnedEndpoint(t *testing.T) { rec, rev, _, err := f.store.getRevision(t.Context(), key) require.NoError(t, err) rec.Pinned = true - require.NoError(t, f.store.CompareAndSet(t.Context(), key, &rec, rev)) + _, casErr := f.store.CompareAndSet(t.Context(), key, &rec, rev) + require.NoError(t, casErr) require.NoError(t, f.reaper.sweepOnce(t.Context())) @@ -292,7 +294,8 @@ func TestReaper_PinnedEndpointNotReapedOnScrapeFailure(t *testing.T) { require.NoError(t, err) require.True(t, found) rec.Pinned = true - require.NoError(t, f.store.CompareAndSet(t.Context(), key, &rec, rev)) + _, casErr := f.store.CompareAndSet(t.Context(), key, &rec, rev) + require.NoError(t, casErr) f.transport.breakScrapes() for range maxScrapeFailures + 2 { diff --git a/spinifex/handlers/bedrock/service.go b/spinifex/handlers/bedrock/service.go index f873b5964..b0ff6a0eb 100644 --- a/spinifex/handlers/bedrock/service.go +++ b/spinifex/handlers/bedrock/service.go @@ -351,7 +351,7 @@ func (s *Service) runLaunch(ctx context.Context, key string, rec EndpointRecord, slog.ErrorContext(ctx, "bedrock: re-read endpoint before READY write failed", "group", spec.GroupID, "err", gerr) return } - if err := s.store.CompareAndSet(ctx, key, &rec, rev); err != nil { + if _, err := s.store.CompareAndSet(ctx, key, &rec, rev); err != nil { slog.ErrorContext(ctx, "bedrock: CAS write of READY state failed", "group", spec.GroupID, "err", err) } } @@ -461,7 +461,7 @@ func (s *Service) Delete(ctx context.Context, in *DeleteEndpointInput, _ string) } rec.State = StateDraining rec.Generation++ - if err := s.store.CompareAndSet(ctx, key, &rec, rev); err != nil { + if _, err := s.store.CompareAndSet(ctx, key, &rec, rev); err != nil { return nil, fmt.Errorf("bedrock: mark endpoint %s draining: %w", in.ModelID, err) } } diff --git a/spinifex/handlers/bedrock/store_test.go b/spinifex/handlers/bedrock/store_test.go index 6fb65735e..bb0849571 100644 --- a/spinifex/handlers/bedrock/store_test.go +++ b/spinifex/handlers/bedrock/store_test.go @@ -47,11 +47,12 @@ func TestStore_CreateGetUpdateDelete(t *testing.T) { got.State = StateReady got.Generation = 2 - require.NoError(t, store.CompareAndSet(t.Context(), key, &got, gotRev)) + _, casErr := store.CompareAndSet(t.Context(), key, &got, gotRev) + require.NoError(t, casErr) // The CAS update must fail against the now-stale revision — this is what // stops a launch goroutine and a concurrent writer stomping each other. - err = store.CompareAndSet(t.Context(), key, &got, gotRev) + _, err = store.CompareAndSet(t.Context(), key, &got, gotRev) assert.ErrorIs(t, err, kvstore.ErrConflict) require.NoError(t, store.Purge(t.Context(), key)) diff --git a/spinifex/handlers/ec2/snapshot/service_impl.go b/spinifex/handlers/ec2/snapshot/service_impl.go index 23fb9a676..c3bd9cfbf 100644 --- a/spinifex/handlers/ec2/snapshot/service_impl.go +++ b/spinifex/handlers/ec2/snapshot/service_impl.go @@ -2,7 +2,6 @@ package handlers_ec2_snapshot import ( "context" - "encoding/json" "errors" "fmt" "log/slog" @@ -23,7 +22,7 @@ import ( "github.com/mulgadc/spinifex/spinifex/ebsprovider" "github.com/mulgadc/spinifex/spinifex/filterutil" "github.com/mulgadc/spinifex/spinifex/handlers/ec2/volumestate" - "github.com/mulgadc/spinifex/spinifex/kvutil" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/migrate" "github.com/mulgadc/spinifex/spinifex/objectstore" "github.com/mulgadc/spinifex/spinifex/types" @@ -38,7 +37,13 @@ var _ SnapshotService = (*SnapshotServiceImpl)(nil) const ( KVBucketVolumeSnapshots = "spinifex-volume-snapshots" KVBucketVolumeSnapshotsVersion = 1 - snapshotCleanupTimeout = 5 * time.Second + KVBucketVolumeSnapshotsHistory = 10 + + // snapshotRefCASAttempts bounds one index write's retry on a revision + // conflict. Every snapshot of a volume shares one record, so simultaneous + // creates on one volume contend once each; the default five sits below that. + snapshotRefCASAttempts = 25 + snapshotCleanupTimeout = 5 * time.Second ) // SnapshotServiceImpl implements SnapshotService with S3-backed storage. @@ -46,7 +51,7 @@ type SnapshotServiceImpl struct { config *config.Config store objectstore.ObjectStore natsConn *nats.Conn - snapKV jetstream.KeyValue + snapRefs *kvstore.Store[[]string] mutex sync.RWMutex metadata *ebsmetadata.Store provider ebsprovider.EBSProvider @@ -84,13 +89,13 @@ func NewSnapshotServiceImplWithNATS(ctx context.Context, cfg *config.Config, nat return nil, nil, fmt.Errorf("create JetStream: %w", err) } - kv, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketVolumeSnapshots, 10) + refs := kvstore.New[[]string](js, snapshotRefsConfig()) + // Opened eagerly: the handle is returned to the caller, which wires it into + // the volume service, so it has to exist before this function returns. + kv, err := refs.KV(ctx) if err != nil { return nil, nil, fmt.Errorf("create KV bucket %s: %w", KVBucketVolumeSnapshots, err) } - if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketVolumeSnapshots, kv, KVBucketVolumeSnapshotsVersion); err != nil { - return nil, nil, fmt.Errorf("migrate %s: %w", KVBucketVolumeSnapshots, err) - } slog.Info("Snapshot service initialized with JetStream KV", "bucket", KVBucketVolumeSnapshots) @@ -98,7 +103,7 @@ func NewSnapshotServiceImplWithNATS(ctx context.Context, cfg *config.Config, nat config: cfg, store: store, natsConn: natsConn, - snapKV: kv, + snapRefs: refs, metadata: ebsmetadata.NewStore(store, cfg.Predastore.Bucket), }, kv, nil } @@ -113,11 +118,25 @@ func NewSnapshotServiceImplWithStore(cfg *config.Config, store objectstore.Objec metadata: ebsmetadata.NewStore(store, cfg.Predastore.Bucket), } if len(snapshotKV) > 0 { - svc.snapKV = snapshotKV[0] + svc.snapRefs = kvstore.Over[[]string](nil, snapshotKV[0], snapshotRefsConfig()) } return svc } +// snapshotRefsConfig describes the volume-to-snapshot index bucket. The value +// is the list of snapshot IDs taken from one volume, keyed by volume ID. +func snapshotRefsConfig() kvstore.Config { + return kvstore.Config{ + Name: KVBucketVolumeSnapshots, + History: KVBucketVolumeSnapshotsHistory, + Missing: "volume snapshot KV bucket not initialized", + Attempts: snapshotRefCASAttempts, + OnOpen: func(ctx context.Context, kv jetstream.KeyValue) error { + return migrate.DefaultRegistry.RunKV(ctx, KVBucketVolumeSnapshots, kv, KVBucketVolumeSnapshotsVersion) + }, + } +} + // getSnapshotConfig reads the caller's own snapshot document. A snapshot // outside the caller's prefix is absent, so ownership needs no second check. func (s *SnapshotServiceImpl) getSnapshotConfig(ctx context.Context, accountID, snapshotID string) (ebsmetadata.Snapshot, error) { @@ -781,58 +800,21 @@ func (s *SnapshotServiceImpl) CopySnapshot(ctx context.Context, input *ec2.CopyS // addSnapshotRef adds snapshotID to the volume's snapshot list in KV. // Uses CAS (Create/Update with revision) to prevent lost updates under concurrency. func (s *SnapshotServiceImpl) addSnapshotRef(ctx context.Context, volumeID, snapshotID string) error { - if s.snapKV == nil { + if s.snapRefs == nil { slog.Debug("addSnapshotRef: snapshotKV is nil, skipping", "volumeId", volumeID, "snapshotId", snapshotID) return nil } - const maxRetries = 5 - for attempt := range maxRetries { - entry, err := s.snapKV.Get(ctx, volumeID) - var snapshots []string - - if err != nil { - if !errors.Is(err, jetstream.ErrKeyNotFound) { - return fmt.Errorf("addSnapshotRef: failed to get KV key %s: %w", volumeID, err) - } - // Key doesn't exist yet — create with just this snapshot - data, err := json.Marshal([]string{snapshotID}) - if err != nil { - return fmt.Errorf("addSnapshotRef: failed to marshal snapshot list: %w", err) - } - if _, err := s.snapKV.Create(ctx, volumeID, data); err != nil { - if attempt < maxRetries-1 { - continue // concurrent Create/Update — retry - } - return fmt.Errorf("addSnapshotRef: failed to create KV key %s: %w", volumeID, err) - } - slog.Info("addSnapshotRef: added snapshot ref", "volumeId", volumeID, "snapshotId", snapshotID) - return nil - } - - if err := json.Unmarshal(entry.Value(), &snapshots); err != nil { - return fmt.Errorf("addSnapshotRef: failed to unmarshal KV value for %s: %w", volumeID, err) - } - - snapshots = append(snapshots, snapshotID) - - data, err := json.Marshal(snapshots) - if err != nil { - return fmt.Errorf("addSnapshotRef: failed to marshal snapshot list: %w", err) - } - - if _, err := s.snapKV.Update(ctx, volumeID, data, entry.Revision()); err != nil { - if attempt < maxRetries-1 { - continue // concurrent update — retry - } - return fmt.Errorf("addSnapshotRef: failed to update KV key %s: %w", volumeID, err) - } - - slog.Info("addSnapshotRef: added snapshot ref", "volumeId", volumeID, "snapshotId", snapshotID) - return nil + err := s.snapRefs.Upsert(ctx, volumeID, func(snapshots *[]string) (bool, error) { + *snapshots = append(*snapshots, snapshotID) + return true, nil + }) + if err != nil { + return fmt.Errorf("addSnapshotRef: failed to record %s against volume %s: %w", snapshotID, volumeID, err) } - return fmt.Errorf("addSnapshotRef: exhausted retries for KV key %s", volumeID) + slog.Info("addSnapshotRef: added snapshot ref", "volumeId", volumeID, "snapshotId", snapshotID) + return nil } // removeSnapshotRefForCleanup removes a reference with a bounded context that @@ -847,77 +829,89 @@ func (s *SnapshotServiceImpl) removeSnapshotRefForCleanup(ctx context.Context, v // Deletes the key if the list becomes empty. // Uses CAS (Update with revision) to prevent lost updates under concurrency. func (s *SnapshotServiceImpl) removeSnapshotRef(ctx context.Context, volumeID, snapshotID string) error { - if s.snapKV == nil { + if s.snapRefs == nil { slog.Debug("removeSnapshotRef: snapshotKV is nil, skipping", "volumeId", volumeID, "snapshotId", snapshotID) return nil } - const maxRetries = 5 - for attempt := range maxRetries { - entry, err := s.snapKV.Get(ctx, volumeID) - if err != nil { - if errors.Is(err, jetstream.ErrKeyNotFound) { - return nil - } - return fmt.Errorf("removeSnapshotRef: failed to get KV key %s: %w", volumeID, err) - } - - var snapshots []string - if err := json.Unmarshal(entry.Value(), &snapshots); err != nil { - return fmt.Errorf("removeSnapshotRef: failed to unmarshal KV value for %s: %w", volumeID, err) - } - - filtered := snapshots[:0] - for _, snap := range snapshots { + // A removal that finds nothing commits nothing: an identical record would + // still bump the revision, waking every watcher and costing the concurrent + // writers a CAS attempt each. + var removed, nowEmpty bool + err := s.snapRefs.Mutate(ctx, volumeID, func(snapshots *[]string) (bool, error) { + before := len(*snapshots) + filtered := (*snapshots)[:0] + for _, snap := range *snapshots { if snap != snapshotID { filtered = append(filtered, snap) } } - - if len(filtered) == 0 { - if err := s.snapKV.Delete(ctx, volumeID); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { - return fmt.Errorf("removeSnapshotRef: failed to delete KV key %s: %w", volumeID, err) - } - } else { - data, err := json.Marshal(filtered) - if err != nil { - return fmt.Errorf("removeSnapshotRef: failed to marshal snapshot list: %w", err) - } - if _, err := s.snapKV.Update(ctx, volumeID, data, entry.Revision()); err != nil { - if attempt < maxRetries-1 { - continue // concurrent update — retry - } - return fmt.Errorf("removeSnapshotRef: failed to update KV key %s: %w", volumeID, err) - } - } - - slog.Info("removeSnapshotRef: removed snapshot ref", "volumeId", volumeID, "snapshotId", snapshotID) + *snapshots = filtered + removed = len(filtered) != before + nowEmpty = len(filtered) == 0 + return removed, nil + }) + if errors.Is(err, kvstore.ErrNotFound) { return nil } + if err != nil { + return fmt.Errorf("removeSnapshotRef: failed to drop %s from volume %s: %w", snapshotID, volumeID, err) + } - return fmt.Errorf("removeSnapshotRef: exhausted retries for KV key %s", volumeID) + // Still runs when nothing was removed, so a key left empty by a lost drop + // race is not stranded until the volume gains another snapshot. + if nowEmpty { + if err := s.dropEmptySnapshotRefs(ctx, volumeID); err != nil { + return err + } + } + + if removed { + slog.Info("removeSnapshotRef: removed snapshot ref", "volumeId", volumeID, "snapshotId", snapshotID) + } + return nil +} + +// dropEmptySnapshotRefs deletes a volume's index key once nothing references it, +// so a volume that has lost its last snapshot leaves no key behind. +// +// The delete is guarded on the revision the emptiness was read at, so a +// CreateSnapshot landing in between keeps its key rather than having it deleted +// out from under it. +func (s *SnapshotServiceImpl) dropEmptySnapshotRefs(ctx context.Context, volumeID string) error { + snapshots, rev, err := s.snapRefs.Get(ctx, volumeID) + if err != nil { + if errors.Is(err, kvstore.ErrNotFound) { + return nil + } + return fmt.Errorf("dropEmptySnapshotRefs: failed to get KV key %s: %w", volumeID, err) + } + if len(*snapshots) > 0 { + return nil + } + // A lost race means a snapshot was added back, so the key has to stay. + if err := s.snapRefs.CompareAndDelete(ctx, volumeID, rev); err != nil && + !errors.Is(err, kvstore.ErrConflict) { + return fmt.Errorf("dropEmptySnapshotRefs: failed to delete KV key %s: %w", volumeID, err) + } + return nil } // volumeHasSnapshots returns true if the volume has any snapshots in KV. func (s *SnapshotServiceImpl) volumeHasSnapshots(ctx context.Context, volumeID string) (bool, error) { - if s.snapKV == nil { + if s.snapRefs == nil { return false, nil } - entry, err := s.snapKV.Get(ctx, volumeID) + snapshots, _, err := s.snapRefs.Get(ctx, volumeID) if err != nil { - if errors.Is(err, jetstream.ErrKeyNotFound) { + if errors.Is(err, kvstore.ErrNotFound) { return false, nil } return false, fmt.Errorf("volumeHasSnapshots: failed to get KV key %s: %w", volumeID, err) } - var snapshots []string - if err := json.Unmarshal(entry.Value(), &snapshots); err != nil { - return false, fmt.Errorf("volumeHasSnapshots: failed to unmarshal KV value for %s: %w", volumeID, err) - } - - return len(snapshots) > 0, nil + return len(*snapshots) > 0, nil } // ApplyRecordTags mirrors CreateTags into the owning snapshot metadata so diff --git a/spinifex/handlers/ec2/snapshot/service_impl_test.go b/spinifex/handlers/ec2/snapshot/service_impl_test.go index 5b5921000..6be52a85f 100644 --- a/spinifex/handlers/ec2/snapshot/service_impl_test.go +++ b/spinifex/handlers/ec2/snapshot/service_impl_test.go @@ -5,7 +5,9 @@ import ( "encoding/json" "errors" "path" + "strconv" "strings" + "sync" "testing" "github.com/aws/aws-sdk-go/aws" @@ -15,6 +17,7 @@ import ( "github.com/mulgadc/spinifex/spinifex/config" "github.com/mulgadc/spinifex/spinifex/ebsmetadata" "github.com/mulgadc/spinifex/spinifex/ebsprovider" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/objectstore" "github.com/mulgadc/spinifex/spinifex/testutil" "github.com/nats-io/nats.go/jetstream" @@ -594,9 +597,18 @@ func setupTestNATSKV(t *testing.T) jetstream.KeyValue { return kv } +// newRefService builds a service whose only wiring is the volume-snapshot +// index, which is all the ref tests touch. A nil kv is the disabled case. +func newRefService(kv jetstream.KeyValue) *SnapshotServiceImpl { + if kv == nil { + return &SnapshotServiceImpl{} + } + return &SnapshotServiceImpl{snapRefs: kvstore.Over[[]string](nil, kv, snapshotRefsConfig())} +} + func TestAddSnapshotRef(t *testing.T) { kv := setupTestNATSKV(t) - svc := &SnapshotServiceImpl{snapKV: kv} + svc := newRefService(kv) require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-a")) require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-b")) @@ -610,7 +622,7 @@ func TestAddSnapshotRef(t *testing.T) { func TestRemoveSnapshotRef(t *testing.T) { kv := setupTestNATSKV(t) - svc := &SnapshotServiceImpl{snapKV: kv} + svc := newRefService(kv) require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-a")) require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-b")) @@ -633,7 +645,7 @@ func TestRemoveSnapshotRef(t *testing.T) { func TestRemoveSnapshotRef_NonExistentKey(t *testing.T) { kv := setupTestNATSKV(t) - svc := &SnapshotServiceImpl{snapKV: kv} + svc := newRefService(kv) // Should not error on non-existent key require.NoError(t, svc.removeSnapshotRef(t.Context(), "vol-nonexistent", "snap-x")) @@ -641,7 +653,7 @@ func TestRemoveSnapshotRef_NonExistentKey(t *testing.T) { func TestRemoveSnapshotRefForCleanupSurvivesCancellation(t *testing.T) { kv := setupTestNATSKV(t) - svc := &SnapshotServiceImpl{snapKV: kv} + svc := newRefService(kv) require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-a")) ctx, cancel := context.WithCancel(t.Context()) @@ -654,7 +666,7 @@ func TestRemoveSnapshotRefForCleanupSurvivesCancellation(t *testing.T) { func TestVolumeHasSnapshots(t *testing.T) { kv := setupTestNATSKV(t) - svc := &SnapshotServiceImpl{snapKV: kv} + svc := newRefService(kv) // No entry → false has, err := svc.volumeHasSnapshots(t.Context(), "vol-1") @@ -675,7 +687,7 @@ func TestVolumeHasSnapshots(t *testing.T) { } func TestKVNilFallback(t *testing.T) { - svc := &SnapshotServiceImpl{snapKV: nil} + svc := newRefService(nil) // All methods should be no-ops when KV is nil require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-a")) @@ -1273,3 +1285,96 @@ type copyFailingProvider struct { func (p *copyFailingProvider) CopySnapshot(context.Context, ebsprovider.CopySnapshotRequest) (*ebsprovider.Snapshot, error) { return nil, errors.New("provider copy failed") } + +// Concurrent adds against one volume must not lose a reference, which is what +// the CAS loop exists to prevent. +func TestAddSnapshotRefConcurrentNoLostUpdates(t *testing.T) { + kv := setupTestNATSKV(t) + svc := newRefService(kv) + + const refs = 16 + var wg sync.WaitGroup + errs := make([]error, refs) + for i := range refs { + wg.Go(func() { + errs[i] = svc.addSnapshotRef(t.Context(), "vol-1", "snap-"+strconv.Itoa(i)) + }) + } + wg.Wait() + require.NoError(t, errors.Join(errs...)) + + entry, err := kv.Get(t.Context(), "vol-1") + require.NoError(t, err) + var snapshots []string + require.NoError(t, json.Unmarshal(entry.Value(), &snapshots)) + assert.Len(t, snapshots, refs) +} + +// Removing the last reference takes the key with it, so a volume that has lost +// every snapshot leaves nothing behind for volumeHasSnapshots to read. +func TestRemoveSnapshotRefDropsTheEmptiedKey(t *testing.T) { + kv := setupTestNATSKV(t) + svc := newRefService(kv) + + require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-a")) + require.NoError(t, svc.removeSnapshotRef(t.Context(), "vol-1", "snap-a")) + + _, err := kv.Get(t.Context(), "vol-1") + require.ErrorIs(t, err, jetstream.ErrKeyNotFound, "the emptied index key must be gone") + + has, err := svc.volumeHasSnapshots(t.Context(), "vol-1") + require.NoError(t, err) + assert.False(t, has) +} + +// dropEmptySnapshotRefs leaves a key that is no longer empty alone, so a volume +// that gained a snapshot after the check keeps its index. +func TestDropEmptySnapshotRefsLeavesANonEmptyKey(t *testing.T) { + kv := setupTestNATSKV(t) + svc := newRefService(kv) + + require.NoError(t, svc.addSnapshotRef(t.Context(), "vol-1", "snap-b")) + require.NoError(t, svc.dropEmptySnapshotRefs(t.Context(), "vol-1")) + + entry, err := kv.Get(t.Context(), "vol-1") + require.NoError(t, err) + var snapshots []string + require.NoError(t, json.Unmarshal(entry.Value(), &snapshots)) + assert.Equal(t, []string{"snap-b"}, snapshots) +} + +// The empty-key delete is guarded on the revision the emptiness was read at, so +// an add landing between that read and the delete keeps its reference. An +// unguarded delete drops it, and volumeHasSnapshots then reports no snapshots +// for a volume that has one — which is what DeleteVolume gates on. +func TestDropEmptySnapshotRefsDoesNotDropAConcurrentAdd(t *testing.T) { + kv := setupTestNATSKV(t) + svc := newRefService(kv) + + // Repeated because the losing interleaving is a window, not a certainty. + for i := range 200 { + volume := "vol-" + strconv.Itoa(i) + _, err := kv.Put(t.Context(), volume, []byte("[]")) + require.NoError(t, err) + + start := make(chan struct{}) + var wg sync.WaitGroup + var dropErr, addErr error + wg.Go(func() { + <-start + dropErr = svc.dropEmptySnapshotRefs(t.Context(), volume) + }) + wg.Go(func() { + <-start + addErr = svc.addSnapshotRef(t.Context(), volume, "snap-b") + }) + close(start) + wg.Wait() + require.NoError(t, dropErr) + require.NoError(t, addErr) + + has, err := svc.volumeHasSnapshots(t.Context(), volume) + require.NoError(t, err) + assert.True(t, has, "iteration %d: the concurrent add was dropped by the empty-key delete", i) + } +} diff --git a/spinifex/handlers/ec2/vpc/external_ipam.go b/spinifex/handlers/ec2/vpc/external_ipam.go index bececbc3c..330b0d6f6 100644 --- a/spinifex/handlers/ec2/vpc/external_ipam.go +++ b/spinifex/handlers/ec2/vpc/external_ipam.go @@ -50,7 +50,9 @@ func NewExternalIPAM(ctx context.Context, js jetstream.JetStream, pools []extern if err != nil { return nil, err } - kv = alloc.KV() + if kv, err = alloc.KV(ctx); err != nil { + return nil, err + } } return &ExternalIPAM{kv: kv, pools: pools, static: alloc, perPool: map[string]external.Allocator{}}, nil } diff --git a/spinifex/handlers/ec2/vpc/ipam.go b/spinifex/handlers/ec2/vpc/ipam.go index 2bde94360..ca58c9dd4 100644 --- a/spinifex/handlers/ec2/vpc/ipam.go +++ b/spinifex/handlers/ec2/vpc/ipam.go @@ -2,13 +2,12 @@ package handlers_ec2_vpc import ( "context" - "encoding/json" "errors" "fmt" "log/slog" "net/netip" - "github.com/mulgadc/spinifex/spinifex/kvutil" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/migrate" "github.com/nats-io/nats.go/jetstream" ) @@ -16,8 +15,15 @@ import ( const ( KVBucketIPAM = "spinifex-vpc-ipam" KVBucketIPAMVersion = 2 + KVBucketIPAMHistory = 5 ) +// ipamCASAttempts bounds one allocation's retry on a revision conflict. Every +// address in a subnet shares one record, so a bulk RunInstances contends with +// itself once per instance; the default five is below a single realistic launch +// burst and returns exhaustion where the caller expects an address. +const ipamCASAttempts = 50 + // ErrIPOutOfRange means the requested address is not a usable host address in // the subnet: outside the CIDR, in the reserved head, or the broadcast address. var ErrIPOutOfRange = errors.New("ip address not usable in subnet") @@ -42,72 +48,74 @@ type IPAMRecord struct { // IPAM manages IP address allocation for VPC subnets using NATS KV with CAS. type IPAM struct { - kv jetstream.KeyValue + store *kvstore.Store[IPAMRecord] +} + +// ipamConfig describes the IPAM bucket. js is captured for the migration, which +// resolves the stream itself and so needs the client rather than the handle; a +// nil js is a caller supplying its own already-migrated handle. +func ipamConfig(js jetstream.JetStream) kvstore.Config { + cfg := kvstore.Config{ + Name: KVBucketIPAM, + History: KVBucketIPAMHistory, + Missing: "IPAM KV bucket not initialized", + Attempts: ipamCASAttempts, + Exhausted: func(key string, attempts int) error { + return fmt.Errorf("IPAM CAS exhausted for subnet %s after %d attempts", key, attempts) + }, + } + if js != nil { + cfg.OnOpen = func(ctx context.Context, kv jetstream.KeyValue) error { + return migrate.DefaultRegistry.RunKVWithJetStream(ctx, KVBucketIPAM, kv, js, KVBucketIPAMVersion) + } + } + return cfg } // NewIPAM creates a new IPAM instance backed by NATS JetStream KV. func NewIPAM(ctx context.Context, js jetstream.JetStream) (*IPAM, error) { - kv, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketIPAM, 5) - if err != nil { + store := kvstore.New[IPAMRecord](js, ipamConfig(js)) + // Opened eagerly: a bucket that cannot be created must fail construction + // rather than the first allocation that needs it. + if _, err := store.KV(ctx); err != nil { return nil, fmt.Errorf("failed to create IPAM KV bucket: %w", err) } - if err := migrate.DefaultRegistry.RunKVWithJetStream(ctx, KVBucketIPAM, kv, js, KVBucketIPAMVersion); err != nil { - return nil, fmt.Errorf("migrate %s: %w", KVBucketIPAM, err) - } - return &IPAM{kv: kv}, nil + return &IPAM{store: store}, nil } // NewIPAMWithKV creates an IPAM with an existing KV bucket (for testing). func NewIPAMWithKV(kv jetstream.KeyValue) *IPAM { - return &IPAM{kv: kv} + return &IPAM{store: kvstore.Over[IPAMRecord](nil, kv, ipamConfig(nil))} } // AllocateIP allocates an IP from the subnet, reserving the first 4 and last // addresses per AWS convention. Uses CAS for conflict-free multi-node allocation. func (m *IPAM) AllocateIP(ctx context.Context, subnetId, cidrBlock, purpose, ownerID string) (string, error) { - for attempt := range 5 { - record, revision, err := m.getRecord(ctx, subnetId) - if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { - return "", fmt.Errorf("get IPAM record: %w", err) + var ip string + err := m.store.Upsert(ctx, subnetId, func(record *IPAMRecord) (bool, error) { + // Only an absent record takes the caller's CIDR. A stored one keeps its + // own, so a caller passing a stale subnet CIDR cannot re-base a pool + // that already has addresses handed out of it. + if record.CidrBlock == "" { + record.SubnetId = subnetId + record.CidrBlock = cidrBlock } - if record == nil { - record = &IPAMRecord{ - SubnetId: subnetId, - CidrBlock: cidrBlock, - } - } - - // Find next available IP - ip, err := m.nextAvailableIP(record) + next, err := m.nextAvailableIP(record) if err != nil { - return "", err + return false, err } - record.Allocated = append(record.Allocated, IPEntry{IP: ip, Purpose: purpose, OwnerID: ownerID}) - - // CAS write - data, err := json.Marshal(record) - if err != nil { - return "", fmt.Errorf("marshal IPAM record: %w", err) - } - - if revision == 0 { - _, err = m.kv.Create(ctx, subnetId, data) - } else { - _, err = m.kv.Update(ctx, subnetId, data, revision) - } - - if err != nil { - slog.Debug("IPAM CAS conflict, retrying", "subnet", subnetId, "attempt", attempt) - continue // CAS conflict, retry - } - - slog.Info("IPAM allocated IP", "subnet", subnetId, "ip", ip, "purpose", purpose, "owner", ownerID) - return ip, nil + record.Allocated = append(record.Allocated, IPEntry{IP: next, Purpose: purpose, OwnerID: ownerID}) + ip = next + return true, nil + }) + if err != nil { + return "", err } - return "", fmt.Errorf("IPAM allocation failed after CAS retries for subnet %s", subnetId) + slog.Info("IPAM allocated IP", "subnet", subnetId, "ip", ip, "purpose", purpose, "owner", ownerID) + return ip, nil } // ClaimIP reserves a caller-specified IP in the subnet, for a pinned address @@ -132,48 +140,30 @@ func (m *IPAM) ClaimIP(ctx context.Context, subnetId, cidrBlock, purpose, ownerI ip := addr.String() - for attempt := range 5 { - record, revision, err := m.getRecord(ctx, subnetId) - if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { - return "", fmt.Errorf("get IPAM record: %w", err) + err = m.store.Upsert(ctx, subnetId, func(record *IPAMRecord) (bool, error) { + // Only an absent record takes the caller's CIDR, matching AllocateIP: a + // stale subnet CIDR must not re-base a pool with addresses already out. + if record.CidrBlock == "" { + record.SubnetId = subnetId + record.CidrBlock = cidrBlock } - - if record == nil { - record = &IPAMRecord{ - SubnetId: subnetId, - CidrBlock: cidrBlock, - } - } - + // Re-checked on every attempt rather than once up front, so a claim that + // lost the race to the same address reports it in use instead of + // appending a duplicate entry. for _, entry := range record.Allocated { if entry.IP == ip { - return "", fmt.Errorf("%w: %s already allocated in subnet %s", ErrIPInUse, ip, subnetId) + return false, fmt.Errorf("%w: %s already allocated in subnet %s", ErrIPInUse, ip, subnetId) } } - record.Allocated = append(record.Allocated, IPEntry{IP: ip, Purpose: purpose, OwnerID: ownerID}) - - data, err := json.Marshal(record) - if err != nil { - return "", fmt.Errorf("marshal IPAM record: %w", err) - } - - if revision == 0 { - _, err = m.kv.Create(ctx, subnetId, data) - } else { - _, err = m.kv.Update(ctx, subnetId, data, revision) - } - - if err != nil { - slog.Debug("IPAM claim CAS conflict, retrying", "subnet", subnetId, "attempt", attempt) - continue // CAS conflict, retry: a concurrent writer may have taken this IP - } - - slog.Info("IPAM claimed IP", "subnet", subnetId, "ip", ip, "purpose", purpose, "owner", ownerID) - return ip, nil + return true, nil + }) + if err != nil { + return "", err } - return "", fmt.Errorf("IPAM claim failed after CAS retries for subnet %s", subnetId) + slog.Info("IPAM claimed IP", "subnet", subnetId, "ip", ip, "purpose", purpose, "owner", ownerID) + return ip, nil } // isClaimableAddr reports whether addr is a usable host address in prefix: @@ -198,76 +188,38 @@ func isClaimableAddr(prefix netip.Prefix, addr netip.Addr) bool { // ReleaseIP releases a previously allocated IP address back to the subnet pool. func (m *IPAM) ReleaseIP(ctx context.Context, subnetId, ip string) error { - for attempt := range 5 { - record, revision, err := m.getRecord(ctx, subnetId) - if err != nil { - return fmt.Errorf("get IPAM record for release: %w", err) - } - if record == nil { - return fmt.Errorf("no IPAM record for subnet %s", subnetId) - } - - found := false + err := m.store.Mutate(ctx, subnetId, func(record *IPAMRecord) (bool, error) { for i, entry := range record.Allocated { if entry.IP == ip { record.Allocated = append(record.Allocated[:i], record.Allocated[i+1:]...) - found = true - break + return true, nil } } - if !found { - return fmt.Errorf("IP %s not allocated in subnet %s", ip, subnetId) - } - - data, err := json.Marshal(record) - if err != nil { - return fmt.Errorf("marshal IPAM record: %w", err) - } - - if _, err := m.kv.Update(ctx, subnetId, data, revision); err != nil { - slog.Debug("IPAM release CAS conflict, retrying", "subnet", subnetId, "attempt", attempt) - continue - } - - slog.Info("IPAM released IP", "subnet", subnetId, "ip", ip) - return nil + return false, fmt.Errorf("IP %s not allocated in subnet %s", ip, subnetId) + }) + if errors.Is(err, kvstore.ErrNotFound) { + return fmt.Errorf("no IPAM record for subnet %s", subnetId) + } + if err != nil { + return err } - return fmt.Errorf("IPAM release failed after CAS retries for subnet %s", subnetId) + slog.Info("IPAM released IP", "subnet", subnetId, "ip", ip) + return nil } // AllocatedIPs returns the list of allocated IP entries for a subnet. func (m *IPAM) AllocatedIPs(ctx context.Context, subnetId string) ([]IPEntry, error) { - record, _, err := m.getRecord(ctx, subnetId) + record, _, err := m.store.Get(ctx, subnetId) if err != nil { - if errors.Is(err, jetstream.ErrKeyNotFound) { + if errors.Is(err, kvstore.ErrNotFound) { return nil, nil } return nil, err } - if record == nil { - return nil, nil - } return record.Allocated, nil } -func (m *IPAM) getRecord(ctx context.Context, subnetId string) (*IPAMRecord, uint64, error) { - entry, err := m.kv.Get(ctx, subnetId) - if err != nil { - if errors.Is(err, jetstream.ErrKeyNotFound) { - return nil, 0, jetstream.ErrKeyNotFound - } - return nil, 0, err - } - - var record IPAMRecord - if err := json.Unmarshal(entry.Value(), &record); err != nil { - return nil, 0, fmt.Errorf("unmarshal IPAM record: %w", err) - } - - return &record, entry.Revision(), nil -} - // nextAvailableIP finds the next available IP in the subnet, skipping reserved addresses. func (m *IPAM) nextAvailableIP(record *IPAMRecord) (string, error) { prefix, err := netip.ParsePrefix(record.CidrBlock) diff --git a/spinifex/handlers/ec2/vpc/ipam_test.go b/spinifex/handlers/ec2/vpc/ipam_test.go index e6422b18a..e51865133 100644 --- a/spinifex/handlers/ec2/vpc/ipam_test.go +++ b/spinifex/handlers/ec2/vpc/ipam_test.go @@ -347,3 +347,39 @@ func TestCompareIPs_NilSortsFirst(t *testing.T) { func itoa(i int) string { return strconv.Itoa(i) } + +// Concurrent allocations against one subnet must not hand the same address to +// two callers, which is what the CAS loop exists to prevent. +func TestIPAM_AllocateConcurrentNoDuplicates(t *testing.T) { + t.Parallel() + ipam := setupTestIPAM(t) + + const allocations = 16 + var ( + wg sync.WaitGroup + mu sync.Mutex + ips = make(map[string]bool, allocations) + errs []error + ) + for i := range allocations { + wg.Go(func() { + ip, err := ipam.AllocateIP(t.Context(), "subnet-1", "10.0.1.0/24", + PurposeENIPrimary, "eni-"+strconv.Itoa(i)) + mu.Lock() + defer mu.Unlock() + if err != nil { + errs = append(errs, err) + return + } + ips[ip] = true + }) + } + wg.Wait() + + require.Empty(t, errs) + assert.Len(t, ips, allocations) + + entries, err := ipam.AllocatedIPs(t.Context(), "subnet-1") + require.NoError(t, err) + assert.Len(t, entries, allocations) +} diff --git a/spinifex/handlers/ecs/store.go b/spinifex/handlers/ecs/store.go index 10978ea57..20d00ea4c 100644 --- a/spinifex/handlers/ecs/store.go +++ b/spinifex/handlers/ecs/store.go @@ -4,12 +4,11 @@ import ( "context" "errors" "fmt" - "log/slog" "time" + "github.com/mulgadc/spinifex/spinifex/kvlease" "github.com/mulgadc/spinifex/spinifex/kvutil" "github.com/mulgadc/spinifex/spinifex/migrate" - "github.com/mulgadc/spinifex/spinifex/otelsetup" "github.com/nats-io/nats.go" "github.com/nats-io/nats.go/jetstream" ) @@ -218,28 +217,11 @@ func GetOrCreateAccountBucket(ctx context.Context, js jetstream.JetStream, accou } // InitLeaderBucket creates (or attaches to) the shared spinifex-ecs-leader bucket -// used for per-cluster scheduler leader-lease CAS locks. The bucket is configured -// with History=1 and a 60s TTL so stale leases expire on their own when a leader -// dies mid-cycle. kvutil.GetOrCreateBucket doesn't expose a TTL knob, so this -// attaches directly and creates only when absent. +// used for per-cluster scheduler leader-lease CAS locks. func InitLeaderBucket(ctx context.Context, js jetstream.JetStream) (jetstream.KeyValue, error) { - // Attach before create. This runs on every scheduler tick, and CreateKeyValue - // against a bucket that exists with any other config is a STREAM.CREATE the - // meta leader answers with an error, so creating first bills one per tick. - kv, err := js.KeyValue(ctx, KVBucketECSLeader) - if errors.Is(err, jetstream.ErrBucketNotFound) { - kv, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{ - Bucket: KVBucketECSLeader, - History: 1, - TTL: KVBucketECSLeaderTTL, - }) - } - if err != nil { - return nil, fmt.Errorf("failed to create or open ECS leader bucket %s: %w", KVBucketECSLeader, err) - } - if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketECSLeader, kv, KVBucketECSLeaderVersion); err != nil { - return nil, fmt.Errorf("migrate %s: %w", KVBucketECSLeader, err) - } - slog.Info("ECS leader bucket initialized", "bucket", KVBucketECSLeader, "ttl_ms", otelsetup.Millis(KVBucketECSLeaderTTL)) - return kv, nil + return kvlease.OpenBucket(ctx, js, kvlease.BucketConfig{ + Name: KVBucketECSLeader, + TTL: KVBucketECSLeaderTTL, + Version: KVBucketECSLeaderVersion, + }) } diff --git a/spinifex/handlers/eks/store.go b/spinifex/handlers/eks/store.go index 30364541d..731a2769e 100644 --- a/spinifex/handlers/eks/store.go +++ b/spinifex/handlers/eks/store.go @@ -6,14 +6,13 @@ import ( "encoding/hex" "errors" "fmt" - "log/slog" "strings" "time" + "github.com/mulgadc/spinifex/spinifex/kvlease" "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/kvutil" "github.com/mulgadc/spinifex/spinifex/migrate" - "github.com/mulgadc/spinifex/spinifex/otelsetup" "github.com/nats-io/nats.go" "github.com/nats-io/nats.go/jetstream" ) @@ -247,30 +246,12 @@ func GetOrCreateAccountBucket(ctx context.Context, js jetstream.JetStream, accou // InitLeaderBucket creates (or attaches to) the shared spinifex-eks-leader // bucket used for per-cluster reconciler leader-lease CAS locks, at the given -// replica count (clamped to a minimum of 1). The bucket is configured with -// History=1 and a 60s TTL so stale leases expire on their own when a leader -// dies mid-cycle. kvutil.GetOrCreateBucketWithReplicas doesn't expose a TTL -// knob, so this function sets Replicas on its own create call, which runs only -// when the bucket is absent. +// replica count (clamped to a minimum of 1). func InitLeaderBucket(ctx context.Context, js jetstream.JetStream, replicas int) (jetstream.KeyValue, error) { - // Attach before create, so replicas applies to a genuine first creation only. - // CreateKeyValue against a bucket that exists with any other config is a - // STREAM.CREATE the meta leader answers with an error rather than a no-op. - kv, err := js.KeyValue(ctx, KVBucketEKSLeader) - if errors.Is(err, jetstream.ErrBucketNotFound) { - kv, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{ - Bucket: KVBucketEKSLeader, - History: 1, - TTL: KVBucketEKSLeaderTTL, - Replicas: max(replicas, 1), - }) - } - if err != nil { - return nil, fmt.Errorf("failed to create or open EKS leader bucket %s: %w", KVBucketEKSLeader, err) - } - if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketEKSLeader, kv, KVBucketEKSLeaderVersion); err != nil { - return nil, fmt.Errorf("migrate %s: %w", KVBucketEKSLeader, err) - } - slog.Info("EKS leader bucket initialized", "bucket", KVBucketEKSLeader, "ttl_ms", otelsetup.Millis(KVBucketEKSLeaderTTL)) - return kv, nil + return kvlease.OpenBucket(ctx, js, kvlease.BucketConfig{ + Name: KVBucketEKSLeader, + TTL: KVBucketEKSLeaderTTL, + Replicas: max(replicas, 1), + Version: KVBucketEKSLeaderVersion, + }) } diff --git a/spinifex/handlers/ochrevector/appliance.go b/spinifex/handlers/ochrevector/appliance.go index 714fc6ad7..50f272c63 100644 --- a/spinifex/handlers/ochrevector/appliance.go +++ b/spinifex/handlers/ochrevector/appliance.go @@ -367,7 +367,7 @@ func (a *Appliance) promote(ctx context.Context, rec ApplianceRecord, rev uint64 rec.Endpoint = endpoint rec.Port = port rec.UpdatedAt = time.Now().UTC() - if err := a.store.CompareAndSet(ctx, appliancePostgresKey, &rec, rev); err != nil { + if _, err := a.store.CompareAndSet(ctx, appliancePostgresKey, &rec, rev); err != nil { if current, _, getErr := a.getRecord(ctx); getErr == nil && current != nil && current.State == ApplianceStateAvailable { return connInfo(*current), nil } diff --git a/spinifex/handlers/quota/quota.go b/spinifex/handlers/quota/quota.go index ef0267980..508399fa6 100644 --- a/spinifex/handlers/quota/quota.go +++ b/spinifex/handlers/quota/quota.go @@ -5,8 +5,10 @@ package handlers_quota import ( + "fmt" "time" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/utils" "github.com/nats-io/nats.go/jetstream" ) @@ -69,10 +71,10 @@ type Limits struct { // derive their usage from account-filtered Describe* calls and need no state. type Service struct { limits Limits - // usage holds the per-account vCPU counters (key {accountID}). Nil when - // quotas are disabled, in which case Exempt short-circuits every check - // before the counter is touched. - usage jetstream.KeyValue + // usage holds the per-account vCPU counters (key {accountID}). Backed by a + // nil bucket when quotas are disabled, in which case Exempt short-circuits + // every check before the counter is touched. + usage *kvstore.Store[int] // bedrockUsage resolves the stream-fed Bedrock token counter // Nil until SetBedrockUsage is called, in which case @@ -93,8 +95,25 @@ type Service struct { // New constructs a quota Service from the configured limits and the gateway-owned // account-usage KV bucket. usage may be nil when quotas are disabled; Exempt then // short-circuits every check before the counter is read. +// +// The store is built over the handle rather than over a JetStream client, so a +// bucket the gateway could not open stays unopenable here instead of being +// quietly recreated under a counter that enforces a cap. func New(limits Limits, usage jetstream.KeyValue) *Service { - return &Service{limits: limits, usage: usage, rpm: newRPMLimiter()} + return &Service{limits: limits, usage: kvstore.Over[int](nil, usage, usageConfig()), rpm: newRPMLimiter()} +} + +// usageConfig describes the account-usage bucket as the vCPU counter reads it. +// The bucket itself is created by the gateway, so only the CAS policy matters. +func usageConfig() kvstore.Config { + return kvstore.Config{ + Name: KVBucketAccountUsage, + Attempts: vcpuCASRetries, + Missing: "account usage KV bucket not configured", + Exhausted: func(key string, attempts int) error { + return fmt.Errorf("vcpu counter CAS exhausted for %s after %d attempts", key, attempts) + }, + } } // SetOverrides attaches the per-account override bucket. Called once at wiring diff --git a/spinifex/handlers/quota/reconcile_test.go b/spinifex/handlers/quota/reconcile_test.go index d6a5f5532..a683f9f18 100644 --- a/spinifex/handlers/quota/reconcile_test.go +++ b/spinifex/handlers/quota/reconcile_test.go @@ -81,9 +81,7 @@ func TestReconcileSkipsSystemAccount(t *testing.T) { s := newVCPUService(t, Limits{Enabled: true, VCPUs: 100}) // Park a non-zero value directly under the system key; AddVCPU would no-op it. - if _, err := s.usage.PutString(t.Context(), utils.GlobalAccountID, "42"); err != nil { - t.Fatalf("seed system counter: %v", err) - } + seedCounter(t, s, utils.GlobalAccountID, 42) if err := s.Reconcile(context.Background(), accountList(utils.GlobalAccountID, testAccount), staticTotals(map[string]int{utils.GlobalAccountID: 99, testAccount: 2}, true, nil)); err != nil { @@ -188,9 +186,7 @@ func TestReconcileAccountTouchesOnlyThatAccount(t *testing.T) { func TestReconcileAccountSkipsSystemAccount(t *testing.T) { s := newVCPUService(t, Limits{Enabled: true, VCPUs: 100}) - if _, err := s.usage.PutString(t.Context(), utils.GlobalAccountID, "42"); err != nil { - t.Fatalf("seed system counter: %v", err) - } + seedCounter(t, s, utils.GlobalAccountID, 42) if err := s.ReconcileAccount(context.Background(), utils.GlobalAccountID, staticTotals(map[string]int{utils.GlobalAccountID: 0}, true, nil)); err != nil { t.Fatalf("ReconcileAccount: %v", err) @@ -217,6 +213,15 @@ func TestReconcileDisabledNoop(t *testing.T) { } } +// seedCounter writes a counter straight to the bucket, bypassing the CAS path +// under test. +func seedCounter(t *testing.T, s *Service, accountID string, value int) { + t.Helper() + if err := s.usage.Set(t.Context(), accountID, &value); err != nil { + t.Fatalf("seed counter %s: %v", accountID, err) + } +} + func assertCounter(t *testing.T, s *Service, accountID string, want int) { t.Helper() got, _, err := s.readVCPU(t.Context(), accountID) diff --git a/spinifex/handlers/quota/vcpu.go b/spinifex/handlers/quota/vcpu.go index fae7ce736..149aaff10 100644 --- a/spinifex/handlers/quota/vcpu.go +++ b/spinifex/handlers/quota/vcpu.go @@ -2,25 +2,26 @@ package handlers_quota import ( "context" - "encoding/json" "errors" - "fmt" "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/ec2" "github.com/mulgadc/spinifex/spinifex/awserrors" gateway_ec2_instance "github.com/mulgadc/spinifex/spinifex/gateway/ec2/instance" "github.com/mulgadc/spinifex/spinifex/instancetypes" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/nats-io/nats.go" - "github.com/nats-io/nats.go/jetstream" ) // vcpuCASRetries bounds AddVCPU's retry on a revision conflict. Each retry // implies another writer committed, so a single grow can conflict at most as -// many times as there are concurrent grows on the account; the bound sits well -// above any realistic in-flight launch burst and trips only as a circuit -// breaker, never on genuine contention. -const vcpuCASRetries = 100 +// many times as there are concurrent grows on the account, and this covers a +// realistic in-flight launch burst. +// +// Sized against the backoff, not just the contending set: the retry pauses +// between attempts, so the budget is a wall-clock one and the counter write +// sits on the RunInstances response path. +const vcpuCASRetries = 25 // CheckVCPU rejects with ResourceLimitExceeded when charging want more vCPUs to // accountID would exceed the configured cap. It only reads the counter; the @@ -43,35 +44,14 @@ func (s *Service) CheckVCPU(ctx context.Context, accountID string, want int) err // value and a skip flag; skip leaves the counter untouched and never creates a // key. A revision conflict is retried; exhausting the bound is a hard error. func (s *Service) casVCPU(ctx context.Context, accountID string, next func(current int) (value int, skip bool)) error { - for range vcpuCASRetries { - current, revision, err := s.readVCPU(ctx, accountID) - if err != nil { - return err - } - value, skip := next(current) + return s.usage.Upsert(ctx, accountID, func(current *int) (bool, error) { + value, skip := next(*current) if skip { - return nil + return false, nil } - data, err := json.Marshal(value) - if err != nil { - return err - } - if revision == 0 { - _, err = s.usage.Create(ctx, accountID, data) - } else { - _, err = s.usage.Update(ctx, accountID, data, revision) - } - if err == nil { - return nil - } - // Create reports a lost race as ErrKeyExists, Update as - // ErrKeyRevisionMismatch — and only the latter holds on a replicated - // bucket, where the conflict carries a different API error code. - if !errors.Is(err, jetstream.ErrKeyExists) && !errors.Is(err, jetstream.ErrKeyRevisionMismatch) { - return err - } - } - return fmt.Errorf("vcpu counter CAS exhausted for %s after %d attempts", accountID, vcpuCASRetries) + *current = value + return true, nil + }) } // AddVCPU adds delta vCPUs to accountID's counter under CAS so concurrent grows @@ -111,18 +91,17 @@ func (s *Service) setVCPUAt(ctx context.Context, accountID string, value int, si if revision != sinceRevision || current == value { return nil } - data, err := json.Marshal(value) - if err != nil { - return err - } + // Single-shot rather than Upsert: the guard is the revision the caller + // observed, so a retry against the winner's revision is the write this + // deliberately drops. if revision == 0 { - _, err = s.usage.Create(ctx, accountID, data) + _, err = s.usage.Create(ctx, accountID, &value) } else { - _, err = s.usage.Update(ctx, accountID, data, revision) + _, err = s.usage.CompareAndSet(ctx, accountID, &value, revision) } // Losing the race is the guard working, not a failure: a charge got there // first and the next pass recomputes against it. - if errors.Is(err, jetstream.ErrKeyExists) || errors.Is(err, jetstream.ErrKeyRevisionMismatch) { + if errors.Is(err, kvstore.ErrExists) || errors.Is(err, kvstore.ErrConflict) { return nil } return err @@ -272,15 +251,12 @@ func instanceTypeFromReservations(reservations []*ec2.Reservation, instanceID st // it was read at. A missing key is the zero counter at revision 0, which AddVCPU // treats as a create. func (s *Service) readVCPU(ctx context.Context, accountID string) (count int, revision uint64, err error) { - entry, err := s.usage.Get(ctx, accountID) + value, revision, err := s.usage.Get(ctx, accountID) if err != nil { - if errors.Is(err, jetstream.ErrKeyNotFound) { + if errors.Is(err, kvstore.ErrNotFound) { return 0, 0, nil } return 0, 0, err } - if err := json.Unmarshal(entry.Value(), &count); err != nil { - return 0, 0, fmt.Errorf("unmarshal vcpu counter for %s: %w", accountID, err) - } - return count, entry.Revision(), nil + return *value, revision, nil } diff --git a/spinifex/handlers/rds/agent.go b/spinifex/handlers/rds/agent.go index 25d938e7e..4a2286f25 100644 --- a/spinifex/handlers/rds/agent.go +++ b/spinifex/handlers/rds/agent.go @@ -7,9 +7,8 @@ import ( "log/slog" "time" - "github.com/nats-io/nats.go/jetstream" - "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/vm" ) @@ -416,7 +415,7 @@ func (s *Service) GetDBBootstrapConfig(ctx context.Context, input *GetDBBootstra // A failure to persist the backfill is not a failure to serve it: the value is // re-derivable and the next fetch records it, whereas failing here would hold a // bootable instance down for a bookkeeping write. -func (s *Service) resolveClientCIDR(ctx context.Context, kv jetstream.KeyValue, rec *DBInstanceRecord) (string, error) { +func (s *Service) resolveClientCIDR(ctx context.Context, kv *kvstore.Bucket, rec *DBInstanceRecord) (string, error) { if rec.VpcCIDR != "" { return rec.VpcCIDR, nil } @@ -453,7 +452,7 @@ func (s *Service) resolveClientCIDR(ctx context.Context, kv jetstream.KeyValue, // rather than returned. Propagating it would discard an attach response this // call has already built and reinstate the retry loop it exists to avoid; the // verdict is idempotent, so the next fetch records it. -func (s *Service) applyStagedBootstrap(ctx context.Context, kv jetstream.KeyValue, accountID string, +func (s *Service) applyStagedBootstrap(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord, generationMatches bool, out *GetDBBootstrapConfigOutput) error { envelope, _, err := readBootstrapPayload(ctx, kv, rec.DBInstanceIdentifier) if err != nil { @@ -500,7 +499,7 @@ func (s *Service) applyStagedBootstrap(ctx context.Context, kv jetstream.KeyValu // was never consumed also never bootstrapped, so there is no datadir to save: // the password is removed on sight and the instance is marked for recreation // rather than left holding a secret at rest for an operator who may never act. -func (s *Service) scrubLegacyBootstrap(ctx context.Context, kv jetstream.KeyValue, +func (s *Service) scrubLegacyBootstrap(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord) error { if rec.Bootstrap.MasterUserPassword == "" { return nil @@ -579,7 +578,7 @@ func (s *Service) AcknowledgeDBBootstrap(ctx context.Context, input *Acknowledge }); err != nil { return nil, err } - if err := deleteBootstrapPayload(ctx, kv, input.DBInstanceIdentifier, jetstream.LastRevision(rev)); err != nil { + if err := deleteBootstrapPayloadAt(ctx, kv, input.DBInstanceIdentifier, rev); err != nil { return nil, err } diff --git a/spinifex/handlers/rds/agent_test.go b/spinifex/handlers/rds/agent_test.go index 16c138367..2d4c48da6 100644 --- a/spinifex/handlers/rds/agent_test.go +++ b/spinifex/handlers/rds/agent_test.go @@ -128,7 +128,9 @@ func readRecord(t *testing.T, svc *Service) (DBInstanceRecord, string) { t.Helper() kv, err := svc.bucket(context.Background(), testAccountID) require.NoError(t, err) - entry, err := kv.Get(context.Background(), DBInstanceKey(testDBID)) + raw, err := kv.KV(context.Background()) + require.NoError(t, err) + entry, err := raw.Get(context.Background(), DBInstanceKey(testDBID)) require.NoError(t, err) var rec DBInstanceRecord require.NoError(t, json.Unmarshal(entry.Value(), &rec)) diff --git a/spinifex/handlers/rds/backup.go b/spinifex/handlers/rds/backup.go index ac5e94bd2..bf4f3f8f8 100644 --- a/spinifex/handlers/rds/backup.go +++ b/spinifex/handlers/rds/backup.go @@ -12,7 +12,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/rds" "github.com/mulgadc/spinifex/spinifex/awserrors" - "github.com/nats-io/nats.go/jetstream" + "github.com/mulgadc/spinifex/spinifex/kvstore" ) // Automated backups: one snapshot per instance per day inside its @@ -203,7 +203,7 @@ var backupCapableStatuses = []Status{StatusAvailable, StatusStopped} // succeeded since the window opened, reporting whether it took one. Called from // the leader-elected reconciler pass, which is what makes it cluster-singular; // the per-instance in-flight guard is what makes it safe anyway. -func (s *Service) runBackupWindow(ctx context.Context, kv jetstream.KeyValue, rev uint64, +func (s *Service) runBackupWindow(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord) (bool, error) { now := time.Now().UTC() window, err := s.resolvedBackupWindow(rec) @@ -267,7 +267,7 @@ func backupRetryDelay(failures int) time.Duration { // costs one duplicate snapshot at worst and is swept on schedule, while a stamped // backup with no index would be invisible to retention for the life of the // instance. -func (s *Service) takeAutomatedBackup(ctx context.Context, kv jetstream.KeyValue, rev uint64, +func (s *Service) takeAutomatedBackup(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord, now time.Time) error { record, err := s.snapshotDBInstance(ctx, kv, rev, accountID, rec, &validatedSnapshot{ DBSnapshotIdentifier: AutomatedSnapshotIdentifier(rec.DBInstanceIdentifier, now), @@ -298,7 +298,7 @@ func (s *Service) takeAutomatedBackup(ctx context.Context, kv jetstream.KeyValue // Counts the failure and reports it against the DB instance, leaving its status // alone. The stamp is what paces the retry inside the window; the count is what // makes a backup that has been failing for days visible and testable. -func (s *Service) recordBackupFailure(ctx context.Context, kv jetstream.KeyValue, accountID string, +func (s *Service) recordBackupFailure(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord, cause error) error { now := time.Now().UTC() slog.WarnContext(ctx, "rds: an automated backup failed; the DB instance is unaffected", @@ -316,7 +316,7 @@ func (s *Service) recordBackupFailure(ctx context.Context, kv jetstream.KeyValue // A window this instance could not be backed up in, evented once and then paced // by the same backoff a failure is. Reported without incrementing the failure // count: nothing about the backup failed, the instance was busy elsewhere. -func (s *Service) recordBackupDeferred(ctx context.Context, kv jetstream.KeyValue, accountID string, +func (s *Service) recordBackupDeferred(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord, opened time.Time, message string) error { now := time.Now().UTC() if rec.LastAutomatedBackupFailureAt == nil || rec.LastAutomatedBackupFailureAt.Before(opened) { @@ -337,7 +337,7 @@ func (s *Service) recordBackupDeferred(ctx context.Context, kv jetstream.KeyValu // single drain through applyPendingModifications — so a deferred change is // applied by exactly the code an immediate one uses, and one instance's VM // replace does not hold up the whole fleet's pass. -func (s *Service) runMaintenanceWindow(ctx context.Context, kv jetstream.KeyValue, rev uint64, +func (s *Service) runMaintenanceWindow(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord) (bool, error) { pending := rec.PendingModifiedValues if pending.empty() || pending.growingFilesystem() || rec.Status != StatusAvailable { @@ -364,7 +364,7 @@ func (s *Service) runMaintenanceWindow(ctx context.Context, kv jetstream.KeyValu rec.Status = StatusModifying rec.UpdatedAt = now if err := updateJSON(ctx, kv, DBInstanceKey(rec.DBInstanceIdentifier), rev, rec); err != nil { - if errors.Is(err, jetstream.ErrKeyRevisionMismatch) { + if errors.Is(err, kvstore.ErrConflict) { // Something else moved the record between the read and here; the next // pass re-reads and the window is still open. return false, nil diff --git a/spinifex/handlers/rds/backup_sweep.go b/spinifex/handlers/rds/backup_sweep.go index 15a07674e..083bea25a 100644 --- a/spinifex/handlers/rds/backup_sweep.go +++ b/spinifex/handlers/rds/backup_sweep.go @@ -10,8 +10,8 @@ import ( "time" "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/vm" - "github.com/nats-io/nats.go/jetstream" ) // The retention sweep is a cluster-wide vm.Reaper rather than a reconciler tick, @@ -49,7 +49,7 @@ func (r *BackupRetentionReaper) Sweep(ctx context.Context) (int, error) { } // A truncated listing must never read as "no accounts": that would leave every // tenant's retention unenforced while looking like a clean pass. - buckets, err := AccountBucketNames(ctx, js) + buckets, err := AccountBuckets(ctx, js) if err != nil { return 0, fmt.Errorf("rds: enumerate account buckets: %w", err) } @@ -63,16 +63,10 @@ func (r *BackupRetentionReaper) Sweep(ctx context.Context) (int, error) { if err := ctx.Err(); err != nil { return reaped, err } - kv, err := js.KeyValue(ctx, bucket) - if err != nil { - failures = append(failures, fmt.Errorf("open %s: %w", bucket, err)) - continue - } - accountID := AccountIDFromBucketName(bucket) - swept, err := r.sweepAccount(ctx, kv, accountID, r.limit-reaped) + swept, err := r.sweepAccount(ctx, bucket, AccountIDFromBucketName(bucket.Name()), r.limit-reaped) reaped += swept if err != nil { - failures = append(failures, fmt.Errorf("sweep %s: %w", bucket, err)) + failures = append(failures, fmt.Errorf("sweep %s: %w", bucket.Name(), err)) } } return reaped, errors.Join(failures...) @@ -80,7 +74,7 @@ func (r *BackupRetentionReaper) Sweep(ctx context.Context) (int, error) { // One account: the over-retention automated snapshots of every instance, then the // retained volumes whose last snapshot is gone. -func (r *BackupRetentionReaper) sweepAccount(ctx context.Context, kv jetstream.KeyValue, +func (r *BackupRetentionReaper) sweepAccount(ctx context.Context, kv *kvstore.Bucket, accountID string, budget int) (int, error) { indexed, err := ListAutomatedBackups(ctx, kv) if err != nil { @@ -115,7 +109,7 @@ func (r *BackupRetentionReaper) sweepAccount(ctx context.Context, kv jetstream.K // available snapshot is never deleted while automated backups are on — if backup // creation has been failing for a week, strict retention would delete the whole // backup set at the exact moment a backup matters most. -func (s *Service) sweepInstanceBackups(ctx context.Context, kv jetstream.KeyValue, +func (s *Service) sweepInstanceBackups(ctx context.Context, kv *kvstore.Bucket, accountID, dbInstanceIdentifier string, stamps []string, budget int) (int, error) { var rec DBInstanceRecord found, err := getJSON(ctx, kv, DBInstanceKey(dbInstanceIdentifier), &rec) @@ -170,7 +164,7 @@ func (s *Service) sweepInstanceBackups(ctx context.Context, kv jetstream.KeyValu // replica reads the same way — and that branch removes a live instance's whole // backup set, newest included. The bucket listing is a second, independent read // path, so an identifier it still names fails this instance's sweep instead. -func confirmDBInstanceGone(ctx context.Context, kv jetstream.KeyValue, dbInstanceIdentifier string) error { +func confirmDBInstanceGone(ctx context.Context, kv *kvstore.Bucket, dbInstanceIdentifier string) error { ids, err := ListDBInstanceIDs(ctx, kv) if err != nil { return err @@ -187,7 +181,7 @@ func confirmDBInstanceGone(ctx context.Context, kv jetstream.KeyValue, dbInstanc // The index entries of one instance, oldest first. An entry whose snapshot record // is gone is a half-finished delete: the entry is removed so the index stops // naming it, and it counts as nothing reaped because no data went with it. -func (s *Service) loadAutomatedBackups(ctx context.Context, kv jetstream.KeyValue, +func (s *Service) loadAutomatedBackups(ctx context.Context, kv *kvstore.Bucket, dbInstanceIdentifier string, stamps []string) ([]automatedBackup, error) { slices.Sort(stamps) entries := make([]automatedBackup, 0, len(stamps)) @@ -248,7 +242,7 @@ func newestAvailableBackup(entries []automatedBackup) string { // snapshot is still being read from — an instance restored from it is still // running — which is a skip rather than a failure: the next pass retries, and the // entry stays so it is not lost from the index. -func (s *Service) deleteAutomatedBackup(ctx context.Context, kv jetstream.KeyValue, +func (s *Service) deleteAutomatedBackup(ctx context.Context, kv *kvstore.Bucket, accountID string, entry automatedBackup) (bool, error) { err := s.removeDBSnapshot(ctx, kv, accountID, entry.snapshot) switch { @@ -276,10 +270,10 @@ func (s *Service) deleteAutomatedBackup(ctx context.Context, kv jetstream.KeyVal // Last, always: while the entry exists the snapshot is still findable by the // sweep, and every step above tolerates work it has already done. -func (s *Service) dropAutomatedBackupIndex(ctx context.Context, kv jetstream.KeyValue, +func (s *Service) dropAutomatedBackupIndex(ctx context.Context, kv *kvstore.Bucket, dbInstanceIdentifier, stamp string) error { key := AutomatedBackupKey(dbInstanceIdentifier, stamp) - if err := kv.Delete(ctx, key); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { + if err := kv.Delete(ctx, key); err != nil { return fmt.Errorf("rds: clear the automated-backup index entry %s: %w", key, err) } return nil @@ -290,7 +284,7 @@ func (s *Service) dropAutomatedBackupIndex(ctx context.Context, kv jetstream.Key // BackupRetentionPeriod=0 modify, which turns the feature off, and the teardown // of the instance itself, whose automated backups would otherwise pin its data // volume after the instance is gone. -func (s *Service) purgeAutomatedBackups(ctx context.Context, kv jetstream.KeyValue, +func (s *Service) purgeAutomatedBackups(ctx context.Context, kv *kvstore.Bucket, accountID, dbInstanceIdentifier string) error { stamps, err := listNames(ctx, kv, AutomatedBackupsPrefix(dbInstanceIdentifier)) if err != nil { @@ -313,7 +307,7 @@ func (s *Service) purgeAutomatedBackups(ctx context.Context, kv jetstream.KeyVal // deletes the volume inline on the last DeleteDBSnapshot, so this only ever fires // for a crash between those two steps — which is precisely what a KV-health-gated // cluster-wide sweep is for. -func (s *Service) reclaimOrphanedVolumes(ctx context.Context, kv jetstream.KeyValue) (int, error) { +func (s *Service) reclaimOrphanedVolumes(ctx context.Context, kv *kvstore.Bucket) (int, error) { ids, err := listNames(ctx, kv, RetainedVolumesPrefix()) if err != nil { return 0, err diff --git a/spinifex/handlers/rds/backup_sweep_test.go b/spinifex/handlers/rds/backup_sweep_test.go index 1c65c0547..ab11f2ffa 100644 --- a/spinifex/handlers/rds/backup_sweep_test.go +++ b/spinifex/handlers/rds/backup_sweep_test.go @@ -8,6 +8,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/rds" "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/vm" "github.com/nats-io/nats.go/jetstream" "github.com/stretchr/testify/assert" @@ -191,7 +192,10 @@ func TestSweep_LeavesTheBackupSetWhenTheRecordCannotBeRead(t *testing.T) { kv, err := h.svc.bucket(t.Context(), testAccountID) require.NoError(t, err) - stale := missingKey{KeyValue: kv, key: DBInstanceKey(testDBID)} + raw, err := kv.KV(t.Context()) + require.NoError(t, err) + stale := kvstore.NewOpenBucket(nil, missingKey{KeyValue: raw, key: DBInstanceKey(testDBID)}, + AccountBucketConfig(testAccountID)) reaped, err := h.svc.sweepInstanceBackups(t.Context(), stale, testAccountID, testDBID, h.automatedStamps(t, testDBID), defaultSweepDeleteLimit) diff --git a/spinifex/handlers/rds/bootstrap.go b/spinifex/handlers/rds/bootstrap.go index eb5dbbac1..c46d3f12d 100644 --- a/spinifex/handlers/rds/bootstrap.go +++ b/spinifex/handlers/rds/bootstrap.go @@ -11,8 +11,8 @@ import ( "github.com/mulgadc/spinifex/spinifex/awserrors" handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/utils" - "github.com/nats-io/nats.go/jetstream" ) // The staged master password is encrypted at rest under bootstrap-payloads/{id} @@ -79,7 +79,7 @@ func bootstrapKeyID(key []byte) string { // Stages the master password for rec's current VM generation. The key's presence // is what makes the bootstrap pending, so it is written once and only ever // deleted by an acknowledgement, a delete, or a generation bump. -func (s *Service) writeBootstrapPayload(ctx context.Context, kv jetstream.KeyValue, +func (s *Service) writeBootstrapPayload(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord, masterPassword string) (string, error) { key, err := s.masterKey() if err != nil { @@ -126,7 +126,7 @@ func (s *Service) writeBootstrapPayload(ctx context.Context, kv jetstream.KeyVal // Returns (nil, 0, nil) when nothing is staged, which is what an acknowledged, // restored or legacy instance reads as. -func readBootstrapPayload(ctx context.Context, kv jetstream.KeyValue, +func readBootstrapPayload(ctx context.Context, kv *kvstore.Bucket, dbInstanceIdentifier string) (*BootstrapPayloadEnvelope, uint64, error) { var envelope BootstrapPayloadEnvelope rev, found, err := getJSONRevision(ctx, kv, BootstrapPayloadKey(dbInstanceIdentifier), &envelope) @@ -138,10 +138,17 @@ func readBootstrapPayload(ctx context.Context, kv jetstream.KeyValue, // An already-absent key is the state this is trying to reach, so it is not an // error. -func deleteBootstrapPayload(ctx context.Context, kv jetstream.KeyValue, - dbInstanceIdentifier string, opts ...jetstream.KVDeleteOpt) error { - err := kv.Delete(ctx, BootstrapPayloadKey(dbInstanceIdentifier), opts...) - if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { +func deleteBootstrapPayload(ctx context.Context, kv *kvstore.Bucket, dbInstanceIdentifier string) error { + if err := kv.Delete(ctx, BootstrapPayloadKey(dbInstanceIdentifier)); err != nil { + return fmt.Errorf("rds bootstrap: delete the payload for %s: %w", dbInstanceIdentifier, err) + } + return nil +} + +// deleteBootstrapPayloadAt is deleteBootstrapPayload guarded on rev, for the +// consumer that must remove only the payload it read. +func deleteBootstrapPayloadAt(ctx context.Context, kv *kvstore.Bucket, dbInstanceIdentifier string, rev uint64) error { + if err := kv.CompareAndDelete(ctx, BootstrapPayloadKey(dbInstanceIdentifier), rev); err != nil { return fmt.Errorf("rds bootstrap: delete the payload for %s: %w", dbInstanceIdentifier, err) } return nil @@ -150,7 +157,7 @@ func deleteBootstrapPayload(ctx context.Context, kv jetstream.KeyValue, // Whether a staged payload is still waiting to be applied. Callers that only // need the yes/no — the modify gate, the reconciler's stale-payload sweep — use // this rather than decrypting. -func bootstrapPending(ctx context.Context, kv jetstream.KeyValue, dbInstanceIdentifier string) (bool, error) { +func bootstrapPending(ctx context.Context, kv *kvstore.Bucket, dbInstanceIdentifier string) (bool, error) { envelope, _, err := readBootstrapPayload(ctx, kv, dbInstanceIdentifier) return envelope != nil, err } @@ -222,7 +229,7 @@ func (s *Service) openBootstrapPayload(envelope *BootstrapPayloadEnvelope, accou // Records that the staged password can no longer reach this datadir. Written // rather than inferred so the reason rides DescribeDBInstances, and idempotent // so a repeated fetch does not churn the record. -func (s *Service) markBootstrapUnrecoverable(ctx context.Context, kv jetstream.KeyValue, +func (s *Service) markBootstrapUnrecoverable(ctx context.Context, kv *kvstore.Bucket, accountID, dbInstanceIdentifier, reason string) error { changed := false if err := s.updateInstance(ctx, kv, dbInstanceIdentifier, func(stored *DBInstanceRecord) { @@ -250,7 +257,7 @@ func (s *Service) markBootstrapUnrecoverable(ctx context.Context, kv jetstream.K // password is not what stops the replacement working — replaceInstanceVM // revokes the data-volume format grant, and only the initial create can hold // one — so rebinding would move the failure rather than fix it. -func (s *Service) discardPendingBootstrap(ctx context.Context, kv jetstream.KeyValue, +func (s *Service) discardPendingBootstrap(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord) error { pending, err := bootstrapPending(ctx, kv, rec.DBInstanceIdentifier) if err != nil || !pending { diff --git a/spinifex/handlers/rds/bootstrap_test.go b/spinifex/handlers/rds/bootstrap_test.go index f3101d47a..303d96979 100644 --- a/spinifex/handlers/rds/bootstrap_test.go +++ b/spinifex/handlers/rds/bootstrap_test.go @@ -45,7 +45,9 @@ func storedPayload(t *testing.T, svc *Service, id string) (*BootstrapPayloadEnve t.Helper() kv, err := svc.bucket(t.Context(), testAccountID) require.NoError(t, err) - entry, err := kv.Get(t.Context(), BootstrapPayloadKey(id)) + raw, err := kv.KV(t.Context()) + require.NoError(t, err) + entry, err := raw.Get(t.Context(), BootstrapPayloadKey(id)) if err != nil { return nil, "" } @@ -417,16 +419,18 @@ func TestCreateDBInstance_StagesThePasswordEncrypted(t *testing.T) { keys, err := bucketKeys(t.Context(), kv) require.NoError(t, err) require.Contains(t, keys, BootstrapPayloadKey(testDBInstanceID)) + raw, err := kv.KV(t.Context()) + require.NoError(t, err) // Nothing anywhere in the account bucket may hold the supplied password. for _, key := range keys { - entry, err := kv.Get(t.Context(), key) + entry, err := raw.Get(t.Context(), key) require.NoError(t, err) assert.NotContains(t, string(entry.Value()), "Sup3rSecret!", "cleartext at rest under %s", key) } var envelope BootstrapPayloadEnvelope - entry, err := kv.Get(t.Context(), BootstrapPayloadKey(testDBInstanceID)) + entry, err := raw.Get(t.Context(), BootstrapPayloadKey(testDBInstanceID)) require.NoError(t, err) require.NoError(t, json.Unmarshal(entry.Value(), &envelope)) plaintext, err := handlers_iam.DecryptSecret(envelope.EncryptedPayload, testMasterKey) diff --git a/spinifex/handlers/rds/create.go b/spinifex/handlers/rds/create.go index ee7f0925c..c6a3bf0b5 100644 --- a/spinifex/handlers/rds/create.go +++ b/spinifex/handlers/rds/create.go @@ -9,9 +9,9 @@ import ( "github.com/aws/aws-sdk-go/service/rds" "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/utils" "github.com/mulgadc/spinifex/spinifex/vm" - "github.com/nats-io/nats.go/jetstream" ) // The first VM behind a DB instance. Replacement and recovery increment it, @@ -61,7 +61,7 @@ func (s *Service) CreateDBInstance(ctx context.Context, input *rds.CreateDBInsta rec := newDBInstanceRecord(accountID, req, placement, parameters) rollbackRev, createErr := createJSONRevision(ctx, kv, key, &rec) if createErr != nil { - if errors.Is(createErr, jetstream.ErrKeyExists) { + if errors.Is(createErr, kvstore.ErrExists) { return nil, awserrors.Errorf(awserrors.ErrorDBInstanceAlreadyExists, "DB instance %s already exists", req.Identifier) } @@ -192,7 +192,7 @@ 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, +func (s *Service) recordLaunch(ctx context.Context, kv *kvstore.Bucket, key, accountID, expectedResourceID string, launched *LaunchOutput, initialCreate bool) (*DBInstanceRecord, uint64, error) { var rec DBInstanceRecord rev, found, err := getJSONRevision(ctx, kv, key, &rec) @@ -243,7 +243,7 @@ func (s *Service) recordLaunch(ctx context.Context, kv jetstream.KeyValue, key, return &rec, updatedRev, nil } -func (s *Service) rollbackDBInstanceReservation(ctx context.Context, kv jetstream.KeyValue, +func (s *Service) rollbackDBInstanceReservation(ctx context.Context, kv *kvstore.Bucket, key, identifier, resourceID string, rev uint64) { rbCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), rollbackTimeout) defer cancel() @@ -264,11 +264,11 @@ func (s *Service) rollbackDBInstanceReservation(ctx context.Context, kv jetstrea var rollbackErr error for range rollbackDeleteAttempts { - rollbackErr = kv.Delete(rbCtx, key, jetstream.LastRevision(rev)) - if rollbackErr == nil || errors.Is(rollbackErr, jetstream.ErrKeyNotFound) { + rollbackErr = kv.CompareAndDelete(rbCtx, key, rev) + if rollbackErr == nil { return } - if !errors.Is(rollbackErr, jetstream.ErrKeyRevisionMismatch) { + if !errors.Is(rollbackErr, kvstore.ErrConflict) { break } diff --git a/spinifex/handlers/rds/create_test.go b/spinifex/handlers/rds/create_test.go index 11ba8a1f3..ae3202c9e 100644 --- a/spinifex/handlers/rds/create_test.go +++ b/spinifex/handlers/rds/create_test.go @@ -12,7 +12,7 @@ import ( "github.com/aws/aws-sdk-go/service/rds" "github.com/mulgadc/spinifex/spinifex/awserrors" handlers_dns "github.com/mulgadc/spinifex/spinifex/handlers/dns" - "github.com/mulgadc/spinifex/spinifex/handlers/ec2/vpc" + handlers_ec2_vpc "github.com/mulgadc/spinifex/spinifex/handlers/ec2/vpc" iammock "github.com/mulgadc/spinifex/spinifex/handlers/iam/mock" "github.com/mulgadc/spinifex/spinifex/testutil" "github.com/mulgadc/spinifex/spinifex/utils" diff --git a/spinifex/handlers/rds/delete.go b/spinifex/handlers/rds/delete.go index 829495c84..20b98dc47 100644 --- a/spinifex/handlers/rds/delete.go +++ b/spinifex/handlers/rds/delete.go @@ -12,9 +12,9 @@ import ( "github.com/aws/aws-sdk-go/service/rds" "github.com/mulgadc/spinifex/spinifex/awserrors" "github.com/mulgadc/spinifex/spinifex/handlers/sysinstance" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/tags" "github.com/mulgadc/spinifex/spinifex/utils" - "github.com/nats-io/nats.go/jetstream" ) var errFinalSnapshotInProgress = errors.New("rds: final DB snapshot creation is still in progress") @@ -83,7 +83,7 @@ func (s *Service) DeleteDBInstance(ctx context.Context, input *rds.DeleteDBInsta if reservation.creator { s.rollbackFinalSnapshotReservation(ctx, kv, finalSnapshot, reservation.revision) } - if errors.Is(err, jetstream.ErrKeyRevisionMismatch) { + if errors.Is(err, kvstore.ErrConflict) { return nil, awserrors.Errorf(awserrors.ErrorDBInstanceInvalidState, "DB instance %s changed state concurrently; retry the delete", id) } @@ -153,7 +153,7 @@ func finalSnapshotIsOurs(existing *DBSnapshotRecord, rec *DBInstanceRecord) bool // Rejects a taken identifier before anything is torn down, as AWS does. Leaving // it to the snapshot step would terminate the VM first and then fail on an // error no retry can clear. -func (s *Service) checkFinalSnapshotAvailable(ctx context.Context, kv jetstream.KeyValue, rec *DBInstanceRecord, identifier string) error { +func (s *Service) checkFinalSnapshotAvailable(ctx context.Context, kv *kvstore.Bucket, rec *DBInstanceRecord, identifier string) error { if identifier == "" { return nil } @@ -175,7 +175,7 @@ func (s *Service) checkFinalSnapshotAvailable(ctx context.Context, kv jetstream. // // Every step treats a missing resource as done, so this is safe to re-run from // any point it stopped at. -func (s *Service) teardownDBInstance(ctx context.Context, kv jetstream.KeyValue, accountID string, +func (s *Service) teardownDBInstance(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord, finalSnapshotCreator bool) error { // Asked to stop cleanly even though the VM is about to go: it is what makes // the final snapshot a clean checkpoint rather than one needing WAL replay. @@ -203,7 +203,7 @@ func (s *Service) teardownDBInstance(ctx context.Context, kv jetstream.KeyValue, s.deleteInstanceENIs(ctx, accountID, rec) if rec.InstanceID != "" { - if err := s.DeleteInstanceIndex(ctx, rec.InstanceID); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { + if err := s.DeleteInstanceIndex(ctx, rec.InstanceID); err != nil { return fmt.Errorf("rds: delete the instance index entry for %s: %w", rec.DBInstanceIdentifier, err) } } @@ -214,7 +214,7 @@ func (s *Service) teardownDBInstance(ctx context.Context, kv jetstream.KeyValue, } // Last: while it exists the instance is still nameable, and everything above // is reachable only through it. - if err := kv.Delete(ctx, DBInstanceKey(rec.DBInstanceIdentifier)); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { + if err := kv.Delete(ctx, DBInstanceKey(rec.DBInstanceIdentifier)); err != nil { return fmt.Errorf("rds: delete the DB instance record for %s: %w", rec.DBInstanceIdentifier, err) } @@ -244,7 +244,7 @@ func (s *Service) terminateInstanceVM(ctx context.Context, instanceID string) er // Reserves the final snapshot name before deletion becomes destructive. The // creator is the only worker allowed to cut new data; retries may adopt data it // already cut, but cannot race it into making a duplicate snapshot. -func (s *Service) reserveFinalSnapshot(ctx context.Context, kv jetstream.KeyValue, accountID string, +func (s *Service) reserveFinalSnapshot(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord, identifier string) (finalSnapshotReservation, error) { if identifier == "" || rec.DataVolumeID == "" { return finalSnapshotReservation{}, nil @@ -262,7 +262,7 @@ func (s *Service) reserveFinalSnapshot(ctx context.Context, kv jetstream.KeyValu if err == nil { return finalSnapshotReservation{creator: true, revision: rev}, nil } - if !errors.Is(err, jetstream.ErrKeyExists) { + if !errors.Is(err, kvstore.ErrExists) { return finalSnapshotReservation{}, err } @@ -281,15 +281,14 @@ func (s *Service) reserveFinalSnapshot(ctx context.Context, kv jetstream.KeyValu // Removes a reservation when the instance never entered deleting or the EC2 // create failed. The revision guard prevents cleanup from deleting a record // another worker has already completed. -func (s *Service) rollbackFinalSnapshotReservation(ctx context.Context, kv jetstream.KeyValue, +func (s *Service) rollbackFinalSnapshotReservation(ctx context.Context, kv *kvstore.Bucket, identifier string, rev uint64) { if identifier == "" || rev == 0 { return } rbCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), rollbackTimeout) defer cancel() - if err := kv.Delete(rbCtx, DBSnapshotKey(identifier), jetstream.LastRevision(rev)); err != nil && - !errors.Is(err, jetstream.ErrKeyNotFound) { + if err := kv.CompareAndDelete(rbCtx, DBSnapshotKey(identifier), rev); err != nil { slog.WarnContext(rbCtx, "rds: rollback of a final snapshot reservation failed", "dbSnapshot", identifier, "err", err) } @@ -298,7 +297,7 @@ func (s *Service) rollbackFinalSnapshotReservation(ctx context.Context, kv jetst // Records the final snapshot under the reservation written before teardown. A // retry first adopts matching EC2 data left by a dead creator; it never cuts a // second snapshot while that creator may still be running. -func (s *Service) takeFinalSnapshot(ctx context.Context, kv jetstream.KeyValue, accountID string, +func (s *Service) takeFinalSnapshot(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord, creator bool) error { if rec.FinalSnapshotIdentifier == "" || rec.DataVolumeID == "" { return nil @@ -393,7 +392,7 @@ func (s *Service) takeFinalSnapshot(ctx context.Context, kv jetstream.KeyValue, // Completes only the snapshot fields under CAS, preserving tags changed while // the EC2 snapshot was being cut. -func (s *Service) completeFinalSnapshot(ctx context.Context, kv jetstream.KeyValue, +func (s *Service) completeFinalSnapshot(ctx context.Context, kv *kvstore.Bucket, source *DBInstanceRecord, snapshotID string) (bool, error) { key := DBSnapshotKey(source.FinalSnapshotIdentifier) for range tagWriteAttempts { @@ -426,7 +425,7 @@ func (s *Service) completeFinalSnapshot(ctx context.Context, kv jetstream.KeyVal record.Status = SnapshotStatusAvailable if err := updateJSON(ctx, kv, key, rev, &record); err == nil { return true, nil - } else if !errors.Is(err, jetstream.ErrKeyRevisionMismatch) { + } else if !errors.Is(err, kvstore.ErrConflict) { return false, err } } @@ -447,7 +446,7 @@ func (s *Service) recordFinalSnapshotCreated(ctx context.Context, accountID stri // copying them, so the volume cannot be deleted while any snapshot survives — // including the final one just taken. It is retained instead, recorded with the // snapshots holding it so the last DeleteDBSnapshot can release it. -func (s *Service) releaseDataVolume(ctx context.Context, kv jetstream.KeyValue, accountID string, rec *DBInstanceRecord) error { +func (s *Service) releaseDataVolume(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord) error { if rec.DataVolumeID == "" { return nil } @@ -482,7 +481,7 @@ func (s *Service) releaseDataVolume(ctx context.Context, kv jetstream.KeyValue, // the last DeleteDBSnapshot can release it. holdersUnresolved marks the // case where the volume store refused the delete but named nobody, so a release // has to re-check rather than trust an empty list. -func (s *Service) retainDataVolume(ctx context.Context, kv jetstream.KeyValue, accountID string, +func (s *Service) retainDataVolume(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord, holders []string, holdersUnresolved bool) error { retained := RetainedVolumeRecord{ VolumeID: rec.DataVolumeID, diff --git a/spinifex/handlers/rds/describe.go b/spinifex/handlers/rds/describe.go index 63f3959d5..6e9a376bb 100644 --- a/spinifex/handlers/rds/describe.go +++ b/spinifex/handlers/rds/describe.go @@ -8,7 +8,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/rds" "github.com/mulgadc/spinifex/spinifex/awserrors" - "github.com/nats-io/nats.go/jetstream" + "github.com/mulgadc/spinifex/spinifex/kvstore" ) // The prefix AWS gives a DB instance's immutable resource ID, and the two filter @@ -116,7 +116,7 @@ func dbInstanceFilterMatcher(filters []*rds.Filter) (func(*DBInstanceRecord) boo } // Returns the record plus its revision, for callers that follow with a CAS. -func (s *Service) getDBInstance(ctx context.Context, kv jetstream.KeyValue, id string) (*DBInstanceRecord, uint64, error) { +func (s *Service) getDBInstance(ctx context.Context, kv *kvstore.Bucket, id string) (*DBInstanceRecord, uint64, error) { var rec DBInstanceRecord rev, found, err := getJSONRevision(ctx, kv, DBInstanceKey(id), &rec) if err != nil { diff --git a/spinifex/handlers/rds/dns.go b/spinifex/handlers/rds/dns.go index e6ddac796..9a2ee21b6 100644 --- a/spinifex/handlers/rds/dns.go +++ b/spinifex/handlers/rds/dns.go @@ -5,7 +5,7 @@ import ( "log/slog" handlers_dns "github.com/mulgadc/spinifex/spinifex/handlers/dns" - "github.com/nats-io/nats.go/jetstream" + "github.com/mulgadc/spinifex/spinifex/kvstore" ) // The vanity hostname for a DB instance, or "" on a deployment with no base @@ -31,20 +31,15 @@ func (s *Service) DesiredDNSChanges() (changes []handlers_dns.Change, ok bool) { if err != nil { return nil, false } - buckets, err := AccountBucketNames(ctx, js) + buckets, err := AccountBuckets(ctx, js) if err != nil { slog.Warn("rds DesiredDNSChanges: enumerate account buckets", "err", err) return nil, false } for _, bucket := range buckets { - kv, err := js.KeyValue(ctx, bucket) + bucketChanges, err := desiredBucketDNSChanges(ctx, bucket, s.deps.BaseDomain) if err != nil { - slog.Warn("rds DesiredDNSChanges: open account bucket", "bucket", bucket, "err", err) - return nil, false - } - bucketChanges, err := desiredBucketDNSChanges(ctx, kv, s.deps.BaseDomain) - if err != nil { - slog.Warn("rds DesiredDNSChanges: read DB instances", "bucket", bucket, "err", err) + slog.Warn("rds DesiredDNSChanges: read DB instances", "bucket", bucket.Name(), "err", err) return nil, false } changes = append(changes, bucketChanges...) @@ -55,7 +50,7 @@ func (s *Service) DesiredDNSChanges() (changes []handlers_dns.Change, ok bool) { // A deleted instance's record is gone, so it contributes nothing and the // reconcile prunes its record. Anything still holding an ENI IP keeps its name // resolvable, including a failed instance an operator is still investigating. -func desiredBucketDNSChanges(ctx context.Context, kv jetstream.KeyValue, baseDomain string) ([]handlers_dns.Change, error) { +func desiredBucketDNSChanges(ctx context.Context, kv *kvstore.Bucket, baseDomain string) ([]handlers_dns.Change, error) { ids, err := ListDBInstanceIDs(ctx, kv) if err != nil { return nil, err diff --git a/spinifex/handlers/rds/engine_postgres.go b/spinifex/handlers/rds/engine_postgres.go index 34914a268..f1fb1c393 100644 --- a/spinifex/handlers/rds/engine_postgres.go +++ b/spinifex/handlers/rds/engine_postgres.go @@ -4,6 +4,7 @@ import ( "strconv" "strings" "time" + // The zone database Go's own LoadLocation reads, vendored into the binary so // the timezone parameter validates identically wherever the daemon runs. _ "time/tzdata" diff --git a/spinifex/handlers/rds/events.go b/spinifex/handlers/rds/events.go index 35b5e42d2..8a4da1d72 100644 --- a/spinifex/handlers/rds/events.go +++ b/spinifex/handlers/rds/events.go @@ -12,7 +12,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/rds" "github.com/mulgadc/spinifex/spinifex/awserrors" - "github.com/nats-io/nats.go/jetstream" + "github.com/mulgadc/spinifex/spinifex/kvstore" ) // The event ring is the only channel some facts have to the customer — a backup @@ -118,7 +118,7 @@ func (s *Service) appendEvent(ctx context.Context, accountID, sourceType, source if !found { if err := createJSON(ctx, kv, key, &ring); err == nil { return nil - } else if !errors.Is(err, jetstream.ErrKeyExists) { + } else if !errors.Is(err, kvstore.ErrExists) { return err } // Another writer created the ring first; re-read and append to it. @@ -130,7 +130,7 @@ func (s *Service) appendEvent(ctx context.Context, accountID, sourceType, source } // A revision mismatch is a lost race with another appender, not a // duplicate: re-read the ring and append to what they wrote. - if !errors.Is(err, jetstream.ErrKeyRevisionMismatch) { + if !errors.Is(err, kvstore.ErrConflict) { return err } } @@ -277,7 +277,7 @@ func eventRecordLimit(input *rds.DescribeEventsInput) int { // A fully qualified read is one Get. Anything broader has to enumerate, because // a ring outlives the record that would otherwise name it. -func eventRingKeys(ctx context.Context, kv jetstream.KeyValue, sourceType, sourceIdentifier string) ([]string, error) { +func eventRingKeys(ctx context.Context, kv *kvstore.Bucket, sourceType, sourceIdentifier string) ([]string, error) { if sourceType != "" && sourceIdentifier != "" { return []string{EventRingKey(sourceType, sourceIdentifier)}, nil } @@ -285,11 +285,8 @@ func eventRingKeys(ctx context.Context, kv jetstream.KeyValue, sourceType, sourc if sourceType != "" { prefix += sourceType + "/" } - keys, err := kv.Keys(ctx) + keys, err := bucketKeys(ctx, kv) if err != nil { - if errors.Is(err, jetstream.ErrNoKeysFound) { - return nil, nil - } return nil, fmt.Errorf("rds: list event rings: %w", err) } matched := make([]string, 0, len(keys)) diff --git a/spinifex/handlers/rds/lifecycle.go b/spinifex/handlers/rds/lifecycle.go index 1dc5bb3ca..a63f0d94d 100644 --- a/spinifex/handlers/rds/lifecycle.go +++ b/spinifex/handlers/rds/lifecycle.go @@ -13,10 +13,10 @@ import ( "github.com/aws/aws-sdk-go/service/rds" "github.com/mulgadc/spinifex/spinifex/awserrors" handlers_ec2_instance "github.com/mulgadc/spinifex/spinifex/handlers/ec2/instance" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/types" "github.com/mulgadc/spinifex/spinifex/utils" "github.com/nats-io/nats.go" - "github.com/nats-io/nats.go/jetstream" ) // The DB VM's power state is driven over the per-instance EC2 command path @@ -272,7 +272,7 @@ func uncleanStopMessage(ctx context.Context, engineName, operation string) strin // as written. from lists the statuses the operation is legal from, so a stop of // an already-stopping instance is rejected as a state error rather than racing // the stop already running. -func (s *Service) beginTransition(ctx context.Context, accountID, id string, to Status, from ...Status) (*DBInstanceRecord, jetstream.KeyValue, error) { +func (s *Service) beginTransition(ctx context.Context, accountID, id string, to Status, from ...Status) (*DBInstanceRecord, *kvstore.Bucket, error) { if id == "" { return nil, nil, awserrors.Errorf(awserrors.ErrorInvalidParameterValue, "DBInstanceIdentifier is required") } @@ -310,7 +310,7 @@ func (s *Service) beginTransition(ctx context.Context, accountID, id string, to rec.TransitionStartedAt = &now rec.UpdatedAt = now if err := updateJSON(ctx, kv, DBInstanceKey(id), rev, rec); err != nil { - if errors.Is(err, jetstream.ErrKeyRevisionMismatch) { + if errors.Is(err, kvstore.ErrConflict) { return nil, nil, awserrors.Errorf(awserrors.ErrorDBInstanceInvalidState, "DB instance %s changed state concurrently; retry the operation", id) } @@ -321,7 +321,7 @@ func (s *Service) beginTransition(ctx context.Context, accountID, id string, to // Lands the instance in its post-transition state and clears the transition // stamp, returning the record as written. -func (s *Service) completeTransition(ctx context.Context, kv jetstream.KeyValue, id string, to Status) (*DBInstanceRecord, error) { +func (s *Service) completeTransition(ctx context.Context, kv *kvstore.Bucket, id string, to Status) (*DBInstanceRecord, error) { var stored *DBInstanceRecord err := s.updateInstance(ctx, kv, id, func(rec *DBInstanceRecord) { rec.Status = to @@ -336,7 +336,7 @@ func (s *Service) completeTransition(ctx context.Context, kv jetstream.KeyValue, // Records why an operation failed and returns the error the caller surfaces, so // a failed transition cannot be left looking like one still in progress. -func (s *Service) failTransition(ctx context.Context, kv jetstream.KeyValue, accountID string, rec *DBInstanceRecord, reason string) error { +func (s *Service) failTransition(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord, reason string) error { if err := s.updateInstance(ctx, kv, rec.DBInstanceIdentifier, func(stored *DBInstanceRecord) { stored.Status = StatusFailed stored.FailureReason = reason @@ -351,7 +351,7 @@ func (s *Service) failTransition(ctx context.Context, kv jetstream.KeyValue, acc // A read-modify-write under CAS, replayed on contention so a concurrent agent // heartbeat cannot make a lifecycle write disappear. -func (s *Service) updateInstance(ctx context.Context, kv jetstream.KeyValue, id string, mutate func(*DBInstanceRecord)) error { +func (s *Service) updateInstance(ctx context.Context, kv *kvstore.Bucket, id string, mutate func(*DBInstanceRecord)) error { _, err := s.updateInstanceIf(ctx, kv, id, func(rec *DBInstanceRecord) bool { mutate(rec) return true @@ -363,7 +363,7 @@ func (s *Service) updateInstance(ctx context.Context, kv jetstream.KeyValue, id // seen the record it would overwrite — a lease claim, whose whole question is // whether someone else holds it. Reports whether the write happened; a mutate // that returns false leaves the record untouched and is not an error. -func (s *Service) updateInstanceIf(ctx context.Context, kv jetstream.KeyValue, id string, mutate func(*DBInstanceRecord) bool) (bool, error) { +func (s *Service) updateInstanceIf(ctx context.Context, kv *kvstore.Bucket, id string, mutate func(*DBInstanceRecord) bool) (bool, error) { key := DBInstanceKey(id) for range tagWriteAttempts { rec, rev, err := s.getDBInstance(ctx, kv, id) @@ -379,7 +379,7 @@ func (s *Service) updateInstanceIf(ctx context.Context, kv jetstream.KeyValue, i if err == nil { return true, nil } - if !errors.Is(err, jetstream.ErrKeyRevisionMismatch) { + if !errors.Is(err, kvstore.ErrConflict) { return false, err } } diff --git a/spinifex/handlers/rds/modify.go b/spinifex/handlers/rds/modify.go index 7f5ad3966..6965a3e6b 100644 --- a/spinifex/handlers/rds/modify.go +++ b/spinifex/handlers/rds/modify.go @@ -13,7 +13,7 @@ import ( "github.com/aws/aws-sdk-go/service/ec2" "github.com/aws/aws-sdk-go/service/rds" "github.com/mulgadc/spinifex/spinifex/awserrors" - "github.com/nats-io/nats.go/jetstream" + "github.com/mulgadc/spinifex/spinifex/kvstore" ) // The resolved effect of a ModifyDBInstance request: what changes now and what @@ -353,7 +353,7 @@ func (s *Service) planBackupSettings(input *rds.ModifyDBInstanceInput, rec *DBIn // rotation, an ENI security-group re-association, and record settings. Each is // applied before any disruptive work, so a modify that carries both does not // lose the cheap half to a failure in the expensive one. -func (s *Service) applyImmediateModify(ctx context.Context, kv jetstream.KeyValue, accountID string, rec *DBInstanceRecord, plan *modifyPlan) error { +func (s *Service) applyImmediateModify(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord, plan *modifyPlan) error { if !plan.immediate() { return nil } @@ -427,7 +427,7 @@ func (s *Service) applyImmediateModify(ctx context.Context, kv jetstream.KeyValu // The retention change itself has already landed, so a sweep that cannot finish // is reported rather than allowed to fail the modify: the retention reaper reads // the same zero retention and removes whatever is left, including the newest. -func (s *Service) disableAutomatedBackups(ctx context.Context, kv jetstream.KeyValue, accountID, id string) { +func (s *Service) disableAutomatedBackups(ctx context.Context, kv *kvstore.Bucket, accountID, id string) { if err := s.purgeAutomatedBackups(ctx, kv, accountID, id); err != nil { slog.WarnContext(ctx, "rds: sweeping the automated backups of a disabled instance failed; the retention reaper will finish it", "dbInstance", id, "accountId", accountID, "err", err) diff --git a/spinifex/handlers/rds/modify_lease.go b/spinifex/handlers/rds/modify_lease.go index 0cb484d5b..395e269b0 100644 --- a/spinifex/handlers/rds/modify_lease.go +++ b/spinifex/handlers/rds/modify_lease.go @@ -9,7 +9,7 @@ import ( "sync" "time" - "github.com/nats-io/nats.go/jetstream" + "github.com/mulgadc/spinifex/spinifex/kvstore" ) // Applying PendingModifiedValues is the one control-plane operation with two @@ -31,7 +31,7 @@ var errModifyLeaseLost = errors.New("rds: modify lease lost while applying pendi // this change and the caller has nothing to do. func (s *Service) withModifyLease( ctx context.Context, - kv jetstream.KeyValue, + kv *kvstore.Bucket, id string, apply func(context.Context) error, ) (bool, error) { @@ -92,7 +92,7 @@ func (s *Service) newModifyLeaseHolder() string { // is allowed, so a caller that already holds it is not deadlocked by itself. func (s *Service) claimModifyLease( ctx context.Context, - kv jetstream.KeyValue, + kv *kvstore.Bucket, id, holder string, ) (bool, time.Time, error) { var expiresAt time.Time @@ -113,7 +113,7 @@ func (s *Service) claimModifyLease( // work continue after its last successful renewal expires. func (s *Service) renewModifyLease( ctx context.Context, - kv jetstream.KeyValue, + kv *kvstore.Bucket, id, holder string, expiresAt time.Time, cancelApply context.CancelCauseFunc, diff --git a/spinifex/handlers/rds/modify_lease_test.go b/spinifex/handlers/rds/modify_lease_test.go index 1fb2433af..821208c37 100644 --- a/spinifex/handlers/rds/modify_lease_test.go +++ b/spinifex/handlers/rds/modify_lease_test.go @@ -218,7 +218,7 @@ func TestWithModifyLease_CancelsWhenRenewalsFailUntilExpiry(t *testing.T) { finished := make(chan error, 1) go func() { _, applyErr := h.svc.withModifyLease(t.Context(), kv, testDBID, func(ctx context.Context) error { - bucketDeleted <- js.DeleteKeyValue(ctx, kv.Bucket()) + bucketDeleted <- js.DeleteKeyValue(ctx, kv.Name()) <-ctx.Done() return context.Cause(ctx) }) diff --git a/spinifex/handlers/rds/modify_test.go b/spinifex/handlers/rds/modify_test.go index b5bc9bebf..1d6828037 100644 --- a/spinifex/handlers/rds/modify_test.go +++ b/spinifex/handlers/rds/modify_test.go @@ -9,9 +9,9 @@ import ( "github.com/aws/aws-sdk-go/service/rds" "github.com/mulgadc/spinifex/spinifex/awserrors" iammock "github.com/mulgadc/spinifex/spinifex/handlers/iam/mock" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/testutil" "github.com/nats-io/nats.go" - "github.com/nats-io/nats.go/jetstream" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -95,7 +95,7 @@ func (h *modifyHarness) record(t *testing.T) DBInstanceRecord { return rec } -func (h *modifyHarness) kv(t *testing.T) jetstream.KeyValue { +func (h *modifyHarness) kv(t *testing.T) *kvstore.Bucket { t.Helper() kv, err := h.svc.bucket(t.Context(), testAccountID) require.NoError(t, err) diff --git a/spinifex/handlers/rds/network.go b/spinifex/handlers/rds/network.go index f825531b7..f3386d31a 100644 --- a/spinifex/handlers/rds/network.go +++ b/spinifex/handlers/rds/network.go @@ -8,7 +8,7 @@ import ( "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.go/jetstream" + "github.com/mulgadc/spinifex/spinifex/kvstore" ) // The customer-account VPC surface CreateDBInstance reads to place the endpoint @@ -33,7 +33,7 @@ type endpointPlacement struct { // back to the account's default VPC, mirroring AWS's own behaviour. The security // groups are validated against whichever VPC that resolved to, because an ENI // cannot carry a group from another one. -func (s *Service) resolvePlacement(ctx context.Context, kv jetstream.KeyValue, accountID string, req *validatedCreate) (*endpointPlacement, error) { +func (s *Service) resolvePlacement(ctx context.Context, kv *kvstore.Bucket, accountID string, req *validatedCreate) (*endpointPlacement, error) { if s.deps.Network == nil { return nil, awserrors.Errorf(awserrors.ErrorServerInternal, "RDS networking is not wired on this node") } diff --git a/spinifex/handlers/rds/network_test.go b/spinifex/handlers/rds/network_test.go index d55c76500..25d8a7a9a 100644 --- a/spinifex/handlers/rds/network_test.go +++ b/spinifex/handlers/rds/network_test.go @@ -4,7 +4,7 @@ import ( "testing" "github.com/aws/aws-sdk-go/aws" - "github.com/mulgadc/spinifex/spinifex/handlers/ec2/vpc" + handlers_ec2_vpc "github.com/mulgadc/spinifex/spinifex/handlers/ec2/vpc" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) diff --git a/spinifex/handlers/rds/parametergroup.go b/spinifex/handlers/rds/parametergroup.go index ca615d067..20f7a2fd2 100644 --- a/spinifex/handlers/rds/parametergroup.go +++ b/spinifex/handlers/rds/parametergroup.go @@ -13,7 +13,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/rds" "github.com/mulgadc/spinifex/spinifex/awserrors" - "github.com/nats-io/nats.go/jetstream" + "github.com/mulgadc/spinifex/spinifex/kvstore" ) // AWS's cap on one ModifyDBParameterGroup call. @@ -68,7 +68,7 @@ func (s *Service) CreateDBParameterGroup(ctx context.Context, input *rds.CreateD UpdatedAt: now, } if err := createJSON(ctx, kv, DBParameterGroupMetaKey(name), &rec); err != nil { - if errors.Is(err, jetstream.ErrKeyExists) { + if errors.Is(err, kvstore.ErrExists) { return nil, awserrors.Errorf(awserrors.ErrorDBParameterGroupAlreadyExists, "DB parameter group %s already exists", name) } @@ -198,7 +198,7 @@ func (s *Service) ModifyDBParameterGroup(ctx context.Context, input *rds.ModifyD // Applies the complete effective set to every instance currently attached to // the group. Each record is re-read before the command so a concurrent detach // or delete does not receive parameters for a group it no longer uses. -func (s *Service) propagateParameterGroup(ctx context.Context, kv jetstream.KeyValue, accountID, name string) error { +func (s *Service) propagateParameterGroup(ctx context.Context, kv *kvstore.Bucket, accountID, name string) error { ids, err := instancesUsingGroup(ctx, kv, func(rec *DBInstanceRecord) bool { return rec.DBParameterGroupName == name }) @@ -403,7 +403,7 @@ func resolveApplyMethod(spec ParameterSpec, requested string) (string, error) { // is synthesised rather than written on read: the record carries nothing a write // would preserve, and materialising it on a describe would make a read path a // writer for no gain. -func getDBParameterGroup(ctx context.Context, kv jetstream.KeyValue, accountID, name string) (*DBParameterGroupRecord, uint64, error) { +func getDBParameterGroup(ctx context.Context, kv *kvstore.Bucket, accountID, name string) (*DBParameterGroupRecord, uint64, error) { var rec DBParameterGroupRecord rev, found, err := getJSONRevision(ctx, kv, DBParameterGroupMetaKey(name), &rec) if err != nil { @@ -513,7 +513,7 @@ func (s *Service) projectParameterGroupRecord(rec *DBParameterGroupRecord) *rds. // Every path that binds a group to an instance comes through here — create, // modify, restore, the deferred apply and group propagation — so the // cross-engine refusal is one check rather than five. -func (s *Service) resolveGroupParameters(ctx context.Context, kv jetstream.KeyValue, accountID string, engine Engine, group, instanceClass string) ([]Parameter, error) { +func (s *Service) resolveGroupParameters(ctx context.Context, kv *kvstore.Bucket, accountID string, engine Engine, group, instanceClass string) ([]Parameter, error) { rec, _, err := getDBParameterGroup(ctx, kv, accountID, group) if err != nil { return nil, err diff --git a/spinifex/handlers/rds/pending.go b/spinifex/handlers/rds/pending.go index ad9e20eb9..8e6cd296e 100644 --- a/spinifex/handlers/rds/pending.go +++ b/spinifex/handlers/rds/pending.go @@ -7,7 +7,7 @@ import ( "log/slog" "github.com/aws/aws-sdk-go/aws" - "github.com/nats-io/nats.go/jetstream" + "github.com/mulgadc/spinifex/spinifex/kvstore" ) // The single drain for everything a modify recorded but has not delivered. Both @@ -25,7 +25,7 @@ import ( // The caller has already moved the instance into modifying; the record is left // there, because the engine has to come back and report healthy before the // reconciler calls it available. -func (s *Service) applyPendingModifications(ctx context.Context, kv jetstream.KeyValue, accountID string, rec *DBInstanceRecord) error { +func (s *Service) applyPendingModifications(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord) error { pending := rec.PendingModifiedValues if pending.empty() { return nil @@ -146,7 +146,7 @@ func (s *Service) applyPendingModifications(ctx context.Context, kv jetstream.Ke // A re-resolve, never a merge: the whole set is recomputed from the catalog and // the new group's overrides, so a parameter the old group set and the new one // does not reverts to its default rather than lingering. -func (s *Service) applyParameterGroup(ctx context.Context, kv jetstream.KeyValue, accountID string, rec *DBInstanceRecord, group, instanceClass string, tolerateUnreachableAgent bool) error { +func (s *Service) applyParameterGroup(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord, group, instanceClass string, tolerateUnreachableAgent bool) error { engine, err := LookupEngine(rec.Engine) if err != nil { return err @@ -187,7 +187,7 @@ func (s *Service) applyParameterGroup(ctx context.Context, kv jetstream.KeyValue // here in place of the dead agent's reply, with nothing pending a reboot. // Persistence uses a detached ctx: the budget that just found the agent // unreachable must not also fail the write that unblocks recovery from it. -func (s *Service) deferParameterApplyToReplacement(ctx context.Context, kv jetstream.KeyValue, rec *DBInstanceRecord, resolved []Parameter) error { +func (s *Service) deferParameterApplyToReplacement(ctx context.Context, kv *kvstore.Bucket, rec *DBInstanceRecord, resolved []Parameter) error { slog.WarnContext(ctx, "rds: the instance agent is unreachable; deferring the parameter apply to the replacement VM's fresh agent", "dbInstance", rec.DBInstanceIdentifier) rec.Bootstrap.ResolvedParameters = resolved @@ -211,7 +211,7 @@ func (s *Service) deferParameterApplyToReplacement(ctx context.Context, kv jetst // The event is deduped off the stored flag because the reconciler retries a // failed modify every pass. A failure to persist joins the cause rather than // replacing it — the apply is what the caller asked about. -func (s *Service) recordParameterApplyFailure(ctx context.Context, kv jetstream.KeyValue, accountID string, +func (s *Service) recordParameterApplyFailure(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord, group string, cause error, ) error { first := false @@ -237,7 +237,7 @@ func (s *Service) recordParameterApplyFailure(ctx context.Context, kv jetstream. // the control plane has already grown the volume, and this extends the guest's // filesystem onto the capacity that is now there. Both ext4 and XFS grow while // mounted, so it needs no ordering against the engine start. -func (s *Service) finishFilesystemGrow(ctx context.Context, kv jetstream.KeyValue, accountID string, rec *DBInstanceRecord) error { +func (s *Service) finishFilesystemGrow(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord) error { if err := s.growFilesystem(ctx, accountID, rec.DBInstanceIdentifier); err != nil { return fmt.Errorf("extend the filesystem of %s onto its grown volume: %w", rec.DBInstanceIdentifier, err) } diff --git a/spinifex/handlers/rds/reconciler.go b/spinifex/handlers/rds/reconciler.go index f2139a8b2..af4890c21 100644 --- a/spinifex/handlers/rds/reconciler.go +++ b/spinifex/handlers/rds/reconciler.go @@ -195,28 +195,23 @@ func (r *Reconciler) reconcileOnce(ctx context.Context) (time.Duration, error) { if err != nil { return 0, err } - buckets, err := AccountBucketNames(ctx, js) + buckets, err := AccountBuckets(ctx, js) if err != nil { return 0, fmt.Errorf("rds reconciler: enumerate account buckets: %w", err) } var revisit time.Duration var failures []error for _, bucket := range buckets { - kv, err := js.KeyValue(ctx, bucket) - if err != nil { - failures = append(failures, fmt.Errorf("open %s: %w", bucket, err)) - continue - } - due, err := r.reconcileAccount(ctx, kv, AccountIDFromBucketName(bucket)) + due, err := r.reconcileAccount(ctx, bucket, AccountIDFromBucketName(bucket.Name())) revisit = reconciler.Earliest(revisit, due) if err != nil { - failures = append(failures, fmt.Errorf("reconcile %s: %w", bucket, err)) + failures = append(failures, fmt.Errorf("reconcile %s: %w", bucket.Name(), err)) } } return revisit, errors.Join(failures...) } -func (r *Reconciler) reconcileAccount(ctx context.Context, kv jetstream.KeyValue, accountID string) (time.Duration, error) { +func (r *Reconciler) reconcileAccount(ctx context.Context, kv *kvstore.Bucket, accountID string) (time.Duration, error) { ids, err := ListDBInstanceIDs(ctx, kv) if err != nil { return 0, err @@ -244,7 +239,7 @@ func (r *Reconciler) reconcileAccount(ctx context.Context, kv jetstream.KeyValue // The reconciler owns every transitional state that no single API call can // finish: the one it drives itself (creating), and the ones whose caller may // have died partway through. A settled instance is left alone. -func (r *Reconciler) reconcileInstance(ctx context.Context, kv jetstream.KeyValue, accountID, id string) (time.Duration, error) { +func (r *Reconciler) reconcileInstance(ctx context.Context, kv *kvstore.Bucket, accountID, id string) (time.Duration, error) { var rec DBInstanceRecord rev, found, err := getJSONRevision(ctx, kv, DBInstanceKey(id), &rec) if err != nil || !found { @@ -350,7 +345,7 @@ func (r *Reconciler) failureRevisit(rec *DBInstanceRecord) time.Duration { // vpcd applies the requested group list declaratively, so this removes the ENI // from the system VPC's default group rather than merely adding the new one // alongside it — and that default group is the whole of the exposure. -func (r *Reconciler) remediateSystemENISG(ctx context.Context, kv jetstream.KeyValue, +func (r *Reconciler) remediateSystemENISG(ctx context.Context, kv *kvstore.Bucket, rev uint64, rec *DBInstanceRecord) (bool, error) { if rec.SystemENIID == "" || rec.Status == StatusDeleting { return false, nil @@ -409,7 +404,7 @@ func (r *Reconciler) ensuredSystemSG(ctx context.Context) (string, error) { // operator acts. Only reported: deleting the payload on the grounds that a // healthy engine implies a completed bootstrap is exactly the inference this // protocol refuses to make. -func (r *Reconciler) reportStalePendingBootstrap(ctx context.Context, kv jetstream.KeyValue, +func (r *Reconciler) reportStalePendingBootstrap(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord) error { key := accountID + "/" + rec.DBInstanceIdentifier stale := rec.Status == StatusAvailable && time.Since(rec.CreatedAt) > r.svc.bootstrapTimeout() @@ -454,7 +449,7 @@ func (r *Reconciler) forgetPendingReport(key string) { delete(r.reportedPending, key) } -func (r *Reconciler) reconcileCreating(ctx context.Context, kv jetstream.KeyValue, rev uint64, accountID string, rec *DBInstanceRecord) error { +func (r *Reconciler) reconcileCreating(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord) error { // No lower bound on the heartbeat: the VM is new, so any beat naming it is // necessarily this instance's. healthy, err := r.engineReady(ctx, accountID, rec, time.Time{}) @@ -485,7 +480,7 @@ func withAgentReason(rec *DBInstanceRecord, reason string) string { // Reboot and start both end the same way: the engine comes back and says so. // The API call that began them returns before that happens, so this is what // actually lands the instance in available. -func (r *Reconciler) reconcileRestarting(ctx context.Context, kv jetstream.KeyValue, rev uint64, accountID string, rec *DBInstanceRecord) error { +func (r *Reconciler) reconcileRestarting(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord) error { // The VM keeps its instance ID across a restart, so only a beat sent after // the transition began proves the engine came back rather than that it was // up before it went down. @@ -516,7 +511,7 @@ func (r *Reconciler) reconcileRestarting(ctx context.Context, kv jetstream.KeyVa // half-applied, and the in-guest filesystem grow, which can only run once the // agent is up again. Both are driven from here so a customer's modify completes // without them, not just without them watching. -func (r *Reconciler) reconcileModifying(ctx context.Context, kv jetstream.KeyValue, rev uint64, accountID string, rec *DBInstanceRecord) error { +func (r *Reconciler) reconcileModifying(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord) error { pending := rec.PendingModifiedValues overrun := time.Since(transitionStarted(rec)) > transitionTimeout @@ -590,7 +585,7 @@ func (r *Reconciler) reconcileModifying(ctx context.Context, kv jetstream.KeyVal // A stop whose caller died leaves the VM possibly still running, so the stop is // re-issued rather than assumed: it is idempotent, and a VM no node holds is // confirmed down against the fleet before the record calls it stopped. -func (r *Reconciler) reconcileStopping(ctx context.Context, kv jetstream.KeyValue, rev uint64, accountID string, rec *DBInstanceRecord) error { +func (r *Reconciler) reconcileStopping(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord) error { if r.svc.deps.Instances == nil { return errors.New("rds reconciler: no instance command path configured") } @@ -617,7 +612,7 @@ func (r *Reconciler) reconcileStopping(ctx context.Context, kv jetstream.KeyValu // resource, so replaying it converges rather than failing on what it already // did; only a teardown still stuck past the bound is called failed, which is // what lets the customer retry the delete. -func (r *Reconciler) reconcileDeleting(ctx context.Context, kv jetstream.KeyValue, rev uint64, accountID string, rec *DBInstanceRecord) error { +func (r *Reconciler) reconcileDeleting(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord) error { err := r.svc.teardownDBInstance(ctx, kv, accountID, rec, false) if err == nil || errors.Is(err, errFinalSnapshotInProgress) { return nil @@ -636,7 +631,7 @@ func (r *Reconciler) reconcileDeleting(ctx context.Context, kv jetstream.KeyValu // died. It is returned to where the snapshot found it; the quiesce needs no // undoing, because the agent's own hold deadline is far shorter than this bound // and has already released the engine. -func (r *Reconciler) reconcileBackingUp(ctx context.Context, kv jetstream.KeyValue, rev uint64, accountID string, rec *DBInstanceRecord) error { +func (r *Reconciler) reconcileBackingUp(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord) error { if time.Since(transitionStarted(rec)) <= transitionTimeout { return nil } @@ -660,7 +655,7 @@ func (r *Reconciler) reconcileBackingUp(ctx context.Context, kv jetstream.KeyVal // The record is re-read rather than carried over from the status pass above, // which may have transitioned the instance — a backup must not be started against // a status that has since moved. -func (r *Reconciler) reconcileWindows(ctx context.Context, kv jetstream.KeyValue, accountID, id string) error { +func (r *Reconciler) reconcileWindows(ctx context.Context, kv *kvstore.Bucket, accountID, id string) error { var rec DBInstanceRecord rev, found, err := getJSONRevision(ctx, kv, DBInstanceKey(id), &rec) if err != nil || !found { @@ -680,7 +675,7 @@ func (r *Reconciler) reconcileWindows(ctx context.Context, kv jetstream.KeyValue // snapshot is the authority: it was tagged with the DB snapshot identifier before // the record was flipped, so its presence says the data exists and its absence // says the create died before cutting it. -func (r *Reconciler) reconcileSnapshots(ctx context.Context, kv jetstream.KeyValue, accountID string) (time.Duration, error) { +func (r *Reconciler) reconcileSnapshots(ctx context.Context, kv *kvstore.Bucket, accountID string) (time.Duration, error) { ids, err := ListDBSnapshotIDs(ctx, kv) if err != nil { return 0, fmt.Errorf("list DB snapshots: %w", err) @@ -713,16 +708,16 @@ func (r *Reconciler) reconcileSnapshots(ctx context.Context, kv jetstream.KeyVal // Either adopts the EC2 snapshot the dead worker cut, or withdraws the record so // its identifier is usable again. A record left in creating would otherwise hold // the name forever while naming nothing a customer can restore. -func (r *Reconciler) resolveCreatingSnapshot(ctx context.Context, kv jetstream.KeyValue, rev uint64, +func (r *Reconciler) resolveCreatingSnapshot(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBSnapshotRecord) error { snapshotID, err := r.svc.findEC2SnapshotFor(ctx, accountID, rec.DBSnapshotIdentifier) if err != nil { return err } if snapshotID == "" { - if err := kv.Delete(ctx, DBSnapshotKey(rec.DBSnapshotIdentifier), jetstream.LastRevision(rev)); err != nil { + if err := kv.CompareAndDelete(ctx, DBSnapshotKey(rec.DBSnapshotIdentifier), rev); err != nil { switch { - case errors.Is(err, jetstream.ErrKeyNotFound), errors.Is(err, jetstream.ErrKeyRevisionMismatch): + case errors.Is(err, kvstore.ErrConflict): // A concurrent completion or delete owns the newer revision. return nil default: @@ -743,7 +738,7 @@ func (r *Reconciler) resolveCreatingSnapshot(ctx context.Context, kv jetstream.K // conservative reading is the one that never overstates the snapshot. rec.CrashConsistent = true if err := updateJSON(ctx, kv, DBSnapshotKey(rec.DBSnapshotIdentifier), rev, rec); err != nil { - if errors.Is(err, jetstream.ErrKeyRevisionMismatch) { + if errors.Is(err, kvstore.ErrConflict) { return nil } return err @@ -829,7 +824,7 @@ func (r *Reconciler) engineReady(ctx context.Context, accountID string, rec *DBI // A CAS write, so a transition raced by an agent report or a lifecycle op is // dropped rather than clobbering the newer state; the next pass re-reads. -func (r *Reconciler) transition(ctx context.Context, kv jetstream.KeyValue, rev uint64, rec *DBInstanceRecord, to Status, reason string) error { +func (r *Reconciler) transition(ctx context.Context, kv *kvstore.Bucket, rev uint64, rec *DBInstanceRecord, to Status, reason string) error { if !CanTransition(rec.Status, to) { return fmt.Errorf("illegal transition %s -> %s", rec.Status, to) } @@ -845,7 +840,7 @@ func (r *Reconciler) transition(ctx context.Context, kv jetstream.KeyValue, rev rec.UpdatedAt = time.Now().UTC() if err := updateJSON(ctx, kv, DBInstanceKey(rec.DBInstanceIdentifier), rev, rec); err != nil { - if errors.Is(err, jetstream.ErrKeyRevisionMismatch) { + if errors.Is(err, kvstore.ErrConflict) { slog.DebugContext(ctx, "rds reconciler: transition lost a revision race; retrying next pass", "dbInstance", rec.DBInstanceIdentifier, "to", to) return nil @@ -855,7 +850,7 @@ func (r *Reconciler) transition(ctx context.Context, kv jetstream.KeyValue, rev slog.InfoContext(ctx, "rds reconciler: DB instance transitioned", "dbInstance", rec.DBInstanceIdentifier, "from", from, "to", to, "reason", reason) if from == StatusCreating && to == StatusAvailable { - r.svc.RecordEvent(ctx, AccountIDFromBucketName(kv.Bucket()), EventSourceTypeDBInstance, + r.svc.RecordEvent(ctx, AccountIDFromBucketName(kv.Name()), EventSourceTypeDBInstance, rec.DBInstanceIdentifier, "DB instance is available.", EventCategoryAvailability) } return nil @@ -863,7 +858,7 @@ func (r *Reconciler) transition(ctx context.Context, kv jetstream.KeyValue, rev // transition against a freshly read revision, for a caller that has written the // record itself since the pass opened and so cannot use the revision it read. -func (r *Reconciler) transitionFresh(ctx context.Context, kv jetstream.KeyValue, id string, to Status, reason string) error { +func (r *Reconciler) transitionFresh(ctx context.Context, kv *kvstore.Bucket, id string, to Status, reason string) error { var rec DBInstanceRecord rev, found, err := getJSONRevision(ctx, kv, DBInstanceKey(id), &rec) if err != nil || !found { diff --git a/spinifex/handlers/rds/recovery.go b/spinifex/handlers/rds/recovery.go index 0f052eb49..1668111b7 100644 --- a/spinifex/handlers/rds/recovery.go +++ b/spinifex/handlers/rds/recovery.go @@ -7,7 +7,7 @@ import ( "log/slog" "time" - "github.com/nats-io/nats.go/jetstream" + "github.com/mulgadc/spinifex/spinifex/kvstore" ) // How long an instance may be observed dark before it is called failed, unless @@ -73,7 +73,7 @@ func classifyHealth(obs healthObservation) healthVerdict { // engine reports healthy again — the recovery the AMI's in-guest // Restart=on-failure and EC2's VM auto-restart provide underneath. v1.0 does no // repair of its own. -func (r *Reconciler) reconcileHealth(ctx context.Context, kv jetstream.KeyValue, rev uint64, accountID string, rec *DBInstanceRecord) error { +func (r *Reconciler) reconcileHealth(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord) error { obs := r.observeAgent(accountID, rec) // The VM lookup is a fleet-wide describe fan-out, so it is issued only where @@ -101,7 +101,7 @@ func (r *Reconciler) reconcileHealth(ctx context.Context, kv jetstream.KeyValue, // A healthy heartbeat is the only thing that resets the failure clock. VM state // alone must not: a VM that boots and immediately wedges would otherwise reset // it every pass and mask a persistent fault indefinitely. -func (r *Reconciler) clearFailure(ctx context.Context, kv jetstream.KeyValue, rev uint64, accountID string, rec *DBInstanceRecord) error { +func (r *Reconciler) clearFailure(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord) error { if rec.Status == StatusFailed { rec.UnhealthySince = nil if err := r.transition(ctx, kv, rev, rec, StatusAvailable, ""); err != nil { @@ -122,7 +122,7 @@ func (r *Reconciler) clearFailure(ctx context.Context, kv jetstream.KeyValue, re // Starts the failure clock on the first dark pass and fails the instance once // the grace window has elapsed. Both writes carry the clock, so the timestamp a // later leader measures against is the one the original observation stamped. -func (r *Reconciler) recordFailure(ctx context.Context, kv jetstream.KeyValue, rev uint64, accountID string, rec *DBInstanceRecord, obs healthObservation) error { +func (r *Reconciler) recordFailure(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord, obs healthObservation) error { // Terminal for the control plane in v1.0: nothing retries, so re-recording a // failure that is already recorded would only churn the record. if rec.Status == StatusFailed { @@ -151,10 +151,10 @@ func (r *Reconciler) recordFailure(ctx context.Context, kv jetstream.KeyValue, r // Writes the clock without a status change. A lost revision race is dropped // rather than retried: the record moved under this pass, and the next one // re-reads and re-stamps against whatever it moved to. -func (r *Reconciler) persistFailureClock(ctx context.Context, kv jetstream.KeyValue, rev uint64, rec *DBInstanceRecord) error { +func (r *Reconciler) persistFailureClock(ctx context.Context, kv *kvstore.Bucket, rev uint64, rec *DBInstanceRecord) error { rec.UpdatedAt = time.Now().UTC() err := updateJSON(ctx, kv, DBInstanceKey(rec.DBInstanceIdentifier), rev, rec) - if errors.Is(err, jetstream.ErrKeyRevisionMismatch) { + if errors.Is(err, kvstore.ErrConflict) { slog.DebugContext(ctx, "rds reconciler: failure clock lost a revision race; retrying next pass", "dbInstance", rec.DBInstanceIdentifier) return nil diff --git a/spinifex/handlers/rds/replace.go b/spinifex/handlers/rds/replace.go index ac90cf8b5..fd67d592b 100644 --- a/spinifex/handlers/rds/replace.go +++ b/spinifex/handlers/rds/replace.go @@ -6,8 +6,8 @@ import ( "fmt" "log/slog" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/utils" - "github.com/nats-io/nats.go/jetstream" ) // The one primitive a class change, a storage grow that rides the same outage, @@ -40,7 +40,7 @@ type replaceInput struct { // Only the VM's identity is written here — the class, the storage size and the // pending values belong to the caller that asked for the change, so a failure // part-way through leaves the request still recorded and retryable. -func (s *Service) replaceInstanceVM(ctx context.Context, kv jetstream.KeyValue, accountID string, rec *DBInstanceRecord, in replaceInput) error { +func (s *Service) replaceInstanceVM(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBInstanceRecord, in replaceInput) error { if rec.ENIID == "" || rec.DataVolumeID == "" { return fmt.Errorf("rds: DB instance %s has no persisted ENI and data volume to replace onto", rec.DBInstanceIdentifier) @@ -199,7 +199,7 @@ func (s *Service) launchReplacementVM(ctx context.Context, accountID string, rec // agent's IMDS credentials still resolve to this DB instance. func (s *Service) rewriteInstanceIndex(ctx context.Context, accountID string, rec *DBInstanceRecord, oldInstanceID, newInstanceID string) error { if oldInstanceID != "" && oldInstanceID != newInstanceID { - if err := s.DeleteInstanceIndex(ctx, oldInstanceID); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { + if err := s.DeleteInstanceIndex(ctx, oldInstanceID); err != nil { return fmt.Errorf("rds: drop the superseded instance index entry for %s: %w", rec.DBInstanceIdentifier, err) } } @@ -226,7 +226,7 @@ func (s *Service) rollbackReplacementLaunch( defer cancel() var cleanupErrs []error - if err := s.DeleteInstanceIndex(cleanupCtx, launched.InstanceID); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { + if err := s.DeleteInstanceIndex(cleanupCtx, launched.InstanceID); err != nil { cleanupErrs = append(cleanupErrs, fmt.Errorf("delete replacement instance index %s: %w", launched.InstanceID, err)) } if oldInstanceID != "" { diff --git a/spinifex/handlers/rds/restore.go b/spinifex/handlers/rds/restore.go index 71fbe2853..3b8b723c0 100644 --- a/spinifex/handlers/rds/restore.go +++ b/spinifex/handlers/rds/restore.go @@ -11,9 +11,9 @@ import ( "github.com/aws/aws-sdk-go/service/ec2" "github.com/aws/aws-sdk-go/service/rds" "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/tags" "github.com/mulgadc/spinifex/spinifex/utils" - "github.com/nats-io/nats.go/jetstream" ) // Builds a new DB instance on a volume created from the snapshot. It is a fresh @@ -76,7 +76,7 @@ func (s *Service) RestoreDBInstanceFromDBSnapshot(ctx context.Context, input *rd rec := newRestoredDBInstanceRecord(accountID, req, placement, parameters, snapshot) rollbackRev, createErr := createJSONRevision(ctx, kv, key, &rec) if createErr != nil { - if errors.Is(createErr, jetstream.ErrKeyExists) { + if errors.Is(createErr, kvstore.ErrExists) { return nil, awserrors.Errorf(awserrors.ErrorDBInstanceAlreadyExists, "DB instance %s already exists", req.Identifier) } diff --git a/spinifex/handlers/rds/service.go b/spinifex/handlers/rds/service.go index 91b1ff8fa..a64833e3d 100644 --- a/spinifex/handlers/rds/service.go +++ b/spinifex/handlers/rds/service.go @@ -11,6 +11,7 @@ import ( "time" "github.com/mulgadc/spinifex/spinifex/admin" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/nats-io/nats.go" "github.com/nats-io/nats.go/jetstream" ) @@ -86,6 +87,13 @@ type Service struct { // persisted only on change or on the slower floor. livenessMu sync.Mutex liveness map[string]*agentLiveness + + // buckets memoises one Bucket per account, plus the single system bucket. + // Resolution is per request, so without this every RDS call pays a bucket + // create round trip before it does any work. + bucketsMu sync.Mutex + buckets map[string]*kvstore.Bucket + systemB *kvstore.Bucket } type agentLiveness struct { @@ -158,20 +166,48 @@ func (s *Service) js() (jetstream.JetStream, error) { return jetstream.New(s.nc) } -func (s *Service) bucket(ctx context.Context, accountID string) (jetstream.KeyValue, error) { +// bucket resolves accountID's bucket, memoised: the resolution used to bill a +// stream create per RDS call, and a Bucket also carries the recovery a raw +// handle cannot — a lost stream is reopened under the next read rather than +// failing every call until the process restarts. +func (s *Service) bucket(ctx context.Context, accountID string) (*kvstore.Bucket, error) { js, err := s.js() if err != nil { return nil, err } - return GetOrCreateAccountBucket(ctx, js, accountID) + s.bucketsMu.Lock() + b, ok := s.buckets[accountID] + if !ok { + b = kvstore.NewBucket(js, AccountBucketConfig(accountID)) + if s.buckets == nil { + s.buckets = make(map[string]*kvstore.Bucket) + } + s.buckets[accountID] = b + } + s.bucketsMu.Unlock() + // Opened outside the lock, under the Bucket's own, so one account's slow + // first open does not hold every other account's resolution behind it. + if _, err := b.KV(ctx); err != nil { + return nil, err + } + return b, nil } -func (s *Service) systemBucket(ctx context.Context) (jetstream.KeyValue, error) { +func (s *Service) systemBucket(ctx context.Context) (*kvstore.Bucket, error) { js, err := s.js() if err != nil { return nil, err } - return GetOrCreateSystemBucket(ctx, js) + s.bucketsMu.Lock() + if s.systemB == nil { + s.systemB = kvstore.NewBucket(js, SystemBucketConfig()) + } + b := s.systemB + s.bucketsMu.Unlock() + if _, err := b.KV(ctx); err != nil { + return nil, err + } + return b, nil } // Both paths empty deliberately disables TLS; a partial configuration is an @@ -200,83 +236,79 @@ func (s *Service) tlsAvailable() (bool, error) { return caCert != nil && caKey != nil, nil } +// The eight helpers below are the account bucket's record layer. They are +// adapters over kvstore rather than a codec of their own: the JSON, the error +// wrapping and the stream recovery belong to kvstore, and what stays here is +// the not-found convention its callers are written against. + // Returns (false, nil) when the key is absent. -func getJSON(ctx context.Context, kv jetstream.KeyValue, key string, out any) (bool, error) { - entry, err := kv.Get(ctx, key) - if errors.Is(err, jetstream.ErrKeyNotFound) { +func getJSON[T any](ctx context.Context, b *kvstore.Bucket, key string, out *T) (bool, error) { + v, _, err := kvstore.On[T](b).Get(ctx, key) + if errors.Is(err, kvstore.ErrNotFound) { return false, nil } if err != nil { return false, err } - if err := json.Unmarshal(entry.Value(), out); err != nil { - return false, fmt.Errorf("unmarshal %s: %w", key, err) - } + *out = *v return true, nil } // getJSON plus the entry revision, for callers that follow with a CAS update. -func getJSONRevision(ctx context.Context, kv jetstream.KeyValue, key string, out any) (uint64, bool, error) { - entry, err := kv.Get(ctx, key) - if errors.Is(err, jetstream.ErrKeyNotFound) { +func getJSONRevision[T any](ctx context.Context, b *kvstore.Bucket, key string, out *T) (uint64, bool, error) { + v, rev, err := kvstore.On[T](b).Get(ctx, key) + if errors.Is(err, kvstore.ErrNotFound) { return 0, false, nil } if err != nil { return 0, false, err } - if err := json.Unmarshal(entry.Value(), out); err != nil { - return 0, false, fmt.Errorf("unmarshal %s: %w", key, err) - } - return entry.Revision(), true, nil + *out = *v + return rev, true, nil } -func putJSON(ctx context.Context, kv jetstream.KeyValue, key string, v any) error { - data, err := json.Marshal(v) +// getJSONRevisionAny is getJSONRevision for the one caller whose record type is +// chosen at runtime from the ARN, so no type parameter can name it. The bytes +// still come through kvstore; only the decode is dynamic. +func getJSONRevisionAny(ctx context.Context, b *kvstore.Bucket, key string, out any) (uint64, bool, error) { + raw, rev, err := kvstore.On[json.RawMessage](b).Get(ctx, key) + if errors.Is(err, kvstore.ErrNotFound) { + return 0, false, nil + } if err != nil { - return err + return 0, false, err } - if _, err := kv.Put(ctx, key, data); err != nil { - return fmt.Errorf("put %s: %w", key, err) + if err := json.Unmarshal(*raw, out); err != nil { + return 0, false, fmt.Errorf("unmarshal %s: %w", key, err) } - return nil + return rev, true, nil +} + +func putJSON[T any](ctx context.Context, b *kvstore.Bucket, key string, v T) error { + return kvstore.On[T](b).Set(ctx, key, &v) } // Writes v at key only if nothing is stored there, so the key doubles as a -// cluster-wide reservation. Returns jetstream.ErrKeyExists when it is taken. -func createJSON(ctx context.Context, kv jetstream.KeyValue, key string, v any) error { - _, err := createJSONRevision(ctx, kv, key, v) +// cluster-wide reservation. Returns kvstore.ErrExists when it is taken. +func createJSON[T any](ctx context.Context, b *kvstore.Bucket, key string, v T) error { + _, err := createJSONRevision(ctx, b, key, v) return err } // createJSON plus the created entry's revision, for callers that must undo only // the reservation they created rather than a concurrent replacement. -func createJSONRevision(ctx context.Context, kv jetstream.KeyValue, key string, v any) (uint64, error) { - data, err := json.Marshal(v) - if err != nil { - return 0, err - } - rev, err := kv.Create(ctx, key, data) - if err != nil { - return 0, err - } - return rev, nil +func createJSONRevision[T any](ctx context.Context, b *kvstore.Bucket, key string, v T) (uint64, error) { + return kvstore.On[T](b).Create(ctx, key, &v) } -// 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) +// Writes v at key only if the stored entry is still at rev. Returns +// kvstore.ErrConflict when another writer got there first. +func updateJSON[T any](ctx context.Context, b *kvstore.Bucket, key string, rev uint64, v T) error { + _, err := updateJSONRevision(ctx, b, 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 0, err - } - updatedRev, err := kv.Update(ctx, key, data, rev) - if err != nil { - return 0, fmt.Errorf("update %s: %w", key, err) - } - return updatedRev, nil +func updateJSONRevision[T any](ctx context.Context, b *kvstore.Bucket, key string, rev uint64, v T) (uint64, error) { + return kvstore.On[T](b).CompareAndSet(ctx, key, &v, rev) } diff --git a/spinifex/handlers/rds/snapshot.go b/spinifex/handlers/rds/snapshot.go index c00dc28b9..f3810e10c 100644 --- a/spinifex/handlers/rds/snapshot.go +++ b/spinifex/handlers/rds/snapshot.go @@ -13,9 +13,9 @@ import ( "github.com/aws/aws-sdk-go/service/ec2" "github.com/aws/aws-sdk-go/service/rds" "github.com/mulgadc/spinifex/spinifex/awserrors" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/tags" "github.com/mulgadc/spinifex/spinifex/utils" - "github.com/nats-io/nats.go/jetstream" ) // The EC2 snapshot surface the RDS control plane drives. A DB snapshot is an @@ -78,7 +78,7 @@ func (s *Service) CreateDBSnapshot(ctx context.Context, input *rds.CreateDBSnaps // rev is the revision the caller read the DB instance at: the CAS that moves it // into backing-up is what serialises the two kinds of snapshot against each other // and against a lifecycle op. -func (s *Service) snapshotDBInstance(ctx context.Context, kv jetstream.KeyValue, rev uint64, +func (s *Service) snapshotDBInstance(ctx context.Context, kv *kvstore.Bucket, rev uint64, accountID string, rec *DBInstanceRecord, req *validatedSnapshot) (*DBSnapshotRecord, error) { if s.deps.Snapshots == nil { return nil, errors.New("rds: no snapshot service configured") @@ -107,7 +107,7 @@ func (s *Service) snapshotDBInstance(ctx context.Context, kv jetstream.KeyValue, record := newDBSnapshotRecord(accountID, rec, req) key := DBSnapshotKey(req.DBSnapshotIdentifier) if err := createJSON(ctx, kv, key, &record); err != nil { - if errors.Is(err, jetstream.ErrKeyExists) { + if errors.Is(err, kvstore.ErrExists) { return nil, awserrors.Errorf(awserrors.ErrorDBSnapshotAlreadyExists, "DB snapshot %s already exists", req.DBSnapshotIdentifier) } @@ -308,7 +308,7 @@ func (s *Service) DeleteDBSnapshot(ctx context.Context, input *rds.DeleteDBSnaps // data volume behind it. Shared with the retention sweep, which cannot go // through DeleteDBSnapshot: that rejects the rds: namespace automated snapshots // live in, so a customer can never delete one by hand. -func (s *Service) removeDBSnapshot(ctx context.Context, kv jetstream.KeyValue, accountID string, rec *DBSnapshotRecord) error { +func (s *Service) removeDBSnapshot(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBSnapshotRecord) error { if s.deps.Snapshots == nil { return errors.New("rds: no snapshot service configured") } @@ -330,7 +330,7 @@ func (s *Service) removeDBSnapshot(ctx context.Context, kv jetstream.KeyValue, a } // Last: while it exists the snapshot is still nameable, and a retry re-runs // the steps above, each of which tolerates work it has already done. - if err := kv.Delete(ctx, DBSnapshotKey(id)); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { + if err := kv.Delete(ctx, DBSnapshotKey(id)); err != nil { return fmt.Errorf("rds: delete the DB snapshot record for %s: %w", id, err) } @@ -343,7 +343,7 @@ func (s *Service) removeDBSnapshot(ctx context.Context, kv jetstream.KeyValue, a // the case where an instance restored from it is still running. That is // translated into an RDS fault naming the instance rather than surfaced as the // raw EC2 code, which says nothing a customer can act on. -func (s *Service) deleteEC2Snapshot(ctx context.Context, kv jetstream.KeyValue, accountID string, rec *DBSnapshotRecord) error { +func (s *Service) deleteEC2Snapshot(ctx context.Context, kv *kvstore.Bucket, accountID string, rec *DBSnapshotRecord) error { if rec.SnapshotID == "" { return nil } @@ -365,7 +365,7 @@ func (s *Service) deleteEC2Snapshot(ctx context.Context, kv jetstream.KeyValue, // What the customer has to remove before the snapshot can go. A restored // instance is named; anything else is a volume outside RDS's own bookkeeping, // which is reported as such rather than guessed at. -func (s *Service) describeSnapshotDependents(ctx context.Context, kv jetstream.KeyValue, dbSnapshotIdentifier string) string { +func (s *Service) describeSnapshotDependents(ctx context.Context, kv *kvstore.Bucket, dbSnapshotIdentifier string) string { ids, err := ListDBInstanceIDs(ctx, kv) if err != nil { slog.WarnContext(ctx, "rds: listing the instances restored from a snapshot failed", @@ -392,7 +392,7 @@ func (s *Service) describeSnapshotDependents(ctx context.Context, kv jetstream.K // rather than taken from the retained record: that index is what DeleteVolume // enforces against, and a record written before another snapshot was taken would // otherwise read as "nothing holds it". -func (s *Service) releaseRetainedVolume(ctx context.Context, kv jetstream.KeyValue, rec *DBSnapshotRecord) error { +func (s *Service) releaseRetainedVolume(ctx context.Context, kv *kvstore.Bucket, rec *DBSnapshotRecord) error { if rec.SourceVolumeID == "" { return nil } @@ -408,7 +408,7 @@ func (s *Service) releaseRetainedVolume(ctx context.Context, kv jetstream.KeyVal // Deletes a retained data volume once nothing holds it, or records the holders it // still has. Reports whether the volume actually went, which is what lets the // reaper count the ones it reclaimed. -func (s *Service) reclaimRetainedVolume(ctx context.Context, kv jetstream.KeyValue, +func (s *Service) reclaimRetainedVolume(ctx context.Context, kv *kvstore.Bucket, retained *RetainedVolumeRecord) (bool, error) { holders, err := s.snapshotsHolding(ctx, retained.VolumeID) if err != nil { @@ -438,7 +438,7 @@ func (s *Service) reclaimRetainedVolume(ctx context.Context, kv jetstream.KeyVal default: return false, fmt.Errorf("rds: delete the retained data volume %s: %w", retained.VolumeID, err) } - if err := kv.Delete(ctx, RetainedVolumeKey(retained.VolumeID)); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { + if err := kv.Delete(ctx, RetainedVolumeKey(retained.VolumeID)); err != nil { return false, fmt.Errorf("rds: clear the retained-volume record for %s: %w", retained.VolumeID, err) } slog.InfoContext(ctx, "rds: retained data volume reclaimed; its last snapshot is gone", @@ -586,7 +586,7 @@ func newDBSnapshotRecord(accountID string, rec *DBInstanceRecord, req *validated } // Returns the record plus its revision, for callers that follow with a CAS. -func (s *Service) getDBSnapshot(ctx context.Context, kv jetstream.KeyValue, id string) (*DBSnapshotRecord, uint64, error) { +func (s *Service) getDBSnapshot(ctx context.Context, kv *kvstore.Bucket, id string) (*DBSnapshotRecord, uint64, error) { var rec DBSnapshotRecord rev, found, err := getJSONRevision(ctx, kv, DBSnapshotKey(id), &rec) if err != nil { @@ -601,7 +601,7 @@ func (s *Service) getDBSnapshot(ctx context.Context, kv jetstream.KeyValue, id s // Rejects a taken identifier before any work starts. A snapshot record outlives // the instance it came from, so a name can be held by a snapshot of an instance // that no longer exists. -func (s *Service) checkDBSnapshotAvailable(ctx context.Context, kv jetstream.KeyValue, id string) error { +func (s *Service) checkDBSnapshotAvailable(ctx context.Context, kv *kvstore.Bucket, id string) error { var existing DBSnapshotRecord found, err := getJSON(ctx, kv, DBSnapshotKey(id), &existing) if err != nil { @@ -615,10 +615,10 @@ func (s *Service) checkDBSnapshotAvailable(ctx context.Context, kv jetstream.Key // Withdraws a creating record whose snapshot never happened, on a context // detached from the caller's so an expired request deadline still clears it. -func (s *Service) discardSnapshotRecord(ctx context.Context, kv jetstream.KeyValue, id string) { +func (s *Service) discardSnapshotRecord(ctx context.Context, kv *kvstore.Bucket, id string) { rbCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), rollbackTimeout) defer cancel() - if err := kv.Delete(rbCtx, DBSnapshotKey(id)); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { + if err := kv.Delete(rbCtx, DBSnapshotKey(id)); err != nil { slog.WarnContext(rbCtx, "rds: withdrawing the record of a snapshot that failed left it behind", "dbSnapshot", id, "err", err) } @@ -628,7 +628,7 @@ func (s *Service) discardSnapshotRecord(ctx context.Context, kv jetstream.KeyVal // CAS, returning the status the instance goes back to. That write is the // per-instance guard: a second snapshot, an automated one or a lifecycle // op finds the instance already backing-up and is rejected rather than queued. -func (s *Service) beginSnapshotOperation(ctx context.Context, kv jetstream.KeyValue, rev uint64, +func (s *Service) beginSnapshotOperation(ctx context.Context, kv *kvstore.Bucket, rev uint64, rec *DBInstanceRecord, dbSnapshotIdentifier string) (Status, error) { legal := []Status{StatusAvailable, StatusStopped} if !slices.Contains(legal, rec.Status) { @@ -649,7 +649,7 @@ func (s *Service) beginSnapshotOperation(ctx context.Context, kv jetstream.KeyVa rec.UpdatedAt = now if err := updateJSON(ctx, kv, DBInstanceKey(rec.DBInstanceIdentifier), rev, rec); err != nil { - if errors.Is(err, jetstream.ErrKeyRevisionMismatch) { + if errors.Is(err, kvstore.ErrConflict) { return "", awserrors.Errorf(awserrors.ErrorDBInstanceInvalidState, "DB instance %s changed state concurrently; retry the snapshot", rec.DBInstanceIdentifier) } @@ -662,7 +662,7 @@ func (s *Service) beginSnapshotOperation(ctx context.Context, kv jetstream.KeyVa // leaves a usable instance rather than one stuck in backing-up. A delete // accepted while the snapshot ran owns the record by now, so the resume only // applies to an instance still in backing-up. -func (s *Service) endSnapshotOperation(ctx context.Context, kv jetstream.KeyValue, id string, resume Status) { +func (s *Service) endSnapshotOperation(ctx context.Context, kv *kvstore.Bucket, id string, resume Status) { endCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), defaultTimeout) defer cancel() err := s.updateInstance(endCtx, kv, id, func(stored *DBInstanceRecord) { diff --git a/spinifex/handlers/rds/store.go b/spinifex/handlers/rds/store.go index 9ab3ce773..1924401f6 100644 --- a/spinifex/handlers/rds/store.go +++ b/spinifex/handlers/rds/store.go @@ -11,6 +11,7 @@ import ( "strings" "time" + "github.com/mulgadc/spinifex/spinifex/kvlease" "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/kvutil" "github.com/mulgadc/spinifex/spinifex/migrate" @@ -190,15 +191,51 @@ func AccountBucketName(accountID string) string { return KVBucketRDSAccountPrefix + accountID } +// AccountBucketConfig describes one tenant's bucket. The name carries the +// account, so a Bucket is per-tenant and the Service holds one per account +// rather than one for the package. +func AccountBucketConfig(accountID string) kvstore.Config { + bucket := AccountBucketName(accountID) + return kvstore.Config{ + Name: bucket, + History: KVBucketRDSAccountHistory, + Missing: "rds service: nil nats connection", + OnOpen: func(ctx context.Context, kv jetstream.KeyValue) error { + return migrate.DefaultRegistry.RunKV(ctx, bucket, kv, KVBucketRDSAccountVersion) + }, + } +} + +// AccountBucketReadConfig describes one tenant's bucket for a cross-tenant pass +// that only reads it. No OnOpen: a reconcile pass is not where a schema +// migration should be discovered, and it would bill a version read per account. +func AccountBucketReadConfig(accountID string) kvstore.Config { + return kvstore.Config{ + Name: AccountBucketName(accountID), + History: KVBucketRDSAccountHistory, + Missing: "rds service: nil nats connection", + AttachOnly: true, + } +} + +// SystemBucketConfig describes the cross-tenant bucket holding the instanceID → +// DB instance reverse index. +func SystemBucketConfig() kvstore.Config { + return kvstore.Config{ + Name: KVBucketRDSSystem, + History: KVBucketRDSSystemHistory, + Missing: "rds service: nil nats connection", + OnOpen: func(ctx context.Context, kv jetstream.KeyValue) error { + return migrate.DefaultRegistry.RunKV(ctx, KVBucketRDSSystem, kv, KVBucketRDSSystemVersion) + }, + } +} + // Creates the bucket on first use; subsequent calls return the existing handle. func GetOrCreateAccountBucket(ctx context.Context, js jetstream.JetStream, accountID string) (jetstream.KeyValue, error) { - bucket := AccountBucketName(accountID) - kv, err := kvutil.GetOrCreateBucket(ctx, js, bucket, KVBucketRDSAccountHistory) + kv, err := kvstore.NewBucket(js, AccountBucketConfig(accountID)).KV(ctx) if err != nil { - return nil, fmt.Errorf("failed to create RDS per-account KV bucket %s: %w", bucket, err) - } - if err := migrate.DefaultRegistry.RunKV(ctx, bucket, kv, KVBucketRDSAccountVersion); err != nil { - return nil, fmt.Errorf("migrate %s: %w", bucket, err) + return nil, fmt.Errorf("failed to create RDS per-account KV bucket %s: %w", AccountBucketName(accountID), err) } return kv, nil } @@ -206,37 +243,21 @@ func GetOrCreateAccountBucket(ctx context.Context, js jetstream.JetStream, accou // Created lazily alongside the first DB instance rather than at daemon boot, so // a cluster with no RDS usage carries no bucket. func GetOrCreateSystemBucket(ctx context.Context, js jetstream.JetStream) (jetstream.KeyValue, error) { - kv, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketRDSSystem, KVBucketRDSSystemHistory) + kv, err := kvstore.NewBucket(js, SystemBucketConfig()).KV(ctx) if err != nil { return nil, fmt.Errorf("failed to create RDS system KV bucket %s: %w", KVBucketRDSSystem, err) } - if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketRDSSystem, kv, KVBucketRDSSystemVersion); err != nil { - return nil, fmt.Errorf("migrate %s: %w", KVBucketRDSSystem, err) - } return kv, nil } -// kvutil.GetOrCreateBucket exposes no TTL knob, so the lease bucket attaches -// directly and creates only when absent. +// InitLeaderBucket creates (or attaches to) the shared spinifex-rds-leader +// bucket used for the reconciler's leader-lease CAS lock. func InitLeaderBucket(ctx context.Context, js jetstream.JetStream) (jetstream.KeyValue, error) { - // Attach before create. This runs on every reconcile tick, and CreateKeyValue - // against a bucket that exists with any other config is a STREAM.CREATE the - // meta leader answers with an error, so creating first bills one per tick. - kv, err := js.KeyValue(ctx, KVBucketRDSLeader) - if errors.Is(err, jetstream.ErrBucketNotFound) { - kv, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{ - Bucket: KVBucketRDSLeader, - History: 1, - TTL: KVBucketRDSLeaderTTL, - }) - } - if err != nil { - return nil, fmt.Errorf("failed to create or open RDS leader bucket %s: %w", KVBucketRDSLeader, err) - } - if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketRDSLeader, kv, KVBucketRDSLeaderVersion); err != nil { - return nil, fmt.Errorf("migrate %s: %w", KVBucketRDSLeader, err) - } - return kv, nil + return kvlease.OpenBucket(ctx, js, kvlease.BucketConfig{ + Name: KVBucketRDSLeader, + TTL: KVBucketRDSLeaderTTL, + Version: KVBucketRDSLeaderVersion, + }) } // Every RDS per-account bucket in the cluster. The reconciler and the DNS @@ -255,22 +276,33 @@ func AccountBucketNames(ctx context.Context, js jetstream.JetStream) ([]string, return names, nil } -// AccountWatchBuckets returns every per-account bucket as a watchable handle, -// for a reconciler that must be woken by a cluster change rather than poll for -// it. The set is re-read on each call because a new account's bucket appears +// AccountBuckets returns every per-account bucket as a typed handle, for the +// cross-tenant passes — the reconciler, the backup sweep and the DNS desired +// set. The set is re-read on each call because a new account's bucket appears // without notice: JetStream publishes no bucket-created event. -func AccountWatchBuckets(ctx context.Context, js jetstream.JetStream) ([]*kvstore.Bucket, error) { +// +// The handles are unopened, so a bucket that cannot be reached surfaces from +// the pass that reads it rather than costing every caller an open it may not +// need. They attach and never create: these passes enumerate the buckets they +// read, so an absent one means it was deleted under them, not that it is new. +func AccountBuckets(ctx context.Context, js jetstream.JetStream) ([]*kvstore.Bucket, error) { names, err := AccountBucketNames(ctx, js) if err != nil { return nil, err } buckets := make([]*kvstore.Bucket, 0, len(names)) for _, name := range names { - buckets = append(buckets, kvstore.NewBucket(js, kvstore.Config{Name: name, History: 1})) + buckets = append(buckets, kvstore.NewBucket(js, AccountBucketReadConfig(AccountIDFromBucketName(name)))) } return buckets, nil } +// AccountWatchBuckets is AccountBuckets for a reconciler that must be woken by +// a cluster change rather than poll for it. +func AccountWatchBuckets(ctx context.Context, js jetstream.JetStream) ([]*kvstore.Bucket, error) { + return AccountBuckets(ctx, js) +} + // The account a per-account bucket belongs to, for callers that enumerated // buckets rather than starting from an account. func AccountIDFromBucketName(bucket string) string { @@ -279,13 +311,13 @@ func AccountIDFromBucketName(bucket string) string { // The DB instance identifiers held in one account bucket. An empty bucket // yields no names rather than an error. -func ListDBInstanceIDs(ctx context.Context, kv jetstream.KeyValue) ([]string, error) { +func ListDBInstanceIDs(ctx context.Context, kv *kvstore.Bucket) ([]string, error) { return listNames(ctx, kv, DBInstancesPrefix()) } // Manual and automated snapshots alike: the type is on the record, so a listing // filtered by it still has to read every one. -func ListDBSnapshotIDs(ctx context.Context, kv jetstream.KeyValue) ([]string, error) { +func ListDBSnapshotIDs(ctx context.Context, kv *kvstore.Bucket) ([]string, error) { names, err := listNames(ctx, kv, DBSnapshotsPrefix()) if err != nil { return nil, err @@ -298,7 +330,7 @@ func ListDBSnapshotIDs(ctx context.Context, kv jetstream.KeyValue) ([]string, er // Walks the .../meta keys, which is what makes a group's own record findable // among the per-parameter keys hanging off the same prefix. -func ListDBParameterGroupNames(ctx context.Context, kv jetstream.KeyValue) ([]string, error) { +func ListDBParameterGroupNames(ctx context.Context, kv *kvstore.Bucket) ([]string, error) { keys, err := bucketKeys(ctx, kv) if err != nil { return nil, err @@ -320,7 +352,7 @@ func ListDBParameterGroupNames(ctx context.Context, kv jetstream.KeyValue) ([]st } // The stored overrides of one parameter group, keyed by parameter name. -func ListDBParameterOverrides(ctx context.Context, kv jetstream.KeyValue, group string) (map[string]DBParameterRecord, error) { +func ListDBParameterOverrides(ctx context.Context, kv *kvstore.Bucket, group string) (map[string]DBParameterRecord, error) { prefix := DBParameterGroupParamsPrefix(group) names, err := listNames(ctx, kv, prefix) if err != nil { @@ -345,7 +377,7 @@ func ListDBParameterOverrides(ctx context.Context, kv jetstream.KeyValue, group // Every account's automated-backup index, grouped by DB instance, from one // bucket listing: the retention sweep needs every instance's set, and a Keys call // per instance would cost one listing each. -func ListAutomatedBackups(ctx context.Context, kv jetstream.KeyValue) (map[string][]string, error) { +func ListAutomatedBackups(ctx context.Context, kv *kvstore.Bucket) (map[string][]string, error) { keys, err := bucketKeys(ctx, kv) if err != nil { return nil, err @@ -382,7 +414,7 @@ func splitAutomatedBackupKey(key string) (string, string, bool) { // The leaf names directly under prefix. A nested key belongs to a sub-space and // is skipped rather than reported as a malformed name. -func listNames(ctx context.Context, kv jetstream.KeyValue, prefix string) ([]string, error) { +func listNames(ctx context.Context, kv *kvstore.Bucket, prefix string) ([]string, error) { keys, err := bucketKeys(ctx, kv) if err != nil { return nil, err @@ -403,7 +435,11 @@ func listNames(ctx context.Context, kv jetstream.KeyValue, prefix string) ([]str // An empty bucket yields no keys rather than an error, so a first-ever describe // answers with an empty list instead of failing. -func bucketKeys(ctx context.Context, kv jetstream.KeyValue) ([]string, error) { +func bucketKeys(ctx context.Context, b *kvstore.Bucket) ([]string, error) { + kv, err := b.KV(ctx) + if err != nil { + return nil, err + } keys, err := kv.Keys(ctx) if err != nil { if errors.Is(err, jetstream.ErrNoKeysFound) { diff --git a/spinifex/handlers/rds/subnetgroup.go b/spinifex/handlers/rds/subnetgroup.go index ce61a5087..5390a460d 100644 --- a/spinifex/handlers/rds/subnetgroup.go +++ b/spinifex/handlers/rds/subnetgroup.go @@ -13,7 +13,7 @@ import ( "github.com/aws/aws-sdk-go/service/ec2" "github.com/aws/aws-sdk-go/service/rds" "github.com/mulgadc/spinifex/spinifex/awserrors" - "github.com/nats-io/nats.go/jetstream" + "github.com/mulgadc/spinifex/spinifex/kvstore" ) // AWS's own name rules. The name is a KV key rather than a DNS label, so the @@ -76,7 +76,7 @@ func (s *Service) CreateDBSubnetGroup(ctx context.Context, input *rds.CreateDBSu // Create rather than Put, so two concurrent creates of one name have exactly // one winner instead of both reporting success over each other's subnets. if err := createJSON(ctx, kv, DBSubnetGroupKey(name), &rec); err != nil { - if errors.Is(err, jetstream.ErrKeyExists) { + if errors.Is(err, kvstore.ErrExists) { return nil, awserrors.Errorf(awserrors.ErrorDBSubnetGroupAlreadyExists, "DB subnet group %s already exists", name) } @@ -234,7 +234,7 @@ func (s *Service) resolveGroupSubnets(ctx context.Context, accountID string, req // The record plus its revision. A missing group raises AWS's own fault, so a // well-formed name that resolves to nothing is distinguishable from a bad one. -func getDBSubnetGroup(ctx context.Context, kv jetstream.KeyValue, name string) (*DBSubnetGroupRecord, uint64, error) { +func getDBSubnetGroup(ctx context.Context, kv *kvstore.Bucket, name string) (*DBSubnetGroupRecord, uint64, error) { var rec DBSubnetGroupRecord rev, found, err := getJSONRevision(ctx, kv, DBSubnetGroupKey(name), &rec) if err != nil { @@ -249,7 +249,7 @@ func getDBSubnetGroup(ctx context.Context, kv jetstream.KeyValue, name string) ( // The identifiers of every instance in the account matching uses, sorted. Both // group deletes share it: an in-use guard that missed one instance would strand // a live database's configuration. -func instancesUsingGroup(ctx context.Context, kv jetstream.KeyValue, uses func(*DBInstanceRecord) bool) ([]string, error) { +func instancesUsingGroup(ctx context.Context, kv *kvstore.Bucket, uses func(*DBInstanceRecord) bool) ([]string, error) { ids, err := ListDBInstanceIDs(ctx, kv) if err != nil { return nil, err diff --git a/spinifex/handlers/rds/tags.go b/spinifex/handlers/rds/tags.go index c5ce5a2b2..869da4de6 100644 --- a/spinifex/handlers/rds/tags.go +++ b/spinifex/handlers/rds/tags.go @@ -11,7 +11,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/rds" "github.com/mulgadc/spinifex/spinifex/awserrors" - "github.com/nats-io/nats.go/jetstream" + "github.com/mulgadc/spinifex/spinifex/kvstore" ) // AWS's own tag limits. The reserved prefix is refused rather than dropped: a @@ -180,9 +180,9 @@ func (s *Service) resolveTaggable(resourceName, accountID string) (taggableResou // The record plus its revision. A missing record raises the resource's own // not-found fault, so an ARN that parses but names nothing is distinguishable // from one that does not parse. -func readTagged(ctx context.Context, kv jetstream.KeyValue, resource taggableResource, identifier string) (TaggedRecord, uint64, error) { +func readTagged(ctx context.Context, kv *kvstore.Bucket, resource taggableResource, identifier string) (TaggedRecord, uint64, error) { rec := resource.newRecord() - rev, found, err := getJSONRevision(ctx, kv, resource.key(identifier), rec) + rev, found, err := getJSONRevisionAny(ctx, kv, resource.key(identifier), rec) if err != nil { return nil, 0, err } @@ -195,7 +195,7 @@ func readTagged(ctx context.Context, kv jetstream.KeyValue, resource taggableRes // Applies mutate to the record's tags under CAS. Losing the CAS means another // tag write landed first, so the mutation is replayed against its result rather // than overwriting it. -func mutateTags(ctx context.Context, kv jetstream.KeyValue, resource taggableResource, identifier string, mutate func(map[string]string) (map[string]string, error)) error { +func mutateTags(ctx context.Context, kv *kvstore.Bucket, resource taggableResource, identifier string, mutate func(map[string]string) (map[string]string, error)) error { key := resource.key(identifier) for range tagWriteAttempts { rec, rev, err := readTagged(ctx, kv, resource, identifier) @@ -214,7 +214,7 @@ func mutateTags(ctx context.Context, kv jetstream.KeyValue, resource taggableRes } // A revision mismatch is a lost race with another writer, not a // duplicate: re-read the record and re-apply the tag mutation. - if !errors.Is(err, jetstream.ErrKeyRevisionMismatch) { + if !errors.Is(err, kvstore.ErrConflict) { return err } } diff --git a/spinifex/handlers/sts/assume_role.go b/spinifex/handlers/sts/assume_role.go index e69843f65..467380c9a 100644 --- a/spinifex/handlers/sts/assume_role.go +++ b/spinifex/handlers/sts/assume_role.go @@ -19,10 +19,10 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/iam" "github.com/aws/aws-sdk-go/service/sts" - "github.com/nats-io/nats.go/jetstream" "github.com/mulgadc/spinifex/spinifex/awserrors" handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/utils" ) @@ -553,8 +553,8 @@ func (s *STSServiceImpl) mintSession(ctx context.Context, env sessionEnvelope, d ExpiresAt: expiresAt, CreatedAt: now, } - if err := putSessionCredential(ctx, s.sessionsBucket, cred); err != nil { - if errors.Is(err, jetstream.ErrKeyExists) { + if err := putSessionCredential(ctx, s.sessions, cred); err != nil { + if errors.Is(err, kvstore.ErrExists) { continue } return nil, "", "", fmt.Errorf("persist session credential: %w", err) diff --git a/spinifex/handlers/sts/service_impl.go b/spinifex/handlers/sts/service_impl.go index dcaede670..fe25b12d0 100644 --- a/spinifex/handlers/sts/service_impl.go +++ b/spinifex/handlers/sts/service_impl.go @@ -7,6 +7,7 @@ import ( "log/slog" handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/nats-io/nats.go" "github.com/nats-io/nats.go/jetstream" ) @@ -26,11 +27,11 @@ const masterKeySize = 32 // AES-256, must match handlers_iam. // change. RunJanitor is the exception: it already owns a real context and // threads it into the sweep. type STSServiceImpl struct { - natsConn *nats.Conn - js jetstream.JetStream - sessionsBucket jetstream.KeyValue - iamSvc handlers_iam.IAMService - masterKey []byte + natsConn *nats.Conn + js jetstream.JetStream + sessions *kvstore.Store[SessionCredential] + iamSvc handlers_iam.IAMService + masterKey []byte } var _ STSService = (*STSServiceImpl)(nil) @@ -57,7 +58,7 @@ func NewSTSServiceImpl(ctx context.Context, natsConn *nats.Conn, iamSvc handlers return nil, fmt.Errorf("get JetStream context: %w", err) } - sessionsBucket, err := initSessionCredentialsBucket(ctx, js, replicas) + sessions, err := initSessionCredentialsStore(ctx, js, replicas) if err != nil { return nil, fmt.Errorf("init session credentials bucket: %w", err) } @@ -67,10 +68,10 @@ func NewSTSServiceImpl(ctx context.Context, natsConn *nats.Conn, iamSvc handlers "replicas", replicas) return &STSServiceImpl{ - natsConn: natsConn, - js: js, - sessionsBucket: sessionsBucket, - iamSvc: iamSvc, - masterKey: masterKey, + natsConn: natsConn, + js: js, + sessions: sessions, + iamSvc: iamSvc, + masterKey: masterKey, }, nil } diff --git a/spinifex/handlers/sts/service_impl_test.go b/spinifex/handlers/sts/service_impl_test.go index a15e84b0b..1b8ba311d 100644 --- a/spinifex/handlers/sts/service_impl_test.go +++ b/spinifex/handlers/sts/service_impl_test.go @@ -60,8 +60,8 @@ func TestNewSTSServiceImpl_RejectsWrongMasterKeySize(t *testing.T) { func TestNewSTSServiceImpl_InitializesBucket(t *testing.T) { svc, _ := newTestSetup(t) - require.NotNil(t, svc.sessionsBucket) - assert.Equal(t, KVBucketSessionCredentials, svc.sessionsBucket.Bucket()) + require.NotNil(t, svc.sessions) + assert.Equal(t, KVBucketSessionCredentials, svc.sessions.Name()) } func TestNewSTSServiceImpl_NormalisesNegativeClusterSize(t *testing.T) { @@ -116,7 +116,7 @@ func TestLookupSessionCredential_HitRoundTrips(t *testing.T) { ExpiresAt: now.Add(time.Hour), CreatedAt: now, } - require.NoError(t, putSessionCredential(t.Context(), svc.sessionsBucket, cred)) + require.NoError(t, putSessionCredential(t.Context(), svc.sessions, cred)) got, err := svc.LookupSessionCredential(cred.AccessKeyID) require.NoError(t, err) @@ -134,13 +134,13 @@ func TestLookupSessionCredential_UnmarshalFailureSurfacesError(t *testing.T) { // the prefix is valid (so the lookup reaches the bucket) but the JSON // body is garbage. This guards against the "lookup returns nil silently // on parse failure" silent-failure mode. - _, err := svc.sessionsBucket.Put(t.Context(), akid, []byte("not json")) + _, err := sessionsKV(t, svc).Put(t.Context(), akid, []byte("not json")) require.NoError(t, err) got, err := svc.LookupSessionCredential(akid) require.Error(t, err) assert.Nil(t, got) - assert.Contains(t, err.Error(), "unmarshal session credential") + assert.Contains(t, err.Error(), "decode") } // Sanity check: the production marshaller round-trips through the lookup diff --git a/spinifex/handlers/sts/session_credentials.go b/spinifex/handlers/sts/session_credentials.go index 5f58e4d82..e3bcd26c9 100644 --- a/spinifex/handlers/sts/session_credentials.go +++ b/spinifex/handlers/sts/session_credentials.go @@ -11,6 +11,7 @@ import ( "strings" "time" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/kvutil" "github.com/mulgadc/spinifex/spinifex/migrate" "github.com/mulgadc/spinifex/spinifex/otelsetup" @@ -63,27 +64,35 @@ type SessionCredential struct { CreatedAt time.Time `json:"created_at"` } -// initSessionCredentialsBucket opens (or creates) the session-credentials KV bucket. -// History is fixed at 1: credentials are write-once at mint and delete-once at expiry. -func initSessionCredentialsBucket(ctx context.Context, js jetstream.JetStream, replicas int) (jetstream.KeyValue, error) { - // kvutil clamps replicas to a minimum of 1, so a zero clusterSize still creates. - kv, err := kvutil.GetOrCreateBucketWithReplicas(ctx, js, KVBucketSessionCredentials, 1, replicas) - if err != nil { +// sessionCredentialsConfig describes the session-credentials bucket. History is +// fixed at 1: credentials are write-once at mint and delete-once at expiry. +func sessionCredentialsConfig(replicas int) kvstore.Config { + return kvstore.Config{ + Name: KVBucketSessionCredentials, + History: 1, + Replicas: replicas, + Missing: "session credentials KV bucket not initialized", + OnOpen: func(ctx context.Context, kv jetstream.KeyValue) error { + return migrate.DefaultRegistry.RunKV(ctx, KVBucketSessionCredentials, kv, KVBucketSessionCredentialsVersion) + }, + } +} + +// initSessionCredentialsStore opens (or creates) the session-credentials bucket. +func initSessionCredentialsStore(ctx context.Context, js jetstream.JetStream, replicas int) (*kvstore.Store[SessionCredential], error) { + // kvstore clamps replicas to a minimum of 1, so a zero clusterSize still creates. + store := kvstore.New[SessionCredential](js, sessionCredentialsConfig(replicas)) + // Opened eagerly: a bucket that cannot be created must fail service + // construction, not the first AssumeRole that needs to mint into it. + if _, err := store.KV(ctx); err != nil { return nil, fmt.Errorf("open session credentials bucket: %w", err) } - - if err := migrate.DefaultRegistry.RunKV( - ctx, KVBucketSessionCredentials, kv, KVBucketSessionCredentialsVersion, - ); err != nil { - return nil, fmt.Errorf("migrate %s: %w", KVBucketSessionCredentials, err) - } - - return kv, nil + return store, nil } // putSessionCredential persists a SessionCredential via CAS create, enforcing the ASIA-prefix -// invariant. Returns jetstream.ErrKeyExists on collision so callers can retry. -func putSessionCredential(ctx context.Context, bucket jetstream.KeyValue, cred *SessionCredential) error { +// invariant. Returns kvstore.ErrExists on collision so callers can retry. +func putSessionCredential(ctx context.Context, store *kvstore.Store[SessionCredential], cred *SessionCredential) error { if cred == nil { return errors.New("nil session credential") } @@ -91,11 +100,7 @@ func putSessionCredential(ctx context.Context, bucket jetstream.KeyValue, cred * return fmt.Errorf("session AKID must start with %q, got %q", SessionAccessKeyIDPrefix, cred.AccessKeyID) } - data, err := json.Marshal(cred) - if err != nil { - return fmt.Errorf("marshal session credential: %w", err) - } - if _, err := bucket.Create(ctx, cred.AccessKeyID, data); err != nil { + if _, err := store.Create(ctx, cred.AccessKeyID, cred); err != nil { return fmt.Errorf("store session credential: %w", err) } return nil @@ -127,18 +132,14 @@ func (s *STSServiceImpl) LookupSessionCredential(accessKeyID string) (*SessionCr if !strings.HasPrefix(accessKeyID, SessionAccessKeyIDPrefix) { return nil, nil } - entry, err := s.sessionsBucket.Get(ctx, accessKeyID) + cred, _, err := s.sessions.Get(ctx, accessKeyID) if err != nil { - if errors.Is(err, jetstream.ErrKeyNotFound) { + if errors.Is(err, kvstore.ErrNotFound) { return nil, nil } return nil, fmt.Errorf("get session credential: %w", err) } - var cred SessionCredential - if err := json.Unmarshal(entry.Value(), &cred); err != nil { - return nil, fmt.Errorf("unmarshal session credential: %w", err) - } - return &cred, nil + return cred, nil } // RunJanitor periodically removes expired session credentials. Idempotent; multiple @@ -164,10 +165,20 @@ func (s *STSServiceImpl) RunJanitor(ctx context.Context) { // sweepExpired deletes all records whose ExpiresAt is past the grace period. // Per-key errors are logged and skipped; returns the delete count. +// +// Deliberately on the raw handle rather than Store.List: the sweep needs the key +// to delete it, and one undecodable record must not stop it expiring every +// other credential. func (s *STSServiceImpl) sweepExpired(ctx context.Context, now time.Time) int { cutoff := now.Add(-janitorGracePeriod) - keys, err := kvutil.Keys(ctx, s.sessionsBucket) + bucket, err := s.sessions.KV(ctx) + if err != nil { + slog.Warn("STS janitor: open session credential bucket failed", "err", err) + return 0 + } + + keys, err := kvutil.Keys(ctx, bucket) if err != nil { if errors.Is(err, jetstream.ErrNoKeysFound) { return 0 @@ -181,7 +192,7 @@ func (s *STSServiceImpl) sweepExpired(ctx context.Context, now time.Time) int { if key == utils.VersionKey { continue } - entry, err := s.sessionsBucket.Get(ctx, key) + entry, err := bucket.Get(ctx, key) if err != nil { if errors.Is(err, jetstream.ErrKeyNotFound) { continue @@ -202,7 +213,7 @@ func (s *STSServiceImpl) sweepExpired(ctx context.Context, now time.Time) int { continue } - if err := s.sessionsBucket.Delete(ctx, key); err != nil { + if err := bucket.Delete(ctx, key); err != nil { slog.Warn("STS janitor: delete expired session credential failed", "key", key, "err", err) continue diff --git a/spinifex/handlers/sts/session_credentials_test.go b/spinifex/handlers/sts/session_credentials_test.go index 7b6e004a5..9df45beca 100644 --- a/spinifex/handlers/sts/session_credentials_test.go +++ b/spinifex/handlers/sts/session_credentials_test.go @@ -2,11 +2,11 @@ package handlers_sts import ( "context" - "encoding/json" "errors" "testing" "time" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/kvutil" "github.com/mulgadc/spinifex/spinifex/testutil" "github.com/mulgadc/spinifex/spinifex/utils" @@ -15,10 +15,19 @@ import ( "github.com/stretchr/testify/require" ) -func setupBucket(t *testing.T) jetstream.KeyValue { +func setupBucket(t *testing.T) *kvstore.Store[SessionCredential] { t.Helper() _, nc, _ := testutil.StartTestJetStream(t) - kv, err := initSessionCredentialsBucket(t.Context(), testutil.NewJetStream(t, nc), 1) + store, err := initSessionCredentialsStore(t.Context(), testutil.NewJetStream(t, nc), 1) + require.NoError(t, err) + return store +} + +// sessionsKV reaches the raw handle under the store, for the tests that seed a +// record the codec would refuse or assert on a key without decoding it. +func sessionsKV(t *testing.T, svc *STSServiceImpl) jetstream.KeyValue { + t.Helper() + kv, err := svc.sessions.KV(t.Context()) require.NoError(t, err) return kv } @@ -44,7 +53,9 @@ func TestInitSessionCredentialsBucket_StampsVersion(t *testing.T) { _, nc, _ := testutil.StartTestJetStream(t) js := testutil.NewJetStream(t, nc) - kv, err := initSessionCredentialsBucket(t.Context(), js, 1) + store, err := initSessionCredentialsStore(t.Context(), js, 1) + require.NoError(t, err) + kv, err := store.KV(t.Context()) require.NoError(t, err) version, err := kvutil.ReadVersion(t.Context(), kv) @@ -56,13 +67,13 @@ func TestInitSessionCredentialsBucket_Idempotent(t *testing.T) { _, nc, _ := testutil.StartTestJetStream(t) js := testutil.NewJetStream(t, nc) - kv1, err := initSessionCredentialsBucket(t.Context(), js, 1) + store1, err := initSessionCredentialsStore(t.Context(), js, 1) require.NoError(t, err) // Reopen — must return the same bucket without error. - kv2, err := initSessionCredentialsBucket(t.Context(), js, 1) + store2, err := initSessionCredentialsStore(t.Context(), js, 1) require.NoError(t, err) - assert.Equal(t, kv1.Bucket(), kv2.Bucket()) + assert.Equal(t, store1.Name(), store2.Name()) } func TestPutSessionCredential_RoundTrip(t *testing.T) { @@ -71,11 +82,9 @@ func TestPutSessionCredential_RoundTrip(t *testing.T) { require.NoError(t, putSessionCredential(t.Context(), bucket, cred)) - entry, err := bucket.Get(t.Context(), cred.AccessKeyID) + got, _, err := bucket.Get(t.Context(), cred.AccessKeyID) require.NoError(t, err) - var got SessionCredential - require.NoError(t, json.Unmarshal(entry.Value(), &got)) assert.Equal(t, cred.AccessKeyID, got.AccessKeyID) assert.Equal(t, cred.AssumedRoleARN, got.AssumedRoleARN) assert.Equal(t, cred.SessionTokenHMAC, got.SessionTokenHMAC) @@ -108,7 +117,9 @@ func TestPutSessionCredential_RejectsNonASIAPrefix(t *testing.T) { // Bucket must contain no credentials — every put failed before reaching // bucket.Create, so no AKID should be present. - keys, err := bucket.Keys(t.Context()) + kv, err := bucket.KV(t.Context()) + require.NoError(t, err) + keys, err := kv.Keys(t.Context()) if !errors.Is(err, jetstream.ErrNoKeysFound) { require.NoError(t, err) for _, k := range keys { @@ -132,11 +143,11 @@ func TestPutSessionCredential_CollisionReturnsKeyExists(t *testing.T) { require.NoError(t, putSessionCredential(t.Context(), bucket, cred)) - // Second create with the same AKID must surface jetstream.ErrKeyExists so the + // Second create with the same AKID must surface kvstore.ErrExists so the // mint helper can retry with a freshly generated AKID. err := putSessionCredential(t.Context(), bucket, cred) require.Error(t, err) - assert.ErrorIs(t, err, jetstream.ErrKeyExists) + assert.ErrorIs(t, err, kvstore.ErrExists) } func TestVerifySessionToken_MatchAndMismatch(t *testing.T) { @@ -176,7 +187,7 @@ func putCredWithExpiry(t *testing.T, svc *STSServiceImpl, akid string, expiresAt t.Helper() cred := newTestSessionCredential(akid) cred.ExpiresAt = expiresAt - require.NoError(t, putSessionCredential(t.Context(), svc.sessionsBucket, cred)) + require.NoError(t, putSessionCredential(t.Context(), svc.sessions, cred)) } func TestSweepExpired_DeletesPastGraceOnly(t *testing.T) { @@ -192,14 +203,14 @@ func TestSweepExpired_DeletesPastGraceOnly(t *testing.T) { assert.Equal(t, 2, deleted) // Live and within-grace must still exist; past-grace must be gone. - _, err := svc.sessionsBucket.Get(t.Context(), "ASIALIVE000000000001") + _, err := sessionsKV(t, svc).Get(t.Context(), "ASIALIVE000000000001") require.NoError(t, err) - _, err = svc.sessionsBucket.Get(t.Context(), "ASIAJUSTEXPIRED00002") + _, err = sessionsKV(t, svc).Get(t.Context(), "ASIAJUSTEXPIRED00002") require.NoError(t, err) - _, err = svc.sessionsBucket.Get(t.Context(), "ASIAPASTGRACE0000003") + _, err = sessionsKV(t, svc).Get(t.Context(), "ASIAPASTGRACE0000003") require.ErrorIs(t, err, jetstream.ErrKeyNotFound) - _, err = svc.sessionsBucket.Get(t.Context(), "ASIAANCIENT000000004") + _, err = sessionsKV(t, svc).Get(t.Context(), "ASIAANCIENT000000004") require.ErrorIs(t, err, jetstream.ErrKeyNotFound) } @@ -215,18 +226,18 @@ func TestSweepExpired_SkipsCorruptRecord(t *testing.T) { svc, _ := newTestSetup(t) now := time.Now().UTC() - _, err := svc.sessionsBucket.Put(t.Context(), "ASIACORRUPT000000001", []byte("not json")) + _, err := sessionsKV(t, svc).Put(t.Context(), "ASIACORRUPT000000001", []byte("not json")) require.NoError(t, err) putCredWithExpiry(t, svc, "ASIAEXPIRED000000002", now.Add(-24*time.Hour)) deleted := svc.sweepExpired(t.Context(), now) assert.Equal(t, 1, deleted) - _, err = svc.sessionsBucket.Get(t.Context(), "ASIAEXPIRED000000002") + _, err = sessionsKV(t, svc).Get(t.Context(), "ASIAEXPIRED000000002") require.ErrorIs(t, err, jetstream.ErrKeyNotFound) // Corrupt record is left in place — janitor is not authorised to delete // data it cannot interpret; an operator must inspect it. - _, err = svc.sessionsBucket.Get(t.Context(), "ASIACORRUPT000000001") + _, err = sessionsKV(t, svc).Get(t.Context(), "ASIACORRUPT000000001") require.NoError(t, err) } @@ -237,7 +248,7 @@ func TestSweepExpired_IgnoresVersionKey(t *testing.T) { svc, _ := newTestSetup(t) assert.Equal(t, 0, svc.sweepExpired(t.Context(), time.Now().UTC())) - _, err := svc.sessionsBucket.Get(t.Context(), utils.VersionKey) + _, err := sessionsKV(t, svc).Get(t.Context(), utils.VersionKey) require.NoError(t, err, "version key must survive the sweep") } diff --git a/spinifex/kvlease/bucket.go b/spinifex/kvlease/bucket.go index c2d79c7c0..a3bbd9454 100644 --- a/spinifex/kvlease/bucket.go +++ b/spinifex/kvlease/bucket.go @@ -4,8 +4,11 @@ import ( "context" "errors" "fmt" + "log/slog" "time" + "github.com/mulgadc/spinifex/spinifex/migrate" + "github.com/mulgadc/spinifex/spinifex/otelsetup" "github.com/nats-io/nats.go" "github.com/nats-io/nats.go/jetstream" ) @@ -16,9 +19,57 @@ var ( bucketRetryStep = 1 * time.Second ) -// NATSBucket attaches to the named bucket, creating it only when genuinely absent, and -// retrying while JetStream is still forming quorum. Callers with their own bucket initialiser -// should pass that instead. +// BucketConfig describes a lease bucket. History is not a field: a lease key has +// no history worth keeping, so every lease bucket is created with one. +type BucketConfig struct { + Name string + + // TTL expires a lease its holder died without releasing, so a crashed leader + // does not park the key forever. It must match the lease's own Config.TTL. + TTL time.Duration + + // Replicas is the replica count applied on a genuine first creation only. + // Zero leaves it to the server, which means one. + Replicas int + + // Version runs the registered schema migrations for the bucket after it is + // opened. Zero skips them, for a bucket holding nothing but lease keys. + Version int +} + +// OpenBucket attaches to a lease bucket, creating it only when genuinely absent, +// and runs any schema migration registered for it. +// +// Attach before create, because callers re-open this on every reconcile tick: +// CreateKeyValue against a bucket that already exists is a STREAM.CREATE the +// meta leader answers with an error, so creating first bills one per tick. It is +// also what makes Replicas mean "on first creation" rather than a config change +// against a live bucket. +func OpenBucket(ctx context.Context, js jetstream.JetStream, cfg BucketConfig) (jetstream.KeyValue, error) { + kv, err := js.KeyValue(ctx, cfg.Name) + if errors.Is(err, jetstream.ErrBucketNotFound) { + kv, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{ + Bucket: cfg.Name, + History: 1, + TTL: cfg.TTL, + Replicas: cfg.Replicas, + }) + } + if err != nil { + return nil, fmt.Errorf("kvlease: open or create lease bucket %s: %w", cfg.Name, err) + } + if cfg.Version > 0 { + if err := migrate.DefaultRegistry.RunKV(ctx, cfg.Name, kv, cfg.Version); err != nil { + return nil, fmt.Errorf("migrate %s: %w", cfg.Name, err) + } + } + slog.DebugContext(ctx, "kvlease: lease bucket ready", "bucket", cfg.Name, "ttl_ms", otelsetup.Millis(cfg.TTL)) + return kv, nil +} + +// NATSBucket is OpenBucket as a BucketFunc, retrying while JetStream is still +// forming quorum on a cold multi-node start. Callers with their own bucket +// initialiser should pass that instead. func NATSBucket(nc *nats.Conn, bucket string, ttl time.Duration) BucketFunc { return func(ctx context.Context) (jetstream.KeyValue, error) { js, err := jetstream.New(nc) @@ -27,16 +78,7 @@ func NATSBucket(nc *nats.Conn, bucket string, ttl time.Duration) BucketFunc { } deadline := time.Now().Add(bucketRetryFor) for { - // CreateKeyValue reports "stream name already in use" when the bucket exists, so - // attach first and create only when genuinely absent. - kv, err := js.KeyValue(ctx, bucket) - if errors.Is(err, jetstream.ErrBucketNotFound) { - kv, err = js.CreateKeyValue(ctx, jetstream.KeyValueConfig{ - Bucket: bucket, - History: 1, - TTL: ttl, - }) - } + kv, err := OpenBucket(ctx, js, BucketConfig{Name: bucket, TTL: ttl}) if err == nil { return kv, nil } diff --git a/spinifex/kvstore/bucket.go b/spinifex/kvstore/bucket.go index 31094b240..a911e1270 100644 --- a/spinifex/kvstore/bucket.go +++ b/spinifex/kvstore/bucket.go @@ -40,6 +40,11 @@ type Config struct { // as instance state is by the owning node's next write. RecreateIfMissing bool + // AttachOnly opens an existing bucket and never creates one, for a caller + // reading buckets it did not make: one deleted between a listing and the + // read is then an error rather than a silently recreated empty bucket. + AttachOnly bool + // OnOpen runs after every successful open, including a recovery reopen, so // a recreated bucket is re-stamped rather than left unversioned. Owners // use it to run their schema migrations. @@ -175,13 +180,16 @@ func (b *Bucket) Reopen(ctx context.Context) (jetstream.KeyValue, error) { b.mu.Lock() defer b.mu.Unlock() - b.kv = nil + // Checked before the handle is discarded. A Bucket over a fixed handle has + // nothing to resolve against, so dropping it would turn one transient + // outage into a permanent failure for every later call. if b.js == nil { if b.cfg.Missing == "" { return nil, errors.New("kvstore: cannot reopen without a JetStream client") } return nil, errors.New(b.cfg.Missing) } + b.kv = nil // Another goroutine may already have repaired it, so reconnecting comes // first and is the only step that is always safe. @@ -231,6 +239,82 @@ func (b *Bucket) withKV(ctx context.Context, op func(jetstream.KeyValue) error) return op(kv) } +// Delete removes a key. Idempotent: an already-absent key is success. +// +// On Bucket rather than Store because it decodes nothing, so every typed view +// over one bucket shares it and a codec-free caller needs no type parameter. +func (b *Bucket) Delete(ctx context.Context, key string) error { + err := b.withKV(ctx, func(kv jetstream.KeyValue) error { + if err := kv.Delete(ctx, key); err != nil { + return fmt.Errorf("kvstore: delete %s: %w", key, err) + } + return nil + }) + if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { + return err + } + return nil +} + +// CompareAndDelete removes a key only if it is still at rev, returning +// ErrConflict when it is not. +// +// Unlike Delete this is not idempotent: a guarded delete carries an expected +// last sequence, and an absent, tombstoned or purged key fails that check, so +// an already-gone key reports ErrConflict rather than success. +// +// This is the delete half of CompareAndSet, for a caller undoing a reservation +// it made — it must remove its own write and not a replacement that landed on +// top of it. +func (b *Bucket) CompareAndDelete(ctx context.Context, key string, rev uint64) error { + kv, err := b.KV(ctx) + if err != nil { + return err + } + // Not through withKV: the revision is the caller's, so a replay against a + // reopened bucket would be guarded on a revision from the bucket it lost. + if err := kv.Delete(ctx, key, jetstream.LastRevision(rev)); err != nil { + if errors.Is(err, jetstream.ErrKeyRevisionMismatch) { + return fmt.Errorf("%w: %s", ErrConflict, key) + } + return fmt.Errorf("kvstore: delete %s: %w", key, err) + } + return nil +} + +// Purge is Delete for a key whose history must go with it, so a later Create +// sees a key that never existed rather than one with a delete marker on top. +func (b *Bucket) Purge(ctx context.Context, key string) error { + err := b.withKV(ctx, func(kv jetstream.KeyValue) error { + if err := kv.Purge(ctx, key); err != nil { + return fmt.Errorf("kvstore: purge %s: %w", key, err) + } + return nil + }) + if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { + return err + } + return nil +} + +// Exists reports whether a key is present without decoding its value, so a +// record that cannot be unmarshalled is still reported as present. +func (b *Bucket) Exists(ctx context.Context, key string) (bool, error) { + err := b.withKV(ctx, func(kv jetstream.KeyValue) error { + if _, err := kv.Get(ctx, key); err != nil { + return fmt.Errorf("kvstore: get %s: %w", key, err) + } + return nil + }) + if err != nil { + if errors.Is(err, jetstream.ErrKeyNotFound) { + return false, nil + } + return false, err + } + return true, nil +} + // onOpen runs the configured open hook, naming the bucket on failure so a // migration error is not mistaken for a connection one. func (b *Bucket) onOpen(ctx context.Context, kv jetstream.KeyValue) error { @@ -298,6 +382,12 @@ func (b *Bucket) WatchFrom(ctx context.Context, filter string, revision uint64) // open picks the kvutil helper matching the configured TTL and replica count. func (b *Bucket) open(ctx context.Context) (jetstream.KeyValue, error) { switch { + case b.cfg.AttachOnly: + kv, err := b.js.KeyValue(ctx, b.cfg.Name) + if err != nil { + return nil, fmt.Errorf("open KV bucket %s: %w", b.cfg.Name, err) + } + return kv, nil case b.cfg.TTL > 0: return kvutil.GetOrCreateBucketWithOptions(ctx, b.js, kvutil.BucketOptions{ Name: b.cfg.Name, diff --git a/spinifex/kvstore/recovery_test.go b/spinifex/kvstore/recovery_test.go index f1e08f9ea..33cb52315 100644 --- a/spinifex/kvstore/recovery_test.go +++ b/spinifex/kvstore/recovery_test.go @@ -191,11 +191,25 @@ func TestStore_CompareAndSetDoesNotReRun(t *testing.T) { require.NoError(t, err) loseStream(t, js) - err = store.CompareAndSet(t.Context(), "acct-a/one", &record{Name: "two"}, rev) + _, err = store.CompareAndSet(t.Context(), "acct-a/one", &record{Name: "two"}, rev) require.Error(t, err, "a revision-guarded write must not be replayed onto a reopened bucket") assert.ErrorContains(t, err, "acct-a/one") } +// CompareAndDelete guards on the caller's revision, so replaying it against a +// reopened bucket would delete under a revision from the bucket it lost. +func TestBucket_CompareAndDeleteDoesNotReRun(t *testing.T) { + t.Parallel() + _, js, store := newRecoverableStore(t, kvstore.Config{RecreateIfMissing: true}) + rev, err := store.Create(t.Context(), "acct-a/one", &record{Name: "one"}) + require.NoError(t, err) + loseStream(t, js) + + err = store.CompareAndDelete(t.Context(), "acct-a/one", rev) + require.Error(t, err, "a revision-guarded delete must not be replayed onto a reopened bucket") + assert.ErrorContains(t, err, "acct-a/one") +} + // TestBucket_OnOpenRunsOnEveryOpen covers the hook a recreated bucket depends // on: an unstamped, unmigrated bucket is not a recovered one. func TestBucket_OnOpenRunsOnEveryOpen(t *testing.T) { @@ -403,3 +417,31 @@ func TestBucket_DescriptionReachesTheCreatedBucket(t *testing.T) { require.NoError(t, err) assert.Equal(t, "a bucket that says what it is", status.(*jetstream.KeyValueBucketStatus).StreamInfo().Config.Description) } + +// A Store over a fixed handle cannot reopen, so a transient outage must leave +// the handle it already has. Discarding it would end every later call with the +// configured Missing reason, long after the stream came back. +func TestBucket_FixedHandleSurvivesATransientOutage(t *testing.T) { + t.Parallel() + _, nc, _ := testutil.StartTestJetStream(t) + js := testutil.NewJetStream(t, nc) + + const bucket = "kvstore-fixed-handle-test" + cfg := jetstream.KeyValueConfig{Bucket: bucket, History: 1} + kv, err := js.CreateKeyValue(t.Context(), cfg) + require.NoError(t, err) + + store := kvstore.Over[record](nil, kv, kvstore.Config{Name: bucket, Missing: "test: no JetStream client"}) + require.NoError(t, store.Set(t.Context(), "acct-a/one", &record{Name: "one"})) + + loseStream := func() { require.NoError(t, js.DeleteKeyValue(t.Context(), bucket)) } + loseStream() + _, _, err = store.Get(t.Context(), "acct-a/one") + require.Error(t, err, "the call meeting the lost stream still fails") + + _, err = js.CreateKeyValue(t.Context(), cfg) + require.NoError(t, err) + + err = store.Set(t.Context(), "acct-a/one", &record{Name: "two"}) + require.NoError(t, err, "the store must work again once the stream is back") +} diff --git a/spinifex/kvstore/store.go b/spinifex/kvstore/store.go index 70bf1ef83..c7f9d6e99 100644 --- a/spinifex/kvstore/store.go +++ b/spinifex/kvstore/store.go @@ -99,35 +99,6 @@ func (s *Store[T]) Upsert(ctx context.Context, key string, fn func(*T) (bool, er }) } -// Delete removes a key. Idempotent: an already-absent key is success. -func (s *Store[T]) Delete(ctx context.Context, key string) error { - err := s.withKV(ctx, func(kv jetstream.KeyValue) error { - if err := kv.Delete(ctx, key); err != nil { - return fmt.Errorf("kvstore: delete %s: %w", key, err) - } - return nil - }) - if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { - return err - } - return nil -} - -// Purge is Delete for a key whose history must go with it, so a later Create -// sees a key that never existed rather than one with a delete marker on top. -func (s *Store[T]) Purge(ctx context.Context, key string) error { - err := s.withKV(ctx, func(kv jetstream.KeyValue) error { - if err := kv.Purge(ctx, key); err != nil { - return fmt.Errorf("kvstore: purge %s: %w", key, err) - } - return nil - }) - if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { - return err - } - return nil -} - // List decodes every record whose key carries the given prefix. An empty // prefix matches everything and an empty bucket is not an error. A key that // disappears between the listing and the read is skipped. @@ -316,43 +287,29 @@ func (s *Store[T]) Set(ctx context.Context, key string, v *T) error { }) } -// Exists reports whether a key is present without decoding its value, so a -// record that cannot be unmarshalled is still reported as present. -func (s *Store[T]) Exists(ctx context.Context, key string) (bool, error) { - err := s.withKV(ctx, func(kv jetstream.KeyValue) error { - if _, err := kv.Get(ctx, key); err != nil { - return fmt.Errorf("kvstore: get %s: %w", key, err) - } - return nil - }) - if err != nil { - if errors.Is(err, jetstream.ErrKeyNotFound) { - return false, nil - } - return false, err - } - return true, nil -} - // CompareAndSet writes v only if key is still at rev, returning ErrConflict // when it is not. For callers whose lost race is a decision rather than a // failure; callers that should retry want Mutate. -func (s *Store[T]) CompareAndSet(ctx context.Context, key string, v *T, rev uint64) error { +// +// The revision the write landed at is returned, for a caller chaining a second +// CAS onto its own first one. +func (s *Store[T]) CompareAndSet(ctx context.Context, key string, v *T, rev uint64) (uint64, error) { kv, err := s.KV(ctx) if err != nil { - return err + return 0, err } data, err := json.Marshal(v) if err != nil { - return fmt.Errorf("kvstore: encode %s: %w", key, err) + return 0, fmt.Errorf("kvstore: encode %s: %w", key, err) } - if _, err := kv.Update(ctx, key, data, rev); err != nil { + next, err := kv.Update(ctx, key, data, rev) + if err != nil { if errors.Is(err, jetstream.ErrKeyRevisionMismatch) || errors.Is(err, jetstream.ErrKeyExists) { - return fmt.Errorf("%w: %s", ErrConflict, key) + return 0, fmt.Errorf("%w: %s", ErrConflict, key) } - return fmt.Errorf("kvstore: update %s: %w", key, err) + return 0, fmt.Errorf("kvstore: update %s: %w", key, err) } - return nil + return next, nil } // keys lists an already-opened bucket, treating an empty bucket as an empty diff --git a/spinifex/kvstore/store_test.go b/spinifex/kvstore/store_test.go index 73330e980..1a99a51d2 100644 --- a/spinifex/kvstore/store_test.go +++ b/spinifex/kvstore/store_test.go @@ -316,19 +316,58 @@ func TestStore_CompareAndSetRejectsAStaleRevision(t *testing.T) { // Another writer commits first, moving the revision on. seedRaw(t, store, "acct-a/one", record{Name: "one", Count: 2}) - err := store.CompareAndSet(t.Context(), "acct-a/one", &record{Name: "one", Count: 99}, stale) + _, err := store.CompareAndSet(t.Context(), "acct-a/one", &record{Name: "one", Count: 99}, stale) require.ErrorIs(t, err, kvstore.ErrConflict) got, fresh, err := store.Get(t.Context(), "acct-a/one") require.NoError(t, err) assert.Equal(t, 2, got.Count, "a rejected write must leave the winner's value in place") - require.NoError(t, store.CompareAndSet(t.Context(), "acct-a/one", &record{Name: "one", Count: 3}, fresh)) + _, err = store.CompareAndSet(t.Context(), "acct-a/one", &record{Name: "one", Count: 3}, fresh) + require.NoError(t, err) got, _, err = store.Get(t.Context(), "acct-a/one") require.NoError(t, err) assert.Equal(t, 3, got.Count, "the same call commits at the current revision") } +// TestBucket_CompareAndDeleteWithdrawsOnlyItsOwnWrite pins the reservation +// rollback: the holder removes the record it wrote, and a replacement that +// landed on top of it survives. +func TestBucket_CompareAndDeleteWithdrawsOnlyItsOwnWrite(t *testing.T) { + t.Parallel() + store := newStore(t) + own := mustCreate(t, store, "acct-a/mine", record{Name: "mine", Count: 1}) + + require.NoError(t, store.CompareAndDelete(t.Context(), "acct-a/mine", own)) + _, _, err := store.Get(t.Context(), "acct-a/mine") + require.ErrorIs(t, err, kvstore.ErrNotFound, "the holder's own write is withdrawn") + + stale := mustCreate(t, store, "acct-a/theirs", record{Name: "theirs", Count: 1}) + // A replacement lands on top of the reservation. + seedRaw(t, store, "acct-a/theirs", record{Name: "theirs", Count: 2}) + + err = store.CompareAndDelete(t.Context(), "acct-a/theirs", stale) + require.ErrorIs(t, err, kvstore.ErrConflict) + got, _, err := store.Get(t.Context(), "acct-a/theirs") + require.NoError(t, err, "a lost guard must leave the replacement in place") + assert.Equal(t, 2, got.Count) +} + +// Unlike Delete, a guarded delete carries an expected last sequence, so a key +// that is already gone fails the guard rather than reporting success. +func TestBucket_CompareAndDeleteReportsAnAbsentKeyAsConflict(t *testing.T) { + t.Parallel() + store := newStore(t) + rev := mustCreate(t, store, "acct-a/one", record{Name: "one"}) + require.NoError(t, store.CompareAndDelete(t.Context(), "acct-a/one", rev)) + + err := store.CompareAndDelete(t.Context(), "acct-a/one", rev) + require.ErrorIs(t, err, kvstore.ErrConflict, "the spent revision does not delete a second time") + + err = store.CompareAndDelete(t.Context(), "never-existed", 1) + require.ErrorIs(t, err, kvstore.ErrConflict) +} + // TestStore_CreateReturnsTheClaimRevision pins the claim-then-CAS sequence a // single-writer state machine runs: the winner of the create holds a revision // good enough to commit its own first update, with no read in between. @@ -344,11 +383,11 @@ func TestStore_CreateReturnsTheClaimRevision(t *testing.T) { require.NoError(t, err) assert.Equal(t, read, claimed, "the claim revision is the record's current revision") - require.NoError(t, store.CompareAndSet(t.Context(), "singleton", &record{Name: "available"}, claimed), - "the winner promotes its own claim without re-reading") + _, err = store.CompareAndSet(t.Context(), "singleton", &record{Name: "available"}, claimed) + require.NoError(t, err, "the winner promotes its own claim without re-reading") // The spent revision must not commit a second time. - err = store.CompareAndSet(t.Context(), "singleton", &record{Name: "impostor"}, claimed) + _, err = store.CompareAndSet(t.Context(), "singleton", &record{Name: "impostor"}, claimed) require.ErrorIs(t, err, kvstore.ErrConflict) got, _, err := store.Get(t.Context(), "singleton") diff --git a/spinifex/kvutil/cas.go b/spinifex/kvutil/cas.go index a2cd7ac92..633f22aa4 100644 --- a/spinifex/kvutil/cas.go +++ b/spinifex/kvutil/cas.go @@ -70,8 +70,11 @@ func retryCAS(ctx context.Context, cfg CASConfig, op, key string, attempt func() return ctx.Err() } } + // The caller's error is wrapped around the last conflict rather than + // replacing it: the API error naming the condition is what tells a stale + // replica read apart from a genuine race. if cfg.Exhausted != nil { - return cfg.Exhausted(key, attempts) + return fmt.Errorf("%w: %w", cfg.Exhausted(key, attempts), lastErr) } return fmt.Errorf("CAS %s exhausted %d attempts for key %s: %w", op, attempts, key, lastErr) } diff --git a/spinifex/network/external/dhcp/manager.go b/spinifex/network/external/dhcp/manager.go index faf8720fd..cf8022027 100644 --- a/spinifex/network/external/dhcp/manager.go +++ b/spinifex/network/external/dhcp/manager.go @@ -5,7 +5,6 @@ import ( "encoding/json" "errors" "fmt" - "github.com/mulgadc/spinifex/spinifex/otelsetup" "log/slog" "math/rand/v2" "net" @@ -13,6 +12,8 @@ import ( "sync" "time" + "github.com/mulgadc/spinifex/spinifex/otelsetup" + "github.com/nats-io/nats.go" "github.com/nats-io/nats.go/jetstream" "golang.org/x/sync/singleflight" diff --git a/spinifex/network/external/static_pool.go b/spinifex/network/external/static_pool.go index 921459882..a8cc2b4f9 100644 --- a/spinifex/network/external/static_pool.go +++ b/spinifex/network/external/static_pool.go @@ -2,7 +2,6 @@ package external import ( "context" - "encoding/json" "errors" "fmt" "log/slog" @@ -10,7 +9,7 @@ import ( "net/netip" "github.com/mulgadc/spinifex/spinifex/awserrors" - "github.com/mulgadc/spinifex/spinifex/kvutil" + "github.com/mulgadc/spinifex/spinifex/kvstore" "github.com/mulgadc/spinifex/spinifex/migrate" "github.com/nats-io/nats.go/jetstream" ) @@ -21,8 +20,12 @@ const ( // carries over with no manual migration. KVBucketStaticPool = "spinifex-external-ipam" KVBucketStaticPoolVersion = 2 + KVBucketStaticPoolHistory = 5 - staticPoolCASRetries = 5 + // One pool record is one key, so every concurrent allocation in a pool + // contends on it. A launch asking for sixteen public IPs is sixteen + // writers, and each round has exactly one winner. + staticPoolCASRetries = 25 // purposeIGWLRP duplicates handlers/ec2/vpc.PurposeIGWLRP so the // allocator can reserve the gateway slot without importing handlers. @@ -58,11 +61,34 @@ type PoolRecord struct { Cursor string `json:"cursor,omitempty"` } +// errPoolContended means a pool key stayed contended for the whole retry +// budget, which the boot-time drift reconcile tolerates and a caller asking for +// an address does not. +var errPoolContended = errors.New("external IPAM: pool contended") + +// staticPoolConfig describes the pool bucket. The name and history are +// unchanged from the pre-Q1 ExternalIPAM so existing cluster data carries over, +// and the migration runs on every open rather than only the first. +func staticPoolConfig() kvstore.Config { + return kvstore.Config{ + Name: KVBucketStaticPool, + History: KVBucketStaticPoolHistory, + Attempts: staticPoolCASRetries, + Missing: "external IPAM: no JetStream client configured", + OnOpen: func(ctx context.Context, kv jetstream.KeyValue) error { + return migrate.DefaultRegistry.RunKV(ctx, KVBucketStaticPool, kv, KVBucketStaticPoolVersion) + }, + Exhausted: func(key string, attempts int) error { + return fmt.Errorf("%w: %s after %d attempts", errPoolContended, key, attempts) + }, + } +} + // StaticPoolAllocator implements Allocator backed by a NATS JetStream KV // bucket. Bucket schema and CAS semantics are unchanged from the pre-Q1 // ExternalIPAM. type StaticPoolAllocator struct { - kv jetstream.KeyValue + store *kvstore.Store[PoolRecord] pools []ExternalPoolConfig } @@ -71,14 +97,7 @@ var _ Allocator = (*StaticPoolAllocator)(nil) // NewStaticPoolAllocator creates the KV bucket (if missing), runs pending // migrations, and seeds each pool's record. func NewStaticPoolAllocator(ctx context.Context, js jetstream.JetStream, pools []ExternalPoolConfig) (*StaticPoolAllocator, error) { - kv, err := kvutil.GetOrCreateBucket(ctx, js, KVBucketStaticPool, 5) - if err != nil { - return nil, fmt.Errorf("create external IPAM KV bucket: %w", err) - } - if err := migrate.DefaultRegistry.RunKV(ctx, KVBucketStaticPool, kv, KVBucketStaticPoolVersion); err != nil { - return nil, fmt.Errorf("migrate %s: %w", KVBucketStaticPool, err) - } - a := &StaticPoolAllocator{kv: kv, pools: pools} + a := &StaticPoolAllocator{store: kvstore.New[PoolRecord](js, staticPoolConfig()), pools: pools} if err := a.initPools(ctx); err != nil { return nil, fmt.Errorf("init external IPAM pools: %w", err) } @@ -88,12 +107,14 @@ func NewStaticPoolAllocator(ctx context.Context, js jetstream.JetStream, pools [ // NewStaticPoolAllocatorWithKV is the test constructor — skips bucket // creation and migrations. func NewStaticPoolAllocatorWithKV(kv jetstream.KeyValue, pools []ExternalPoolConfig) *StaticPoolAllocator { - return &StaticPoolAllocator{kv: kv, pools: pools} + return &StaticPoolAllocator{store: kvstore.Over[PoolRecord](nil, kv, staticPoolConfig()), pools: pools} } -// KV exposes the underlying bucket so callers that share the bucket -// (notably ExternalIPAM facade tests) can construct sibling allocators. -func (a *StaticPoolAllocator) KV() jetstream.KeyValue { return a.kv } +// KV exposes the underlying bucket so callers that share it (the ExternalIPAM +// facade and its tests) can construct sibling allocators over the same handle. +func (a *StaticPoolAllocator) KV(ctx context.Context) (jetstream.KeyValue, error) { + return a.store.KV(ctx) +} // Pools returns the allocator's pool list. Used by the ExternalIPAM // facade to satisfy pool-by-region lookups. @@ -101,97 +122,86 @@ func (a *StaticPoolAllocator) Pools() []ExternalPoolConfig { return a.pools } // Allocate implements Allocator. func (a *StaticPoolAllocator) Allocate(ctx context.Context, req AllocateRequest) (netip.Addr, error) { - for attempt := range staticPoolCASRetries { - record, revision, err := a.getRecord(ctx, req.PoolName) - if err != nil { - return netip.Addr{}, fmt.Errorf("get external IPAM record: %w", err) + var ip string + err := a.store.Mutate(ctx, req.PoolName, func(record *PoolRecord) (bool, error) { + // A record written before the field existed decodes to a nil map, which + // the assignment below would panic on. + if record.Allocated == nil { + record.Allocated = map[string]ExternalIPAllocation{} } - - ip, err := nextAvailableIP(record) + next, err := nextAvailableIP(record) if err != nil { - return netip.Addr{}, err + return false, err } - - record.Allocated[ip] = ExternalIPAllocation{ + record.Allocated[next] = ExternalIPAllocation{ Purpose: req.Purpose, AllocationID: req.AllocationID, ENIId: req.ENIID, InstanceId: req.InstanceID, } - record.Cursor = ip - - data, err := json.Marshal(record) - if err != nil { - return netip.Addr{}, fmt.Errorf("marshal external IPAM record: %w", err) - } - - if _, err := a.kv.Update(ctx, req.PoolName, data, revision); err != nil { - slog.Debug("external IPAM CAS conflict, retrying", "pool", req.PoolName, "attempt", attempt) - continue - } - - addr, ok := netip.AddrFromSlice(net.ParseIP(ip).To4()) - if !ok { - return netip.Addr{}, fmt.Errorf("parse allocated ip %q", ip) - } - slog.Info("external IPAM allocated IP", "pool", req.PoolName, "ip", ip, "purpose", req.Purpose) - return addr.Unmap(), nil + record.Cursor = next + ip = next + return true, nil + }) + if err != nil { + return netip.Addr{}, err } - return netip.Addr{}, fmt.Errorf("external IPAM allocation failed after CAS retries for pool %s", req.PoolName) + + addr, ok := netip.AddrFromSlice(net.ParseIP(ip).To4()) + if !ok { + return netip.Addr{}, fmt.Errorf("parse allocated ip %q", ip) + } + slog.InfoContext(ctx, "external IPAM allocated IP", "pool", req.PoolName, "ip", ip, "purpose", req.Purpose) + return addr.Unmap(), nil } // Release implements Allocator. func (a *StaticPoolAllocator) Release(ctx context.Context, poolName string, ip netip.Addr, ownerENIID string) error { target := ip.String() - for attempt := range staticPoolCASRetries { - record, revision, err := a.getRecord(ctx, poolName) - if err != nil { - return fmt.Errorf("get external IPAM record for release: %w", err) - } + var released bool + err := a.store.Mutate(ctx, poolName, func(record *PoolRecord) (bool, error) { + // Reset per attempt: a retry that finds the IP already gone must not + // inherit the previous attempt's verdict. + released = false alloc, ok := record.Allocated[target] if !ok { // An owner-scoped release of an already-freed IP is an idempotent // no-op (a duplicated or stale teardown), not an error. if ownerENIID != "" { - return nil + return false, nil } - return fmt.Errorf("IP %s not allocated in pool %s", target, poolName) + return false, fmt.Errorf("IP %s not allocated in pool %s", target, poolName) } // Ownership guard: when a caller names an owner, only free the lease if // it still belongs to that ENI. External IPs are recycled and GC/teardown // sweeps re-emit releases, so a release for a prior owner must not free an // IP since reassigned to a live instance (that would double-allocate). if ownerENIID != "" && alloc.ENIId != "" && alloc.ENIId != ownerENIID { - slog.Info("external IPAM release skip — IP reassigned to a different ENI (stale release)", + slog.InfoContext(ctx, "external IPAM release skip — IP reassigned to a different ENI (stale release)", "pool", poolName, "ip", target, "stale_owner_eni", ownerENIID, "current_owner_eni", alloc.ENIId) - return nil + return false, nil } if alloc.Purpose == purposeIGWLRP { - return fmt.Errorf("cannot release gateway IP %s in pool %s", target, poolName) + return false, fmt.Errorf("cannot release gateway IP %s in pool %s", target, poolName) } delete(record.Allocated, target) - - data, err := json.Marshal(record) - if err != nil { - return fmt.Errorf("marshal external IPAM record: %w", err) - } - - if _, err := a.kv.Update(ctx, poolName, data, revision); err != nil { - slog.Debug("external IPAM release CAS conflict, retrying", "pool", poolName, "attempt", attempt) - continue - } - - slog.Info("external IPAM released IP", "pool", poolName, "ip", target) - return nil + released = true + return true, nil + }) + if err != nil { + return err } - return fmt.Errorf("external IPAM release failed after CAS retries for pool %s", poolName) + if released { + slog.InfoContext(ctx, "external IPAM released IP", "pool", poolName, "ip", target) + } + return nil } // GetPoolRecord returns the current pool record. func (a *StaticPoolAllocator) GetPoolRecord(ctx context.Context, poolName string) (*PoolRecord, error) { - rec, _, err := a.getRecord(ctx, poolName) + rec, _, err := a.store.Get(ctx, poolName) return rec, err } @@ -204,43 +214,63 @@ func (a *StaticPoolAllocator) initPools(ctx context.Context) error { return nil } -func (a *StaticPoolAllocator) initPool(ctx context.Context, pool ExternalPoolConfig) error { - chk, revision, err := a.getRecord(ctx, pool.Name) +// rangesMatch reports whether rec already carries the ranges pool configures. +func rangesMatch(rec *PoolRecord, pool ExternalPoolConfig) bool { + return rec.RangeStart == pool.RangeStart && rec.RangeEnd == pool.RangeEnd && + rec.GwLrpRangeStart == pool.GwLrpRangeStart && rec.GwLrpRangeEnd == pool.GwLrpRangeEnd +} - if err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) { +// reconcileDrift commits the corrected ranges under the store's retry budget. A +// pool is one key, so an ordinary Allocate or Release wins this race routinely; +// conceding to it would leave the drift unfixed until some later boot won. +func (a *StaticPoolAllocator) reconcileDrift(ctx context.Context, pool ExternalPoolConfig) error { + return a.store.Mutate(ctx, pool.Name, func(rec *PoolRecord) (bool, error) { + // Re-checked per attempt: the winner of a lost race may already have + // stored the ranges this node wanted, which is nothing left to do. + if rangesMatch(rec, pool) { + return false, nil + } + rec.RangeStart = pool.RangeStart + rec.RangeEnd = pool.RangeEnd + rec.GwLrpRangeStart = pool.GwLrpRangeStart + rec.GwLrpRangeEnd = pool.GwLrpRangeEnd + return true, nil + }) +} + +func (a *StaticPoolAllocator) initPool(ctx context.Context, pool ExternalPoolConfig) error { + chk, _, err := a.store.Get(ctx, pool.Name) + + if err != nil && !errors.Is(err, kvstore.ErrNotFound) { return err } if err == nil { - if chk.RangeStart != pool.RangeStart || chk.RangeEnd != pool.RangeEnd || - chk.GwLrpRangeStart != pool.GwLrpRangeStart || chk.GwLrpRangeEnd != pool.GwLrpRangeEnd { + if !rangesMatch(chk, pool) { slog.Info("external IPAM pool config drift, reconciling KV", "pool", pool.Name, "old_range", chk.RangeStart+"-"+chk.RangeEnd, "new_range", pool.RangeStart+"-"+pool.RangeEnd, "old_gw_lrp_range", chk.GwLrpRangeStart+"-"+chk.GwLrpRangeEnd, "new_gw_lrp_range", pool.GwLrpRangeStart+"-"+pool.GwLrpRangeEnd) - chk.RangeStart = pool.RangeStart - chk.RangeEnd = pool.RangeEnd - chk.GwLrpRangeStart = pool.GwLrpRangeStart - chk.GwLrpRangeEnd = pool.GwLrpRangeEnd - - data, err := json.Marshal(chk) - if err != nil { - return fmt.Errorf("marshal external IPAM record: %w", err) - } - - if _, err := a.kv.Update(ctx, pool.Name, data, revision); err != nil { - slog.Warn("external IPAM update failed", "pool", pool.Name, "err", err) - return err + if err := a.reconcileDrift(ctx, pool); err != nil { + if !errors.Is(err, errPoolContended) { + return err + } + // Failing the boot over a key that stayed contended would take + // this node down, so an unreconciled drift is reported instead. + slog.WarnContext(ctx, "external IPAM drift reconcile stayed contended; the stored ranges still differ", + "pool", pool.Name, + "want_range", pool.RangeStart+"-"+pool.RangeEnd, + "want_gw_lrp_range", pool.GwLrpRangeStart+"-"+pool.GwLrpRangeEnd, "err", err) } return nil } - slog.Debug("external IPAM pool already initialized", "pool", pool.Name) + slog.DebugContext(ctx, "external IPAM pool already initialized", "pool", pool.Name) return nil } - slog.Info("external IPAM pool not found, creating", "pool", pool.Name) + slog.InfoContext(ctx, "external IPAM pool not found, creating", "pool", pool.Name) gwIP := pool.GatewayIP if gwIP == "" { @@ -263,35 +293,17 @@ func (a *StaticPoolAllocator) initPool(ctx context.Context, pool ExternalPoolCon }, } - data, err := json.Marshal(record) - if err != nil { - return fmt.Errorf("marshal pool record: %w", err) - } - - if _, err := a.kv.Create(ctx, pool.Name, data); err != nil { - if errors.Is(err, jetstream.ErrKeyExists) { + if _, err := a.store.Create(ctx, pool.Name, record); err != nil { + if errors.Is(err, kvstore.ErrExists) { return nil } return fmt.Errorf("create pool KV entry: %w", err) } - slog.Info("external IPAM pool initialized", "pool", pool.Name, "gateway_ip", gwIP) + slog.InfoContext(ctx, "external IPAM pool initialized", "pool", pool.Name, "gateway_ip", gwIP) return nil } -func (a *StaticPoolAllocator) getRecord(ctx context.Context, poolName string) (*PoolRecord, uint64, error) { - entry, err := a.kv.Get(ctx, poolName) - if err != nil { - return nil, 0, err - } - - var record PoolRecord - if err := json.Unmarshal(entry.Value(), &record); err != nil { - return nil, 0, fmt.Errorf("unmarshal external IPAM record: %w", err) - } - return &record, entry.Revision(), nil -} - // nextAvailableIP picks the next unallocated address, skipping [GwLrpRangeStart, // GwLrpRangeEnd]. Resumes one past Cursor, wrapping at RangeEnd. func nextAvailableIP(record *PoolRecord) (string, error) { diff --git a/spinifex/network/external/static_pool_test.go b/spinifex/network/external/static_pool_test.go index 462e6b276..c391e27ec 100644 --- a/spinifex/network/external/static_pool_test.go +++ b/spinifex/network/external/static_pool_test.go @@ -144,6 +144,71 @@ func TestStaticPool_CASConflict(t *testing.T) { } } +// A launch asking for many public IPs at once is that many writers on one pool +// key. The retry bound has to cover the whole contending set, not a guess at +// timing: at the inherited five attempts this drops callers under load. +func TestStaticPool_AllocateUnderHeavyContention(t *testing.T) { + pool := wanPool() + pool.RangeEnd = "192.168.1.200" + a := newStaticAllocator(t, []ExternalPoolConfig{pool}) + + const writers = 16 + results := make([]netip.Addr, writers) + errs := make([]error, writers) + var wg sync.WaitGroup + for i := range writers { + wg.Go(func() { + results[i], errs[i] = a.Allocate(t.Context(), AllocateRequest{PoolName: "wan"}) + }) + } + wg.Wait() + + seen := make(map[string]bool, writers) + for i, err := range errs { + require.NoError(t, err, "concurrent allocation %d must not be dropped", i) + ip := results[i].String() + assert.False(t, seen[ip], "duplicate IP %s handed to two callers", ip) + seen[ip] = true + } +} + +// A pool is one key, so the boot-time drift reconcile contends with every live +// allocation on it. It must retry onto the winner rather than concede, which +// would leave the configured ranges unstored until some later boot won the race. +func TestStaticPool_DriftReconcileConvergesUnderContention(t *testing.T) { + narrow := wanPool() + narrow.RangeEnd = "192.168.1.180" + a := newStaticAllocator(t, []ExternalPoolConfig{narrow}) + + widened := narrow + widened.RangeEnd = "192.168.1.200" + + const writers = 16 + var wg sync.WaitGroup + for range writers { + wg.Go(func() { + _, _ = a.Allocate(t.Context(), AllocateRequest{PoolName: "wan"}) + }) + } + var driftErr error + wg.Go(func() { driftErr = a.reconcileDrift(t.Context(), widened) }) + wg.Wait() + require.NoError(t, driftErr) + + rec, err := a.GetPoolRecord(t.Context(), "wan") + require.NoError(t, err) + assert.Equal(t, "192.168.1.200", rec.RangeEnd, "the reconcile must land the configured range on the winner") +} + +// A write that is not a lost race must surface with its cause rather than being +// retried into a generic exhaustion, which is what the old loop did. +func TestStaticPool_ReleaseSurfacesANonConflictError(t *testing.T) { + a := newStaticAllocator(t, []ExternalPoolConfig{wanPool()}) + err := a.Release(t.Context(), "no-such-pool", mustAddr(t, "192.168.1.151"), "") + require.Error(t, err) + assert.NotErrorIs(t, err, errPoolContended, "an absent pool is not contention") +} + func TestStaticPool_ReleaseUnknown(t *testing.T) { a := newStaticAllocator(t, []ExternalPoolConfig{wanPool()}) err := a.Release(context.Background(), "wan", mustAddr(t, "192.168.1.200"), "")