Session timeout 404/1002 HA fixes (Azure#18594)

simplynaveen20 · web-flow · commit 89fb1fc010a8 · 2021-01-20T16:27:20.000-08:00
* HA fixes

* removing extra onBefore on read request

* incremental change in test case

* incremental change in test case

* ignoring warning from test case

* fixing null pointer in test case on document producer

* adding another test case for non preferred list

* closing clients

* using reflection utils methods

* adding 0th index for single master explicity to go to master region, improving code readablity

* Removing InternalObjectNode from the test case

* Adding seperate CI for multi region and single master

* stopping the ci

* merge with master and fixing new test pipeline
diff --git a/sdk/cosmos/azure-cosmos/pom.xml b/sdk/cosmos/azure-cosmos/pom.xml
@@ -366,6 +366,27 @@ Licensed under the MIT License.
         </plugins>
       </build>
     </profile>
+    <profile>
+      <!-- integration tests, requires Cosmos DB endpoint with multi region support -->
+      <id>multi-region</id>
+      <properties>
+        <test.groups>multi-region</test.groups>
+      </properties>
+      <build>
+        <plugins>
+          <plugin>
+            <groupId>org.apache.maven.plugins</groupId>
+            <artifactId>maven-failsafe-plugin</artifactId>
+            <version>2.22.0</version> <!-- {x-version-update;org.apache.maven.plugins:maven-failsafe-plugin;external_dependency} -->
+            <configuration>
+              <suiteXmlFiles>
+                <suiteXmlFile>src/test/resources/multi-region-testng.xml</suiteXmlFile>
+              </suiteXmlFiles>
+            </configuration>
+          </plugin>
+        </plugins>
+      </build>
+    </profile>
     <profile>
       <!-- integration tests, requires Cosmos DB endpoint -->
       <id>examples</id>
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientRetryPolicy.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ClientRetryPolicy.java
@@ -168,7 +168,7 @@ private ShouldRetryResult shouldRetryOnSessionNotAvailable() {
                     // on all locations, then don't retry the request
                     return ShouldRetryResult.noRetry();
                 } else {
-                    this.retryContext = new RetryContext(this.sessionTokenRetryCount - 1, this.sessionTokenRetryCount > 1);
+                    this.retryContext = new RetryContext(this.sessionTokenRetryCount , true);
                     return ShouldRetryResult.retryAfter(Duration.ZERO);
                 }
             } else {
@@ -177,7 +177,7 @@ private ShouldRetryResult shouldRetryOnSessionNotAvailable() {
                     // we have already tried this request on the write location
                     return ShouldRetryResult.noRetry();
                 } else {
-                    this.retryContext = new RetryContext(this.sessionTokenRetryCount - 1, false);
+                    this.retryContext = new RetryContext(0, false);
                     return ShouldRetryResult.retryAfter(Duration.ZERO);
                 }
             }
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
@@ -1856,9 +1856,6 @@ private Mono<ResourceResponse<Document>> readDocumentInternal(String documentLin
             Mono<RxDocumentServiceRequest> requestObs = addPartitionKeyInformation(request, null, null, options, collectionObs);
 
             return requestObs.flatMap(req -> {
-                if (retryPolicyInstance != null) {
-                    retryPolicyInstance.onBeforeSendRequest(request);
-                }
                 return this.read(request, retryPolicyInstance).map(serviceResponse -> toResourceResponse(serviceResponse, Document.class));
             });
 
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentProducer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentProducer.java
@@ -136,10 +136,15 @@ public DocumentProducer(
                 } catch (Exception e) {
                     return Mono.error(e);
                 }
-                retryPolicy.onBeforeSendRequest(request);
             }
+
+            DocumentClientRetryPolicy finalRetryPolicy = retryPolicy;
             return ObservableHelper.inlineIfPossibleAsObs(
                     () -> {
+                        if(finalRetryPolicy != null) {
+                            finalRetryPolicy.onBeforeSendRequest(request);
+                        }
+
                         ++retries;
                         return executeRequestFunc.apply(request);
                     }, retryPolicy);
diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/SessionNotAvailableRetryTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/SessionNotAvailableRetryTest.java
diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/TestSuiteBase.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/TestSuiteBase.java
@@ -134,7 +134,7 @@ public Mono<ResourceResponse<Database>> deleteDatabase(String id) {
         }
     }
 
-    @BeforeSuite(groups = {"simple", "long", "direct", "multi-master", "emulator", "non-emulator"}, timeOut = SUITE_SETUP_TIMEOUT)
+    @BeforeSuite(groups = {"simple", "long", "direct", "multi-region", "multi-master", "emulator", "non-emulator"}, timeOut = SUITE_SETUP_TIMEOUT)
     public static void beforeSuite() {
         logger.info("beforeSuite Started");
         AsyncDocumentClient houseKeepingClient = createGatewayHouseKeepingDocumentClient().build();
@@ -151,7 +151,7 @@ public static void beforeSuite() {
         }
     }
 
-    @AfterSuite(groups = {"simple", "long", "direct", "multi-master", "emulator", "non-emulator"}, timeOut = SUITE_SHUTDOWN_TIMEOUT)
+    @AfterSuite(groups = {"simple", "long", "direct", "multi-region", "multi-master", "emulator", "non-emulator"}, timeOut = SUITE_SHUTDOWN_TIMEOUT)
     public static void afterSuite() {
         logger.info("afterSuite Started");
         AsyncDocumentClient houseKeepingClient = createGatewayHouseKeepingDocumentClient().build();
diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java
@@ -9,7 +9,6 @@
 import com.azure.cosmos.CosmosClientBuilder;
 import com.azure.cosmos.implementation.AsyncDocumentClient;
 import com.azure.cosmos.implementation.ConnectionPolicy;
-import com.azure.cosmos.implementation.DatabaseAccount;
 import com.azure.cosmos.implementation.GlobalEndpointManager;
 import com.azure.cosmos.implementation.RxDocumentClientImpl;
 import com.azure.cosmos.implementation.RxStoreModel;
@@ -191,4 +190,27 @@ public static void setServerStoreModel (RxDocumentClientImpl client, RxStoreMode
         set(client, storeModel, "storeModel");
     }
 
+    public static ReplicatedResourceClient getReplicatedResourceClient(StoreClient storeClient) {
+        return get(ReplicatedResourceClient.class, storeClient, "replicatedResourceClient");
+    }
+
+    public static ConsistencyReader getConsistencyReader(ReplicatedResourceClient replicatedResourceClient) {
+        return get(ConsistencyReader.class, replicatedResourceClient, "consistencyReader");
+    }
+
+    public static ConsistencyWriter getConsistencyWriter(ReplicatedResourceClient replicatedResourceClient) {
+        return get(ConsistencyWriter.class, replicatedResourceClient, "consistencyWriter");
+    }
+
+    public static StoreReader getStoreReader(ConsistencyReader consistencyReader) {
+        return get(StoreReader.class, consistencyReader, "storeReader");
+    }
+
+    public static void setTransportClient(StoreReader storeReader, TransportClient transportClient) {
+        set(storeReader, transportClient, "transportClient");
+    }
+
+    public static void setTransportClient(ConsistencyWriter consistencyWriter, TransportClient transportClient) {
+        set(consistencyWriter, transportClient, "transportClient");
+    }
 }
diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java
@@ -195,7 +195,7 @@ public CosmosAsyncDatabase getDatabase(String id) {
         }
     }
 
