diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RenameCollectionAwareClientRetryPolicyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RenameCollectionAwareClientRetryPolicyTest.java index 1f1f0a33a7dd..8e6ca063a736 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RenameCollectionAwareClientRetryPolicyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RenameCollectionAwareClientRetryPolicyTest.java @@ -121,6 +121,7 @@ public void shouldRetryWithNotFoundStatusCodeAndReadSessionNotAvailableSubStatus OperationType.Create, "/dbs/db/colls/col/docs/docId", ResourceType.Document); request.requestContext = new DocumentServiceRequestContext(); request.requestContext.resolvedCollectionRid = "rid_0"; + request.getHeaders().put(HttpConstants.HttpHeaders.INTENDED_COLLECTION_RID_HEADER, "rid_0"); renameCollectionAwareClientRetryPolicy.onBeforeSendRequest(request); NotFoundException notFoundException = new NotFoundException(); @@ -138,6 +139,10 @@ public void shouldRetryWithNotFoundStatusCodeAndReadSessionNotAvailableSubStatus .nullException() .shouldRetry(true) .build()); + + // Verify stale intended-collection-rid header was removed so it gets + // re-populated with the new collection rid on retry + assertThat(request.getHeaders().get(HttpConstants.HttpHeaders.INTENDED_COLLECTION_RID_HEADER)).isNull(); } /** diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/StaleResourceExceptionRetryPolicyTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/StaleResourceExceptionRetryPolicyTest.java index 76a5d3f6bc08..2b1f26de1098 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/StaleResourceExceptionRetryPolicyTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/StaleResourceExceptionRetryPolicyTest.java @@ -98,6 +98,75 @@ public void suppressRetryForExternalCollectionRid() { assertThat(shouldRetryResult.shouldRetry).isFalse(); } + @DataProvider(name = "staleContainerExceptionProvider") + public Object[][] staleContainerExceptionProvider() { + return new Object[][] { + //status code, subStatusCode + { HttpConstants.StatusCodes.GONE, HttpConstants.SubStatusCodes.NAME_CACHE_IS_STALE }, + { HttpConstants.StatusCodes.BADREQUEST, HttpConstants.SubStatusCodes.INCORRECT_CONTAINER_RID_SUB_STATUS }, + }; + } + + @Test(groups = "unit", dataProvider = "staleContainerExceptionProvider") + public void requestContextResetOnRetry(int statusCode, int subStatusCode) { + // Validates that when StaleResourceRetryPolicy retries a stale-container + // exception, it resets the request context so the retry re-resolves the + // collection and sends the updated intended-collection-rid header. + String testCollectionLink = "/dbs/test/colls/contextResetTest"; + + DocumentCollection documentCollection = new DocumentCollection(); + documentCollection.setId("contextResetTest"); + documentCollection.setResourceId("oldRid"); + + DocumentCollection documentCollectionAfterRefresh = new DocumentCollection(); + documentCollectionAfterRefresh.setId("contextResetTest"); + documentCollectionAfterRefresh.setResourceId("newRid"); + + RxCollectionCache rxCollectionCache = Mockito.mock(RxCollectionCache.class); + Mockito + .when(rxCollectionCache.resolveByNameAsync(Mockito.any(), Mockito.any(), Mockito.isNull(), Mockito.isNull(), Mockito.isNull())) + .thenReturn(Mono.just(documentCollection)) + .thenReturn(Mono.just(documentCollectionAfterRefresh)); + doNothing().when(rxCollectionCache).refresh(Mockito.any(), Mockito.any(), Mockito.isNull()); + + ISessionContainer sessionContainer = Mockito.mock(ISessionContainer.class); + doNothing().when(sessionContainer).clearTokenByResourceId(documentCollection.getResourceId()); + + StaleResourceRetryPolicy staleResourceRetryPolicy = new StaleResourceRetryPolicy( + rxCollectionCache, + null, + testCollectionLink, + null, + null, + sessionContainer, + TestUtils.mockDiagnosticsClientContext(), + null + ); + + // Simulate a request with a stale resolvedCollectionRid and intended header + RxDocumentServiceRequest request = RxDocumentServiceRequest.createFromName( + TestUtils.mockDiagnosticsClientContext(), + OperationType.Read, "/dbs/test/colls/contextResetTest/docs/doc1", ResourceType.Document); + request.requestContext = new DocumentServiceRequestContext(); + request.requestContext.resolvedCollectionRid = "oldRid"; + request.getHeaders().put(HttpConstants.HttpHeaders.INTENDED_COLLECTION_RID_HEADER, "oldRid"); + + staleResourceRetryPolicy.onBeforeSendRequest(request); + + CosmosException exception = BridgeInternal.createCosmosException(statusCode); + BridgeInternal.setSubStatusCode(exception, subStatusCode); + + ShouldRetryResult shouldRetryResult = staleResourceRetryPolicy.shouldRetry(exception).block(); + assertThat(shouldRetryResult.shouldRetry).isTrue(); + + // Verify request context was updated for a clean retry + assertThat(request.requestContext.resolvedCollectionRid).isEqualTo("newRid"); + assertThat(request.getHeaders().get(HttpConstants.HttpHeaders.INTENDED_COLLECTION_RID_HEADER)).isNull(); + + // Verify session token was cleaned up for the old rid + verify(sessionContainer, Mockito.times(1)).clearTokenByResourceId("oldRid"); + } + @Test(groups = "unit") public void cleanSessionToken() { String testCollectionLink = "/dbs/test/colls/staledExceptionTest"; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RenameCollectionAwareClientRetryPolicy.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RenameCollectionAwareClientRetryPolicy.java index ad1858673a37..d376a74fb99e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RenameCollectionAwareClientRetryPolicy.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RenameCollectionAwareClientRetryPolicy.java @@ -70,6 +70,7 @@ public Mono shouldRetry(Exception e) { request.forceNameCacheRefresh = true; request.requestContext.resolvedCollectionRid = null; + request.getHeaders().remove(HttpConstants.HttpHeaders.INTENDED_COLLECTION_RID_HEADER); Mono> collectionObs = this.collectionCache.resolveCollectionAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), request); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/StaleResourceRetryPolicy.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/StaleResourceRetryPolicy.java index b3e8d0274a91..cb4edc8d2373 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/StaleResourceRetryPolicy.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/StaleResourceRetryPolicy.java @@ -124,6 +124,14 @@ public Mono shouldRetry(Exception e) { this.sessionContainer.clearTokenByResourceId(oldCollectionRid.get()); } + // Update request context with the refreshed collection rid + // and remove the stale intended-collection-rid header so it + // gets re-populated on retry. + if (this.request != null && this.request.requestContext != null) { + this.request.requestContext.resolvedCollectionRid = refreshedCollectionRid; + this.request.getHeaders().remove(HttpConstants.HttpHeaders.INTENDED_COLLECTION_RID_HEADER); + } + this.retried = true; if (this.shouldSuppressRetry.get()) { return Mono.just(ShouldRetryResult.error(e));