diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientProbeWiringTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientProbeWiringTests.java index 8864bd15b262..09dc16b00b7b 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientProbeWiringTests.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientProbeWiringTests.java @@ -14,15 +14,20 @@ import org.mockito.Mockito; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; +import reactor.core.Disposable; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import java.lang.reflect.Field; import java.net.URI; import java.time.Duration; +import java.util.Arrays; import java.util.HashMap; +import java.util.LinkedHashSet; import java.util.Map; import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; @@ -249,8 +254,356 @@ public void refreshWithNewRegion_probesOnlyNewlyAddedRegion() throws Exception { } } + @Test(groups = { "unit" }, timeOut = TIMEOUT) + public void thinClientDisabledAtRuntime_stopsProbeTraffic() throws Exception { + // Dynamic disable: a wired-and-probing orchestrator must stop issuing probe traffic once + // COSMOS.THINCLIENT_ENABLED is set to an explicit false at runtime, and the probe cycle + // must be skipped on subsequent forced refreshes (routing separately short-circuits to V1). + AtomicInteger probeCallCount = new AtomicInteger(0); + Map statusByEndpoint = new HashMap<>(); + statusByEndpoint.put(URI.create("https://testaccount-eastus.documents.azure.com:10250/connectivity-probe"), 200); + statusByEndpoint.put(URI.create("https://testaccount-eastasia.documents.azure.com:10250/connectivity-probe"), 200); + HttpClient httpClient = stubHttpClient(statusByEndpoint, probeCallCount); + + DatabaseAccount databaseAccount = new DatabaseAccount(DB_ACCOUNT_WITH_THINCLIENT_LOCATIONS); + Mockito.when(databaseAccountManagerInternal.getDatabaseAccountFromEndpoint(any())).thenReturn(Flux.just(databaseAccount)); + Mockito.when(databaseAccountManagerInternal.getServiceEndpoint()).thenReturn(new URI("https://testaccount.documents.azure.com:443")); + + ConnectionPolicy connectionPolicy = new ConnectionPolicy(DirectConnectionConfig.getDefaultConfig()); + connectionPolicy.setEndpointDiscoveryEnabled(true); + connectionPolicy.setMultipleWriteRegionsEnabled(true); + + String previous = System.getProperty(THINCLIENT_ENABLED_PROPERTY); + GlobalEndpointManager gem = new GlobalEndpointManager(databaseAccountManagerInternal, connectionPolicy, new Configs()); + try { + // Phase 1: default (unset) -> probes fire, gate goes healthy. + System.clearProperty(THINCLIENT_ENABLED_PROPERTY); + gem.setThinClientHttpClient(httpClient); + gem.init(); + waitForProxyDecision(gem, Boolean.TRUE, Duration.ofSeconds(5)); + int callsBeforeDisable = probeCallCount.get(); + assertThat(callsBeforeDisable).as("probes fired while enabled").isGreaterThanOrEqualTo(2); + + // Phase 2: explicit runtime opt-out -> forced refresh must NOT issue any new probe traffic. + System.setProperty(THINCLIENT_ENABLED_PROPERTY, "false"); + gem.refreshLocationAsync(null, true).block(); + Thread.sleep(500); // allow any (incorrectly) scheduled probe to run + + assertThat(probeCallCount.get()) + .as("no new probe traffic is issued while COSMOS.THINCLIENT_ENABLED=false") + .isEqualTo(callsBeforeDisable); + } finally { + restoreProperty(THINCLIENT_ENABLED_PROPERTY, previous); + LifeCycleUtils.closeQuietly(gem); + } + } + + @Test(groups = { "unit" }, timeOut = TIMEOUT) + public void thinClientOptedOutAtInit_probeResumesAfterOptOutDropped() throws Exception { + // Lazy-wiring: an orchestrator wired while COSMOS.THINCLIENT_ENABLED=false must issue NO probe + // traffic, yet must NOT be permanently disabled -- dropping the opt-out (unset) at runtime and + // forcing a refresh must resume probing and let the gate flip healthy. + AtomicInteger probeCallCount = new AtomicInteger(0); + Map statusByEndpoint = new HashMap<>(); + statusByEndpoint.put(URI.create("https://testaccount-eastus.documents.azure.com:10250/connectivity-probe"), 200); + statusByEndpoint.put(URI.create("https://testaccount-eastasia.documents.azure.com:10250/connectivity-probe"), 200); + HttpClient httpClient = stubHttpClient(statusByEndpoint, probeCallCount); + + DatabaseAccount databaseAccount = new DatabaseAccount(DB_ACCOUNT_WITH_THINCLIENT_LOCATIONS); + Mockito.when(databaseAccountManagerInternal.getDatabaseAccountFromEndpoint(any())).thenReturn(Flux.just(databaseAccount)); + Mockito.when(databaseAccountManagerInternal.getServiceEndpoint()).thenReturn(new URI("https://testaccount.documents.azure.com:443")); + + ConnectionPolicy connectionPolicy = new ConnectionPolicy(DirectConnectionConfig.getDefaultConfig()); + connectionPolicy.setEndpointDiscoveryEnabled(true); + connectionPolicy.setMultipleWriteRegionsEnabled(true); + + String previous = System.getProperty(THINCLIENT_ENABLED_PROPERTY); + GlobalEndpointManager gem = new GlobalEndpointManager(databaseAccountManagerInternal, connectionPolicy, new Configs()); + try { + // Phase 1: opt out BEFORE wiring/init -> probe is wired but never fires while disabled. + System.setProperty(THINCLIENT_ENABLED_PROPERTY, "false"); + gem.setThinClientHttpClient(httpClient); + gem.init(); + Thread.sleep(500); // give any (incorrectly) scheduled probe a chance to run + assertThat(probeCallCount.get()) + .as("no probe traffic is issued while opted out at init") + .isEqualTo(0); + + // Phase 2: drop the opt-out and force a refresh -> probing resumes and the gate flips healthy. + System.clearProperty(THINCLIENT_ENABLED_PROPERTY); + gem.refreshLocationAsync(null, true).block(); + waitForProxyDecision(gem, Boolean.TRUE, Duration.ofSeconds(5)); + + assertThat(probeCallCount.get()) + .as("probes resume for each thin-client region after the opt-out is dropped") + .isGreaterThanOrEqualTo(2); + assertThat(gem.getProxyProbeDecision()) + .as("gate flips healthy once probing resumes and all regions are proven") + .isEqualTo(Boolean.TRUE); + } finally { + restoreProperty(THINCLIENT_ENABLED_PROPERTY, previous); + LifeCycleUtils.closeQuietly(gem); + } + } + + @Test(groups = { "unit" }, timeOut = TIMEOUT) + public void overlappingProbeFire_doesNotCancelActiveCycle() throws Exception { + // New-A regression: a second (overlapping) probe fire must NOT cancel the sole in-flight cycle. + // Pre-fix, fireThinClientProbeCycle() did getAndSet(newDisposable) + dispose(old) on the CALLING + // thread, so a rapid overlapping fire disposed the only cycle doing real HTTP probes -> the gate + // stayed FALSE forever. The fix never disposes on fire (close() owns cancellation; overlapping + // fires are deduped by the probe client's single-flight CAS), so the in-flight cycle survives and + // the gate still flips healthy. + AtomicInteger probeCallCount = new AtomicInteger(0); + CountDownLatch probeEntered = new CountDownLatch(2); // both regions probed concurrently (flatMap) + CountDownLatch releaseProbe = new CountDownLatch(1); // holds both region sends in-flight + Map statusByEndpoint = new HashMap<>(); + statusByEndpoint.put(URI.create("https://testaccount-eastus.documents.azure.com:10250/connectivity-probe"), 200); + statusByEndpoint.put(URI.create("https://testaccount-eastasia.documents.azure.com:10250/connectivity-probe"), 200); + + // A blocking stub: each probe registers itself and then parks until releaseProbe is counted down, + // simulating a real in-flight HTTP probe. subscribeOn(boundedElastic) gives each concurrent probe + // its own thread so both can be in-flight at once (mirrors non-blocking Netty sends under flatMap). + HttpClient blockingClient = Mockito.mock(HttpClient.class); + org.mockito.stubbing.Answer> blockingAnswer = invocation -> { + HttpRequest req = invocation.getArgument(0); + Integer status = statusByEndpoint.get(req.uri()); + if (status == null) { + return Mono.error(new RuntimeException("Unexpected probe URI: " + req.uri())); + } + return Mono.fromCallable(() -> { + probeCallCount.incrementAndGet(); + probeEntered.countDown(); + releaseProbe.await(); + return stubResponse(req, status); + }) + .subscribeOn(reactor.core.scheduler.Schedulers.boundedElastic()); + }; + Mockito.when(blockingClient.send(any(HttpRequest.class), any(Duration.class))).thenAnswer(blockingAnswer); + Mockito.when(blockingClient.send(any(HttpRequest.class))).thenAnswer(blockingAnswer); + + DatabaseAccount databaseAccount = new DatabaseAccount(DB_ACCOUNT_WITH_THINCLIENT_LOCATIONS); + Mockito.when(databaseAccountManagerInternal.getDatabaseAccountFromEndpoint(any())).thenReturn(Flux.just(databaseAccount)); + Mockito.when(databaseAccountManagerInternal.getServiceEndpoint()).thenReturn(new URI("https://testaccount.documents.azure.com:443")); + + ConnectionPolicy connectionPolicy = new ConnectionPolicy(DirectConnectionConfig.getDefaultConfig()); + connectionPolicy.setEndpointDiscoveryEnabled(true); + connectionPolicy.setMultipleWriteRegionsEnabled(true); + + String previous = System.getProperty(THINCLIENT_ENABLED_PROPERTY); + GlobalEndpointManager gem = new GlobalEndpointManager(databaseAccountManagerInternal, connectionPolicy, new Configs()); + try { + System.clearProperty(THINCLIENT_ENABLED_PROPERTY); // tri-state unset -> probe-gated + gem.setThinClientHttpClient(blockingClient); + gem.init(); // cycle D1 fires; both region sends park on releaseProbe + + // D1 has won single-flight and both regions are in-flight concurrently. + assertThat(probeEntered.await(5, TimeUnit.SECONDS)) + .as("the initial probe cycle reached in-flight HTTP for both thin-client regions") + .isTrue(); + int callsInFlight = probeCallCount.get(); + + // Fire an overlapping cycle while D1 is still in-flight. Pre-fix, this disposed D1 here. + gem.refreshLocationAsync(null, true).block(); + + // Single-flight: the overlapping fire is deduped and issues NO new probe traffic. + assertThat(probeCallCount.get()) + .as("overlapping fire is deduped by single-flight and issues no new probe traffic") + .isEqualTo(callsInFlight); + + // Let the still-alive cycle complete. On the fixed code D1 was never disposed. + releaseProbe.countDown(); + + waitForProxyDecision(gem, Boolean.TRUE, Duration.ofSeconds(5)); + assertThat(gem.getProxyProbeDecision()) + .as("the in-flight cycle survives the overlapping fire and flips the gate healthy") + .isEqualTo(Boolean.TRUE); + } finally { + releaseProbe.countDown(); // never leave a probe thread parked + restoreProperty(THINCLIENT_ENABLED_PROPERTY, previous); + LifeCycleUtils.closeQuietly(gem); + } + } + + @Test(groups = { "unit" }, timeOut = TIMEOUT) + public void overlappingProbeFire_closeStillCancelsTheLiveCycle() throws Exception { + // New-A residual regression (PR #49796 comment r3575154107): an overlapping (no-op) probe fire + // must NOT replace the tracked disposable that points at the sole in-flight cycle. close() owns + // cancellation and disposes whatever handle is tracked, so if a no-op fire clobbers the handle, + // close() cancels the no-op and the real in-flight probe I/O keeps running until its per-probe + // timeout -- i.e. close() can no longer promptly stop the live cycle. + // + // Pre-fix, fireThinClientProbeCycle() did set(newDisposable) on EVERY fire, so an overlapping + // fire (single-flight loser -> a no-op that completes immediately) overwrote the tracked handle: + // the tracked disposable was then the disposed no-op, not the live cycle. The fix uses + // getAndUpdate(keep-live-predecessor): while a live cycle is tracked, a no-op fire is discarded + // and the handle keeps pointing at the real cycle, so close() promptly cancels it. + AtomicInteger probeCallCount = new AtomicInteger(0); + CountDownLatch probeEntered = new CountDownLatch(2); // both regions probed concurrently (flatMap) + CountDownLatch releaseProbe = new CountDownLatch(1); // holds both region sends in-flight + Map statusByEndpoint = new HashMap<>(); + statusByEndpoint.put(URI.create("https://testaccount-eastus.documents.azure.com:10250/connectivity-probe"), 200); + statusByEndpoint.put(URI.create("https://testaccount-eastasia.documents.azure.com:10250/connectivity-probe"), 200); + + HttpClient blockingClient = Mockito.mock(HttpClient.class); + org.mockito.stubbing.Answer> blockingAnswer = invocation -> { + HttpRequest req = invocation.getArgument(0); + Integer status = statusByEndpoint.get(req.uri()); + if (status == null) { + return Mono.error(new RuntimeException("Unexpected probe URI: " + req.uri())); + } + return Mono.fromCallable(() -> { + probeCallCount.incrementAndGet(); + probeEntered.countDown(); + releaseProbe.await(); + return stubResponse(req, status); + }) + .subscribeOn(reactor.core.scheduler.Schedulers.boundedElastic()); + }; + Mockito.when(blockingClient.send(any(HttpRequest.class), any(Duration.class))).thenAnswer(blockingAnswer); + Mockito.when(blockingClient.send(any(HttpRequest.class))).thenAnswer(blockingAnswer); + + DatabaseAccount databaseAccount = new DatabaseAccount(DB_ACCOUNT_WITH_THINCLIENT_LOCATIONS); + Mockito.when(databaseAccountManagerInternal.getDatabaseAccountFromEndpoint(any())).thenReturn(Flux.just(databaseAccount)); + Mockito.when(databaseAccountManagerInternal.getServiceEndpoint()).thenReturn(new URI("https://testaccount.documents.azure.com:443")); + + ConnectionPolicy connectionPolicy = new ConnectionPolicy(DirectConnectionConfig.getDefaultConfig()); + connectionPolicy.setEndpointDiscoveryEnabled(true); + connectionPolicy.setMultipleWriteRegionsEnabled(true); + + String previous = System.getProperty(THINCLIENT_ENABLED_PROPERTY); + GlobalEndpointManager gem = new GlobalEndpointManager(databaseAccountManagerInternal, connectionPolicy, new Configs()); + try { + System.clearProperty(THINCLIENT_ENABLED_PROPERTY); // tri-state unset -> probe-gated + gem.setThinClientHttpClient(blockingClient); + gem.init(); // cycle D1 fires; both region sends park on releaseProbe + + assertThat(probeEntered.await(5, TimeUnit.SECONDS)) + .as("the initial probe cycle reached in-flight HTTP for both thin-client regions") + .isTrue(); + + // D1 is the tracked, live cycle. + Disposable liveCycle = getTrackedProbeDisposable(gem); + assertThat(liveCycle).as("the initial in-flight cycle is tracked").isNotNull(); + assertThat(liveCycle.isDisposed()).as("the tracked cycle is still live").isFalse(); + + // Fire an overlapping cycle while D1 is still in-flight (single-flight -> a no-op loser). + gem.refreshLocationAsync(null, true).block(); + + // The crux: the tracked handle must STILL be the live cycle, not the no-op fire. + // Pre-fix (set() on every fire) this is the disposed no-op, so both assertions go red. + assertThat(getTrackedProbeDisposable(gem)) + .as("overlapping no-op fire must not replace the handle to the live cycle") + .isSameAs(liveCycle); + assertThat(liveCycle.isDisposed()) + .as("the live cycle is untouched by the overlapping fire") + .isFalse(); + + // close() disposes the tracked handle: it must promptly cancel the REAL in-flight cycle. + gem.close(); + assertThat(liveCycle.isDisposed()) + .as("close() promptly cancels the live in-flight cycle") + .isTrue(); + } finally { + releaseProbe.countDown(); // never leave a probe thread parked + restoreProperty(THINCLIENT_ENABLED_PROPERTY, previous); + LifeCycleUtils.closeQuietly(gem); + } + } + + @Test(groups = { "unit" }, timeOut = TIMEOUT) + public void staleTopologyGrowth_doesNotFlipGateForUnprobedRegion() throws Exception { + // Issue #1 regression (stale-green): if the topology GROWS while a cycle is in flight, the + // completing cycle must recompute the gate against the LATEST topology, not the snapshot it + // captured at subscribe time. + // + // Pre-fix, applyCycleResult(endpoints, results) recomputed against the winning cycle's stale + // snapshot T1={E1,E2}: after proving E1 and E2 it computed computeGate(T1)=TRUE, clobbering the + // correct FALSE that the overlapping T2={E1,E2,E3} trigger had already published (E3 unproven). + // That globally enabled thin-client routing for the still-unproven E3 region until the next + // refresh -> misrouting. The fix recomputes against latestTopology (T2) so the gate stays FALSE. + URI e1 = URI.create("https://region1.example.com:10250/"); + URI e2 = URI.create("https://region2.example.com:10250/"); + URI e3 = URI.create("https://region3.example.com:10250/"); + + Map statusByEndpoint = new HashMap<>(); + statusByEndpoint.put(URI.create("https://region1.example.com:10250/connectivity-probe"), 200); + statusByEndpoint.put(URI.create("https://region2.example.com:10250/connectivity-probe"), 200); + statusByEndpoint.put(URI.create("https://region3.example.com:10250/connectivity-probe"), 200); + + CountDownLatch probeEntered = new CountDownLatch(2); // E1 + E2 probed by the winning cycle C1 + CountDownLatch releaseProbe = new CountDownLatch(1); // holds C1's two region sends in-flight + + // Blocking stub: each probe registers itself then parks until releaseProbe fires, simulating a + // real in-flight HTTP probe. subscribeOn(boundedElastic) gives each concurrent probe its own + // thread so both can be in-flight at once (mirrors non-blocking Netty sends under flatMap). + HttpClient blockingClient = Mockito.mock(HttpClient.class); + org.mockito.stubbing.Answer> blockingAnswer = invocation -> { + HttpRequest req = invocation.getArgument(0); + Integer status = statusByEndpoint.get(req.uri()); + if (status == null) { + return Mono.error(new RuntimeException("Unexpected probe URI: " + req.uri())); + } + return Mono.fromCallable(() -> { + probeEntered.countDown(); + releaseProbe.await(); + return stubResponse(req, status); + }) + .subscribeOn(reactor.core.scheduler.Schedulers.boundedElastic()); + }; + Mockito.when(blockingClient.send(any(HttpRequest.class), any(Duration.class))).thenAnswer(blockingAnswer); + Mockito.when(blockingClient.send(any(HttpRequest.class))).thenAnswer(blockingAnswer); + + EndpointProbeClient client = new EndpointProbeClient(blockingClient); + try { + Set t1 = new LinkedHashSet<>(Arrays.asList(e1, e2)); + Set t2 = new LinkedHashSet<>(Arrays.asList(e1, e2, e3)); + + // C1: the winning cycle probes T1 and parks both regions in-flight (async subscribe). + CountDownLatch c1Done = new CountDownLatch(1); + client.runProbeCycle(t1).subscribe(v -> { }, e -> c1Done.countDown(), c1Done::countDown); + + assertThat(probeEntered.await(5, TimeUnit.SECONDS)) + .as("C1 reached in-flight HTTP for both T1 regions and holds the single-flight slot") + .isTrue(); + + // C2: overlapping trigger publishing the grown topology T2. Deduped by single-flight, it + // issues NO new probe traffic (E3 is never sent here) and returns the freshly recomputed + // gate = FALSE because E3 is unproven. This also records latestTopology = T2. + Boolean c2Gate = client.runProbeCycle(t2).block(); + assertThat(c2Gate) + .as("overlapping trigger recomputes the gate against the grown topology (E3 unproven)") + .isEqualTo(Boolean.FALSE); + + // Release C1; it proves E1 and E2 and then completes. + releaseProbe.countDown(); + assertThat(c1Done.await(5, TimeUnit.SECONDS)) + .as("C1 completes after release") + .isTrue(); + + // Core assertion: the completing cycle recomputed the gate against the LATEST topology (T2), + // so it stays FALSE because E3 was never probed. Pre-fix this clobbered to a stale-green TRUE. + assertThat(client.isThinClientRoutable()) + .as("gate reflects the latest topology (E3 unproven) and does NOT flip to a stale green") + .isFalse(); + } finally { + releaseProbe.countDown(); // never leave a probe thread parked + client.close(); + } + } + // ---- helpers ---- + private static final String THINCLIENT_ENABLED_PROPERTY = "COSMOS.THINCLIENT_ENABLED"; + + private static void restoreProperty(String key, String previous) { + if (previous == null) { + System.clearProperty(key); + } else { + System.setProperty(key, previous); + } + } + private GlobalEndpointManager newGemWithAccount(String accountJson) throws Exception { DatabaseAccount databaseAccount = new DatabaseAccount(accountJson); Mockito.when(databaseAccountManagerInternal.getDatabaseAccountFromEndpoint(ArgumentMatchers.any())) @@ -273,6 +626,13 @@ private static LocationCache getLocationCache(GlobalEndpointManager gem) throws return (LocationCache) f.get(gem); } + @SuppressWarnings("unchecked") + private static Disposable getTrackedProbeDisposable(GlobalEndpointManager gem) throws Exception { + Field f = GlobalEndpointManager.class.getDeclaredField("thinClientProbeCycleDisposable"); + f.setAccessible(true); + return ((AtomicReference) f.get(gem)).get(); + } + private static HttpClient stubHttpClient(Map statusByEndpoint, AtomicInteger callCount) { HttpClient mock = Mockito.mock(HttpClient.class); Mockito.when(mock.send(any(HttpRequest.class), any(Duration.class))) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/EndpointProbeClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/EndpointProbeClient.java index a0bcc260b40e..6b570ada7017 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/EndpointProbeClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/EndpointProbeClient.java @@ -23,6 +23,7 @@ import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; /** * Drives the thin-client HTTP/2 connectivity probe lifecycle with a per-region, one-shot @@ -69,6 +70,13 @@ public class EndpointProbeClient implements Closeable { private final AtomicBoolean closed = new AtomicBoolean(false); private final AtomicBoolean cycleInProgress = new AtomicBoolean(false); + // The most recently observed topology (published by every runProbeCycle trigger, winner or + // overlapping loser). applyCycleResult recomputes the gate against THIS rather than the winning + // cycle's captured snapshot, so a topology that advances mid-cycle can't be clobbered by a stale + // green verdict when the in-flight cycle completes. + private final AtomicReference> latestTopology = + new AtomicReference<>(Collections.emptySet()); + public EndpointProbeClient(HttpClient httpClient) { this.httpClient = Objects.requireNonNull( httpClient, @@ -99,9 +107,11 @@ public Mono runProbeCycle(Collection regionalEndpoints) { return Mono.just(this.thinClientRoutable); } - // Recompute the cached gate against the latest topology using the current proven set, so - // it reflects the newest endpoints even if this cycle probes nothing or is skipped as an - // overlapping trigger. + // Publish the newest topology so a cycle completing later recomputes the gate against + // this value (never a stale captured snapshot), then recompute the cached gate against it + // using the current proven set — so the gate reflects the newest endpoints even if this + // cycle probes nothing or is skipped as an overlapping trigger. + this.latestTopology.set(endpoints); this.thinClientRoutable = computeGate(endpoints); if (!this.cycleInProgress.compareAndSet(false, true)) { @@ -115,7 +125,7 @@ public Mono runProbeCycle(Collection regionalEndpoints) { .filter(endpoint -> !this.provenHealthyEndpoints.contains(endpoint)) .flatMap(this::probeEndpointOnce) .collectList() - .map(results -> applyCycleResult(endpoints, results)) + .map(this::applyCycleResult) .onErrorResume(t -> { logger.warn( "Thin-client probe cycle threw an unexpected error; leaving failed regions un-cached.", t); @@ -189,7 +199,7 @@ private Mono probeEndpointOnce(URI regionalEndpoint) { }); } - private Boolean applyCycleResult(Collection endpoints, List results) { + private Boolean applyCycleResult(List results) { // Dropped if the client closed mid-cycle so we don't mutate a dead client. if (this.closed.get()) { @@ -204,7 +214,10 @@ private Boolean applyCycleResult(Collection endpoints, List { }, t -> logger.warn("Thin-client probe cycle subscription errored unexpectedly; ignoring.", t)); - Disposable oldDisposable = this.thinClientProbeCycleDisposable.getAndSet(newDisposable); - if (oldDisposable != null && !oldDisposable.isDisposed()) { - oldDisposable.dispose(); - } + // close() cancels the tracked subscription to promptly stop the ACTIVE probe I/O (the probe + // client's own close() only flips a flag; in-flight probes otherwise self-terminate on the + // per-probe timeout). So the tracked handle must always be the real in-flight cycle, never a + // no-op trigger: while a cycle is running, EndpointProbeClient's single-flight CAS forces every + // concurrent fire onto the no-op path, so THIS subscription is the no-op and the already-tracked + // one is the real cycle. Keep the live predecessor and let this no-op self-complete; only + // install the new handle when nothing live is tracked (it won the single-flight and is the real + // cycle). Firing on every refresh still runs runThinClientProbeCycleMono -> runProbeCycle, which + // republishes the latest topology and recomputes the gate even on the no-op path, so a topology + // delta arriving mid-cycle is never lost. We never dispose the predecessor here (only close() + // does), so an active cycle is never aborted. + this.thinClientProbeCycleDisposable.getAndUpdate( + previous -> (previous != null && !previous.isDisposed()) ? previous : newDisposable); } private Mono runThinClientProbeCycleMono() { return Mono.defer(() -> { + // No-op when COSMOS.THINCLIENT_ENABLED is explicitly set (true or false): the flag is then + // a hard contract that decides routing directly, so the probe verdict is irrelevant and we + // skip the traffic. Read live here, so dropping the opt-out/opt-in back to unset resumes + // probing on the next refresh. Only an unset flag lets the probe gate routing. + if (Configs.isThinClientEnabled() != null) { + return Mono.empty(); + } EndpointProbeClient probeClient = this.thinClientProbeClient.get(); if (probeClient == null) { return Mono.empty(); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index c31d590e6d35..85b20c5f976d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -921,29 +921,20 @@ public void init(CosmosClientMetadataCachesSnapshot metadataCachesSnapshot, Func this.reactorHttpClient, this.additionalHeaders); - // Wire thin-client HttpClient into GEM so the connectivity-probe orchestrator - // can fan out probes after every topology refresh. Must happen BEFORE - // globalEndpointManager.init() so the first refresh probes immediately. - // The probe client is wired whenever thin-client is usable — GATEWAY mode + HTTP/2 and - // COSMOS.THINCLIENT_ENABLED not an explicit false (a hard opt-out); see - // ThinClientConnectivityConfig.canThinClientBeUsed(). We deliberately wire even on an - // explicit true: the wiring decision is made once here at init, but the tri-state flag is - // re-read lazily per request, so a runtime transition of COSMOS.THINCLIENT_ENABLED from - // true back to unset (an operator dropping the opt-in to rely on probe-based rollout) - // must still have a live probe to consult instead of silently pinning to Gateway V1. - // Wiring the probe for an explicit opt-in is otherwise free: shouldUseThinClientStoreModel - // returns the explicit verdict directly and never consults the probe. When the probe is - // NOT wired (hard opt-out), GEM's probeClient stays null and `getProxyProbeDecision()` - // renders no decision (null). Wiring itself is guarded inside GEM so any failure cannot - // trip client init. - if (this.thinClientConnectivityConfig.canThinClientBeUsed()) { - try { - this.globalEndpointManager.setThinClientHttpClient(this.reactorHttpClient); - } catch (Throwable t) { - // Defense in depth: GEM already swallows wiring failures, but if anything - // does escape we must not fail CosmosClient construction over a probe. - logger.warn("Failed to wire thin-client connectivity-probe HttpClient; continuing without probe gating.", t); - } + // Wire thin-client HttpClient into GEM so the connectivity-probe orchestrator can fan out + // probes after every topology refresh. Must happen BEFORE globalEndpointManager.init() so + // the first refresh probes immediately. We always wire the probe client and do NOT gate on + // COSMOS.THINCLIENT_ENABLED here: the flag is runtime-mutable and re-read lazily, so gating + // wiring on it would make an init-time hard opt-out (false) permanent. Instead the probe + // cycle itself is a no-op whenever the flag is explicitly set (true or false); it only + // probes when the flag is unset (the case where the probe verdict actually gates routing). + // Wiring is guarded inside GEM so any failure cannot trip client init. + try { + this.globalEndpointManager.setThinClientHttpClient(this.reactorHttpClient); + } catch (Throwable t) { + // Defense in depth: GEM already swallows wiring failures, but if anything + // does escape we must not fail CosmosClient construction over a probe. + logger.warn("Failed to wire thin-client connectivity-probe HttpClient; continuing without probe gating.", t); } this.perPartitionFailoverConfigModifier