-    @BeforeSuite(groups = {"simple", "long", "direct", "multi-master", "emulator", "non-emulator"}, timeOut = SUITE_SETUP_TIMEOUT)
+    @BeforeSuite(groups = {"simple", "long", "direct", "multi-region", "multi-master", "emulator", "non-emulator"}, timeOut = SUITE_SETUP_TIMEOUT)
     public static void beforeSuite() {
 
         logger.info("beforeSuite Started");
@@ -211,7 +211,7 @@ public static void beforeSuite() {
         }
     }
 
-    @AfterSuite(groups = {"simple", "long", "direct", "multi-master", "emulator", "non-emulator"}, timeOut = SUITE_SHUTDOWN_TIMEOUT)
+    @AfterSuite(groups = {"simple", "long", "direct", "multi-region", "multi-master", "emulator", "non-emulator"}, timeOut = SUITE_SHUTDOWN_TIMEOUT)
     public static void afterSuite() {
 
         logger.info("afterSuite Started");
diff --git a/sdk/cosmos/azure-cosmos/src/test/resources/multi-region-testng.xml b/sdk/cosmos/azure-cosmos/src/test/resources/multi-region-testng.xml
@@ -0,0 +1,35 @@
+<!--
+  ~ The MIT License (MIT)
+  ~ Copyright (c) 2018 Microsoft Corporation
+  ~
+  ~ Permission is hereby granted, free of charge, to any person obtaining a copy
+  ~ of this software and associated documentation files (the "Software"), to deal
+  ~ in the Software without restriction, including without limitation the rights
+  ~ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+  ~ copies of the Software, and to permit persons to whom the Software is
+  ~ furnished to do so, subject to the following conditions:
+  ~
+  ~ The above copyright notice and this permission notice shall be included in all
+  ~ copies or substantial portions of the Software.
+  ~
+  ~ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+  ~ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+  ~ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+  ~ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+  ~ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+  ~ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+  ~ SOFTWARE.
+  -->
+<!DOCTYPE suite SYSTEM "http://testng.org/testng-1.0.dtd">
+<suite name="multi-region">
+  <test name="multi-region" group-by-instances="true">
+    <groups>
+      <run>
+        <include name="multi-region"/>
+      </run>
+    </groups>
+    <packages>
+      <package name="com.azure.cosmos.*"/>
+    </packages>
+  </test>
+</suite>
diff --git a/sdk/cosmos/tests.yml b/sdk/cosmos/tests.yml
@@ -248,6 +248,17 @@ jobs:
           ACCOUNT_CONSISTENCY: 'Session'
           DESIRED_CONSISTENCIES: '["Session"]'
           PREFERRED_LOCATIONS: null
+        # 15m 10s, expect passed
+        Session_Single_Master_Multi_Region_Tcp_MultiRegion:
+          ArmTemplateParameters: "@{ enableMultipleWriteLocations = $false; defaultConsistencyLevel = 'Session'; enableMultipleRegions = $true}"
+          DisplayName: Single Master Multi Region Multi-Region
+          Pool: azsdk-pool-mms-ubuntu-1804-general
+          OSVmImage: MMSUbuntu18.04
+          ProfileFlag: '-Pmulti-region'
+          PROTOCOLS: '["Tcp"]'
+          ACCOUNT_CONSISTENCY: 'Session'
+          DESIRED_CONSISTENCIES: '["Session"]'
+          PREFERRED_LOCATIONS: null
 
       PreRunSteps:
         - template: /eng/pipelines/templates/steps/install-reporting-tools.yml

Original file line number	Diff line number	Diff line change
`@@ -168,7 +168,7 @@ private ShouldRetryResult shouldRetryOnSessionNotAvailable() {`
`168`	`168`	`// on all locations, then don't retry the request`
`169`	`169`	`return ShouldRetryResult.noRetry();`
`170`	`170`	`} else {`
`171`		`- this.retryContext = new RetryContext(this.sessionTokenRetryCount - 1, this.sessionTokenRetryCount > 1);`
	`171`	`+ this.retryContext = new RetryContext(this.sessionTokenRetryCount , true);`
`172`	`172`	`return ShouldRetryResult.retryAfter(Duration.ZERO);`
`173`	`173`	`}`
`174`	`174`	`} else {`
`@@ -177,7 +177,7 @@ private ShouldRetryResult shouldRetryOnSessionNotAvailable() {`
`177`	`177`	`// we have already tried this request on the write location`
`178`	`178`	`return ShouldRetryResult.noRetry();`
`179`	`179`	`} else {`
`180`		`- this.retryContext = new RetryContext(this.sessionTokenRetryCount - 1, false);`
	`180`	`+ this.retryContext = new RetryContext(0, false);`
`181`	`181`	`return ShouldRetryResult.retryAfter(Duration.ZERO);`
`182`	`182`	`}`
`183`	`183`	`}`
Original file line number	Diff line number	Diff line change
`@@ -134,7 +134,7 @@ public Mono<ResourceResponse<Database>> deleteDatabase(String id) {`
`134`	`134`	`}`
`135`	`135`	`}`
`136`	`136`
`137`		`- @BeforeSuite(groups = {"simple", "long", "direct", "multi-master", "emulator", "non-emulator"}, timeOut = SUITE_SETUP_TIMEOUT)`
	`137`	`+ @BeforeSuite(groups = {"simple", "long", "direct", "multi-region", "multi-master", "emulator", "non-emulator"}, timeOut = SUITE_SETUP_TIMEOUT)`
`138`	`138`	`public static void beforeSuite() {`
`139`	`139`	`logger.info("beforeSuite Started");`
`140`	`140`	`AsyncDocumentClient houseKeepingClient = createGatewayHouseKeepingDocumentClient().build();`
`@@ -151,7 +151,7 @@ public static void beforeSuite() {`
`151`	`151`	`}`
`152`	`152`	`}`
`153`	`153`
`154`		`- @AfterSuite(groups = {"simple", "long", "direct", "multi-master", "emulator", "non-emulator"}, timeOut = SUITE_SHUTDOWN_TIMEOUT)`
	`154`	`+ @AfterSuite(groups = {"simple", "long", "direct", "multi-region", "multi-master", "emulator", "non-emulator"}, timeOut = SUITE_SHUTDOWN_TIMEOUT)`
`155`	`155`	`public static void afterSuite() {`
`156`	`156`	`logger.info("afterSuite Started");`
`157`	`157`	`AsyncDocumentClient houseKeepingClient = createGatewayHouseKeepingDocumentClient().build();`
Original file line number	Diff line number	Diff line change
`@@ -195,7 +195,7 @@ public CosmosAsyncDatabase getDatabase(String id) {`
`195`	`195`	`}`
`196`	`196`	`}`
`197`	`197`
`198`		`- @BeforeSuite(groups = {"simple", "long", "direct", "multi-master", "emulator", "non-emulator"}, timeOut = SUITE_SETUP_TIMEOUT)`
	`198`	`+ @BeforeSuite(groups = {"simple", "long", "direct", "multi-region", "multi-master", "emulator", "non-emulator"}, timeOut = SUITE_SETUP_TIMEOUT)`
`199`	`199`	`public static void beforeSuite() {`
`200`	`200`
`201`	`201`	`logger.info("beforeSuite Started");`
`@@ -211,7 +211,7 @@ public static void beforeSuite() {`
`211`	`211`	`}`
`212`	`212`	`}`
`213`	`213`
`214`		`- @AfterSuite(groups = {"simple", "long", "direct", "multi-master", "emulator", "non-emulator"}, timeOut = SUITE_SHUTDOWN_TIMEOUT)`
	`214`	`+ @AfterSuite(groups = {"simple", "long", "direct", "multi-region", "multi-master", "emulator", "non-emulator"}, timeOut = SUITE_SHUTDOWN_TIMEOUT)`
`215`	`215`	`public static void afterSuite() {`
`216`	`216`
`217`	`217`	`logger.info("afterSuite Started");`