diff --git a/.agent-docs/concurrent-keys/phase1-design.md b/.agent-docs/concurrent-keys/phase1-design.md index df01c026..1ac645b4 100644 --- a/.agent-docs/concurrent-keys/phase1-design.md +++ b/.agent-docs/concurrent-keys/phase1-design.md @@ -253,6 +253,25 @@ This is the highest-risk piece of Phase 1. The **T0 PoC** validated the swap mec This order is *necessary* — the PoC found that flipping the pointer too early leaves `GetClient()` nil mid-swap (H6) and breaks the env on init failure (H7) — and *sufficient* only with the store-handling approach below. +### Two re-anchor cases — Case A (new key) vs Case B (already-accepted key) + +The sequence above is **Case A**: the new anchor is a key relay has not previously accepted, so no SDK client exists for it. Relay must build one, hand over the store, wait for `Initialized()`, then flip and re-wire. + +**Case B** is the abbreviated path taken when the new anchor's key **already has a live SDK client** — most commonly a *former* anchor that is still inside its grace period (it was demoted on an earlier rotation, kept alive to serve downstream traffic, and is now being promoted back). In that situation: + +1. **No `Build`.** The existing client is reused as-is — there is no second upstream connection to stand up. +2. **No store handover.** The store the existing client created is already populated and initialized; nothing is handed over because nothing new is constructed. +3. **Atomically flip the rotator's anchor pointer** (`CommitAnchor`) — identical to Case A step 3. +4. **Call `ReplaceCredential`** on the event dispatcher + metrics publisher — identical to Case A step 4. +5. Big-segment sync is re-wired (T2.d), identical to Case A. +6. The retiring anchor's client is closed by the existing `removeCredential` path when its own grace period ends — identical to Case A step 6. + +The two paths **converge after the flip**: steps 3–6 are the same. The only difference is the front of the sequence — Case A builds + initializes + hands over the store; Case B reuses what is already there and does none of that. The caller branches on whether a client already exists for the new anchor (`c.clients[newAnchor] != nil`). + +Because Case B does no client build, there is no init-failure rollback to consider for it — the client it reuses was already initialized and serving. Rollback handling (preserve previous anchor, log a structured error) applies to **Case A only**. + +**Reconcile/additions interaction:** so the synchronous re-anchor owns the new anchor's setup end-to-end, `Rotator.Reconcile` does not flip the anchor itself — it returns a `ReconcileResult.AnchorChange` and the caller invokes `CommitAnchor` at the right moment. In **Case A** the new anchor would otherwise appear in the reconcile's `additions` list and `addCredential` would fire an *async* `startSDKClient` that races the synchronous build — so Reconcile strips the new anchor from `additions` in Case A and the synchronous path installs the peripherals (envStreams, handlers, connection mapping) itself. In **Case B** the new anchor was already accepted, so it was never going to appear in `additions` — no stripping is needed there. + ### The data store: hand the existing store over to the new client An earlier version of this design assumed two SDK clients pointed at the same env would feed the *same* data store as a side-effect. The PoC (H1, H5) showed this is **wrong for the in-memory store**: each SDK client construction calls `storeAdapter.Build()`, which atomically swaps in a *new, empty* store, so the new client would otherwise have to re-sync from scratch (an empty-store window). This affects only the in-memory case; with a persistent store (Redis, DynamoDB) the data lives outside the wrapper and survives the swap. @@ -276,17 +295,17 @@ This is the concrete form of decoupling the store's lifecycle from the client's. ### Failure handling -If the new client fails to initialize, the swap **rolls back**: the rotator's anchor pointer stays on the old key, the previous accepted set is preserved, a structured error is logged, and an alarm is raised. The old anchor's client (still alive in its grace period) continues to serve. This is the §8 atomicity principle applied to re-anchor. +If the new client fails to initialize, the swap **rolls back**: the rotator's anchor pointer stays on the old key (the caller simply does not call `CommitAnchor`), the previous accepted set is preserved, and a structured error is logged. The old anchor's client (still alive in its grace period) continues to serve. This is the §8 atomicity principle applied to re-anchor, and it applies to **Case A only** — Case B reuses an already-initialized client and has nothing to fail. Relay has no dedicated alarm infrastructure today; an `Error`-level structured log (`globalLoggers.Errorf`) is the strongest signal available and is sufficient. ### Consolidated specification for T2.c / T2.d | # | Requirement | Source | Owner | |---|---|---|---| -| 1 | Build + initialize the new anchor client *before* flipping the pointer; flip atomically. | H5, H6 | T2.c | -| 2 | On init failure, roll back to old anchor; preserve previous accepted set; log + alarm. | H7 | T2.c | -| 3 | Hand the existing store over to the new client (adapter reuses its store); ensure the retiring client's `Close()` does not tear down the shared store. | H1, H5 | T2.c | -| 4 | Re-wire big-segment sync on re-anchor (recreate or replace-credential). | H3 | T2.d | -| 5 | Call `ReplaceCredential` on event dispatcher + metrics publisher. | §7 | T2.c (already wired in `addCredential`) | +| 1 | **Case A**: build + initialize the new anchor client *before* flipping the pointer; flip atomically via `CommitAnchor`. **Case B** (new anchor already has a live client): skip the build, reuse it, then flip. | H5, H6 | T2.c | +| 2 | **Case A** init failure: roll back (do not `CommitAnchor`); preserve previous accepted set; log a structured error. Not applicable to Case B. | H7 | T2.c | +| 3 | **Case A**: hand the existing store over to the new client (adapter reuses its store, refcounted so the retiring client's `Close()` does not tear it down). **Case B**: no handover — the store is already populated. | H1, H5 | T2.c | +| 4 | Re-wire big-segment sync on re-anchor (recreate or replace-credential). Same for both cases. | H3 | T2.d | +| 5 | Call `ReplaceCredential` on event dispatcher + metrics publisher (synchronously, in the re-anchor sequence — not via `addCredential`, since Reconcile strips the new anchor from `additions`). Same for both cases. | §7 | T2.c | | 6 | Expect duplicate downstream `put`; retain connections for credentials still in the accepted set. | H2 | T2.c (awareness) | | 7 | No `httpconfig` change. | H4 | n/a | diff --git a/internal/credential/rotator.go b/internal/credential/rotator.go index db6881da..1a46529f 100644 --- a/internal/credential/rotator.go +++ b/internal/credential/rotator.go @@ -2,6 +2,7 @@ package credential import ( "maps" + "slices" "sync" "time" @@ -209,6 +210,42 @@ func (r *Rotator) StepTime(now time.Time) (additions []SDKCredential, expiration return additions, expirations } +// ReconcileResult signals state changes that the caller must apply synchronously rather than rely +// on the normal addCredential / removeCredential flow driven by StepTime. +// +// AnchorChange is non-nil when the SDK anchor changed during Reconcile. The rotator does NOT flip +// its anchor pointer in that case — the caller must drive the synchronous re-anchor sequence +// (build the new anchor's SDK client if one does not exist, wait for Initialized, then invoke +// CommitAnchor to atomically move the pointer, then call ReplaceCredential on the event dispatcher +// and metrics publisher, then re-wire big-segment sync). +// +// MobilePrimaryRepoint is non-nil when the primary mobile key changed AND the new primary was +// already in the accepted set. In that case it does not appear in StepTime's additions list and +// addCredential's primary-mobile gate will not fire for it, so the caller must invoke +// eventDispatcher.ReplaceCredential synchronously. When nil, either the primary mobile key did not +// change, or it changed to a newly-accepted key — in which case the normal addCredential path +// handles the ReplaceCredential call via the existing gate. +type ReconcileResult struct { + AnchorChange *AnchorChange + MobilePrimaryRepoint *config.MobileKey +} + +// AnchorChange describes an SDK anchor transition produced by Reconcile. +// +// NewAnchorPreviouslyAccepted distinguishes the two re-anchor paths: +// - false (the anchor is a new key): the new anchor was not previously in the accepted set. The +// synchronous re-anchor must register the credential mappings (envStreams, handlers, connection +// mapping), construct and initialize a new SDK client, then invoke CommitAnchor + ReplaceCredential. +// - true (the anchor is a previously-accepted key): the new anchor was already accepted (typically +// a former anchor still in its grace period). Its credential mappings are already registered and a +// client may already exist; the synchronous re-anchor reuses it (or constructs one only if missing +// — see the re-anchor sequence in env_context_impl.go), then invokes CommitAnchor + ReplaceCredential. +type AnchorChange struct { + PreviousAnchor config.SDKKey + NewAnchor config.SDKKey + NewAnchorPreviouslyAccepted bool +} + // Reconcile updates the rotator to match set. The set names its own anchor (the primary SDK key) and // primary mobile key. It diffs the desired accepted set against the current one and queues additions // and expirations (drained by the next StepTime call); keys newly present are accepted, and keys no @@ -220,13 +257,101 @@ func (r *Rotator) StepTime(now time.Time) (additions []SDKCredential, expiration // The set is assumed well-formed: AcceptedSetBuilder.Build validates that an anchor was designated // (and, because WithAnchor adds the key as it designates it, that the anchor is among the SDK // keys), so Reconcile trusts what it is handed rather than re-validating. -func (r *Rotator) Reconcile(set AcceptedSet, now time.Time) { +// +// Reconcile does NOT flip the SDK anchor pointer when the anchor changes — the returned +// ReconcileResult.AnchorChange signals the change so the caller can drive the synchronous re-anchor +// sequence, then call CommitAnchor to atomically move the pointer. When the new anchor is a new key +// (NewAnchorPreviouslyAccepted == false) it is also stripped from additions so that the async +// startSDKClient invocation in addCredential does not race the synchronous client build. +func (r *Rotator) Reconcile(set AcceptedSet, now time.Time) ReconcileResult { r.mu.Lock() defer r.mu.Unlock() - r.reconcileSDKKeys(set, set.anchor, now) + var result ReconcileResult + + previousAnchor := r.anchorKey + newAnchor := set.anchor + if previousAnchor != newAnchor && newAnchor.Defined() { + _, alreadyAccepted := r.acceptedSDKKeys[newAnchor] + result.AnchorChange = &AnchorChange{ + PreviousAnchor: previousAnchor, + NewAnchor: newAnchor, + NewAnchorPreviouslyAccepted: alreadyAccepted, + } + } + + r.reconcileSDKKeys(set, now) + + if result.AnchorChange != nil && !result.AnchorChange.NewAnchorPreviouslyAccepted { + // The anchor is a new key: reconcileAcceptedKeys just appended it to r.additions. Strip it — + // the synchronous re-anchor sequence in env_context_impl owns the new anchor's setup + // (credential mappings + client build + flip + ReplaceCredential). If addCredential drained this + // addition normally, its async startSDKClient would race the synchronous build. When the anchor + // is a previously-accepted key it was already in acceptedSDKKeys, so reconcileAcceptedKeys did + // not add it — no strip needed. + r.additions = slices.DeleteFunc(r.additions, func(c SDKCredential) bool { return c == newAnchor }) + } + + previousMobile := r.primaryMobileKey + newMobile := set.primaryMobileKey + var newMobileAlreadyAccepted bool + if newMobile.Defined() { + _, newMobileAlreadyAccepted = r.acceptedMobileKeys[newMobile] + } r.reconcileMobileKeys(set, now) + if previousMobile != newMobile && newMobile.Defined() && newMobileAlreadyAccepted { + // Primary mobile key changed to a key already in the accepted set: addCredential's gate will + // not fire for it (it's not in additions), so the caller must call ReplaceCredential itself. + m := newMobile + result.MobilePrimaryRepoint = &m + } + r.reconcileEnvironmentID(set) + + return result +} + +// CommitAnchor atomically moves the rotator's SDK anchor pointer to the given key. The caller +// invokes this once the synchronous re-anchor sequence is ready to flip — i.e. after the new +// anchor's client is built and reports Initialized (when the anchor is a new key) or after confirming +// the existing client will be reused (when the anchor is a previously-accepted key). Until CommitAnchor +// is called, the rotator's anchor stays on the previous key so GetClient() returns the still-serving +// old client and the gate in addCredential does not fire for the pending new anchor. +// +// Aside from Initialize (which establishes the initial anchor), CommitAnchor is the only path that +// moves the anchor pointer: Reconcile deliberately does not flip it (see reconcileSDKKeys). +func (r *Rotator) CommitAnchor(key config.SDKKey) { + r.mu.Lock() + defer r.mu.Unlock() + r.anchorKey = key +} + +// RevertAnchorChange undoes the accepted-set effects of an AnchorChange whose synchronous re-anchor +// failed and rolled back. Because CommitAnchor was never called, the anchor pointer still names the +// previous anchor; this realigns the accepted set with it so the two don't disagree. +// +// - If the previous anchor is a defined key that was revoked in the same reconcile (it is no longer +// accepted — an immediate revocation rather than a grace demotion), re-admit it as a permanent +// key, since it remains the anchor and keeps serving. If it is still accepted (grace demotion), +// leave it and its expiry untouched. An undefined previous anchor (the env's first SDK key) is +// never admitted. +// - Drop the failed new anchor, but only if it was brand new; a previously-accepted key that was +// promoted and failed stays accepted as the non-anchor key it already was. +func (r *Rotator) RevertAnchorChange(change AnchorChange) { + r.mu.Lock() + defer r.mu.Unlock() + + // Only re-admit a defined previous anchor. When an env gains its first SDK key, the previous anchor + // is the empty (undefined) key — there is nothing to re-admit, and inserting "" would put an + // undefined credential into the accepted set (the rotator otherwise only holds defined keys). + if change.PreviousAnchor.Defined() { + if _, stillAccepted := r.acceptedSDKKeys[change.PreviousAnchor]; !stillAccepted { + r.acceptedSDKKeys[change.PreviousAnchor] = AcceptedKey{} + } + } + if !change.NewAnchorPreviouslyAccepted { + delete(r.acceptedSDKKeys, change.NewAnchor) + } } // reconcilableKey constrains the generic reconcile helper to a comparable credential (so it can key a @@ -276,7 +401,11 @@ func reconcileAcceptedKeys[K reconcilableKey]( // reconcileAcceptedKeys. The set is trusted as well-formed: BuildAcceptedSet / the builder guarantee // the anchor is present and permanent (WithAnchor forces a nil expiry), so no special handling is // needed here. The caller must hold the write lock. -func (r *Rotator) reconcileSDKKeys(set AcceptedSet, anchor config.SDKKey, now time.Time) { +// +// NOTE: reconcileSDKKeys does NOT flip r.anchorKey when the anchor changes. The Reconcile caller +// signals the anchor change via ReconcileResult.AnchorChange and invokes CommitAnchor to move the +// pointer once the synchronous re-anchor sequence is ready (see Reconcile + CommitAnchor). +func (r *Rotator) reconcileSDKKeys(set AcceptedSet, now time.Time) { desired := make(map[config.SDKKey]AcceptedKey, len(set.sdkKeys)) for key, info := range set.sdkKeys { if info.Expiry != nil && !now.Before(*info.Expiry) { @@ -285,7 +414,6 @@ func (r *Rotator) reconcileSDKKeys(set AcceptedSet, anchor config.SDKKey, now ti desired[key] = info } reconcileAcceptedKeys(desired, r.acceptedSDKKeys, &r.additions, &r.expirations, r.loggers, "SDK key") - r.anchorKey = anchor } // reconcileMobileKeys mirrors reconcileSDKKeys for mobile keys. The set is trusted as well-formed: diff --git a/internal/credential/rotator_test.go b/internal/credential/rotator_test.go index b89eb56a..d56c4ae7 100644 --- a/internal/credential/rotator_test.go +++ b/internal/credential/rotator_test.go @@ -54,10 +54,13 @@ func TestReconcileAnchorOnly(t *testing.T) { anchor := config.SDKKey("anchor") now := time.Now() - r.Reconcile(mustBuild(t, NewAcceptedSetBuilder().WithAnchor(SDKKeyParams{Value: anchor})), now) + result := r.Reconcile(mustBuild(t, NewAcceptedSetBuilder().WithAnchor(SDKKeyParams{Value: anchor})), now) + require.NotNil(t, result.AnchorChange, "anchor transition from empty to defined is signaled") + r.CommitAnchor(result.AnchorChange.NewAnchor) additions, expirations := r.StepTime(now) - assert.ElementsMatch(t, []SDKCredential{anchor}, additions) + // The anchor is stripped from additions — the synchronous re-anchor sequence owns its setup. + assert.Empty(t, additions) assert.Empty(t, expirations) assert.Equal(t, anchor, r.AnchorKey()) assert.ElementsMatch(t, []SDKCredential{anchor}, r.AllCredentials()) @@ -70,12 +73,15 @@ func TestReconcileMultipleSDKKeys(t *testing.T) { other := config.SDKKey("other") now := time.Now() - r.Reconcile( + result := r.Reconcile( mustBuild(t, NewAcceptedSetBuilder().WithAnchor(SDKKeyParams{Value: anchor}).WithSDKKey(SDKKeyParams{Value: other})), now) + require.NotNil(t, result.AnchorChange) + r.CommitAnchor(result.AnchorChange.NewAnchor) additions, expirations := r.StepTime(now) - // Both server keys are accepted; only the anchor is primary. - assert.ElementsMatch(t, []SDKCredential{anchor, other}, additions) + // Both server keys are accepted; only the non-anchor server key is in additions (the anchor is + // owned by the synchronous re-anchor sequence in env_context_impl). + assert.ElementsMatch(t, []SDKCredential{other}, additions) assert.Empty(t, expirations) assert.Equal(t, anchor, r.AnchorKey()) assert.ElementsMatch(t, []SDKCredential{anchor, other}, r.AllCredentials()) @@ -93,8 +99,9 @@ func TestReconcileMultipleMobileKeys(t *testing.T) { mustBuild(t, NewAcceptedSetBuilder().WithAnchor(SDKKeyParams{Value: anchor}).WithPrimaryMobileKey(MobileKeyParams{Value: mob1}).WithMobileKey(MobileKeyParams{Value: mob2})), now) additions, _ := r.StepTime(now) - // Every mobile key is accepted; the designated one is the primary. - assert.ElementsMatch(t, []SDKCredential{anchor, mob1, mob2}, additions) + // Every mobile key is accepted; the anchor is owned by the synchronous re-anchor (stripped from + // additions). The designated primary mobile key and the other mobile key remain in additions. + assert.ElementsMatch(t, []SDKCredential{mob1, mob2}, additions) assert.Equal(t, mob1, r.MobileKey()) assert.ElementsMatch(t, []SDKCredential{anchor, mob1, mob2}, r.AllCredentials()) } @@ -140,7 +147,8 @@ func TestReconcileAcceptsExpiringKeysAsData(t *testing.T) { now) additions, expirations := r.StepTime(now) - assert.ElementsMatch(t, []SDKCredential{anchor, expiringSDK, mob, expiringMobile}, additions) + // Anchor is stripped from additions (owned by the synchronous re-anchor); other keys flow through. + assert.ElementsMatch(t, []SDKCredential{expiringSDK, mob, expiringMobile}, additions) assert.Empty(t, expirations) // Every key is accepted (still authenticates)... assert.ElementsMatch(t, []SDKCredential{anchor, expiringSDK, mob, expiringMobile}, r.AllCredentials()) @@ -191,7 +199,8 @@ func TestReconcileExpiringKeysAreEvictedByStepTime(t *testing.T) { WithMobileKey(MobileKeyParams{Value: expiringMobile, Expiry: util.PtrOrNil(expiry)})), now) additions, expirations := r.StepTime(now) - require.ElementsMatch(t, []SDKCredential{anchor, expiringSDK, mob, expiringMobile}, additions) + // Anchor is stripped from additions (owned by the synchronous re-anchor); other keys flow through. + require.ElementsMatch(t, []SDKCredential{expiringSDK, mob, expiringMobile}, additions) require.Empty(t, expirations) // At the exact expiry, expiry is strict (now must be strictly after), so nothing is dropped yet. @@ -217,15 +226,18 @@ func TestReconcileAlreadyExpiredKeyIsIgnoredOnAdd(t *testing.T) { now := time.Unix(2000, 0) alreadyExpired := now.Add(-time.Hour) - r.Reconcile( + result := r.Reconcile( mustBuild(t, NewAcceptedSetBuilder(). WithAnchor(SDKKeyParams{Value: anchor}). WithSDKKey(SDKKeyParams{Value: staleKey, Expiry: util.PtrOrNil(alreadyExpired)})), now) + require.NotNil(t, result.AnchorChange) + r.CommitAnchor(result.AnchorChange.NewAnchor) additions, expirations := r.StepTime(now) - // Only the anchor is added; the stale key is never accepted. - assert.ElementsMatch(t, []SDKCredential{anchor}, additions) + // The fresh anchor is stripped from additions (the synchronous re-anchor sequence owns its setup), + // and the already-expired stale key is never accepted — so nothing is added. + assert.Empty(t, additions) assert.Empty(t, expirations) assert.ElementsMatch(t, []SDKCredential{anchor}, r.AllCredentials()) } @@ -273,9 +285,11 @@ func TestReconcileDeExpiryRestoresKey(t *testing.T) { func TestAcceptedKeys(t *testing.T) { t.Run("single anchor plus primary mobile", func(t *testing.T) { r := newTestRotator() - r.Reconcile(mustBuild(t, NewAcceptedSetBuilder(). + result := r.Reconcile(mustBuild(t, NewAcceptedSetBuilder(). WithAnchor(SDKKeyParams{Value: "sdk-anchor", Key: util.PtrOrNil("default")}). WithPrimaryMobileKey(MobileKeyParams{Value: "mob-primary", Key: util.PtrOrNil("mob-1")})), time.Unix(0, 0)) + require.NotNil(t, result.AnchorChange) + r.CommitAnchor(result.AnchorChange.NewAnchor) set := r.AcceptedKeys() require.Len(t, set.Server, 1) @@ -298,11 +312,13 @@ func TestAcceptedKeys(t *testing.T) { t.Run("multiple keys include the anchor; expiry populated", func(t *testing.T) { r := newTestRotator() expiry := time.Date(2099, 6, 1, 0, 0, 0, 0, time.UTC) - r.Reconcile(mustBuild(t, NewAcceptedSetBuilder(). + result := r.Reconcile(mustBuild(t, NewAcceptedSetBuilder(). WithAnchor(SDKKeyParams{Value: "sdk-anchor", Key: util.PtrOrNil("default")}). WithSDKKey(SDKKeyParams{Value: "sdk-b", Key: util.PtrOrNil("b-service")}). WithSDKKey(SDKKeyParams{Value: "sdk-old", Key: util.PtrOrNil("old-key"), Expiry: util.PtrOrNil(expiry)}). WithPrimaryMobileKey(MobileKeyParams{Value: "mob-primary"})), time.Unix(0, 0)) + require.NotNil(t, result.AnchorChange) + r.CommitAnchor(result.AnchorChange.NewAnchor) set := r.AcceptedKeys() require.Len(t, set.Server, 3) // anchor + sdk-b + sdk-old @@ -349,3 +365,164 @@ func TestReconcileClearsStaleKeyIdentifier(t *testing.T) { require.True(t, ok) assert.Nil(t, b.Key, "identifier must be cleared when the new payload carries none") } + +func TestReconcileAnchorChangeToPreviouslyAcceptedKey(t *testing.T) { + // When the anchor moves to a key that was already accepted (a non-anchor server key promoted to + // anchor), the AnchorChange reports NewAnchorPreviouslyAccepted == true, and the key is not queued + // as a new addition (it was already accepted). + r := newTestRotator() + anchor := config.SDKKey("anchor") + other := config.SDKKey("other") + now := time.Now() + + result := r.Reconcile(mustBuild(t, NewAcceptedSetBuilder(). + WithAnchor(SDKKeyParams{Value: anchor}). + WithSDKKey(SDKKeyParams{Value: other})), now) + require.NotNil(t, result.AnchorChange) + r.CommitAnchor(result.AnchorChange.NewAnchor) + r.StepTime(now) + + // Move the anchor to `other`, keeping the old anchor accepted. + result = r.Reconcile(mustBuild(t, NewAcceptedSetBuilder(). + WithAnchor(SDKKeyParams{Value: other}). + WithSDKKey(SDKKeyParams{Value: anchor})), now) + require.NotNil(t, result.AnchorChange) + assert.Equal(t, anchor, result.AnchorChange.PreviousAnchor) + assert.Equal(t, other, result.AnchorChange.NewAnchor) + assert.True(t, result.AnchorChange.NewAnchorPreviouslyAccepted, "other was already in the accepted set") + + additions, _ := r.StepTime(now) + assert.NotContains(t, additions, SDKCredential(other), "an already-accepted new anchor is not a fresh addition") +} + +func TestReconcileMobilePrimaryRepointToAlreadyAcceptedKey(t *testing.T) { + // Switching the primary mobile key to a key that is already accepted must be signaled via + // MobilePrimaryRepoint, because addCredential's gate won't fire for it (it's not in additions). + r := newTestRotator() + anchor := config.SDKKey("anchor") + mob1 := config.MobileKey("mob1") + mob2 := config.MobileKey("mob2") + now := time.Now() + + result := r.Reconcile(mustBuild(t, NewAcceptedSetBuilder(). + WithAnchor(SDKKeyParams{Value: anchor}). + WithPrimaryMobileKey(MobileKeyParams{Value: mob1}). + WithMobileKey(MobileKeyParams{Value: mob2})), now) + require.NotNil(t, result.AnchorChange) + r.CommitAnchor(result.AnchorChange.NewAnchor) + r.StepTime(now) + + // Make mob2 (already accepted) the primary. It is not a new addition, so it must be signaled. + result = r.Reconcile(mustBuild(t, NewAcceptedSetBuilder(). + WithAnchor(SDKKeyParams{Value: anchor}). + WithPrimaryMobileKey(MobileKeyParams{Value: mob2}). + WithMobileKey(MobileKeyParams{Value: mob1})), now) + assert.Nil(t, result.AnchorChange, "the anchor did not change") + require.NotNil(t, result.MobilePrimaryRepoint, "an already-accepted new primary mobile key must be signaled") + assert.Equal(t, mob2, *result.MobilePrimaryRepoint) +} + +func TestReconcileMobilePrimaryToNewKeyDoesNotSignalRepoint(t *testing.T) { + // When the new primary mobile key was NOT already accepted, it arrives via the additions list, so + // addCredential's gate handles the ReplaceCredential and MobilePrimaryRepoint stays nil. + r := newTestRotator() + anchor := config.SDKKey("anchor") + mob1 := config.MobileKey("mob1") + mob2 := config.MobileKey("mob2") + now := time.Now() + + result := r.Reconcile(mustBuild(t, NewAcceptedSetBuilder(). + WithAnchor(SDKKeyParams{Value: anchor}). + WithPrimaryMobileKey(MobileKeyParams{Value: mob1})), now) + require.NotNil(t, result.AnchorChange) + r.CommitAnchor(result.AnchorChange.NewAnchor) + r.StepTime(now) + + result = r.Reconcile(mustBuild(t, NewAcceptedSetBuilder(). + WithAnchor(SDKKeyParams{Value: anchor}). + WithPrimaryMobileKey(MobileKeyParams{Value: mob2})), now) + assert.Nil(t, result.MobilePrimaryRepoint, "a brand-new primary mobile key is handled via additions, not the repoint signal") + + additions, _ := r.StepTime(now) + assert.Contains(t, additions, SDKCredential(mob2), "a brand-new primary mobile key is queued as an addition") +} + +func TestRevertAnchorChangeReadmitsRevokedPreviousAnchor(t *testing.T) { + // When the previous anchor was immediately revoked in the same reconcile (dropped from the accepted + // set) and the re-anchor rolled back, RevertAnchorChange re-admits it and drops the failed new key. + r := newTestRotator() + keyA := config.SDKKey("keyA") + keyB := config.SDKKey("keyB") + now := time.Now() + + res := r.Reconcile(mustBuild(t, NewAcceptedSetBuilder().WithAnchor(SDKKeyParams{Value: keyA})), now) + require.NotNil(t, res.AnchorChange) + r.CommitAnchor(res.AnchorChange.NewAnchor) + r.StepTime(now) + + // Move the anchor to keyB, omitting keyA entirely (immediate revocation). Reconcile drops keyA. + res = r.Reconcile(mustBuild(t, NewAcceptedSetBuilder().WithAnchor(SDKKeyParams{Value: keyB})), now) + require.NotNil(t, res.AnchorChange) + require.False(t, res.AnchorChange.NewAnchorPreviouslyAccepted) + + // Simulate the rollback: the caller did NOT CommitAnchor, so the pointer still names keyA. + r.RevertAnchorChange(*res.AnchorChange) + + assert.Equal(t, keyA, r.AnchorKey()) + creds := r.AllCredentials() + assert.Contains(t, creds, SDKCredential(keyA), "previous anchor re-admitted") + assert.NotContains(t, creds, SDKCredential(keyB), "failed new anchor dropped") +} + +func TestRevertAnchorChangeLeavesGraceDemotedPreviousAnchorUntouched(t *testing.T) { + // When the previous anchor was demoted with a grace expiry (still accepted) rather than revoked, + // RevertAnchorChange must leave it — and its expiry — untouched, and must not make it permanent. + r := newTestRotator() + keyA := config.SDKKey("keyA") + keyB := config.SDKKey("keyB") + now := time.Unix(1000, 0) + expiry := now.Add(time.Hour) + + res := r.Reconcile(mustBuild(t, NewAcceptedSetBuilder().WithAnchor(SDKKeyParams{Value: keyA})), now) + require.NotNil(t, res.AnchorChange) + r.CommitAnchor(res.AnchorChange.NewAnchor) + r.StepTime(now) + + // Move the anchor to keyB while keeping keyA accepted with a grace expiry. + res = r.Reconcile(mustBuild(t, NewAcceptedSetBuilder(). + WithAnchor(SDKKeyParams{Value: keyB}). + WithSDKKey(SDKKeyParams{Value: keyA, Expiry: util.PtrOrNil(expiry)})), now) + require.NotNil(t, res.AnchorChange) + + r.RevertAnchorChange(*res.AnchorChange) + + set := r.AcceptedKeys() + kaInfo, ok := set.Server[keyA] + require.True(t, ok, "grace-demoted previous anchor stays accepted") + require.NotNil(t, kaInfo.Expiry, "RevertAnchorChange must not wipe the grace expiry / make it permanent") + assert.Equal(t, expiry, *kaInfo.Expiry) + _, keyBAccepted := set.Server[keyB] + assert.False(t, keyBAccepted, "failed new anchor dropped") +} + +func TestRevertAnchorChangeDoesNotAdmitUndefinedPreviousAnchor(t *testing.T) { + // When an env gains its first SDK key, the AnchorChange's previous anchor is the empty (undefined) + // key. A rollback must not insert that empty key into the accepted set. + r := newTestRotator() + keyB := config.SDKKey("keyB") + now := time.Now() + + res := r.Reconcile(mustBuild(t, NewAcceptedSetBuilder().WithAnchor(SDKKeyParams{Value: keyB})), now) + require.NotNil(t, res.AnchorChange) + require.False(t, res.AnchorChange.PreviousAnchor.Defined(), "the first SDK key has an undefined previous anchor") + + // Simulate a failed build / rollback. + r.RevertAnchorChange(*res.AnchorChange) + + for _, cred := range r.AllCredentials() { + if sdkKey, ok := cred.(config.SDKKey); ok { + assert.True(t, sdkKey.Defined(), "revert must not insert an undefined SDK key into the accepted set") + } + } + assert.NotContains(t, r.AllCredentials(), SDKCredential(keyB), "failed new anchor dropped") +} diff --git a/internal/relayenv/env_context_impl.go b/internal/relayenv/env_context_impl.go index 7bad18ae..a37658b5 100644 --- a/internal/relayenv/env_context_impl.go +++ b/internal/relayenv/env_context_impl.go @@ -123,6 +123,12 @@ type envContextImpl struct { connectionMapper ConnectionMapper offline bool closed bool + + // reconcileMu serializes reconcileCredentials calls — only one runs at a time, including the + // synchronous re-anchor sequence inside it. Held separately from mu so that GetClient / GetStore / + // GetEvaluator / addCredential continue to run during the (potentially seconds-long) SDK client + // construction when re-anchoring to a new key. + reconcileMu sync.Mutex } // Implementation of the DataStoreQueries interface that the streams package uses as an abstraction of @@ -430,17 +436,13 @@ func (c *envContextImpl) cleanupExpiredCredentials(interval time.Duration) { func (c *envContextImpl) addCredential(newCredential credential.SDKCredential) { c.mu.Lock() defer c.mu.Unlock() - c.envStreams.AddCredential(newCredential) - for streamProvider, handlers := range c.handlers { - if h := streamProvider.Handler(sdkauth.NewScoped(c.filterKey, newCredential)); h != nil { - handlers[newCredential] = h - } - } + + c.registerCredentialMappings(newCredential) // A new SDK key means: // 1. we should start a new SDK client*, but only for the anchor: there is a single upstream - // connection per environment, owned by the anchor key. Non-anchor server keys get envStreams - // + handler bundles above, but no upstream client — matching today's mobile-key behavior. + // connection per environment, owned by the anchor key. Non-anchor server keys get their + // credential mappings registered above, but no upstream client — matching today's mobile-key behavior. // 2. we should tell all event forwarding components that use an SDK key to use the new one, // again only when it is the anchor, since events collapse to the anchor per kind. // A new mobile key does not require starting a new SDK client, but does requiring updating any event forwarding @@ -472,8 +474,6 @@ func (c *envContextImpl) addCredential(newCredential credential.SDKCredential) { } } } - - c.connectionMapper.AddConnectionMapping(sdkauth.NewScoped(c.filterKey, newCredential), c) } func (c *envContextImpl) removeCredential(oldCredential credential.SDKCredential) { @@ -526,21 +526,9 @@ func (c *envContextImpl) startSDKClient(sdkKey config.SDKKey, readyCh chan<- Env } c.clients[sdkKey] = client - // The data store instance is created by the SDK when it creates the client. Now that - // we have a data store, we can finish setting up the Evaluator that we'll use for this - // environment. - store := c.storeAdapter.GetStore() - dataProvider := ldstoreimpl.NewDataStoreEvaluatorDataProvider(store, c.loggers) - evalOptions := []ldeval.EvaluatorOption{ - // We're setting EnableSecondaryKey because we may be doing evaluations for client-side SDKs that - // are sending old-style user data with the "secondary" attribute. This option doesn't affect - // evaluations done for newer client-side SDKs that send contexts. - ldeval.EvaluatorOptionEnableSecondaryKey(true), - } - if c.sdkBigSegments != nil { - evalOptions = append(evalOptions, ldeval.EvaluatorOptionBigSegmentProvider(c.sdkBigSegments)) - } - c.evaluator = ldeval.NewEvaluatorWithOptions(dataProvider, evalOptions...) + // The data store instance is created by the SDK when it creates the client. Now that we have a + // data store, we can finish setting up the Evaluator for this environment. + c.rebuildEvaluator() } c.initErr = err c.mu.Unlock() @@ -605,16 +593,273 @@ func (c *envContextImpl) ReconcileCredentials(newSet credential.AcceptedSet) { // reconcileCredentials is the time-injectable implementation of ReconcileCredentials. now is the // reference time for expiry math; production callers pass time.Now() via ReconcileCredentials. // -// The Rotator owns the diff (add → re-anchor → remove) and queues the resulting additions and -// expirations; triggerCredentialChanges then applies them, draining additions before expirations so -// the accepted set is a superset during the transition. addCredential opens an upstream client only -// for the anchor, so non-anchor server keys are accepted and routed without a second connection. +// Order of operations: add → re-anchor → remove. Additions drain first so credential mappings +// are registered before any synchronous re-anchor runs, the re-anchor swaps the upstream client +// while the old anchor is still serving, and expirations drain last so the old anchor's client +// (and any other revoked keys) are torn down only after the new anchor is fully operational. addCredential +// opens an upstream client only for the anchor — non-anchor server keys are accepted and routed +// without a second connection. +// +// Re-anchor handling: Reconcile defers the SDK anchor flip and strips the new anchor from additions +// so this method owns the new anchor's setup. The synchronous re-anchor sequence — build a new client +// (new anchor key) or reuse an existing one (previously-accepted anchor key), CommitAnchor, +// ReplaceCredential on event dispatcher + metrics publisher, re-wire big-segment sync — happens between the addition +// and expiration phases. The MobilePrimaryRepoint case (primary mobile key changed to a key that +// was already in the accepted set) is handled in the same window: addCredential's gate won't fire +// for it because the key isn't in additions, so ReplaceCredential is called synchronously. +// +// reconcileMu serializes concurrent reconciles. If two reconciles arrive while a synchronous build +// is in flight, the second blocks here until the first completes — matching the all-or-nothing +// atomicity requirement. func (c *envContextImpl) reconcileCredentials(newSet credential.AcceptedSet, now time.Time) { - c.keyRotator.Reconcile(newSet, now) - c.triggerCredentialChanges(now) + c.reconcileMu.Lock() + defer c.reconcileMu.Unlock() + + result := c.keyRotator.Reconcile(newSet, now) + additions, expirations := c.keyRotator.StepTime(now) + + for _, cred := range additions { + c.addCredential(cred) + } + + if result.AnchorChange != nil { + if committed := c.reanchor(result.AnchorChange); !committed { + // The re-anchor rolled back — the new anchor's client never came up. Back out just this + // anchor change (other changes in the payload stand). The env-side undo mirrors + // RevertAnchorChange's accepted-set logic exactly: + // - A brand-new anchor: reanchor registered its mappings this cycle, so undo them here; + // RevertAnchorChange drops it from the accepted set. + // - A previously-accepted anchor (e.g. a non-anchor key promoted to anchor): reanchor did + // NOT register mappings (they predate this reconcile), so we must NOT tear them down; + // RevertAnchorChange keeps it accepted. It reverts to the non-anchor key it already was. + if !result.AnchorChange.NewAnchorPreviouslyAccepted { + c.removeCredential(result.AnchorChange.NewAnchor) + } + c.keyRotator.RevertAnchorChange(*result.AnchorChange) + // Keep the previous anchor's client serving by not expiring it here — even if this payload + // revoked it outright. (A grace-demoted previous anchor isn't in expirations anyway, so this + // only matters for an immediate revocation.) + previousAnchor := result.AnchorChange.PreviousAnchor + expirations = slices.DeleteFunc(expirations, func(cred credential.SDKCredential) bool { + return cred == previousAnchor + }) + } + } + + if result.MobilePrimaryRepoint != nil { + c.mu.RLock() + dispatcher := c.eventDispatcher + c.mu.RUnlock() + if dispatcher != nil { + dispatcher.ReplaceCredential(*result.MobilePrimaryRepoint) + } + } + + for _, cred := range expirations { + c.removeCredential(cred) + } +} + +// reanchor drives the synchronous re-anchor sequence for an SDK anchor change signaled by Reconcile's +// ReconcileResult.AnchorChange. Invoked by reconcileCredentials after additions have been processed +// and before expirations, so the previous anchor's client is still alive while the new client is built +// (or reused). +// +// reanchor holds c.mu for the whole sequence and releases it only around the SDK client build (which +// must not hold the lock — see buildNewAnchorClient). Holding one continuous lock otherwise keeps +// Close() (which also takes c.mu) from tearing down clients or the dispatcher mid-commit, and lets +// commitReanchor assume the lock is held rather than re-acquiring it. +// +// - When there is no existing client for the new anchor and the env is online: register its credential +// mappings if the key is brand new (Reconcile stripped it from additions), build a new SDK client, +// and on Initialized commit the anchor. On init failure, roll back: do not commit, leave the previous +// anchor authoritative (its client keeps serving), and log a structured error. +// - When a client already exists (e.g. a former anchor still in its grace period), or the env is +// offline: no build, just commit. +// +// Returns true if the anchor was committed, false if it rolled back (init failure or the env closed +// mid-build), so reconcileCredentials can back out the anchor change. The old anchor's client is not +// closed here; its grace-period expiration drives removeCredential. +func (c *envContextImpl) reanchor(change *credential.AnchorChange) bool { + newAnchor := change.NewAnchor + previousAnchor := change.PreviousAnchor + + c.mu.Lock() + defer c.mu.Unlock() + + // The two signals below answer two different questions, so they are used independently: + // - NewAnchorPreviouslyAccepted: were this key's credential mappings already registered? If it + // was already accepted (e.g. a non-anchor server key that addCredential registered mappings for + // but never built a client), the mappings exist, so skip re-registering. If it is brand new + // (Reconcile stripped it from additions, so addCredential never ran for it), register them here. + // - the client check below: is there already a client for this key? If so, reuse it; otherwise + // build one. + // These genuinely differ: promoting a previously-accepted non-anchor key to anchor has mappings but + // no client (register: skip, client: build). The one-way invariant "a client exists => + // NewAnchorPreviouslyAccepted" holds (removeCredential deletes a key's client in lockstep with the + // rotator dropping it), which is why gating mapping registration on NewAnchorPreviouslyAccepted is + // safe: a key with a live client is always already accepted, so it never re-registers. + if !change.NewAnchorPreviouslyAccepted { + c.registerCredentialMappings(newAnchor) + } + + // A live client for the new anchor is reused as-is; an offline env has no upstream client to build. + // Either way there is nothing to build, so fall through to commit. + why := "reused existing client" + if c.clients[newAnchor] == nil { + if c.offline { + why = "offline — no client build" + } else { + // Build the new client without the lock: sdkClientFactory can block for up to sdkInitTimeout, + // and holding c.mu that long would stall every GetClient/GetStore caller (see reconcileMu). + // reanchor's deferred Unlock releases the lock we re-acquire here on return. + c.mu.Unlock() + client := c.buildNewAnchorClient(newAnchor, previousAnchor) + c.mu.Lock() + + if client == nil { + // Init failed; buildNewAnchorClient already logged and closed the half-built client. Do not + // commit — leave the previous anchor authoritative. + return false + } + if c.closed { + // The env was torn down while the lock was released for the build (Close() does not hold + // reconcileMu, so it can run concurrently); its client-teardown loop has already finished and + // would never close this one, so discard the freshly-built client rather than install it into + // a closed env (mirrors the guard in startSDKClient). + _ = client.Close() + return false + } + if existing := c.clients[newAnchor]; existing != nil && existing != client { + // Stale-client guard: the lock was released for the build, so re-check and close any client + // installed concurrently for newAnchor. + _ = existing.Close() + } + c.clients[newAnchor] = client + // With store handover, GetStore() returns the SAME wrapper the old client used, so the rebuilt + // evaluator serves the already-populated data immediately (no empty-store window). + c.rebuildEvaluator() + why = "built new client" + } + } + + return c.commitReanchor(newAnchor, previousAnchor, why) +} + +// buildNewAnchorClient constructs the SDK client for a re-anchor to newAnchor. It must run without c.mu +// held: sdkClientFactory can block for up to sdkInitTimeout, and holding the lock that long would stall +// every GetClient/GetStore caller. It touches only fields fixed at construction (sdkClientFactory, +// sdkConfig, sdkInitTimeout, globalLoggers), so it needs no lock — mirroring startSDKClient, which also +// builds before locking. +// +// Returns the initialized client, or nil if the build failed, in which case it has already closed any +// half-built client and logged a structured error. initErr is deliberately left untouched on failure: +// it feeds the request middleware, and setting it to the new anchor's ErrInitializationFailed would 401 +// an env that is serving fine on the previous anchor. +func (c *envContextImpl) buildNewAnchorClient(newAnchor, previousAnchor config.SDKKey) sdks.LDClientContext { + client, err := c.sdkClientFactory(newAnchor, c.sdkConfig, c.sdkInitTimeout) + if err != nil || client == nil || !client.Initialized() { + var initialized bool + if client != nil { + initialized = client.Initialized() + _ = client.Close() + } + c.globalLoggers.Errorf("Re-anchor to SDK key %s failed (err=%v initialized=%v); "+ + "preserving previous anchor %s", + newAnchor.Masked(), err, initialized, previousAnchor.Masked()) + return nil + } + return client } +// commitReanchor is the second half of the re-anchor sequence: atomically move the rotator's anchor +// pointer, clear any stale init error now that a healthy client is current, and repoint downstream +// event/metrics forwarding. The caller must hold c.mu — reanchor holds it across the whole sequence, so +// the commit and Close() (which also takes c.mu) are mutually exclusive and Close can't tear the client +// or dispatcher out mid-commit. This mirrors addCredential, which likewise repoints event forwarding +// and reads rotator state under c.mu. +// +// Returns false without committing if the env was closed first, so callers report the rollback rather +// than a phantom success. +func (c *envContextImpl) commitReanchor(newAnchor, previousAnchor config.SDKKey, why string) bool { + if c.closed { + // Close() ran before we could commit. Don't flip the anchor or touch the (now-closed) dispatcher + // and metrics publisher; the env is being torn down. + return false + } + + c.keyRotator.CommitAnchor(newAnchor) + // The anchor now points at a healthy client (freshly built and Initialized, or a reused live + // client), so clear any init error a prior client left behind — otherwise GetInitError() and the + // request middleware would keep reporting a still-serving env as failed. + c.initErr = nil + + if c.metricsEventPub != nil { + c.metricsEventPub.ReplaceCredential(newAnchor) + } + if c.eventDispatcher != nil { + c.eventDispatcher.ReplaceCredential(newAnchor) + } + + // Big-segment synchronization is intentionally left pointing at the previous anchor key across a + // re-anchor: this matches pre-concurrent-keys behavior (there was no re-anchor, so it never moved) + // and does not regress. When big-segment re-anchor is implemented, its re-wire hook belongs right + // here, after the event/metrics ReplaceCredential calls — either recreate the BigSegmentSynchronizer + // for newAnchor, or add a credential-replacement method to it. + + c.globalLoggers.Infof("Re-anchored SDK from %s to %s (%s)", previousAnchor.Masked(), newAnchor.Masked(), why) + return true +} + +// rebuildEvaluator constructs the environment's Evaluator against the current data store. It is called +// after (re)creating an SDK client, once the store is available, and is shared by the initial client +// startup and the re-anchor path. It reads and writes envContextImpl fields directly, so the caller +// must hold c.mu. +// +// EnableSecondaryKey is set because we may evaluate for client-side SDKs sending old-style user data +// with the "secondary" attribute; it has no effect for newer SDKs that send contexts. +func (c *envContextImpl) rebuildEvaluator() { + store := c.storeAdapter.GetStore() + dataProvider := ldstoreimpl.NewDataStoreEvaluatorDataProvider(store, c.loggers) + evalOptions := []ldeval.EvaluatorOption{ + ldeval.EvaluatorOptionEnableSecondaryKey(true), + } + if c.sdkBigSegments != nil { + evalOptions = append(evalOptions, ldeval.EvaluatorOptionBigSegmentProvider(c.sdkBigSegments)) + } + c.evaluator = ldeval.NewEvaluatorWithOptions(dataProvider, evalOptions...) +} + +// registerCredentialMappings wires relay's downstream-facing routing for cred: it registers the +// credential with the env's stream machinery, builds the per-stream-provider HTTP handlers, and adds +// the connection→env mapping, so incoming SDK/client connections that authenticate with cred are +// served by this env. It does NOT start the upstream SDK client or repoint event/metrics forwarding — +// those are anchor-only concerns owned by the callers (addCredential, and the re-anchor sequence). +// The caller must hold c.mu. +func (c *envContextImpl) registerCredentialMappings(cred credential.SDKCredential) { + c.envStreams.AddCredential(cred) + for streamProvider, handlers := range c.handlers { + if h := streamProvider.Handler(sdkauth.NewScoped(c.filterKey, cred)); h != nil { + handlers[cred] = h + } + } + c.connectionMapper.AddConnectionMapping(sdkauth.NewScoped(c.filterKey, cred), c) +} + +// triggerCredentialChanges drains the rotator's StepTime queue and applies the resulting additions +// and expirations. It runs on the cleanup ticker (cleanupExpiredCredentials), so it can fire at any +// moment — including while a synchronous re-anchor is in flight inside reconcileCredentials. +// +// It takes reconcileMu for the whole StepTime + add/remove pass so the ticker is serialized against +// reconcileCredentials exactly the way concurrent reconciles already are. Without it, a credential +// expiry firing during an in-flight re-anchor would drain the same StepTime queue the reconcile +// relies on (the ticker could steal additions a reconcile just queued) and could removeCredential — +// closing a client — partway through the re-anchor sequence. reconcileCredentials never calls this, +// so taking reconcileMu here introduces no re-entrancy. func (c *envContextImpl) triggerCredentialChanges(now time.Time) { + c.reconcileMu.Lock() + defer c.reconcileMu.Unlock() + additions, expirations := c.keyRotator.StepTime(now) for _, cred := range additions { c.addCredential(cred) diff --git a/internal/relayenv/env_context_impl_test.go b/internal/relayenv/env_context_impl_test.go index e77b7eaa..ebd6a12a 100644 --- a/internal/relayenv/env_context_impl_test.go +++ b/internal/relayenv/env_context_impl_test.go @@ -60,12 +60,17 @@ func requireClientReady(t *testing.T, clientCh chan *testclient.FakeLDClient) *t func makeBasicEnv(t *testing.T, envConfig config.EnvConfig, clientFactory sdks.ClientFactoryFunc, loggers ldlog.Loggers, readyCh chan EnvContext) EnvContext { + return makeBasicEnvWithMapper(t, envConfig, clientFactory, loggers, readyCh, mockConnectionMapper{}) +} + +func makeBasicEnvWithMapper(t *testing.T, envConfig config.EnvConfig, clientFactory sdks.ClientFactoryFunc, + loggers ldlog.Loggers, readyCh chan EnvContext, connMapper ConnectionMapper) EnvContext { env, err := NewEnvContext(EnvContextImplParams{ Identifiers: EnvIdentifiers{ConfiguredName: envName}, EnvConfig: envConfig, ClientFactory: clientFactory, Loggers: loggers, - ConnectionMapper: mockConnectionMapper{}, + ConnectionMapper: connMapper, }, readyCh) require.NoError(t, err) return env @@ -81,6 +86,34 @@ func (m mockConnectionMapper) RemoveConnectionMapping(scopedCredential sdkauth.S } +// recordingConnectionMapper tracks which scoped credentials currently have a connection mapping, so a +// test can assert that a credential's mapping survived (or was torn down). +type recordingConnectionMapper struct { + mu sync.Mutex + mapped map[sdkauth.ScopedCredential]bool +} + +func (m *recordingConnectionMapper) AddConnectionMapping(scopedCredential sdkauth.ScopedCredential, _ EnvContext) { + m.mu.Lock() + defer m.mu.Unlock() + if m.mapped == nil { + m.mapped = map[sdkauth.ScopedCredential]bool{} + } + m.mapped[scopedCredential] = true +} + +func (m *recordingConnectionMapper) RemoveConnectionMapping(scopedCredential sdkauth.ScopedCredential) { + m.mu.Lock() + defer m.mu.Unlock() + delete(m.mapped, scopedCredential) +} + +func (m *recordingConnectionMapper) isMapped(filterKey config.FilterKey, cred credential.SDKCredential) bool { + m.mu.Lock() + defer m.mu.Unlock() + return m.mapped[sdkauth.NewScoped(filterKey, cred)] +} + func TestConstructorBasicProperties(t *testing.T) { envConfig := st.EnvWithAllCredentials.Config envConfig.TTL = configtypes.NewOptDuration(time.Hour) diff --git a/internal/relayenv/env_context_reanchor_synchronous_test.go b/internal/relayenv/env_context_reanchor_synchronous_test.go new file mode 100644 index 00000000..f04fb386 --- /dev/null +++ b/internal/relayenv/env_context_reanchor_synchronous_test.go @@ -0,0 +1,599 @@ +package relayenv + +// Regression tests for the synchronous re-anchor sequence. +// +// These cover: re-anchoring to a new key (build success and init-failure rollback), re-anchoring to a +// previously-accepted key (reuse the existing client), no orphan clients, store-handover survival, and +// the mobile-primary repoint signal (the gap when the new primary mobile key was already in the +// accepted set, so the primary-mobile gate does not fire for it). + +import ( + "errors" + "testing" + "time" + + "github.com/launchdarkly/ld-relay/v8/config" + "github.com/launchdarkly/ld-relay/v8/internal/credential" + "github.com/launchdarkly/ld-relay/v8/internal/sdks" + st "github.com/launchdarkly/ld-relay/v8/internal/sharedtest" + "github.com/launchdarkly/ld-relay/v8/internal/sharedtest/testclient" + + "github.com/launchdarkly/go-sdk-common/v3/ldlog" + "github.com/launchdarkly/go-sdk-common/v3/ldlogtest" + ld "github.com/launchdarkly/go-server-sdk/v7" + "github.com/launchdarkly/go-server-sdk/v7/subsystems/ldstoreimpl" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const reanchorSyncTestKey2 = config.SDKKey("reanchor-sync-new-anchor") +const reanchorSyncExpiringKey = config.SDKKey("reanchor-sync-expiring-key") + +// reanchorViaReconcile drives the production ReconcileCredentials path with a payload that +// designates newKey as the anchor and keeps oldKey accepted with an expiry one hour in the future, +// mirroring the backend's default-rotation behavior (the new anchor is non-expiring; the demoted +// old anchor carries an expiry). extraAcceptedSDK, if defined, is added as a permanent non-anchor +// SDK key — used to set up the reuse path (a key already in the accepted set later becoming the anchor). +func reanchorViaReconcile( + t *testing.T, + env EnvContext, + newKey, oldKey, extraAcceptedSDK config.SDKKey, + primaryMobile config.MobileKey, + envID config.EnvironmentID, + now time.Time, +) { + t.Helper() + builder := credential.NewAcceptedSetBuilder().WithAnchor(credential.SDKKeyParams{Value: newKey}) + if oldKey.Defined() && oldKey != newKey { + expiry := now.Add(time.Hour) + builder = builder.WithSDKKey(credential.SDKKeyParams{Value: oldKey, Expiry: &expiry}) + } + if extraAcceptedSDK.Defined() && extraAcceptedSDK != newKey { + builder = builder.WithSDKKey(credential.SDKKeyParams{Value: extraAcceptedSDK}) + } + if primaryMobile.Defined() { + builder = builder.WithPrimaryMobileKey(credential.MobileKeyParams{Value: primaryMobile}) + } + if envID.Defined() { + builder = builder.WithEnvironmentID(envID) + } + set, err := builder.Build() + require.NoError(t, err) + env.(*envContextImpl).reconcileCredentials(set, now) +} + +// TestReanchorSync_CaseA_BuildsNewClientAndMovesAnchor exercises the happy path where the new anchor +// is a brand-new SDK key with no existing client: the build succeeds, and the anchor commits. The +// store is handed over (no empty-store window) and the new client becomes current. +func TestReanchorSync_CaseA_BuildsNewClientAndMovesAnchor(t *testing.T) { + featureKind := ldstoreimpl.Features() + flagKey := st.Flag1ServerSide.Flag.Key + envConfig := st.EnvMain.Config + + mockLog := ldlogtest.NewMockLog() + defer mockLog.DumpIfTestFailed(t) + + clientCh := make(chan *testclient.FakeLDClient, 10) + readyCh := make(chan EnvContext, 1) + env := makeBasicEnv(t, envConfig, testclient.FakeLDClientFactoryWithChannel(true, clientCh), mockLog.Loggers, readyCh) + defer env.Close() + + require.Equal(t, env, requireEnvReady(t, readyCh)) + originalClient := requireClientReady(t, clientCh) + require.Eventually(t, func() bool { return env.GetClient() == originalClient }, time.Second, 10*time.Millisecond) + + // Populate the store as the original anchor's stream sync would. + require.NoError(t, env.GetStore().Init(st.AllData)) + originalStore := env.GetStore() + + // Re-anchor onto a brand-new key via ReconcileCredentials. + now := time.Unix(2000, 0) + reanchorViaReconcile(t, env, reanchorSyncTestKey2, envConfig.SDKKey, "", envConfig.MobileKey, envConfig.EnvID, now) + + // A new client was built synchronously and is now the current anchor's client. + newClient := requireClientReady(t, clientCh) + assert.NotSame(t, originalClient, newClient, "a new anchor key builds a fresh client") + assert.Same(t, newClient, env.GetClient(), "GetClient returns the new anchor's client after the commit") + assert.Nil(t, env.GetInitError(), "successful re-anchor clears any prior init error") + + // The rotator's primary now names the new key (CommitAnchor ran). + assert.Equal(t, reanchorSyncTestKey2, env.(*envContextImpl).keyRotator.AnchorKey()) + assert.Contains(t, env.GetCredentials(), credential.SDKCredential(reanchorSyncTestKey2)) + + // Store handover: the wrapper is the same instance, still initialized, data intact. + assert.Same(t, originalStore, env.GetStore(), "store handover: the wrapper survives the re-anchor") + got, err := env.GetStore().Get(featureKind, flagKey) + require.NoError(t, err) + assert.NotNil(t, got.Item, "data survives the re-anchor (no empty-store window)") + + // The old client is still alive — its grace period has not elapsed (the old client keeps serving; + // closure happens via removeCredential when the expiry fires). + envImpl := env.(*envContextImpl) + envImpl.mu.RLock() + _, oldStillPresent := envImpl.clients[envConfig.SDKKey] + envImpl.mu.RUnlock() + assert.True(t, oldStillPresent, "old anchor's client retained during its grace period") +} + +// TestReanchorSync_CaseA_InitFailureRollsBack confirms that a failed new-client init does NOT move +// the rotator's anchor, does NOT install the broken client, and does NOT close the previous client. +// The old anchor keeps serving; the failure surfaces as a structured Errorf log (initErr is left +// untouched so the still-healthy env is not marked failed). This is the all-or-nothing atomicity +// requirement applied to re-anchor. +func TestReanchorSync_CaseA_InitFailureRollsBack(t *testing.T) { + envConfig := st.EnvMain.Config + fakeErr := errors.New("re-anchor: new client init refused") + + mockLog := ldlogtest.NewMockLog() + defer mockLog.DumpIfTestFailed(t) + + clientCh := make(chan *testclient.FakeLDClient, 10) + healthyFactory := testclient.FakeLDClientFactoryWithChannel(true, clientCh) + + // Fail only for the new anchor key; succeed for the original. + factory := func(sdkKey config.SDKKey, cfg ld.Config, timeout time.Duration) (sdks.LDClientContext, error) { + if sdkKey == reanchorSyncTestKey2 { + return nil, fakeErr + } + return healthyFactory(sdkKey, cfg, timeout) + } + + readyCh := make(chan EnvContext, 1) + env := makeBasicEnv(t, envConfig, factory, mockLog.Loggers, readyCh) + defer env.Close() + + require.Equal(t, env, requireEnvReady(t, readyCh)) + originalClient := requireClientReady(t, clientCh) + require.Eventually(t, func() bool { return env.GetClient() == originalClient }, time.Second, 10*time.Millisecond) + + now := time.Unix(2000, 0) + reanchorViaReconcile(t, env, reanchorSyncTestKey2, envConfig.SDKKey, "", envConfig.MobileKey, envConfig.EnvID, now) + + // Rollback: the env stays healthy on the previous anchor, so GetInitError stays nil — setting it + // would 401 a still-serving env at the request middleware. The failure surfaces via a structured + // Error log instead. Anchor pointer unchanged, no new client installed, old client still in place. + assert.NoError(t, env.GetInitError(), "a failed re-anchor must not mark the still-serving env as failed") + mockLog.AssertMessageMatch(t, true, ldlog.Error, "Re-anchor to SDK key .* failed") + assert.Same(t, originalClient, env.GetClient(), "GetClient still returns the previous anchor's client") + assert.Equal(t, envConfig.SDKKey, env.(*envContextImpl).keyRotator.AnchorKey(), "anchor pointer stayed on the previous key") + + envImpl := env.(*envContextImpl) + envImpl.mu.RLock() + _, newAnchorClientInstalled := envImpl.clients[reanchorSyncTestKey2] + _, oldStillPresent := envImpl.clients[envConfig.SDKKey] + envImpl.mu.RUnlock() + assert.False(t, newAnchorClientInstalled, "no client installed for the failed new anchor") + assert.True(t, oldStillPresent, "old anchor's client preserved on rollback") +} + +// TestReanchorSync_CaseB_ReusesExistingClient covers re-anchoring onto a key that already has a +// live client. The simplest deterministic setup: re-anchor A→B (B's client built), then re-anchor +// B→A while A is still in its grace period. A's client still exists, so the second re-anchor must +// reuse it and build nothing new. +func TestReanchorSync_CaseB_ReusesExistingClient(t *testing.T) { + envConfig := st.EnvMain.Config + + mockLog := ldlogtest.NewMockLog() + defer mockLog.DumpIfTestFailed(t) + + clientCh := make(chan *testclient.FakeLDClient, 10) + readyCh := make(chan EnvContext, 1) + env := makeBasicEnv(t, envConfig, testclient.FakeLDClientFactoryWithChannel(true, clientCh), mockLog.Loggers, readyCh) + defer env.Close() + + require.Equal(t, env, requireEnvReady(t, readyCh)) + originalClient := requireClientReady(t, clientCh) + require.Eventually(t, func() bool { return env.GetClient() == originalClient }, time.Second, 10*time.Millisecond) + require.NoError(t, env.GetStore().Init(st.AllData)) + + // First re-anchor: original → key2 (new key; key2 client built). Keep the original in grace. + now := time.Unix(2000, 0) + reanchorViaReconcile(t, env, reanchorSyncTestKey2, envConfig.SDKKey, "", envConfig.MobileKey, envConfig.EnvID, now) + key2Client := requireClientReady(t, clientCh) + require.Same(t, key2Client, env.GetClient()) + + // The original anchor's client is still alive in its grace period. + envImpl := env.(*envContextImpl) + envImpl.mu.RLock() + originalStillPresent := envImpl.clients[envConfig.SDKKey] == originalClient + envImpl.mu.RUnlock() + require.True(t, originalStillPresent, "original client retained for reuse") + + // Second re-anchor: key2 → original. The original's client exists, so this is the reuse path: no + // Build, the existing client is reused, the anchor flips, and ReplaceCredential runs. + reanchorViaReconcile(t, env, envConfig.SDKKey, reanchorSyncTestKey2, "", envConfig.MobileKey, envConfig.EnvID, now) + + // No new client was created — clientCh must be empty (every prior client was drained). + select { + case c := <-clientCh: + t.Fatalf("re-anchoring to a key with an existing client must not build a new one, but one was created: %v", c.Key) + case <-time.After(100 * time.Millisecond): + } + + assert.Same(t, originalClient, env.GetClient(), "reuses the existing client for the re-anchored key") + assert.Equal(t, envConfig.SDKKey, envImpl.keyRotator.AnchorKey(), "anchor flipped back to the original key") +} + +// TestReanchorSync_CredentialExpiryDuringReanchorIsSerialized exercises the concurrency gap closed +// by serializing the cleanup ticker against reconcileMu: a credential expiry firing (via the ticker +// path, triggerCredentialChanges) while a synchronous re-anchor is mid-build must not run its +// StepTime + add/remove pass concurrently with the in-flight re-anchor. Both paths drain the same +// StepTime queue and can close clients, so they must be serialized the way concurrent reconciles are. +// +// The test wedges a re-anchor open by blocking the new anchor's client build, then fires the ticker +// from another goroutine and asserts (a) the ticker is blocked while the re-anchor holds reconcileMu, +// (b) it completes once the re-anchor releases it, and (c) the final state is consistent — the new +// anchor committed, the expiring non-anchor key dropped, the demoted old anchor retained in grace. +func TestReanchorSync_CredentialExpiryDuringReanchorIsSerialized(t *testing.T) { + envConfig := st.EnvMain.Config + + mockLog := ldlogtest.NewMockLog() + defer mockLog.DumpIfTestFailed(t) + + now := time.Unix(2000, 0) + expiringExpiry := now.Add(30 * time.Minute) // the non-anchor key the ticker will drop + graceExpiry := now.Add(2 * time.Hour) // the demoted old anchor stays alive in its grace period + tickerTime := now.Add(time.Hour) // between the two expiries: drops only the expiring key + + // The new anchor's client build blocks until releaseBuild is closed, holding the re-anchor (and + // thus reconcileMu) open so the ticker has a window to (try to) run concurrently. + buildEntered := make(chan struct{}) + releaseBuild := make(chan struct{}) + + clientCh := make(chan *testclient.FakeLDClient, 10) + healthy := testclient.FakeLDClientFactoryWithChannel(true, clientCh) + factory := func(sdkKey config.SDKKey, cfg ld.Config, timeout time.Duration) (sdks.LDClientContext, error) { + if sdkKey == reanchorSyncTestKey2 { + close(buildEntered) + <-releaseBuild + } + return healthy(sdkKey, cfg, timeout) + } + + readyCh := make(chan EnvContext, 1) + env := makeBasicEnv(t, envConfig, factory, mockLog.Loggers, readyCh) + defer env.Close() + + require.Equal(t, env, requireEnvReady(t, readyCh)) + originalClient := requireClientReady(t, clientCh) + require.Eventually(t, func() bool { return env.GetClient() == originalClient }, time.Second, 10*time.Millisecond) + require.NoError(t, env.GetStore().Init(st.AllData)) + + envImpl := env.(*envContextImpl) + + // Accept a second SDK key as a non-anchor server key carrying a future expiry. addCredential's + // anchor gate won't build a client for it, so it has no client of its own. + initialSet, err := credential.NewAcceptedSetBuilder(). + WithAnchor(credential.SDKKeyParams{Value: envConfig.SDKKey}). + WithSDKKey(credential.SDKKeyParams{Value: reanchorSyncExpiringKey, Expiry: &expiringExpiry}). + WithPrimaryMobileKey(credential.MobileKeyParams{Value: envConfig.MobileKey}). + WithEnvironmentID(envConfig.EnvID). + Build() + require.NoError(t, err) + envImpl.reconcileCredentials(initialSet, now) + require.Contains(t, env.GetCredentials(), credential.SDKCredential(reanchorSyncExpiringKey)) + + // Re-anchor onto a brand-new key while keeping both the demoted old anchor and the expiring key + // accepted. The build will block, holding reconcileMu. + reanchorSet, err := credential.NewAcceptedSetBuilder(). + WithAnchor(credential.SDKKeyParams{Value: reanchorSyncTestKey2}). + WithSDKKey(credential.SDKKeyParams{Value: envConfig.SDKKey, Expiry: &graceExpiry}). + WithSDKKey(credential.SDKKeyParams{Value: reanchorSyncExpiringKey, Expiry: &expiringExpiry}). + WithPrimaryMobileKey(credential.MobileKeyParams{Value: envConfig.MobileKey}). + WithEnvironmentID(envConfig.EnvID). + Build() + require.NoError(t, err) + + reconcileDone := make(chan struct{}) + go func() { + defer close(reconcileDone) + envImpl.reconcileCredentials(reanchorSet, now) + }() + + // Wait until the re-anchor is wedged open inside the client build (holding reconcileMu). + <-buildEntered + + // Fire the cleanup ticker's work while the re-anchor holds reconcileMu. With the ticker serialized + // against reconcileMu, this must block until the re-anchor releases it. + tickerDone := make(chan struct{}) + go func() { + defer close(tickerDone) + envImpl.triggerCredentialChanges(tickerTime) + }() + + select { + case <-tickerDone: + t.Fatal("the cleanup ticker ran its StepTime + add/remove pass concurrently with an in-flight " + + "re-anchor; reconcileMu did not serialize the ticker against reconcileCredentials") + case <-time.After(100 * time.Millisecond): + // Expected: the ticker is blocked on reconcileMu. + } + + // Release the build; the re-anchor finishes and drops reconcileMu, unblocking the ticker. + close(releaseBuild) + <-reconcileDone + select { + case <-tickerDone: + case <-time.After(time.Second): + t.Fatal("the cleanup ticker did not complete after the re-anchor released reconcileMu") + } + + // Final state is consistent: the new anchor committed and serves its client, the expiring + // non-anchor key was dropped by the ticker, and the demoted old anchor is retained in grace. + newClient := requireClientReady(t, clientCh) + assert.Same(t, newClient, env.GetClient(), "the new anchor's client is current after the re-anchor") + assert.Equal(t, reanchorSyncTestKey2, envImpl.keyRotator.AnchorKey()) + assert.NotContains(t, env.GetCredentials(), credential.SDKCredential(reanchorSyncExpiringKey), + "the expiring non-anchor key was dropped by the ticker") + + envImpl.mu.RLock() + _, oldStillPresent := envImpl.clients[envConfig.SDKKey] + envImpl.mu.RUnlock() + assert.True(t, oldStillPresent, "demoted old anchor's client retained during its grace period") +} + +// TestReanchorSync_MobilePrimaryRepoint_AlreadyAcceptedKey covers the primary-mobile gate's gap: +// when the primary mobile key switches to a key that was ALREADY in the accepted set, that key is +// not in StepTime's additions, so addCredential's gate never fires for it. reconcileCredentials +// must therefore call eventDispatcher.ReplaceCredential synchronously via MobilePrimaryRepoint. +// +// This test asserts the rotator-level signal: after reconciling to make an already-accepted mobile +// key the primary, the env's primary mobile key reflects the change. (The ReplaceCredential side +// effect on the dispatcher is exercised by the events package's own dispatcher tests; here we +// confirm the env drives the primary-mobile transition correctly without spawning a client.) +func TestReanchorSync_MobilePrimaryRepoint_AlreadyAcceptedKey(t *testing.T) { + envConfig := st.EnvMain.Config + mob1 := config.MobileKey("mob-primary-1") + mob2 := config.MobileKey("mob-primary-2") + + mockLog := ldlogtest.NewMockLog() + defer mockLog.DumpIfTestFailed(t) + + clientCh := make(chan *testclient.FakeLDClient, 10) + readyCh := make(chan EnvContext, 1) + env := makeBasicEnv(t, envConfig, testclient.FakeLDClientFactoryWithChannel(true, clientCh), mockLog.Loggers, readyCh) + defer env.Close() + + require.Equal(t, env, requireEnvReady(t, readyCh)) + _ = requireClientReady(t, clientCh) + + now := time.Unix(2000, 0) + envImpl := env.(*envContextImpl) + + // Accept both mobile keys, mob1 primary. mob2 is accepted but not primary. + set1, err := credential.NewAcceptedSetBuilder(). + WithAnchor(credential.SDKKeyParams{Value: envConfig.SDKKey}). + WithPrimaryMobileKey(credential.MobileKeyParams{Value: mob1}). + WithMobileKey(credential.MobileKeyParams{Value: mob2}). + WithEnvironmentID(envConfig.EnvID). + Build() + require.NoError(t, err) + envImpl.reconcileCredentials(set1, now) + require.Equal(t, mob1, envImpl.keyRotator.MobileKey()) + require.Contains(t, env.GetCredentials(), credential.SDKCredential(mob2)) + + // Now make mob2 (already accepted) the primary mobile key. No SDK anchor change here. + set2, err := credential.NewAcceptedSetBuilder(). + WithAnchor(credential.SDKKeyParams{Value: envConfig.SDKKey}). + WithPrimaryMobileKey(credential.MobileKeyParams{Value: mob2}). + WithMobileKey(credential.MobileKeyParams{Value: mob1}). + WithEnvironmentID(envConfig.EnvID). + Build() + require.NoError(t, err) + envImpl.reconcileCredentials(set2, now) + + assert.Equal(t, mob2, envImpl.keyRotator.MobileKey(), "primary mobile key repointed to the already-accepted key") + // No SDK client should have been spawned by a mobile-only change. + select { + case c := <-clientCh: + t.Fatalf("a mobile-primary repoint must not build an SDK client, got: %v", c.Key) + case <-time.After(100 * time.Millisecond): + } +} + +// TestReanchorSync_PreviouslyAcceptedNonAnchorPromotedToAnchor covers the case the two-signal split in +// reanchor exists for: a server SDK key accepted as a NON-anchor has its credential mappings registered +// by addCredential but no client (only the anchor gets one). Promoting it to anchor must NOT re-register +// its mappings (NewAnchorPreviouslyAccepted == true) but MUST build a client (existingClient == nil). +func TestReanchorSync_PreviouslyAcceptedNonAnchorPromotedToAnchor(t *testing.T) { + envConfig := st.EnvMain.Config + nonAnchorKey := config.SDKKey("reanchor-sync-nonanchor") + + mockLog := ldlogtest.NewMockLog() + defer mockLog.DumpIfTestFailed(t) + + clientCh := make(chan *testclient.FakeLDClient, 10) + readyCh := make(chan EnvContext, 1) + env := makeBasicEnv(t, envConfig, testclient.FakeLDClientFactoryWithChannel(true, clientCh), mockLog.Loggers, readyCh) + defer env.Close() + + require.Equal(t, env, requireEnvReady(t, readyCh)) + originalClient := requireClientReady(t, clientCh) + require.Eventually(t, func() bool { return env.GetClient() == originalClient }, time.Second, 10*time.Millisecond) + require.NoError(t, env.GetStore().Init(st.AllData)) + + now := time.Unix(2000, 0) + envImpl := env.(*envContextImpl) + + // Accept a second server SDK key as a NON-anchor (permanent). Its mappings are registered but no + // client is built — only the anchor owns an upstream client. + set1, err := credential.NewAcceptedSetBuilder(). + WithAnchor(credential.SDKKeyParams{Value: envConfig.SDKKey}). + WithSDKKey(credential.SDKKeyParams{Value: nonAnchorKey}). + WithPrimaryMobileKey(credential.MobileKeyParams{Value: envConfig.MobileKey}). + WithEnvironmentID(envConfig.EnvID). + Build() + require.NoError(t, err) + envImpl.reconcileCredentials(set1, now) + require.Contains(t, env.GetCredentials(), credential.SDKCredential(nonAnchorKey)) + select { + case c := <-clientCh: + t.Fatalf("a non-anchor server key must not build a client, got: %v", c.Key) + case <-time.After(100 * time.Millisecond): + } + envImpl.mu.RLock() + _, nonAnchorHasClient := envImpl.clients[nonAnchorKey] + envImpl.mu.RUnlock() + require.False(t, nonAnchorHasClient, "non-anchor key has no client of its own") + + // Promote the previously-accepted non-anchor key to anchor: mappings already exist (skip + // re-registration), but there is no client, so the re-anchor must build one. + reanchorViaReconcile(t, env, nonAnchorKey, envConfig.SDKKey, "", envConfig.MobileKey, envConfig.EnvID, now) + + newClient := requireClientReady(t, clientCh) + assert.NotSame(t, originalClient, newClient, "promoting a client-less accepted key builds a fresh client") + assert.Same(t, newClient, env.GetClient()) + assert.Equal(t, nonAnchorKey, envImpl.keyRotator.AnchorKey(), "anchor committed to the promoted key") + assert.NoError(t, env.GetInitError()) +} + +// TestReanchorSync_Offline_CommitsWithoutBuildingClient covers the offline re-anchor branch: when the +// env is offline, re-anchoring to a new key must commit the anchor WITHOUT building a new upstream +// client. (The initial anchor client is still created at startup; offline only skips the re-anchor +// build.) +func TestReanchorSync_Offline_CommitsWithoutBuildingClient(t *testing.T) { + envConfig := st.EnvMain.Config + envConfig.Offline = true + + mockLog := ldlogtest.NewMockLog() + defer mockLog.DumpIfTestFailed(t) + + clientCh := make(chan *testclient.FakeLDClient, 10) + readyCh := make(chan EnvContext, 1) + env := makeBasicEnv(t, envConfig, testclient.FakeLDClientFactoryWithChannel(true, clientCh), mockLog.Loggers, readyCh) + defer env.Close() + + require.Equal(t, env, requireEnvReady(t, readyCh)) + _ = requireClientReady(t, clientCh) // drain the initial anchor client + + now := time.Unix(2000, 0) + reanchorViaReconcile(t, env, reanchorSyncTestKey2, envConfig.SDKKey, "", envConfig.MobileKey, envConfig.EnvID, now) + + // The anchor commits, but the offline branch builds no new client. + assert.Equal(t, reanchorSyncTestKey2, env.(*envContextImpl).keyRotator.AnchorKey(), "offline re-anchor commits the anchor") + select { + case c := <-clientCh: + t.Fatalf("an offline re-anchor must not build a new SDK client, got: %v", c.Key) + case <-time.After(100 * time.Millisecond): + } + assert.NoError(t, env.GetInitError()) +} + +// TestReanchorSync_RollbackWithImmediateRevocationKeepsOldAnchorServing covers the edge where a +// reconcile both moves the anchor to a new key AND immediately revokes the current anchor (no grace +// expiry), and the new anchor's client fails to build. The re-anchor rolls back, backing out just the +// anchor change: the previous anchor is kept serving and re-admitted to the accepted set, and the +// failed new key is dropped. +func TestReanchorSync_RollbackWithImmediateRevocationKeepsOldAnchorServing(t *testing.T) { + envConfig := st.EnvMain.Config + fakeErr := errors.New("re-anchor: new client init refused") + + mockLog := ldlogtest.NewMockLog() + defer mockLog.DumpIfTestFailed(t) + + clientCh := make(chan *testclient.FakeLDClient, 10) + healthyFactory := testclient.FakeLDClientFactoryWithChannel(true, clientCh) + factory := func(sdkKey config.SDKKey, cfg ld.Config, timeout time.Duration) (sdks.LDClientContext, error) { + if sdkKey == reanchorSyncTestKey2 { + return nil, fakeErr + } + return healthyFactory(sdkKey, cfg, timeout) + } + + readyCh := make(chan EnvContext, 1) + env := makeBasicEnv(t, envConfig, factory, mockLog.Loggers, readyCh) + defer env.Close() + + require.Equal(t, env, requireEnvReady(t, readyCh)) + originalClient := requireClientReady(t, clientCh) + require.Eventually(t, func() bool { return env.GetClient() == originalClient }, time.Second, 10*time.Millisecond) + require.NoError(t, env.GetStore().Init(st.AllData)) + + envImpl := env.(*envContextImpl) + now := time.Unix(2000, 0) + + // Re-anchor to a brand-new key that fails to init, while immediately revoking the current anchor — + // it is omitted from the payload entirely (no grace expiry). + set, err := credential.NewAcceptedSetBuilder(). + WithAnchor(credential.SDKKeyParams{Value: reanchorSyncTestKey2}). + WithPrimaryMobileKey(credential.MobileKeyParams{Value: envConfig.MobileKey}). + WithEnvironmentID(envConfig.EnvID). + Build() + require.NoError(t, err) + envImpl.reconcileCredentials(set, now) + + // The rollback backed out the anchor change: the previous anchor still serves and is still accepted, + // the failed new key is gone, and the env is not marked failed. + assert.NoError(t, env.GetInitError(), "a rolled-back re-anchor must not mark the env failed") + assert.Same(t, originalClient, env.GetClient(), "the previous anchor's client keeps serving despite the revocation") + assert.Equal(t, envConfig.SDKKey, envImpl.keyRotator.AnchorKey(), "anchor pointer stayed on the previous key") + assert.Contains(t, env.GetCredentials(), credential.SDKCredential(envConfig.SDKKey), "previous anchor re-admitted to the accepted set") + assert.NotContains(t, env.GetCredentials(), credential.SDKCredential(reanchorSyncTestKey2), "failed new anchor dropped from the accepted set") + + envImpl.mu.RLock() + _, newHasClient := envImpl.clients[reanchorSyncTestKey2] + _, oldHasClient := envImpl.clients[envConfig.SDKKey] + envImpl.mu.RUnlock() + assert.False(t, newHasClient, "no client for the failed new anchor") + assert.True(t, oldHasClient, "previous anchor's client retained") +} + +// TestReanchorSync_PreviouslyAcceptedAnchorPromotionFailureKeepsItsMappings covers a failed promotion +// of an already-accepted non-anchor key: the rollback must NOT tear down that key's credential +// mappings (they predate this reconcile), since RevertAnchorChange keeps it accepted. It should revert +// cleanly to the non-anchor key it already was. +func TestReanchorSync_PreviouslyAcceptedAnchorPromotionFailureKeepsItsMappings(t *testing.T) { + envConfig := st.EnvMain.Config + nonAnchorKey := config.SDKKey("reanchor-sync-nonanchor-promote-fail") + + mockLog := ldlogtest.NewMockLog() + defer mockLog.DumpIfTestFailed(t) + + clientCh := make(chan *testclient.FakeLDClient, 10) + healthyFactory := testclient.FakeLDClientFactoryWithChannel(true, clientCh) + // Fail the build only when nonAnchorKey is promoted to anchor. + factory := func(sdkKey config.SDKKey, cfg ld.Config, timeout time.Duration) (sdks.LDClientContext, error) { + if sdkKey == nonAnchorKey { + return nil, errors.New("promotion build refused") + } + return healthyFactory(sdkKey, cfg, timeout) + } + + mapper := &recordingConnectionMapper{} + readyCh := make(chan EnvContext, 1) + env := makeBasicEnvWithMapper(t, envConfig, factory, mockLog.Loggers, readyCh, mapper) + defer env.Close() + + require.Equal(t, env, requireEnvReady(t, readyCh)) + originalClient := requireClientReady(t, clientCh) + require.Eventually(t, func() bool { return env.GetClient() == originalClient }, time.Second, 10*time.Millisecond) + require.NoError(t, env.GetStore().Init(st.AllData)) + + envImpl := env.(*envContextImpl) + now := time.Unix(2000, 0) + + // Accept nonAnchorKey as a non-anchor server key: mappings are registered (and it becomes mapped), + // but no client is built for it. + set1, err := credential.NewAcceptedSetBuilder(). + WithAnchor(credential.SDKKeyParams{Value: envConfig.SDKKey}). + WithSDKKey(credential.SDKKeyParams{Value: nonAnchorKey}). + WithPrimaryMobileKey(credential.MobileKeyParams{Value: envConfig.MobileKey}). + WithEnvironmentID(envConfig.EnvID). + Build() + require.NoError(t, err) + envImpl.reconcileCredentials(set1, now) + require.True(t, mapper.isMapped(envConfig.FilterKey, nonAnchorKey), "non-anchor key is mapped once accepted") + + // Promote nonAnchorKey to anchor; its client build fails, so the re-anchor rolls back. + reanchorViaReconcile(t, env, nonAnchorKey, envConfig.SDKKey, "", envConfig.MobileKey, envConfig.EnvID, now) + + // The failed promotion rolled back: nonAnchorKey stays accepted AND keeps its connection mapping — + // the rollback must not strip mappings that existed before this reconcile. + assert.Contains(t, env.GetCredentials(), credential.SDKCredential(nonAnchorKey), "previously-accepted key stays accepted after a failed promotion") + assert.True(t, mapper.isMapped(envConfig.FilterKey, nonAnchorKey), "its connection mapping must survive the rollback") + // The previous anchor still serves. + assert.Equal(t, envConfig.SDKKey, envImpl.keyRotator.AnchorKey()) + assert.Same(t, originalClient, env.GetClient()) + assert.NoError(t, env.GetInitError()) +} diff --git a/internal/relayenv/env_context_reanchor_test.go b/internal/relayenv/env_context_reanchor_test.go index d967cf96..22543955 100644 --- a/internal/relayenv/env_context_reanchor_test.go +++ b/internal/relayenv/env_context_reanchor_test.go @@ -1,19 +1,16 @@ package relayenv -// T0 — Re-anchoring PoC (SDK-2453 / SDK-2530). +// Re-anchoring proof-of-concept tests. // -// These tests validate the upstream SDK-client swap mechanism that T2.c will implement. Each test -// answers one of the seven hypotheses in .agent-docs/concurrent-keys/phase1-design.md §7. They are -// written as durable, executable probes of today's primitives so they survive into T2 as regression -// tests and as the executable spec for the re-anchor implementation. -// -// A written summary of the findings lives in -// .agent-docs/concurrent-keys/phase1-T0-reanchor-poc-findings.md. +// These tests validate the upstream SDK-client swap mechanism that the re-anchor implementation builds +// on. Each test answers one of the seven hypotheses about how re-anchoring should behave. They are +// written as durable, executable probes of today's primitives so they survive as regression tests and +// as the executable spec for the re-anchor implementation. // // Terminology: "re-anchor" = swapping the single upstream SDK client when sdkKey.value changes. // Today there is no dedicated re-anchor method; the closest existing path is ReconcileCredentials with // an expiring (grace-period) key (which rotates the primary SDK key and stands up a new client), so -// several tests drive that path and observe where it falls short of the §7 requirements. +// several tests drive that path and observe where it falls short of the requirements. import ( "errors" @@ -119,12 +116,12 @@ func TestReanchorPoC_H1_SharedStoreAdapterRebuildSemantics(t *testing.T) { featureKind := ldstoreimpl.Features() flagKey := st.Flag1ServerSide.Flag.Key - // The design (§7) assumes "two SDK clients pointed at the same env can feed the same store as a - // side-effect." This sub-test shows that assumption is FALSE for the default in-memory store: each - // client init calls storeAdapter.Build, which constructs a brand-new wrapper around a brand-new - // underlying store and atomically swaps it in. No corruption occurs, but the new client starts from - // an empty store. - t.Run("in-memory factory builds a fresh empty store on each client init", func(t *testing.T) { + // Original PoC finding: each storeAdapter.Build call constructed a fresh wrapper around a fresh + // underlying store, so the new anchor's client would start empty. The store-handover change + // (SSERelayDataStoreAdapter.Build reusing its existing wrapper, with refcounted Close) inverts + // this: the second client receives the SAME wrapper, with its data still in place. This sub-test + // now asserts the post-fix invariant. + t.Run("in-memory factory reuses the existing store on a second client init (store handover)", func(t *testing.T) { rec := &recordingStreamUpdates{} adapter := store.NewSSERelayDataStoreAdapter(ldcomponents.InMemoryDataStore(), rec) @@ -142,20 +139,19 @@ func TestReanchorPoC_H1_SharedStoreAdapterRebuildSemantics(t *testing.T) { s2, err := adapter.Build(subsystems.BasicClientContext{}) require.NoError(t, err) - // FINDING: Build swaps in a new store instance... - assert.NotSame(t, s1, s2, "each Build creates a new store wrapper") - assert.Same(t, s2, adapter.GetStore(), "the adapter now points at the new store") + // Post-fix: Build hands the existing wrapper to the new client and the adapter still points at it. + assert.Same(t, s1, s2, "store handover: Build returns the existing wrapper") + assert.Same(t, s2, adapter.GetStore(), "the adapter still points at the shared wrapper") - // ...and that store is empty + uninitialized. So the two clients do NOT share data through the - // in-memory store; the new anchor must re-sync from scratch. - assert.False(t, adapter.GetStore().IsInitialized(), "the new in-memory store starts uninitialized") + // The wrapper stays initialized and the data survives — no empty-store window for the new anchor. + assert.True(t, adapter.GetStore().IsInitialized(), "the shared store remains initialized") got2, err := adapter.GetStore().Get(featureKind, flagKey) require.NoError(t, err) - assert.Nil(t, got2.Item, "the new in-memory store starts empty") + assert.NotNil(t, got2.Item, "data persists across handover") }) // With a persistent store, the underlying data lives outside the wrapper, so the swap preserves it. - // This is the configuration in which the §7 "shared store" assumption actually holds. + // This is the configuration in which the "shared store" assumption actually holds. t.Run("shared (persistent) underlying store preserves data across client init", func(t *testing.T) { underlying, err := ldcomponents.InMemoryDataStore().Build(subsystems.BasicClientContext{}) require.NoError(t, err) @@ -273,7 +269,7 @@ func TestReanchorPoC_H2_NewClientInitialSyncRebroadcastsPut(t *testing.T) { // FINDING: the new anchor's initial sync produces a second full "put" to every connected downstream // stream. From a downstream SDK's perspective this is a duplicate put. It is tolerable (SDKs apply - // puts idempotently) but T2.c must expect it; it is not a corruption. + // puts idempotently) but the re-anchor implementation must expect it; it is not a corruption. assert.Equal(t, 2, rec.allDataCount(), "the new anchor's initial sync re-broadcasts a full put") } @@ -367,8 +363,8 @@ func TestReanchorPoC_H3_BigSegmentSyncIsNotReWiredOnReAnchor(t *testing.T) { // FINDING: big-segment sync is wired to the SDK key at construction and is NOT re-wired by today's // swap path -- the synchronizer is neither recreated nor told about the new key (the // BigSegmentSynchronizer interface has no credential-replacement method). After re-anchor it keeps - // polling/streaming on the OLD anchor key. T2.d must add a re-wire path (a ReplaceCredential-style - // method) or recreate the synchronizer on each re-anchor. + // polling/streaming on the OLD anchor key. The big-segment re-wire (follow-up work) must add a + // re-wire path (a ReplaceCredential-style method) or recreate the synchronizer on each re-anchor. count, sdkKey = capturing.snapshot() assert.Equal(t, 1, count, "synchronizer was not recreated on re-anchor") assert.Equal(t, envConfig.SDKKey, sdkKey, "synchronizer still references the old anchor key") @@ -417,7 +413,13 @@ func TestReanchorPoC_H4_HTTPConfigIsKeyIndependentExceptAuthHeader(t *testing.T) // Hypothesis 5: Order of operations / the in-memory store window. // ----------------------------------------------------------------------------------------------- -func TestReanchorPoC_H5_InMemoryStoreIsWipedByReAnchor(t *testing.T) { +// TestReanchorPoC_H5_StoreSurvivesReAnchor was originally a proof-of-concept test asserting the +// *broken* pre-handover behavior: that re-anchor caused the env's in-memory store to be replaced with +// a fresh, empty one. Once SSERelayDataStoreAdapter.Build was changed to hand over the existing wrapper +// to the new client (refcounted Close), that breakage is gone. This now asserts the post-fix +// invariant — the data store instance survives the re-anchor and keeps its data — and is preserved +// as the executable proof that handover holds end-to-end through env_context. +func TestReanchorPoC_H5_StoreSurvivesReAnchor(t *testing.T) { featureKind := ldstoreimpl.Features() flagKey := st.Flag1ServerSide.Flag.Key envConfig := st.EnvMain.Config @@ -450,31 +452,27 @@ func TestReanchorPoC_H5_InMemoryStoreIsWipedByReAnchor(t *testing.T) { require.Eventually(t, func() bool { return env.GetClient() == client2 }, time.Second, 10*time.Millisecond, "GetClient should return the new anchor's client once it is registered") - // FINDING: starting the new client replaced the data store with a fresh, empty, uninitialized one. - // This happens regardless of operation order, because building the new client is what rebuilds the - // store. So "start-new -> swap-pointer -> close-old" alone is NOT sufficient with an in-memory store: - // there is a window in which evaluations see an empty store until the new anchor finishes its initial - // sync. T2.c must either (a) keep the old store/anchor authoritative until the new client reports - // Initialized()==true, (b) require a persistent store for graceful re-anchor, or (c) decouple the - // data store lifecycle from the client lifecycle so a new client does not rebuild it. + // Post-fix: the adapter handed the existing wrapper to the new client, so the env's store is the + // same instance, still initialized, and the data is intact. There is no empty-store window for the + // new anchor. newStore := env.GetStore() - assert.NotSame(t, oldStore, newStore, "the data store instance was replaced by the new client") - assert.False(t, newStore.IsInitialized(), "the new store is uninitialized until the new anchor re-syncs") + assert.Same(t, oldStore, newStore, "the data store instance survives re-anchor (store handover)") + assert.True(t, newStore.IsInitialized(), "the store stays initialized across re-anchor") got2, err := newStore.Get(featureKind, flagKey) require.NoError(t, err) - assert.Nil(t, got2.Item, "data is absent in the new store until the new anchor re-syncs") + assert.NotNil(t, got2.Item, "data is preserved across re-anchor") } // TestReanchorPoC_H5_StoreHandoverAvoidsEmptyWindow validates the reviewer suggestion that, because // relay owns the data store implementation (it hands the SDK a single storeAdapter), the re-anchor can // hand the existing store over to the new client instead of letting it build a fresh one. Modeled here // by a DataStoreFactory that returns the same underlying store on every Build; the production change -// (T2.c/T2.d) is to make SSERelayDataStoreAdapter reuse its store across the swap. With handover the -// new anchor's client sees the populated, initialized store immediately -- no empty-store window -// (contrast TestReanchorPoC_H5_InMemoryStoreIsWipedByReAnchor). +// is to make SSERelayDataStoreAdapter reuse its store across the swap. With handover the new anchor's +// client sees the populated, initialized store immediately -- no empty-store window (contrast +// TestReanchorPoC_H5_InMemoryStoreIsWipedByReAnchor). // -// CAVEAT for the implementation (not reproducible with the fake client, so documented here and in the -// findings): streamUpdatesStoreWrapper.Close() closes the underlying store. If the new client wraps the +// CAVEAT for the implementation (not reproducible with the fake client, so documented here): +// streamUpdatesStoreWrapper.Close() closes the underlying store. If the new client wraps the // SAME underlying store, closing the retiring client must NOT close it -- the store's lifecycle has to // be owned by the adapter, not by the client being retired. func TestReanchorPoC_H5_StoreHandoverAvoidsEmptyWindow(t *testing.T) { @@ -530,7 +528,14 @@ func TestReanchorPoC_H5_StoreHandoverAvoidsEmptyWindow(t *testing.T) { // Hypothesis 6: Behavior during the swap window (requests arriving mid-swap). // ----------------------------------------------------------------------------------------------- -func TestReanchorPoC_H6_AnchorPointerFlipsBeforeNewClientIsRegistered(t *testing.T) { +// TestReanchorPoC_H6_AnchorHoldsUntilNewClientReady is the inversion of the original H6 finding. +// The originally-observed broken behavior: reconcileCredentials flipped the rotator's anchor +// synchronously and then built the new client asynchronously, opening a window where GetClient() +// (== clients[AnchorKey()]) returned nil for the not-yet-registered new key. The fix builds the new +// client synchronously and only commits the anchor once it reports Initialized, so no such nil window +// exists: while the new client is still building, the anchor stays on the old key and GetClient() keeps +// returning the old, still-serving client. +func TestReanchorPoC_H6_AnchorHoldsUntilNewClientReady(t *testing.T) { envConfig := st.EnvMain.Config mockLog := ldlogtest.NewMockLog() @@ -540,10 +545,13 @@ func TestReanchorPoC_H6_AnchorPointerFlipsBeforeNewClientIsRegistered(t *testing inner := testclient.FakeLDClientFactoryWithChannel(true, clientCh) // gate blocks construction of the NEW anchor's client so we can observe the swap window - // deterministically (no sleeps / no racing). + // deterministically. entered signals that the synchronous build has reached the factory (and is now + // blocked on gate) — at which point the re-anchor is mid-flight but has not yet committed. gate := make(chan struct{}) + entered := make(chan struct{}, 1) gatedFactory := func(sdkKey config.SDKKey, cfg ld.Config, timeout time.Duration) (sdks.LDClientContext, error) { if sdkKey == reanchorTestKey2 { + entered <- struct{}{} <-gate } return inner(sdkKey, cfg, timeout) @@ -557,30 +565,46 @@ func TestReanchorPoC_H6_AnchorPointerFlipsBeforeNewClientIsRegistered(t *testing client1 := requireClientReady(t, clientCh) require.Eventually(t, func() bool { return env.GetClient() == client1 }, time.Second, 10*time.Millisecond) - // Re-anchor. reconcileCredentials flips the rotator's primary SDK key synchronously, then starts the - // new client on a background goroutine (which blocks in the factory on `gate`). + // The re-anchor is synchronous, so the reconcile blocks in the gated factory until we release it. + // Drive it on a background goroutine and build the set up front (keeping require off that goroutine). start := time.Unix(1000, 0) - reanchor(t, env, reanchorTestKey2, envConfig.SDKKey, start) - - // FINDING: there is a window where the anchor pointer already names the new key but no client exists - // for it yet, so GetClient() returns nil. GetClient() == clients[rotator.AnchorKey()], and the rotator's - // anchor flipped to the new key before startSDKClient registered the client. A request arriving in - // this window gets a nil client. T2.c must not advance the anchor pointer until the new client is - // registered (and ideally Initialized()). - assert.Nil(t, env.GetClient(), "GetClient() is nil during the swap window") + set, err := credential.NewAcceptedSetBuilder(). + WithAnchor(credential.SDKKeyParams{Value: reanchorTestKey2}). + WithSDKKey(credential.SDKKeyParams{Value: envConfig.SDKKey, Expiry: util.PtrOrNil(start.Add(time.Hour))}). + Build() + require.NoError(t, err) + done := make(chan struct{}) + go func() { + defer close(done) + env.(*envContextImpl).reconcileCredentials(set, start) + }() + + // The build has reached the factory and is blocked: the re-anchor is mid-flight, pre-commit. + <-entered + envImpl := env.(*envContextImpl) + assert.Same(t, client1, env.GetClient(), "GetClient() keeps returning the old client during the build — never nil") + assert.Equal(t, envConfig.SDKKey, envImpl.keyRotator.AnchorKey(), "the anchor stays on the old key until the new client is ready") - // Release the gate; the new client registers and GetClient() recovers. + // Release the gate; the new client initializes, the anchor commits, and GetClient() advances. close(gate) + <-done client2 := requireClientReady(t, clientCh) require.Eventually(t, func() bool { return env.GetClient() == client2 }, time.Second, 10*time.Millisecond, - "GetClient() recovers once the new client is registered") + "GetClient() returns the new client once it is registered and the anchor commits") + assert.Equal(t, reanchorTestKey2, envImpl.keyRotator.AnchorKey()) } // ----------------------------------------------------------------------------------------------- // Hypothesis 7: Failure modes — new client init fails. // ----------------------------------------------------------------------------------------------- -func TestReanchorPoC_H7_FailedNewClientLeavesEnvWithoutAnchorClient(t *testing.T) { +// TestReanchorPoC_H7_FailedNewClientRollsBackToOldAnchor is the inversion of the original H7 finding. +// The originally-observed broken behavior: a failed new-client init left the anchor already flipped to +// a key with no client, so GetClient() returned nil and the environment was broken even though the old +// client was still alive. The fix builds and validates the new client before committing: on init +// failure it does NOT commit the anchor, surfaces the error, and leaves the old anchor authoritative +// (its client keeps serving). This is the all-or-nothing atomicity requirement applied to re-anchor. +func TestReanchorPoC_H7_FailedNewClientRollsBackToOldAnchor(t *testing.T) { envConfig := st.EnvMain.Config fakeErr := errors.New("new anchor client failed to initialize") @@ -606,25 +630,25 @@ func TestReanchorPoC_H7_FailedNewClientLeavesEnvWithoutAnchorClient(t *testing.T client1 := requireClientReady(t, clientCh) require.Eventually(t, func() bool { return env.GetClient() == client1 }, time.Second, 10*time.Millisecond) - // Re-anchor onto a key whose client init fails (old key kept valid for a grace hour). + // Re-anchor onto a key whose client init fails (old key kept valid for a grace hour). The re-anchor + // is synchronous, so the init failure and rollback are complete by the time reanchor returns. start := time.Unix(1000, 0) reanchor(t, env, reanchorTestKey2, envConfig.SDKKey, start) - require.Eventually(t, func() bool { return env.GetInitError() != nil }, time.Second, 10*time.Millisecond, - "the failed new-client init should surface as an init error") - assert.Equal(t, fakeErr, env.GetInitError()) - - // FINDING: the rotator already flipped the anchor to the new key, but no client exists for it, so - // GetClient() returns nil -- even though the OLD anchor's client is still alive and valid during its - // grace period. A failed re-anchor breaks the environment with today's code. This is exactly the §8 - // atomicity requirement: T2.c must validate that the new client initializes BEFORE swapping the - // anchor pointer, and roll back to the old anchor on failure (preserving the previous accepted set). - assert.Nil(t, env.GetClient(), "GetClient() is nil after a failed re-anchor") - + // Post-fix: the re-anchor rolled back. The env stays healthy on the old anchor, so GetInitError + // stays nil — setting it would 401 a still-serving env at the request middleware. The failure + // surfaces via a structured Error log instead. The anchor pointer stayed on the old key, whose + // client is still alive and serving, so GetClient() never returns nil and no client is installed + // for the failed new anchor. + assert.NoError(t, env.GetInitError(), "a failed re-anchor must not mark the still-serving env as failed") + mockLog.AssertMessageMatch(t, true, ldlog.Error, "Re-anchor to SDK key .* failed") + assert.Same(t, client1, env.GetClient(), "GetClient() still returns the old anchor's client after rollback") envImpl := env.(*envContextImpl) + assert.Equal(t, envConfig.SDKKey, envImpl.keyRotator.AnchorKey(), "the anchor pointer stays on the old key") envImpl.mu.RLock() _, oldStillPresent := envImpl.clients[envConfig.SDKKey] + _, newInstalled := envImpl.clients[reanchorTestKey2] envImpl.mu.RUnlock() - assert.True(t, oldStillPresent, - "the old anchor's client is still alive -- the data path could have been preserved by rolling back") + assert.True(t, oldStillPresent, "the old anchor's client is preserved") + assert.False(t, newInstalled, "no client is installed for the failed new anchor") } diff --git a/internal/relayenv/store_handover_realclient_test.go b/internal/relayenv/store_handover_realclient_test.go new file mode 100644 index 00000000..016b9ec7 --- /dev/null +++ b/internal/relayenv/store_handover_realclient_test.go @@ -0,0 +1,204 @@ +package relayenv + +// Spike verifying the real ld.LDClient's Close() behavior against the SSERelayDataStoreAdapter / +// streamUpdatesStoreWrapper pair, which the fake-client PoC could not exercise. This is the single +// remaining piece the fake-client PoC could not validate: +// +// > streamUpdatesStoreWrapper.Close() closes the underlying store. With handover the retiring +// > and new clients share one underlying store, so closing the retiring client must NOT close +// > it — the adapter (not the client) must own the store's lifecycle. (Not reproducible with +// > the fake client; verified here against the real client.) +// +// We answer two questions: +// Q1. Does ld.LDClient.Close() invoke Close() on its data store (the wrapper)? +// Q2. After the wrapper's Close() runs, is the underlying store still usable for reads? +// +// Q1 determines whether store handover is at risk at all. Q2 determines whether the remedy needs to +// gate the wrapper's Close (case A: in-memory Close is destructive) or whether it can stay as-is +// (case B: in-memory Close is a no-op and reads still work). + +import ( + "testing" + "time" + + "github.com/launchdarkly/ld-relay/v8/internal/store" + st "github.com/launchdarkly/ld-relay/v8/internal/sharedtest" + + ld "github.com/launchdarkly/go-server-sdk/v7" + "github.com/launchdarkly/go-server-sdk/v7/ldcomponents" + "github.com/launchdarkly/go-server-sdk/v7/subsystems" + "github.com/launchdarkly/go-server-sdk/v7/subsystems/ldstoreimpl" + "github.com/launchdarkly/go-server-sdk/v7/subsystems/ldstoretypes" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// closeObservingStore wraps an in-memory data store factory so we can directly observe when the +// underlying store's Close() is invoked. This is the unambiguous signal that ld.LDClient.Close() +// propagates through streamUpdatesStoreWrapper.Close() to the wrapped store. +type closeObservingStore struct { + inner subsystems.DataStore + closeCount int +} + +func (c *closeObservingStore) Close() error { + c.closeCount++ + return c.inner.Close() +} +func (c *closeObservingStore) Init(d []ldstoretypes.Collection) error { return c.inner.Init(d) } +func (c *closeObservingStore) Get(k ldstoretypes.DataKind, key string) (ldstoretypes.ItemDescriptor, error) { + return c.inner.Get(k, key) +} +func (c *closeObservingStore) GetAll(k ldstoretypes.DataKind) ([]ldstoretypes.KeyedItemDescriptor, error) { + return c.inner.GetAll(k) +} +func (c *closeObservingStore) Upsert(k ldstoretypes.DataKind, key string, item ldstoretypes.ItemDescriptor) (bool, error) { + return c.inner.Upsert(k, key, item) +} +func (c *closeObservingStore) IsInitialized() bool { return c.inner.IsInitialized() } +func (c *closeObservingStore) IsStatusMonitoringEnabled() bool { return c.inner.IsStatusMonitoringEnabled() } + +type closeObservingStoreFactory struct { + observed *closeObservingStore +} + +func (f *closeObservingStoreFactory) Build(ctx subsystems.ClientContext) (subsystems.DataStore, error) { + inner, err := ldcomponents.InMemoryDataStore().Build(ctx) + if err != nil { + return nil, err + } + f.observed = &closeObservingStore{inner: inner} + return f.observed, nil +} + +// realClientUsingAdapter spins up a real ld.LDClient backed by the relay store adapter. Using +// ExternalUpdatesOnly as the data source avoids any network calls (no upstream streaming connection +// is opened), so the test is hermetic. The adapter sees the real client's DataStore.Build() call +// and the real Close() lifecycle on shutdown. +func realClientUsingAdapter(t *testing.T, adapter *store.SSERelayDataStoreAdapter) *ld.LDClient { + t.Helper() + cfg := ld.Config{ + DataStore: adapter, + DataSource: ldcomponents.ExternalUpdatesOnly(), + Events: ldcomponents.NoEvents(), + } + client, err := ld.MakeCustomClient("fake-sdk-key", cfg, 5*time.Second) + require.NoError(t, err) + return client +} + +// TestRealClient_CloseInvokesWrapperClose verifies that closing a real ld.LDClient causes the +// wrapped store's Close() to fire. This is the precondition that makes store handover dangerous — +// if Close did not propagate, there would be no lifecycle hazard to design around. +func TestRealClient_CloseInvokesWrapperClose(t *testing.T) { + factory := &closeObservingStoreFactory{} + rec := &recordingStreamUpdates{} + adapter := store.NewSSERelayDataStoreAdapter(factory, rec) + + client := realClientUsingAdapter(t, adapter) + require.NotNil(t, factory.observed, "the adapter must have built the observed store") + require.Equal(t, 0, factory.observed.closeCount, "no Close yet before client.Close") + + wrapper := adapter.GetStore() + require.NoError(t, wrapper.Init(st.AllData)) + require.True(t, wrapper.IsInitialized()) + + require.NoError(t, client.Close()) + + // The headline finding: closing the real client propagates Close() to the underlying store via + // streamUpdatesStoreWrapper.Close(). If this assertion fails, the lifecycle caveat is not a real + // hazard for this combination and the store-handover fix only needs the Build() reuse, not a + // Close() lifecycle change. + assert.Equal(t, 1, factory.observed.closeCount, + "ld.LDClient.Close should propagate to the underlying data store via the wrapper") +} + +// TestRealClient_ReadsAfterCloseAreStillFunctional asks the second question: after Close runs, is +// the underlying in-memory store still usable for Get? The answer tells us whether the store-handover +// fix needs to actually prevent Close (because reads will fail after it) or whether reads coincidentally +// still work (because the in-memory store's Close is effectively a no-op for read behavior). Even +// if reads happen to work, the fix should still gate Close — relying on undocumented "Close is a +// no-op" behavior is brittle and breaks when persistent stores enter the picture. +func TestRealClient_ReadsAfterCloseAreStillFunctional(t *testing.T) { + factory := &closeObservingStoreFactory{} + rec := &recordingStreamUpdates{} + adapter := store.NewSSERelayDataStoreAdapter(factory, rec) + + client := realClientUsingAdapter(t, adapter) + wrapper := adapter.GetStore() + require.NoError(t, wrapper.Init(st.AllData)) + + featureKind := ldstoreimpl.Features() + flagKey := st.Flag1ServerSide.Flag.Key + + got, err := wrapper.Get(featureKind, flagKey) + require.NoError(t, err) + require.NotNil(t, got.Item, "sanity: data is readable before close") + + require.NoError(t, client.Close()) + + // Read after Close. The outcome here is informational, not a pass/fail design gate: + // - If the read succeeds, the in-memory store's Close is effectively a no-op for queries; the fix + // can still safely gate Close to be defensive (persistent stores may differ). + // - If the read fails, the fix MUST gate Close, since the new anchor would observe a broken store. + gotAfter, errAfter := wrapper.Get(featureKind, flagKey) + t.Logf("Get after client.Close: item=%v err=%v initialized=%v", + gotAfter.Item != nil, errAfter, wrapper.IsInitialized()) +} + +// TestRealClient_HandoverPreservesUnderlyingStore exercises the production store-handover behavior: +// when a second real ld.LDClient is built against the same SSERelayDataStoreAdapter — the re-anchor +// case — the adapter hands the existing wrapper (and underlying store) to the new client rather +// than rebuilding it. Closing the first client must not tear the underlying store down while the +// second client is still holding it; only the final Close releases it. +// +// Sequence: +// 1. Build client1; init data on its wrapper. +// 2. Build client2 — adapter returns the SAME wrapper; no second underlying store is built. +// 3. Close client1 — store stays open because client2 still holds it. +// 4. Close client2 — final release; underlying store closes exactly once. +func TestRealClient_HandoverPreservesUnderlyingStore(t *testing.T) { + factory := &countingStoreFactory{} + rec := &recordingStreamUpdates{} + adapter := store.NewSSERelayDataStoreAdapter(factory, rec) + + client1 := realClientUsingAdapter(t, adapter) + wrapper1 := adapter.GetStore() + require.NoError(t, wrapper1.Init(st.AllData)) + require.Equal(t, 1, factory.buildCount, "one Build for client1") + underlying := factory.lastObserved + + client2 := realClientUsingAdapter(t, adapter) + wrapper2 := adapter.GetStore() + require.Equal(t, 1, factory.buildCount, "client2 reuses the existing wrapper — no second Build") + assert.Same(t, wrapper1, wrapper2, "the adapter hands the same wrapper to both clients") + assert.True(t, wrapper2.IsInitialized(), "the shared store stays initialized across handover") + + require.NoError(t, client1.Close()) + assert.Equal(t, 0, underlying.closeCount, + "client1.Close must NOT tear down the underlying store while client2 still holds it") + + require.NoError(t, client2.Close()) + assert.Equal(t, 1, underlying.closeCount, + "client2.Close is the final release; the underlying store is closed exactly once") +} + +// countingStoreFactory tracks every Build call and exposes the most recently built store so the +// handover test can compare instances across builds. Each build wraps a real in-memory store in a +// closeObservingStore so close events are visible. +type countingStoreFactory struct { + buildCount int + lastObserved *closeObservingStore +} + +func (f *countingStoreFactory) Build(ctx subsystems.ClientContext) (subsystems.DataStore, error) { + inner, err := ldcomponents.InMemoryDataStore().Build(ctx) + if err != nil { + return nil, err + } + observed := &closeObservingStore{inner: inner} + f.buildCount++ + f.lastObserved = observed + return observed, nil +} diff --git a/internal/store/relay_feature_store.go b/internal/store/relay_feature_store.go index 5e3d6a7e..40ca9b25 100644 --- a/internal/store/relay_feature_store.go +++ b/internal/store/relay_feature_store.go @@ -67,15 +67,31 @@ func NewSSERelayDataStoreAdapter( } // Build is called by the SDK when the LDClient is being created. +// +// Store handover (concurrent-keys re-anchor): if the adapter already holds a wrapper from +// a prior client construction, that wrapper is returned again instead of building a fresh one. This +// hands the populated, initialized data store over to the new anchor's client during a re-anchor — +// no empty-store window, no re-sync. The wrapper refcounts its holders so the underlying store is +// only torn down by the final Close (see streamUpdatesStoreWrapper.Close). If the parked wrapper has +// already been fully closed (acquire returns false), a fresh one is built rather than resurrecting a +// wrapper whose underlying store is torn down. func (a *SSERelayDataStoreAdapter) Build( context subsystems.ClientContext, ) (subsystems.DataStore, error) { - var sw *streamUpdatesStoreWrapper + a.mu.Lock() + if existing := a.store; existing != nil { + if sw, ok := existing.(*streamUpdatesStoreWrapper); ok && sw.acquire() { + a.mu.Unlock() + return sw, nil + } + } + a.mu.Unlock() + wrappedStore, err := a.wrappedFactory.Build(context) if err != nil { return nil, err // this will cause client initialization to fail immediately } - sw = newStreamUpdatesStoreWrapper( + sw := newStreamUpdatesStoreWrapper( a.updates, wrappedStore, context.GetLogging().Loggers, @@ -93,6 +109,15 @@ type streamUpdatesStoreWrapper struct { store subsystems.DataStore updates streams.EnvStreamUpdates loggers ldlog.Loggers + + // refCount tracks how many SDK clients hold this wrapper. The first holder is implicit + // (count starts at 1 in newStreamUpdatesStoreWrapper). Each handover (Build reuse) calls + // acquire to bump the count; each client's Close decrements. The underlying store is torn + // down only when the count reaches zero, at which point closed is set so a later acquire + // refuses to hand back a wrapper whose underlying store is gone. Guarded by refMu. + refMu sync.Mutex + refCount int + closed bool } func newStreamUpdatesStoreWrapper( @@ -101,14 +126,42 @@ func newStreamUpdatesStoreWrapper( loggers ldlog.Loggers, ) *streamUpdatesStoreWrapper { relayStore := &streamUpdatesStoreWrapper{ - store: baseFeatureStore, - updates: updates, - loggers: loggers, + store: baseFeatureStore, + updates: updates, + loggers: loggers, + refCount: 1, } return relayStore } +// acquire records an additional holder of the wrapper, used by SSERelayDataStoreAdapter.Build when it +// hands this wrapper to a new client during a concurrent-keys re-anchor. It returns false if the +// wrapper has already been fully closed (refCount reached zero and the underlying store was torn +// down); the caller must then build a fresh wrapper rather than resurrect a dead one. +func (sw *streamUpdatesStoreWrapper) acquire() bool { + sw.refMu.Lock() + defer sw.refMu.Unlock() + if sw.closed { + return false + } + sw.refCount++ + return true +} + func (sw *streamUpdatesStoreWrapper) Close() error { + sw.refMu.Lock() + sw.refCount-- + final := sw.refCount <= 0 + if final { + sw.closed = true + } + sw.refMu.Unlock() + if !final { + // Re-anchor handover in progress: another client is still using this underlying store. + // The retiring client's Close must not tear it down — see SSERelayDataStoreAdapter.Build + // for the other half of this contract. + return nil + } return sw.store.Close() } diff --git a/internal/store/store_rebuild_after_close_test.go b/internal/store/store_rebuild_after_close_test.go new file mode 100644 index 00000000..75efa9c3 --- /dev/null +++ b/internal/store/store_rebuild_after_close_test.go @@ -0,0 +1,66 @@ +package store + +// Regression test for the store-handover refcount contract (concurrent-keys re-anchor). +// +// SSERelayDataStoreAdapter.Build reuses whatever wrapper is parked in a.store so a re-anchor can hand +// the populated store to the new client. The hazard: once the wrapper's refCount reaches zero and its +// underlying store is torn down, a later Build must NOT hand that same, now-closed wrapper back to a +// new client (a use-after-close for a persistent store whose Close releases its connection pool). The +// fix marks a fully-closed wrapper and has acquire refuse it, so Build rebuilds a fresh wrapper. +// +// (The re-anchor flow keeps the anchor client permanent, so refCount doesn't reach zero while the env +// is alive today; this guards the wrapper/adapter contract itself against a future caller.) + +import ( + "testing" + + "github.com/launchdarkly/ld-relay/v8/internal/sharedtest" + + "github.com/launchdarkly/go-server-sdk/v7/subsystems" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// freshStoreFactory builds a new underlying store on every Build, mirroring a real DataStore factory +// (mockStoreFactory returns a single fixed instance, which can't model a rebuild). +type freshStoreFactory struct { + built []*mockStore +} + +func (f *freshStoreFactory) Build(_ subsystems.ClientContext) (subsystems.DataStore, error) { + s := &mockStore{realStore: sharedtest.NewInMemoryStore()} + f.built = append(f.built, s) + return s, nil +} + +func TestStoreAdapterRebuildsAfterFullClose(t *testing.T) { + factory := &freshStoreFactory{} + updates := &mockEnvStreamsUpdates{} + adapter := NewSSERelayDataStoreAdapter(factory, updates) + ctx := subsystems.BasicClientContext{} + + // First client builds the wrapper: refCount = 1. + first, err := adapter.Build(ctx) + require.NoError(t, err) + sw1 := first.(*streamUpdatesStoreWrapper) + require.Equal(t, 1, sw1.refCount) + + // The sole client shuts down: refCount 1 -> 0, wrapper marked closed, underlying store torn down. + require.NoError(t, first.Close()) + require.True(t, factory.built[0].closed, "final Close tears down the underlying store") + + // A subsequent Build (e.g. a later re-anchor) must NOT resurrect the fully-closed wrapper — it + // rebuilds a fresh wrapper backed by a fresh, open store. + second, err := adapter.Build(ctx) + require.NoError(t, err) + sw2 := second.(*streamUpdatesStoreWrapper) + + assert.NotSame(t, sw1, sw2, "adapter must rebuild rather than hand back the torn-down wrapper") + assert.Equal(t, 1, sw2.refCount, "the fresh wrapper starts at refCount 1") + assert.False(t, sw2.store.(*mockStore).closed, "the fresh wrapper's underlying store is open") + assert.Same(t, sw2, adapter.GetStore(), "the adapter now points at the fresh wrapper") + + // acquire on the fully-closed wrapper refuses, so it can never be resurrected. + assert.False(t, sw1.acquire(), "acquire on a fully-closed wrapper must return false") +}