From cd82abbb0c94e41a96d8cf76a07e36efb368a74a Mon Sep 17 00:00:00 2001 From: tiennm99 Date: Sun, 28 Jun 2026 12:52:32 +0700 Subject: [PATCH] refactor(storage): replace value-bytes CAS with version-field optimistic lock Replace content-addressed storage versioning (CAS) with explicit version fields on documents. Versions are managed via optimistic locking pattern: Get returns version, Put increments atomically and fails if version mismatch detected. Callers (portfolio, lolschedule) updated to use GetVersioned/PutVersioned, simplifying concurrency handling and eliminating CAS overhead. --- internal/modules/coin/portfolio.go | 20 +-- internal/modules/coin/portfolio_test.go | 8 +- internal/modules/gold/portfolio.go | 20 +-- internal/modules/gold/portfolio_test.go | 22 ++- internal/modules/lolschedule/cron.go | 42 +++--- internal/storage/dynamodb_kv.go | 36 +---- internal/storage/dynamodb_kv_test.go | 37 +---- internal/storage/kv_store.go | 21 ++- internal/storage/memory_kv.go | 47 ++++-- internal/storage/memory_kv_test.go | 117 ++++++--------- internal/storage/mongodb_kv.go | 134 +++++++++++------ internal/storage/mongodb_kv_test.go | 183 +++++++++++++++++------- internal/storage/prefix.go | 16 ++- internal/storage/prefix_test.go | 36 +++-- 14 files changed, 419 insertions(+), 320 deletions(-) diff --git a/internal/modules/coin/portfolio.go b/internal/modules/coin/portfolio.go index 1c40633..5f4c950 100644 --- a/internal/modules/coin/portfolio.go +++ b/internal/modules/coin/portfolio.go @@ -62,13 +62,13 @@ func SavePortfolio(ctx context.Context, kv storage.KVStore, userID int64, p Port } func UpdatePortfolio(ctx context.Context, kv storage.KVStore, userID int64, now int64, mutate func(*Portfolio) error) (Portfolio, error) { - cas, ok := kv.(storage.CompareAndSwapStore) + vs, ok := kv.(storage.VersionedStore) if !ok { - return Portfolio{}, fmt.Errorf("coin: storage does not support conditional portfolio updates") + return Portfolio{}, fmt.Errorf("coin: storage does not support versioned portfolio updates") } key := portfolioKey(userID) for attempt := 0; attempt < portfolioUpdateAttempts; attempt++ { - p, expected, err := loadPortfolioForUpdate(ctx, kv, key, now) + p, version, err := loadPortfolioForUpdate(ctx, vs, key, now) if err != nil { return Portfolio{}, fmt.Errorf("coin: load portfolio %d: %w", userID, err) } @@ -80,7 +80,7 @@ func UpdatePortfolio(ctx context.Context, kv storage.KVStore, userID int64, now if err != nil { return Portfolio{}, fmt.Errorf("coin: save portfolio %d: json encode: %w", userID, err) } - if err := cas.CompareAndSwap(ctx, key, expected, next); err == nil { + if err := vs.PutVersioned(ctx, key, version, next); err == nil { return p, nil } else if !errors.Is(err, storage.ErrConflict) { return Portfolio{}, fmt.Errorf("coin: save portfolio %d: %w", userID, err) @@ -89,13 +89,13 @@ func UpdatePortfolio(ctx context.Context, kv storage.KVStore, userID int64, now return Portfolio{}, fmt.Errorf("coin: save portfolio %d: %w", userID, storage.ErrConflict) } -func loadPortfolioForUpdate(ctx context.Context, kv storage.KVStore, key string, now int64) (Portfolio, []byte, error) { - raw, err := kv.Get(ctx, key) +func loadPortfolioForUpdate(ctx context.Context, vs storage.VersionedStore, key string, now int64) (Portfolio, int64, error) { + raw, version, err := vs.GetVersioned(ctx, key) switch { case err == nil: var p Portfolio if err := json.Unmarshal(raw, &p); err != nil { - return Portfolio{}, nil, fmt.Errorf("json decode: %w", err) + return Portfolio{}, 0, fmt.Errorf("json decode: %w", err) } p.normalize() if p.Assets == nil { @@ -104,11 +104,11 @@ func loadPortfolioForUpdate(ctx context.Context, kv storage.KVStore, key string, if p.Meta.CreatedAt == 0 { p.Meta.CreatedAt = now } - return p, raw, nil + return p, version, nil case errors.Is(err, storage.ErrNotFound): - return NewPortfolio(now), nil, nil + return NewPortfolio(now), 0, nil default: - return Portfolio{}, nil, err + return Portfolio{}, 0, err } } diff --git a/internal/modules/coin/portfolio_test.go b/internal/modules/coin/portfolio_test.go index aac81de..9b4c8e8 100644 --- a/internal/modules/coin/portfolio_test.go +++ b/internal/modules/coin/portfolio_test.go @@ -72,7 +72,11 @@ type conflictOnceStore struct { conflicted bool } -func (s *conflictOnceStore) CompareAndSwap(ctx context.Context, key string, expected []byte, val []byte) error { +func (s *conflictOnceStore) GetVersioned(ctx context.Context, key string) ([]byte, int64, error) { + return s.KVStore.(storage.VersionedStore).GetVersioned(ctx, key) +} + +func (s *conflictOnceStore) PutVersioned(ctx context.Context, key string, expectedVersion int64, val []byte) error { if !s.conflicted { s.conflicted = true competing := NewPortfolio(1) @@ -82,7 +86,7 @@ func (s *conflictOnceStore) CompareAndSwap(ctx context.Context, key string, expe } return storage.ErrConflict } - return s.KVStore.(storage.CompareAndSwapStore).CompareAndSwap(ctx, key, expected, val) + return s.KVStore.(storage.VersionedStore).PutVersioned(ctx, key, expectedVersion, val) } func TestUpdatePortfolioRetriesAfterWriteConflict(t *testing.T) { diff --git a/internal/modules/gold/portfolio.go b/internal/modules/gold/portfolio.go index 8d3087d..aa70908 100644 --- a/internal/modules/gold/portfolio.go +++ b/internal/modules/gold/portfolio.go @@ -59,13 +59,13 @@ func SavePortfolio(ctx context.Context, kv storage.KVStore, userID int64, p Port } func UpdatePortfolio(ctx context.Context, kv storage.KVStore, userID int64, now int64, mutate func(*Portfolio) error) (Portfolio, error) { - cas, ok := kv.(storage.CompareAndSwapStore) + vs, ok := kv.(storage.VersionedStore) if !ok { - return Portfolio{}, fmt.Errorf("gold: storage does not support conditional portfolio updates") + return Portfolio{}, fmt.Errorf("gold: storage does not support versioned portfolio updates") } key := portfolioKey(userID) for attempt := 0; attempt < portfolioUpdateAttempts; attempt++ { - p, expected, err := loadPortfolioForUpdate(ctx, kv, key, now) + p, version, err := loadPortfolioForUpdate(ctx, vs, key, now) if err != nil { return Portfolio{}, fmt.Errorf("gold: load portfolio %d: %w", userID, err) } @@ -77,7 +77,7 @@ func UpdatePortfolio(ctx context.Context, kv storage.KVStore, userID int64, now if err != nil { return Portfolio{}, fmt.Errorf("gold: save portfolio %d: json encode: %w", userID, err) } - if err := cas.CompareAndSwap(ctx, key, expected, next); err == nil { + if err := vs.PutVersioned(ctx, key, version, next); err == nil { return p, nil } else if !errors.Is(err, storage.ErrConflict) { return Portfolio{}, fmt.Errorf("gold: save portfolio %d: %w", userID, err) @@ -86,23 +86,23 @@ func UpdatePortfolio(ctx context.Context, kv storage.KVStore, userID int64, now return Portfolio{}, fmt.Errorf("gold: save portfolio %d: %w", userID, storage.ErrConflict) } -func loadPortfolioForUpdate(ctx context.Context, kv storage.KVStore, key string, now int64) (Portfolio, []byte, error) { - raw, err := kv.Get(ctx, key) +func loadPortfolioForUpdate(ctx context.Context, vs storage.VersionedStore, key string, now int64) (Portfolio, int64, error) { + raw, version, err := vs.GetVersioned(ctx, key) switch { case err == nil: var p Portfolio if err := json.Unmarshal(raw, &p); err != nil { - return Portfolio{}, nil, fmt.Errorf("json decode: %w", err) + return Portfolio{}, 0, fmt.Errorf("json decode: %w", err) } p.normalize() if p.Meta.CreatedAt == 0 { p.Meta.CreatedAt = now } - return p, raw, nil + return p, version, nil case errors.Is(err, storage.ErrNotFound): - return NewPortfolio(now), nil, nil + return NewPortfolio(now), 0, nil default: - return Portfolio{}, nil, err + return Portfolio{}, 0, err } } diff --git a/internal/modules/gold/portfolio_test.go b/internal/modules/gold/portfolio_test.go index c2ccbc8..28896f6 100644 --- a/internal/modules/gold/portfolio_test.go +++ b/internal/modules/gold/portfolio_test.go @@ -76,7 +76,11 @@ type conflictOnceStore struct { conflicted bool } -func (s *conflictOnceStore) CompareAndSwap(ctx context.Context, key string, expected []byte, val []byte) error { +func (s *conflictOnceStore) GetVersioned(ctx context.Context, key string) ([]byte, int64, error) { + return s.KVStore.(storage.VersionedStore).GetVersioned(ctx, key) +} + +func (s *conflictOnceStore) PutVersioned(ctx context.Context, key string, expectedVersion int64, val []byte) error { if !s.conflicted { s.conflicted = true competing := NewPortfolio(1) @@ -86,7 +90,7 @@ func (s *conflictOnceStore) CompareAndSwap(ctx context.Context, key string, expe } return storage.ErrConflict } - return s.KVStore.(storage.CompareAndSwapStore).CompareAndSwap(ctx, key, expected, val) + return s.KVStore.(storage.VersionedStore).PutVersioned(ctx, key, expectedVersion, val) } func TestUpdatePortfolioRetriesAfterWriteConflict(t *testing.T) { @@ -116,7 +120,11 @@ type alwaysConflictStore struct { attempts int } -func (s *alwaysConflictStore) CompareAndSwap(context.Context, string, []byte, []byte) error { +func (s *alwaysConflictStore) GetVersioned(ctx context.Context, key string) ([]byte, int64, error) { + return s.KVStore.(storage.VersionedStore).GetVersioned(ctx, key) +} + +func (s *alwaysConflictStore) PutVersioned(context.Context, string, int64, []byte) error { s.attempts++ return storage.ErrConflict } @@ -144,9 +152,9 @@ type countingCASStore struct { casCalls int } -func (s *countingCASStore) CompareAndSwap(ctx context.Context, key string, expected, val []byte) error { +func (s *countingCASStore) PutVersioned(ctx context.Context, key string, expectedVersion int64, val []byte) error { s.casCalls++ - return s.MemoryKVStore.CompareAndSwap(ctx, key, expected, val) + return s.MemoryKVStore.PutVersioned(ctx, key, expectedVersion, val) } func TestUpdatePortfolioMutateErrorDoesNotRetryOrPersist(t *testing.T) { @@ -212,8 +220,8 @@ func TestUpdatePortfolioConcurrentIncrementsLoseNoUpdates(t *testing.T) { } } -// plainKVStore hides the embedded store's CompareAndSwap method to model a -// backend without CAS support. +// plainKVStore hides the embedded store's versioned methods to model a backend +// without optimistic-locking support. type plainKVStore struct { storage.KVStore } diff --git a/internal/modules/lolschedule/cron.go b/internal/modules/lolschedule/cron.go index e3690d1..1123d40 100644 --- a/internal/modules/lolschedule/cron.go +++ b/internal/modules/lolschedule/cron.go @@ -134,37 +134,39 @@ func (s *state) dailyPushHandler(ctx context.Context, deps modules.Deps) error { // for the daily push: a winner proceeds to fan out; a loser (another trigger // already claimed today) returns false and sends nothing. // -// The claim is a compare-and-swap on lastPushDateKey so two simultaneous -// triggers cannot both win. Every real KV backend implements -// CompareAndSwapStore; if a bare store without CAS is used (only in narrow -// tests), it falls back to a plain Put — losing the atomic guarantee but -// preserving the "no-op if already today" behaviour. +// The claim is a version-based optimistic write on lastPushDateKey so two +// simultaneous triggers cannot both win. Every real KV backend implements +// VersionedStore; if a bare store without it is used (only in narrow tests), it +// falls back to a plain Put — losing the atomic guarantee but preserving the +// "no-op if already today" behaviour. func claimDailyPush(ctx context.Context, kv storage.KVStore, today string) (bool, error) { - current, err := kv.Get(ctx, lastPushDateKey) + vs, ok := kv.(storage.VersionedStore) + if !ok { + // No versioned support (a bare store in narrow tests): best-effort Put, + // losing the atomic guarantee but preserving "no-op if already today". + current, err := kv.Get(ctx, lastPushDateKey) + if err == nil && string(current) == today { + return false, nil + } + if err != nil && !errors.Is(err, storage.ErrNotFound) { + return false, err + } + return true, kv.Put(ctx, lastPushDateKey, []byte(today)) + } + + current, version, err := vs.GetVersioned(ctx, lastPushDateKey) switch { case err == nil: if string(current) == today { return false, nil // already pushed today } case errors.Is(err, storage.ErrNotFound): - current = nil // never pushed + version = 0 // never pushed default: return false, err } - cas, ok := kv.(storage.CompareAndSwapStore) - if !ok { - if err := kv.Put(ctx, lastPushDateKey, []byte(today)); err != nil { - return false, err - } - return true, nil - } - - var expected []byte - if len(current) > 0 { - expected = current - } - if err := cas.CompareAndSwap(ctx, lastPushDateKey, expected, []byte(today)); err != nil { + if err := vs.PutVersioned(ctx, lastPushDateKey, version, []byte(today)); err != nil { if errors.Is(err, storage.ErrConflict) { return false, nil // another trigger claimed today first } diff --git a/internal/storage/dynamodb_kv.go b/internal/storage/dynamodb_kv.go index 49aa7a4..cabbf0a 100644 --- a/internal/storage/dynamodb_kv.go +++ b/internal/storage/dynamodb_kv.go @@ -3,7 +3,6 @@ package storage import ( "context" "encoding/json" - "errors" "fmt" "strconv" "time" @@ -107,38 +106,9 @@ func (s *DynamoDBKVStore) Put(ctx context.Context, key string, val []byte) error return nil } -func (s *DynamoDBKVStore) CompareAndSwap(ctx context.Context, key string, expected []byte, val []byte) error { - if err := validateKey(key); err != nil { - return err - } - input := &dynamodb.PutItemInput{ - TableName: aws.String(s.table), - Item: map[string]types.AttributeValue{ - dynamoPKAttr: &types.AttributeValueMemberS{Value: s.moduleName}, - dynamoSKAttr: &types.AttributeValueMemberS{Value: key}, - dynamoValueAttr: &types.AttributeValueMemberS{Value: string(val)}, - dynamoUpdatedAtAttr: &types.AttributeValueMemberN{Value: strconv.FormatInt(time.Now().UTC().UnixNano(), 10)}, - }, - } - if expected == nil { - input.ConditionExpression = aws.String("attribute_not_exists(pk) AND attribute_not_exists(sk)") - } else { - input.ConditionExpression = aws.String("#v = :expected") - input.ExpressionAttributeNames = map[string]string{"#v": dynamoValueAttr} - input.ExpressionAttributeValues = map[string]types.AttributeValue{ - ":expected": &types.AttributeValueMemberS{Value: string(expected)}, - } - } - _, err := s.client.PutItem(ctx, input) - if err == nil { - return nil - } - var conflict *types.ConditionalCheckFailedException - if errors.As(err, &conflict) { - return ErrConflict - } - return fmt.Errorf("dynamodb compare-and-swap %s/%s: %w", s.moduleName, key, err) -} +// Note: DynamoDB intentionally does NOT implement VersionedStore. Post-self-host +// it is reachable only via the data migrator (Scan + Put), which never needs +// optimistic locking. The live runtime backend is MongoDB; memory covers tests. // PutJSON marshals val and writes the bytes at key. func (s *DynamoDBKVStore) PutJSON(ctx context.Context, key string, val any) error { diff --git a/internal/storage/dynamodb_kv_test.go b/internal/storage/dynamodb_kv_test.go index bca9b0a..2cf3c51 100644 --- a/internal/storage/dynamodb_kv_test.go +++ b/internal/storage/dynamodb_kv_test.go @@ -152,41 +152,8 @@ func TestDynamoDBKVStore_ListPrefix(t *testing.T) { } } -func TestDynamoDBKVStore_CompareAndSwap(t *testing.T) { - client, table, cleanup := dynamoDBLocalSetup(t) - defer cleanup() - - ctx := context.Background() - s := NewDynamoDBKVStore(client, table, "gold") - - // Create-if-absent succeeds once, then conflicts. - if err := s.CompareAndSwap(ctx, "user:1", nil, []byte("v1")); err != nil { - t.Fatalf("CompareAndSwap create: %v", err) - } - if err := s.CompareAndSwap(ctx, "user:1", nil, []byte("v2")); !errors.Is(err, ErrConflict) { - t.Errorf("CompareAndSwap create over existing: got %v, want ErrConflict", err) - } - - // Swap with matching expected succeeds; stale expected conflicts. - if err := s.CompareAndSwap(ctx, "user:1", []byte("v1"), []byte("v2")); err != nil { - t.Fatalf("CompareAndSwap matching: %v", err) - } - if err := s.CompareAndSwap(ctx, "user:1", []byte("v1"), []byte("v3")); !errors.Is(err, ErrConflict) { - t.Errorf("CompareAndSwap stale: got %v, want ErrConflict", err) - } - got, err := s.Get(ctx, "user:1") - if err != nil { - t.Fatalf("Get: %v", err) - } - if string(got) != "v2" { - t.Errorf("stored value = %q, want %q", got, "v2") - } - - // Non-nil expected on a missing key conflicts (caller reloads and retries). - if err := s.CompareAndSwap(ctx, "user:missing", []byte("v1"), []byte("v2")); !errors.Is(err, ErrConflict) { - t.Errorf("CompareAndSwap missing key: got %v, want ErrConflict", err) - } -} +// DynamoDB intentionally no longer implements optimistic locking — it is +// migrate-only post-self-host (Scan + Put). No CAS/versioned test here. func TestDynamoDBKVStore_DeleteMissingNoError(t *testing.T) { client, table, cleanup := dynamoDBLocalSetup(t) diff --git a/internal/storage/kv_store.go b/internal/storage/kv_store.go index c20ee3f..cca7777 100644 --- a/internal/storage/kv_store.go +++ b/internal/storage/kv_store.go @@ -22,9 +22,20 @@ type KVStore interface { List(ctx context.Context, prefix string) ([]string, error) } -// CompareAndSwapStore is implemented by stores that can conditionally replace -// a value only when it still equals the bytes read by the caller. A nil expected -// value means the key must not exist. -type CompareAndSwapStore interface { - CompareAndSwap(ctx context.Context, key string, expected []byte, val []byte) error +// VersionedStore is implemented by stores that support version-based optimistic +// locking. It replaces value-bytes compare-and-swap: each key carries a +// monotonic version, so a writer swaps only if the version it read is unchanged. +// This decouples concurrency control from the value encoding (so values can be +// stored as native documents rather than exact byte blobs). +// +// Contract: +// - GetVersioned returns the value and its current version, or ErrNotFound. +// A key that exists without a recorded version (written by an older build) +// reports version 0 and no error. +// - PutVersioned writes val only if the stored version still equals +// expectedVersion, then bumps the version. expectedVersion == 0 means +// "create, or adopt a not-yet-versioned key". A mismatch returns ErrConflict. +type VersionedStore interface { + GetVersioned(ctx context.Context, key string) (val []byte, version int64, err error) + PutVersioned(ctx context.Context, key string, expectedVersion int64, val []byte) error } diff --git a/internal/storage/memory_kv.go b/internal/storage/memory_kv.go index d64f5ac..dcdf1cf 100644 --- a/internal/storage/memory_kv.go +++ b/internal/storage/memory_kv.go @@ -9,27 +9,33 @@ import ( "sync" ) +// memEntry is one stored value plus its optimistic-locking version. +type memEntry struct { + val []byte + version int64 +} + // MemoryKVStore is an in-process KVStore for tests and local smoke runs. -// Data is lost on restart; production uses the DynamoDB provider. +// Data is lost on restart; production uses the MongoDB provider. type MemoryKVStore struct { mu sync.RWMutex - m map[string][]byte + m map[string]memEntry } // NewMemoryKVStore returns an empty in-memory store. func NewMemoryKVStore() *MemoryKVStore { - return &MemoryKVStore{m: make(map[string][]byte)} + return &MemoryKVStore{m: make(map[string]memEntry)} } func (s *MemoryKVStore) Get(_ context.Context, key string) ([]byte, error) { s.mu.RLock() defer s.mu.RUnlock() - v, ok := s.m[key] + e, ok := s.m[key] if !ok { return nil, ErrNotFound } - out := make([]byte, len(v)) - copy(out, v) + out := make([]byte, len(e.val)) + copy(out, e.val) return out, nil } @@ -46,24 +52,41 @@ func (s *MemoryKVStore) Put(_ context.Context, key string, val []byte) error { defer s.mu.Unlock() stored := make([]byte, len(val)) copy(stored, val) - s.m[key] = stored + // A plain Put bumps the version too, so a concurrent versioned writer that + // read the old version correctly sees a conflict. + s.m[key] = memEntry{val: stored, version: s.m[key].version + 1} return nil } -func (s *MemoryKVStore) CompareAndSwap(_ context.Context, key string, expected []byte, val []byte) error { +// GetVersioned returns the value and its version, or ErrNotFound. +func (s *MemoryKVStore) GetVersioned(_ context.Context, key string) ([]byte, int64, error) { + s.mu.RLock() + defer s.mu.RUnlock() + e, ok := s.m[key] + if !ok { + return nil, 0, ErrNotFound + } + out := make([]byte, len(e.val)) + copy(out, e.val) + return out, e.version, nil +} + +// PutVersioned writes val only if the stored version equals expectedVersion +// (0 = must not exist yet), then bumps the version. ErrConflict on mismatch. +func (s *MemoryKVStore) PutVersioned(_ context.Context, key string, expectedVersion int64, val []byte) error { s.mu.Lock() defer s.mu.Unlock() - current, ok := s.m[key] - if expected == nil { + e, ok := s.m[key] + if expectedVersion == 0 { if ok { return ErrConflict } - } else if !ok || !bytes.Equal(current, expected) { + } else if !ok || e.version != expectedVersion { return ErrConflict } stored := make([]byte, len(val)) copy(stored, val) - s.m[key] = stored + s.m[key] = memEntry{val: stored, version: e.version + 1} return nil } diff --git a/internal/storage/memory_kv_test.go b/internal/storage/memory_kv_test.go index 8b5d5c3..d38b45d 100644 --- a/internal/storage/memory_kv_test.go +++ b/internal/storage/memory_kv_test.go @@ -6,86 +6,53 @@ import ( "testing" ) -func TestMemoryKVStore_CompareAndSwap(t *testing.T) { +func TestMemoryKVStore_Versioned(t *testing.T) { ctx := context.Background() + s := NewMemoryKVStore() - tests := []struct { - name string - existing []byte // nil = key absent before the call - expected []byte - val []byte - wantErr error - wantVal string // stored value after the call - }{ - { - name: "create when absent", - existing: nil, - expected: nil, - val: []byte("v1"), - wantErr: nil, - wantVal: "v1", - }, - { - name: "create when exists conflicts", - existing: []byte("v1"), - expected: nil, - val: []byte("v2"), - wantErr: ErrConflict, - wantVal: "v1", - }, - { - name: "swap when expected matches", - existing: []byte("v1"), - expected: []byte("v1"), - val: []byte("v2"), - wantErr: nil, - wantVal: "v2", - }, - { - name: "swap when expected stale conflicts", - existing: []byte("v2"), - expected: []byte("v1"), - val: []byte("v3"), - wantErr: ErrConflict, - wantVal: "v2", - }, - { - name: "swap when key missing conflicts", - existing: nil, - expected: []byte("v1"), - val: []byte("v2"), - wantErr: ErrConflict, - wantVal: "", - }, + // Absent key reports ErrNotFound + version 0. + if _, v, err := s.GetVersioned(ctx, "k"); !errors.Is(err, ErrNotFound) || v != 0 { + t.Fatalf("GetVersioned absent: got (v=%d, %v), want (0, ErrNotFound)", v, err) } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - s := NewMemoryKVStore() - if tt.existing != nil { - if err := s.Put(ctx, "k", tt.existing); err != nil { - t.Fatalf("Put: %v", err) - } - } + // Create when absent (expectedVersion 0). + if err := s.PutVersioned(ctx, "k", 0, []byte("v1")); err != nil { + t.Fatalf("PutVersioned create: %v", err) + } + val, ver, err := s.GetVersioned(ctx, "k") + if err != nil || string(val) != "v1" || ver != 1 { + t.Fatalf("after create: got (%q, v=%d, %v), want (v1, 1, nil)", val, ver, err) + } - err := s.CompareAndSwap(ctx, "k", tt.expected, tt.val) - if !errors.Is(err, tt.wantErr) { - t.Fatalf("CompareAndSwap: got %v, want %v", err, tt.wantErr) - } + // Create over existing conflicts. + if err := s.PutVersioned(ctx, "k", 0, []byte("v2")); !errors.Is(err, ErrConflict) { + t.Errorf("create over existing: got %v, want ErrConflict", err) + } - got, err := s.Get(ctx, "k") - if tt.wantVal == "" { - if !errors.Is(err, ErrNotFound) { - t.Fatalf("Get after failed create: got (%q, %v), want ErrNotFound", got, err) - } - return - } - if err != nil { - t.Fatalf("Get: %v", err) - } - if string(got) != tt.wantVal { - t.Errorf("stored value = %q, want %q", got, tt.wantVal) - } - }) + // Swap with matching version succeeds and bumps version. + if err := s.PutVersioned(ctx, "k", 1, []byte("v2")); err != nil { + t.Fatalf("swap matching version: %v", err) + } + if _, ver, _ := s.GetVersioned(ctx, "k"); ver != 2 { + t.Errorf("version after swap = %d, want 2", ver) + } + + // Stale version conflicts. + if err := s.PutVersioned(ctx, "k", 1, []byte("v3")); !errors.Is(err, ErrConflict) { + t.Errorf("stale version: got %v, want ErrConflict", err) + } + + // Swap on a missing key conflicts. + if err := s.PutVersioned(ctx, "missing", 1, []byte("x")); !errors.Is(err, ErrConflict) { + t.Errorf("swap missing key: got %v, want ErrConflict", err) + } + + // A plain Put bumps the version, so a versioned writer holding the old + // version sees a conflict. + if err := s.Put(ctx, "k", []byte("v4")); err != nil { + t.Fatalf("Put: %v", err) + } + if err := s.PutVersioned(ctx, "k", 2, []byte("v5")); !errors.Is(err, ErrConflict) { + t.Errorf("versioned write after plain Put: got %v, want ErrConflict", err) } } diff --git a/internal/storage/mongodb_kv.go b/internal/storage/mongodb_kv.go index df1e592..a5d84db 100644 --- a/internal/storage/mongodb_kv.go +++ b/internal/storage/mongodb_kv.go @@ -18,6 +18,7 @@ const ( mongoIDField = "_id" mongoValueField = "value" mongoUpdatedAtField = "updatedAt" + mongoVersionField = "version" ) // MongoKVStore is a KVStore backed by a single MongoDB collection. The caller @@ -35,12 +36,13 @@ func NewMongoKVStore(coll *mongo.Collection, moduleName string) *MongoKVStore { return &MongoKVStore{coll: coll, moduleName: moduleName} } -// decodeValue extracts the stored value bytes from a decoded document. The -// driver decodes a BSON string into string and a BSON binary into bson.Binary; -// accept both so values written by any path round-trip. String is the current -// encoding (human-readable in the Atlas UI); the binary case is retained for -// backward compatibility with any documents written by an earlier build that -// stored value as BSON binary. +// decodeValue extracts the stored value bytes from a decoded document, +// reconstructing the caller's JSON []byte from whichever representation is on +// disk: +// - native object/array (bson.M / bson.A / bson.D) → re-serialized to JSON; +// - string → the bytes verbatim (bare scalars / non-JSON values, e.g. a date); +// - bson.Binary / []byte → legacy fallback for docs written by an earlier +// build that stored value as a binary blob. func (s *MongoKVStore) decodeValue(key string, doc bson.M) ([]byte, error) { raw, ok := doc[mongoValueField] if !ok { @@ -49,6 +51,8 @@ func (s *MongoKVStore) decodeValue(key string, doc bson.M) ([]byte, error) { switch v := raw.(type) { case string: return []byte(v), nil + case bson.M, bson.A, bson.D: + return nativeToJSON(v) case bson.Binary: return v.Data, nil case []byte: @@ -86,29 +90,22 @@ func (s *MongoKVStore) GetJSON(ctx context.Context, key string, dst any) error { return nil } -// doc builds the persisted document for key/val with a fresh updatedAt stamp. -// value is stored as a BSON string so it is human-readable in the Atlas/Compass -// UI (mirrors DynamoDB's String storage, dynamodb_kv.go). Every current caller -// writes JSON, which is UTF-8 safe; non-UTF-8 callers must encode upstream -// (e.g. base64), same constraint as DynamoDB. updatedAt is int64 unix-nanos to -// match DynamoDB and keep migration faithful. -func (s *MongoKVStore) doc(key string, val []byte) bson.M { - return bson.M{ - mongoIDField: key, - mongoValueField: string(val), - mongoUpdatedAtField: time.Now().UTC().UnixNano(), - } -} - -// Put writes raw bytes at key, creating or overwriting. +// Put writes raw bytes at key, creating or overwriting, and bumps the version +// (so a concurrent versioned writer that read the old version sees a conflict). func (s *MongoKVStore) Put(ctx context.Context, key string, val []byte) error { if err := validateKey(key); err != nil { return err } - _, err := s.coll.ReplaceOne(ctx, + _, err := s.coll.UpdateOne(ctx, bson.M{mongoIDField: key}, - s.doc(key, val), - options.Replace().SetUpsert(true), + bson.M{ + "$set": bson.M{ + mongoValueField: encodeValue(val), + mongoUpdatedAtField: time.Now().UTC().UnixNano(), + }, + "$inc": bson.M{mongoVersionField: 1}, + }, + options.UpdateOne().SetUpsert(true), ) if err != nil { return fmt.Errorf("mongo put %s/%s: %w", s.moduleName, key, err) @@ -125,41 +122,76 @@ func (s *MongoKVStore) PutJSON(ctx context.Context, key string, val any) error { return s.Put(ctx, key, raw) } -// CompareAndSwap conditionally replaces the value only when it still equals -// expected. A nil expected means the key must not yet exist. +// GetVersioned returns the value and its version, or ErrNotFound. A document +// written by an older build that has no version field reports version 0 (and no +// error), so PutVersioned(expectedVersion=0) can adopt it. +func (s *MongoKVStore) GetVersioned(ctx context.Context, key string) ([]byte, int64, error) { + if err := validateKey(key); err != nil { + return nil, 0, err + } + var doc bson.M + err := s.coll.FindOne(ctx, bson.M{mongoIDField: key}).Decode(&doc) + if err != nil { + if errors.Is(err, mongo.ErrNoDocuments) { + return nil, 0, ErrNotFound + } + return nil, 0, fmt.Errorf("mongo get %s/%s: %w", s.moduleName, key, err) + } + val, err := s.decodeValue(key, doc) + if err != nil { + return nil, 0, err + } + return val, decodeVersion(doc), nil +} + +// PutVersioned writes val only if the stored version equals expectedVersion, +// then bumps it. expectedVersion == 0 means "create, or adopt a key that has no +// version field yet" (a legacy doc). ErrConflict on a version mismatch. // -// - expected == nil → InsertOne; the unique _id index makes the absent-insert -// race linearizable (exactly one writer wins, losers get a duplicate-key -// error → ErrConflict). This is a LIVE first-write path (every new -// coin/gold portfolio), not an edge case. -// - expected != nil → UpdateOne filtered on the matching value; MatchedCount -// of 0 means the stored value changed (or the key is absent) → ErrConflict. -func (s *MongoKVStore) CompareAndSwap(ctx context.Context, key string, expected []byte, val []byte) error { +// - expectedVersion == 0 → upsert matching {_id} with version absent or 0; a +// key that already has version ≥ 1 fails the filter, the upsert attempts an +// insert on the same _id, and the unique-_id index returns duplicate-key → +// ErrConflict (linearizable single-winner for the first write). +// - expectedVersion > 0 → UpdateOne on {_id, version:expected}; MatchedCount +// of 0 means the version moved → ErrConflict. +func (s *MongoKVStore) PutVersioned(ctx context.Context, key string, expectedVersion int64, val []byte) error { if err := validateKey(key); err != nil { return err } - if expected == nil { - _, err := s.coll.InsertOne(ctx, s.doc(key, val)) + now := time.Now().UTC().UnixNano() + if expectedVersion == 0 { + _, err := s.coll.UpdateOne(ctx, + bson.M{ + mongoIDField: key, + "$or": bson.A{ + bson.M{mongoVersionField: bson.M{"$exists": false}}, + bson.M{mongoVersionField: int64(0)}, + }, + }, + bson.M{"$set": bson.M{ + mongoValueField: encodeValue(val), + mongoUpdatedAtField: now, + mongoVersionField: int64(1), + }}, + options.UpdateOne().SetUpsert(true), + ) if err != nil { if mongo.IsDuplicateKeyError(err) { return ErrConflict } - return fmt.Errorf("mongo compare-and-swap %s/%s: %w", s.moduleName, key, err) + return fmt.Errorf("mongo put-versioned %s/%s: %w", s.moduleName, key, err) } return nil } res, err := s.coll.UpdateOne(ctx, + bson.M{mongoIDField: key, mongoVersionField: expectedVersion}, bson.M{ - mongoIDField: key, - mongoValueField: string(expected), + "$set": bson.M{mongoValueField: encodeValue(val), mongoUpdatedAtField: now}, + "$inc": bson.M{mongoVersionField: 1}, }, - bson.M{"$set": bson.M{ - mongoValueField: string(val), - mongoUpdatedAtField: time.Now().UTC().UnixNano(), - }}, ) if err != nil { - return fmt.Errorf("mongo compare-and-swap %s/%s: %w", s.moduleName, key, err) + return fmt.Errorf("mongo put-versioned %s/%s: %w", s.moduleName, key, err) } if res.MatchedCount == 0 { return ErrConflict @@ -167,6 +199,22 @@ func (s *MongoKVStore) CompareAndSwap(ctx context.Context, key string, expected return nil } +// decodeVersion reads the version field from a decoded doc, defaulting to 0 for +// legacy docs that predate the version field. The driver may decode a BSON int +// as int32 or int64. +func decodeVersion(doc bson.M) int64 { + switch v := doc[mongoVersionField].(type) { + case int64: + return v + case int32: + return int64(v) + case float64: + return int64(v) + default: + return 0 + } +} + // Delete removes the document at key. Deleting a missing key is not an error // (idempotent) — DeleteOne with a zero match count returns nil. func (s *MongoKVStore) Delete(ctx context.Context, key string) error { diff --git a/internal/storage/mongodb_kv_test.go b/internal/storage/mongodb_kv_test.go index b24803d..0b4f7f7 100644 --- a/internal/storage/mongodb_kv_test.go +++ b/internal/storage/mongodb_kv_test.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "os" + "reflect" "sync" "testing" "time" @@ -89,34 +90,93 @@ func TestMongoKVStore_GetMissing(t *testing.T) { } } -// TestMongoKVStore_ValueStoredAsViewableString verifies value is persisted as a -// BSON string (directly readable in the Atlas/Compass UI), not opaque BinData, -// and that UTF-8 JSON round-trips byte-identically. -func TestMongoKVStore_ValueStoredAsViewableString(t *testing.T) { +// TestMongoKVStore_NativeValueRepresentation verifies a JSON object value is +// persisted as a NATIVE BSON document (expandable/queryable in Atlas), and a +// bare non-JSON value as a string — and that both round-trip. +func TestMongoKVStore_NativeValueRepresentation(t *testing.T) { s, db, cleanup := mongoLocalSetup(t, "wordle") defer cleanup() ctx := context.Background() - payload := []byte(`{"word":"chào","score":42}`) // multi-byte UTF-8 - if err := s.Put(ctx, "u1", payload); err != nil { - t.Fatalf("Put: %v", err) + obj := []byte(`{"word":"chào","score":42}`) // multi-byte UTF-8 object + if err := s.Put(ctx, "obj", obj); err != nil { + t.Fatalf("Put obj: %v", err) } - got, err := s.Get(ctx, "u1") - if err != nil { - t.Fatalf("Get: %v", err) - } - if string(got) != string(payload) { - t.Errorf("round trip: got %q, want %q", got, payload) + // A bare string (lolschedule date-guard style) is not JSON → stored as string. + if err := s.Put(ctx, "bare", []byte("2026-06-28")); err != nil { + t.Fatalf("Put bare: %v", err) } - // Inspect the raw document: the value field must decode as a Go string - // (BSON string type), proving it is viewable directly rather than BinData. - var doc bson.M - if err := db.Collection("wordle").FindOne(ctx, bson.M{"_id": "u1"}).Decode(&doc); err != nil { - t.Fatalf("raw FindOne: %v", err) + var objDoc, bareDoc bson.M + if err := db.Collection("wordle").FindOne(ctx, bson.M{"_id": "obj"}).Decode(&objDoc); err != nil { + t.Fatalf("raw FindOne obj: %v", err) } - if _, ok := doc["value"].(string); !ok { - t.Errorf("value stored as %T, want string (directly viewable)", doc["value"]) + switch objDoc["value"].(type) { + case bson.M, bson.D: + // native object — good + default: + t.Errorf("object value stored as %T, want native bson document", objDoc["value"]) + } + if err := db.Collection("wordle").FindOne(ctx, bson.M{"_id": "bare"}).Decode(&bareDoc); err != nil { + t.Fatalf("raw FindOne bare: %v", err) + } + if _, ok := bareDoc["value"].(string); !ok { + t.Errorf("bare value stored as %T, want string", bareDoc["value"]) + } + + // Round-trips. + if got, _ := s.Get(ctx, "bare"); string(got) != "2026-06-28" { + t.Errorf("bare round trip: got %q", got) + } + var back map[string]any + if err := s.GetJSON(ctx, "obj", &back); err != nil { + t.Fatalf("GetJSON obj: %v", err) + } + if back["word"] != "chào" { + t.Errorf("object round trip lost data: %v", back) + } +} + +// TestMongoKVStore_Int64Fidelity is the codec gate: an int64 field must survive +// the JSON↔BSON round trip as an integer, not collapse to a float (which would +// make GetJSON into an int64 field fail or drift). +func TestMongoKVStore_Int64Fidelity(t *testing.T) { + s, db, cleanup := mongoLocalSetup(t, "coin") + defer cleanup() + + ctx := context.Background() + type meta struct { + CreatedAt int64 `json:"createdAt"` // UnixMilli ~1.7e12 + Invested float64 `json:"invested"` + } + type portfolio struct { + USD float64 `json:"usd"` + Meta meta `json:"meta"` + Qty map[string]int64 `json:"qty"` + } + in := portfolio{USD: 1234.56, Meta: meta{CreatedAt: 1719500000000, Invested: 5000}, Qty: map[string]int64{"BTC": 3}} + if err := s.PutJSON(ctx, "u1", in); err != nil { + t.Fatalf("PutJSON: %v", err) + } + var out portfolio + if err := s.GetJSON(ctx, "u1", &out); err != nil { + t.Fatalf("GetJSON: %v", err) + } + if !reflect.DeepEqual(out, in) { + t.Errorf("fidelity lost: got %+v, want %+v", out, in) + } + // The stored CreatedAt must be a BSON int (int32/int64), not a double. + var doc bson.M + _ = db.Collection("coin").FindOne(ctx, bson.M{"_id": "u1"}).Decode(&doc) + if m, ok := doc["value"].(bson.M); ok { + if meta, ok := m["meta"].(bson.M); ok { + switch meta["createdAt"].(type) { + case int64, int32: + // good — preserved as integer + default: + t.Errorf("createdAt stored as %T, want int64 (json.Number codec)", meta["createdAt"]) + } + } } } @@ -177,45 +237,70 @@ func TestMongoKVStore_ListPrefix(t *testing.T) { } } -func TestMongoKVStore_CompareAndSwap(t *testing.T) { +func TestMongoKVStore_Versioned(t *testing.T) { s, _, cleanup := mongoLocalSetup(t, "gold") defer cleanup() ctx := context.Background() - // Create-if-absent succeeds once, then conflicts. - if err := s.CompareAndSwap(ctx, "user:1", nil, []byte("v1")); err != nil { - t.Fatalf("CompareAndSwap create: %v", err) + // Create-if-absent (version 0) succeeds once, then conflicts. + if err := s.PutVersioned(ctx, "user:1", 0, []byte("v1")); err != nil { + t.Fatalf("PutVersioned create: %v", err) } - if err := s.CompareAndSwap(ctx, "user:1", nil, []byte("v2")); !errors.Is(err, ErrConflict) { - t.Errorf("CompareAndSwap create over existing: got %v, want ErrConflict", err) + if err := s.PutVersioned(ctx, "user:1", 0, []byte("v2")); !errors.Is(err, ErrConflict) { + t.Errorf("create over existing: got %v, want ErrConflict", err) } - // Swap with matching expected succeeds; stale expected conflicts. - if err := s.CompareAndSwap(ctx, "user:1", []byte("v1"), []byte("v2")); err != nil { - t.Fatalf("CompareAndSwap matching: %v", err) + val, ver, err := s.GetVersioned(ctx, "user:1") + if err != nil || string(val) != "v1" || ver != 1 { + t.Fatalf("GetVersioned: got (%q, v=%d, %v), want (v1, 1, nil)", val, ver, err) } - if err := s.CompareAndSwap(ctx, "user:1", []byte("v1"), []byte("v3")); !errors.Is(err, ErrConflict) { - t.Errorf("CompareAndSwap stale: got %v, want ErrConflict", err) + + // Swap with matching version succeeds and bumps; stale version conflicts. + if err := s.PutVersioned(ctx, "user:1", 1, []byte("v2")); err != nil { + t.Fatalf("swap matching version: %v", err) } - got, err := s.Get(ctx, "user:1") - if err != nil { - t.Fatalf("Get: %v", err) + if err := s.PutVersioned(ctx, "user:1", 1, []byte("v3")); !errors.Is(err, ErrConflict) { + t.Errorf("stale version: got %v, want ErrConflict", err) } - if string(got) != "v2" { + if got, _ := s.Get(ctx, "user:1"); string(got) != "v2" { t.Errorf("stored value = %q, want %q", got, "v2") } - // Non-nil expected on a missing key conflicts (caller reloads and retries). - if err := s.CompareAndSwap(ctx, "user:missing", []byte("v1"), []byte("v2")); !errors.Is(err, ErrConflict) { - t.Errorf("CompareAndSwap missing key: got %v, want ErrConflict", err) + // Swap on a missing key conflicts. + if err := s.PutVersioned(ctx, "user:missing", 1, []byte("v2")); !errors.Is(err, ErrConflict) { + t.Errorf("swap missing key: got %v, want ErrConflict", err) } } -// TestMongoKVStore_CompareAndSwap_ConcurrentInsert proves the absent-insert -// race is linearizable: N goroutines racing a nil-expected CAS on the same key -// must yield exactly one winner; every loser gets ErrConflict (never a silent -// overwrite). This is the blocking correctness gate from Phase 1. -func TestMongoKVStore_CompareAndSwap_ConcurrentInsert(t *testing.T) { +// TestMongoKVStore_PutVersioned_AdoptsLegacyDoc proves a document written by an +// older build (no version field) is treated as version 0 and updated cleanly, +// so existing users' portfolios keep working after the rollout. +func TestMongoKVStore_PutVersioned_AdoptsLegacyDoc(t *testing.T) { + s, db, cleanup := mongoLocalSetup(t, "coin") + defer cleanup() + + ctx := context.Background() + // Seed a legacy doc: value string, NO version field. + if _, err := db.Collection("coin").InsertOne(ctx, bson.M{"_id": "user:9", "value": "legacy"}); err != nil { + t.Fatalf("seed legacy doc: %v", err) + } + val, ver, err := s.GetVersioned(ctx, "user:9") + if err != nil || string(val) != "legacy" || ver != 0 { + t.Fatalf("GetVersioned legacy: got (%q, v=%d, %v), want (legacy, 0, nil)", val, ver, err) + } + // Claim it at version 0 — must succeed (no spurious conflict). + if err := s.PutVersioned(ctx, "user:9", 0, []byte("migrated")); err != nil { + t.Fatalf("adopt legacy doc: %v", err) + } + if got, ver, _ := s.GetVersioned(ctx, "user:9"); string(got) != "migrated" || ver != 1 { + t.Errorf("after adopt: got (%q, v=%d), want (migrated, 1)", got, ver) + } +} + +// TestMongoKVStore_PutVersioned_ConcurrentCreate proves the absent-create race +// is linearizable: N goroutines racing PutVersioned(_, 0, _) on one key yield +// exactly one winner; losers get ErrConflict. Blocking correctness gate. +func TestMongoKVStore_PutVersioned_ConcurrentCreate(t *testing.T) { s, _, cleanup := mongoLocalSetup(t, "coin") defer cleanup() @@ -229,8 +314,8 @@ func TestMongoKVStore_CompareAndSwap_ConcurrentInsert(t *testing.T) { wg.Add(1) go func(i int) { defer wg.Done() - <-start // release all goroutines at once to maximize contention - err := s.CompareAndSwap(ctx, "race", nil, []byte(fmt.Sprintf("v%d", i))) + <-start + err := s.PutVersioned(ctx, "race", 0, []byte(fmt.Sprintf("v%d", i))) mu.Lock() defer mu.Unlock() switch { @@ -240,7 +325,7 @@ func TestMongoKVStore_CompareAndSwap_ConcurrentInsert(t *testing.T) { conflicts++ default: others++ - t.Errorf("unexpected CAS error: %v", err) + t.Errorf("unexpected error: %v", err) } }(i) } @@ -248,13 +333,13 @@ func TestMongoKVStore_CompareAndSwap_ConcurrentInsert(t *testing.T) { wg.Wait() if wins != 1 { - t.Errorf("concurrent nil-expected CAS: got %d winners, want exactly 1", wins) + t.Errorf("concurrent create: got %d winners, want exactly 1", wins) } if conflicts != n-1 { - t.Errorf("concurrent nil-expected CAS: got %d conflicts, want %d", conflicts, n-1) + t.Errorf("concurrent create: got %d conflicts, want %d", conflicts, n-1) } if others != 0 { - t.Errorf("concurrent CAS produced %d non-conflict errors", others) + t.Errorf("concurrent create produced %d non-conflict errors", others) } } diff --git a/internal/storage/prefix.go b/internal/storage/prefix.go index 337641d..df7ff2e 100644 --- a/internal/storage/prefix.go +++ b/internal/storage/prefix.go @@ -36,14 +36,22 @@ func (p *prefixedStore) Put(ctx context.Context, key string, val []byte) error { return p.inner.Put(ctx, p.k(key), val) } -func (p *prefixedStore) CompareAndSwap(ctx context.Context, key string, expected []byte, val []byte) error { - cas, ok := p.inner.(CompareAndSwapStore) +func (p *prefixedStore) GetVersioned(ctx context.Context, key string) ([]byte, int64, error) { + vs, ok := p.inner.(VersionedStore) + if !ok { + return nil, 0, fmt.Errorf("storage: inner store does not support versioned reads: %w", errors.ErrUnsupported) + } + return vs.GetVersioned(ctx, p.k(key)) +} + +func (p *prefixedStore) PutVersioned(ctx context.Context, key string, expectedVersion int64, val []byte) error { + vs, ok := p.inner.(VersionedStore) if !ok { // A missing capability is permanent: report it as unsupported so // callers fail fast instead of retrying it as a write conflict. - return fmt.Errorf("storage: inner store does not support compare-and-swap: %w", errors.ErrUnsupported) + return fmt.Errorf("storage: inner store does not support versioned writes: %w", errors.ErrUnsupported) } - return cas.CompareAndSwap(ctx, p.k(key), expected, val) + return vs.PutVersioned(ctx, p.k(key), expectedVersion, val) } func (p *prefixedStore) PutJSON(ctx context.Context, key string, val any) error { diff --git a/internal/storage/prefix_test.go b/internal/storage/prefix_test.go index 5bfc2df..e241c9d 100644 --- a/internal/storage/prefix_test.go +++ b/internal/storage/prefix_test.go @@ -68,17 +68,17 @@ func TestPrefixed_NotFoundPropagates(t *testing.T) { } } -func TestPrefixed_CompareAndSwapDelegatesWithPrefixedKey(t *testing.T) { +func TestPrefixed_VersionedDelegatesWithPrefixedKey(t *testing.T) { ctx := context.Background() base := NewMemoryKVStore() mod := Prefixed(base, "gold") - cas, ok := mod.(CompareAndSwapStore) + vs, ok := mod.(VersionedStore) if !ok { - t.Fatal("Prefixed store does not implement CompareAndSwapStore") + t.Fatal("Prefixed store does not implement VersionedStore") } - if err := cas.CompareAndSwap(ctx, "user:1", nil, []byte("v1")); err != nil { - t.Fatalf("CompareAndSwap create: %v", err) + if err := vs.PutVersioned(ctx, "user:1", 0, []byte("v1")); err != nil { + t.Fatalf("PutVersioned create: %v", err) } // The write must land under the prefixed key on the inner store. @@ -90,31 +90,37 @@ func TestPrefixed_CompareAndSwapDelegatesWithPrefixedKey(t *testing.T) { t.Errorf("inner value = %q, want %q", got, "v1") } - // Stale expected must surface the inner store's conflict. - if err := cas.CompareAndSwap(ctx, "user:1", []byte("stale"), []byte("v2")); !errors.Is(err, ErrConflict) { - t.Errorf("CompareAndSwap stale: got %v, want ErrConflict", err) + // GetVersioned reads through the prefix and returns the inner version. + val, ver, err := vs.GetVersioned(ctx, "user:1") + if err != nil || string(val) != "v1" || ver != 1 { + t.Fatalf("GetVersioned: got (%q, v=%d, %v), want (v1, 1, nil)", val, ver, err) } - if err := cas.CompareAndSwap(ctx, "user:1", []byte("v1"), []byte("v2")); err != nil { - t.Errorf("CompareAndSwap matching: %v", err) + + // Stale version must surface the inner store's conflict. + if err := vs.PutVersioned(ctx, "user:1", 99, []byte("v2")); !errors.Is(err, ErrConflict) { + t.Errorf("stale version: got %v, want ErrConflict", err) + } + if err := vs.PutVersioned(ctx, "user:1", 1, []byte("v2")); err != nil { + t.Errorf("matching version: %v", err) } } -// plainStore hides the inner store's CompareAndSwap method so the wrapper's +// plainStore hides the inner store's versioned methods so the wrapper's // missing-capability path can be exercised. type plainStore struct { KVStore } -func TestPrefixed_CompareAndSwapUnsupportedInner(t *testing.T) { +func TestPrefixed_VersionedUnsupportedInner(t *testing.T) { mod := Prefixed(&plainStore{KVStore: NewMemoryKVStore()}, "gold") - cas := mod.(CompareAndSwapStore) + vs := mod.(VersionedStore) - err := cas.CompareAndSwap(context.Background(), "user:1", nil, []byte("v1")) + err := vs.PutVersioned(context.Background(), "user:1", 0, []byte("v1")) if !errors.Is(err, errors.ErrUnsupported) { t.Errorf("got %v, want errors.ErrUnsupported", err) } if errors.Is(err, ErrConflict) { - t.Error("missing CAS capability must not be reported as a retryable conflict") + t.Error("missing versioned capability must not be reported as a retryable conflict") } }