[Cosmos] Gateway V2 endpoint-probe: topology-safe gate lifecycle (follow-up to #49437)

[jeet1995]

Context

Follow-up to the review comment thread on #49437. Addresses two endpoint-probe gate lifecycle correctness issues raised in review. Both concern how the thin-client (Gateway V2) routing gate is driven by the background connectivity probe, and both are proven with regression tests in ThinClientProbeWiringTests.

---

Issue 1 — Overlapping probe fire could cancel the only active probe cycle

Where: GlobalEndpointManager (probe fire/lifecycle) + EndpointProbeClient (single-flight cycle).

The problem. The probe cycle is fired fire-and-forget on every topology refresh, and the live subscription is tracked in an AtomicReference so close() can promptly cancel in-flight probe I/O. The original code tracked the subscription with .set(newDisposable) on every fire. But EndpointProbeClient runs a single-flight cycle: a fire that arrives while a cycle is already running loses the CAS and returns a fast no-op subscription. .set() then overwrote the handle to the real in-flight cycle with that no-op. On close(), getAndSet(null) returned the no-op — so the real probe I/O was never promptly cancelled and only self-terminated on the per-probe timeout. In the worst concurrent case, the single active cycle's cancel handle was lost entirely.

The fix. Keep the tracked handle pointing at the live cycle:

this.thinClientProbeCycleDisposable.getAndUpdate(
    previous -> (previous != null && !previous.isDisposed()) ? previous : newDisposable);

The live predecessor is retained; the new handle is installed only when nothing live is tracked (this fire won single-flight and is the real cycle) or the predecessor is already disposed. The overlapping no-op loser is fire-and-forget and self-completes — no orphaned subscription. close() disposes the real cycle promptly.

Topology freshness is preserved alongside this: every fire (winner or overlapping loser) publishes the newest topology into latestTopology and recomputes the gate before the single-flight CAS, and applyCycleResult recomputes the final gate against latestTopology rather than the winning cycle's captured snapshot. So a topology that advances mid-cycle (e.g. a new region) can never be clobbered by a stale "green" verdict from an older snapshot.

Regression tests.

• overlappingProbeFire_closeStillCancelsTheLiveCycle — asserts the tracked handle is still the live cycle after an overlapping no-op fire (identity check) and that close() disposes that live cycle. Proven red on the old .set() (fails "overlapping no-op fire must not replace the handle to the live cycle"), green on getAndUpdate.
• overlappingProbeFire_doesNotCancelActiveCycle — the active cycle survives an overlapping trigger and the gate still flips healthy.
• staleTopologyGrowth_doesNotFlipGateForUnprobedRegion — a newly added, unprobed region does not get a stale green verdict.

---

Issue 2 — Init-time opt-out permanently blocked later probe-gated enablement (lazy unset)

Where: RxDocumentClientImpl client-construction wiring of the probe HttpClient into GlobalEndpointManager.

The problem. The probe HttpClient was wired into GEM only when ThinClientConnectivityConfig.canThinClientBeUsed() was true at construction, which consulted COSMOS.IS_THINCLIENT_ENABLED. That flag is tri-state and runtime-mutable (system property / env var), re-read lazily per request: true = force on, false = force off, unset = let the probe verdict gate routing. Gating the wiring on an init-time read meant an explicit false at startup permanently tore out the probe. If an operator later cleared the flag to unset (dropping a hard opt-out to rely on probe-based rollout), there was no live probe to consult, so routing silently pinned to Gateway V1 forever — the dynamic-enablement contract was broken.

The fix. Always wire the probe HttpClient at init and never gate the wiring on the flag:

try {
    this.globalEndpointManager.setThinClientHttpClient(this.reactorHttpClient);
} catch (Throwable t) {
    logger.warn("Failed to wire thin-client connectivity-probe HttpClient; continuing without probe gating.", t);
}

The flag is now honored where it is actually read (lazily, per request): shouldUseThinClientStoreModel returns the explicit verdict directly for true/false and consults the probe decision only when the flag is unset. The probe cycle itself is effectively a no-op whenever the flag is explicitly set — it does real probing only in the unset case where the verdict gates routing. So flipping true/false ↔ unset at runtime now dynamically enables/disables thin-client routing with a live probe already in place, and wiring an explicit opt-in stays free. Wiring is guarded inside GEM so a probe failure can never fail CosmosClient construction.

---

Testing

• ThinClientProbeWiringTests — 11/11 green.
• azure-cosmos build succeeds (revapi NPE trace is known non-fatal noise).

---

Thin-client connectivity-probe drills — live results (PR #49796 @ d7f284f0)

Ran against thinclient-session (multi-write, EUS2+CUS, enableConnectivityProbe=true) from an EUS2 bench VM. All runs on the implicit probe path (GATEWAY + HTTP/2), single CosmosAsyncClient. Routing was read from the authoritative CosmosDiagnostics.endpoint (host:port) of each request and corroborated with socket state (ss). No account topology was mutated; all faults were client-side on the VM only.

1. Issue 2 — init opt-out must not permanently block probe-gated enablement (COSMOS.THINCLIENT_ENABLED false → unset) ✅ PASS

Single client started with -DCOSMOS.THINCLIENT_ENABLED=false, property cleared at runtime (no restart):

phase	COSMOS.THINCLIENT_ENABLED	routing
t=0–86s	false (hard opt-out)	Gateway V1 :443 (eastus2)
t=91s	cleared → unset	still :443 (probe verdict not yet TRUE)
t=309s	unset	elevated to Gateway V2 :10250 (thinclient-session-eastus2.documents.azure.com:10250)

Elevation happened ~218s after unset (one topology-refresh probe cycle), on the same client, 0 errors, 0 non-200. This proves the fix: the probe HttpClient is now always wired at init, so dropping a startup false to unset lets the live probe gate routing dynamically. Pre-fix, the probe was never wired under an init opt-out and routing would have pinned to :443 forever.

2. THINCLIENT_PROBE_NETFAULT — C1 V1→V2 elevation ✅ PASS

iptables -A OUTPUT -p tcp --dport 10250 -j DROP applied before bootstrap:

phase	routing	ss :10250
:10250 blocked	Gateway V1 :443 (writes 201/200, reads 200, eastus2)	0 conns
after iptables -D (t=46s)	—	—
t=300s	elevated to Gateway V2 :10250 (writes+reads 200, eastus2)	8 conns

Probe-gated client under a :10250 block bootstraps on V1; on unblock it elevates to V2 within one probe cycle (~254s), transparently (0 errors).

3. THINCLIENT_PROBE_NETFAULT — C2 regional failover (multi-write) ✅ PASS

Client on V2. DNS-blackholed the write region (thinclient-session-eastus2… → 192.0.2.1) + severed live :10250 conns:

phase	reads	writes	transport
eastus2 blackholed	fail over → central us 200	fail over → central us 200	stays :10250 (V2)

• Transition: ~100s transient 503 (substatus 10001) burst while in-flight requests to the severed eastus2 :10250 timed out; client marked eastus2 unavailable, refreshed, and converged.
• Steady state: 74 requests on central us :10250, ss = 16 conns all to CUS IPs (20.15.133.x).
• 0 :443 downgrades after the fault — failover is regional on the V2 transport, not a V1 fallback (matches the multi-write oracle: reads + writes fail over and succeed).

Full-run endpoint/region tally: 198× :443/eastus2 (V1 blocked) → 22× :10250/eastus2 (V2 elevated) → 74× :10250/centralus (failed over).

[jeet1995]

/azp run java - cosmos - tests

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[Copilot]

Pull request overview

Follow-up to the Gateway V2 (thin-client) endpoint connectivity-probe rollout work, tightening probe lifecycle correctness under overlapping refreshes and runtime flag transitions to prevent stale “green” gates and inadvertent cancellation of active probe cycles.

Changes:

• Ensures probe gate computation uses the latest observed topology (prevents stale-green clobbering after topology growth).
• Adjusts probe-cycle subscription tracking and adds a runtime opt-in/opt-out no-op gate for probe traffic.
• Adds targeted unit tests for overlapping fires and stale-topology growth scenarios.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java	Always wires the thin-client probe HTTP client at init to avoid permanent disablement from init-time opt-out.
sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/GlobalEndpointManager.java	Updates probe fire/disposal behavior and skips probe traffic when COSMOS.THINCLIENT_ENABLED is explicitly set.
sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/EndpointProbeClient.java	Recomputes routing gate against the latest topology reference to avoid stale verdict publication.
sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientProbeWiringTests.java	Adds regression tests covering overlapping probe fires and topology growth during an in-flight cycle, plus runtime opt-out behavior.

[jeet1995]

@sdkReviewAgent

[jeet1995]

/azp run java - cosmos - tests

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[FabianMeiswinkel]

LGTM