[NO REVIEW] Fix 400/1024 error when container is re-created with same name

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();
     }
 
     /**

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";

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RenameCollectionAwareClientRetryPolicy.java

@@ -70,6 +70,7 @@ public Mono<ShouldRetryResult> shouldRetry(Exception e) {
 
                     request.forceNameCacheRefresh = true;
                     request.requestContext.resolvedCollectionRid = null;
+                    request.getHeaders().remove(HttpConstants.HttpHeaders.INTENDED_COLLECTION_RID_HEADER);
 
                     Mono<Utils.ValueHolder<DocumentCollection>> collectionObs = this.collectionCache.resolveCollectionAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), request);
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/StaleResourceRetryPolicy.java

@@ -124,6 +124,14 @@ public Mono<ShouldRetryResult> 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));