Session retry options diagnostic improvement (#36711)

jeet1995 · web-flow · commit 0bfbd9f683d4 · 2023-11-08T09:56:55.000-05:00
* Wired SessionRetryOptions in diagnostics.

* Updated CHANGELOG.md

* Added maxInRegionRetryCount and minInRegionRetryTime to diagnostic string.

* Added maxInRegionRetryCount and minInRegionRetryTime to diagnostic string.

* Updated CHANGELOG.md.

* Updated CHANGELOG.md.

* Refactoring.

* Updated CHANGELOG.md.
diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ClientConfigDiagnosticsTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ClientConfigDiagnosticsTest.java
@@ -8,6 +8,9 @@
 import com.azure.cosmos.CosmosContainerProactiveInitConfigBuilder;
 import com.azure.cosmos.CosmosExcludedRegions;
 import com.azure.cosmos.DirectConnectionConfig;
+import com.azure.cosmos.CosmosRegionSwitchHint;
+import com.azure.cosmos.SessionRetryOptions;
+import com.azure.cosmos.SessionRetryOptionsBuilder;
 import com.azure.cosmos.implementation.directconnectivity.RntbdTransportClient;
 import com.azure.cosmos.implementation.guava25.collect.ImmutableList;
 import com.azure.cosmos.implementation.http.HttpClientConfig;
@@ -40,6 +43,9 @@ public class ClientConfigDiagnosticsTest {
     private static final ImplementationBridgeHelpers.CosmosContainerIdentityHelper.CosmosContainerIdentityAccessor containerIdentityAccessor = ImplementationBridgeHelpers
         .CosmosContainerIdentityHelper
         .getCosmosContainerIdentityAccessor();
+    private static final ImplementationBridgeHelpers.CosmosSessionRetryOptionsHelper.CosmosSessionRetryOptionsAccessor sessionRetryOptionsAccessor = ImplementationBridgeHelpers
+        .CosmosSessionRetryOptionsHelper
+        .getCosmosSessionRetryOptionsAccessor();
 
     @DataProvider(name = "proactiveContainerInitConfigProvider")
     public Object[][] proactiveContainerInitConfigProvider() {
@@ -88,6 +94,48 @@ public Object[][] proactiveContainerInitConfigProvider() {
         };
     }
 
+    @DataProvider(name = "sessionRetryOptionsConfigProvider")
+    public Object[][] sessionRetryOptionsConfigProvider() {
+
+        SessionRetryOptions sessionRetryOptionsWithLocalRegionPreferred = new SessionRetryOptionsBuilder()
+            .regionSwitchHint(CosmosRegionSwitchHint.LOCAL_REGION_PREFERRED)
+            .build();
+
+        SessionRetryOptions sessionRetryOptionsWithRemoteRegionPreferred = new SessionRetryOptionsBuilder()
+            .regionSwitchHint(CosmosRegionSwitchHint.REMOTE_REGION_PREFERRED)
+            .build();
+
+        SessionRetryOptions sessionRetryOptionsWithNoDefaults = new SessionRetryOptionsBuilder()
+            .regionSwitchHint(CosmosRegionSwitchHint.REMOTE_REGION_PREFERRED)
+            .minTimeoutPerRegion(Duration.ofSeconds(2))
+            .maxRetriesPerRegion(7)
+            .build();
+
+        return new Object[][] {
+            {
+                sessionRetryOptionsWithLocalRegionPreferred,
+                reconstructSessionRetryOptionsAsString(
+                    sessionRetryOptionsAccessor.getRegionSwitchHint(sessionRetryOptionsWithLocalRegionPreferred),
+                    sessionRetryOptionsAccessor.getMinInRegionRetryTime(sessionRetryOptionsWithLocalRegionPreferred),
+                    sessionRetryOptionsAccessor.getMaxInRegionRetryCount(sessionRetryOptionsWithLocalRegionPreferred))
+            },
+            {
+                sessionRetryOptionsWithRemoteRegionPreferred,
+                reconstructSessionRetryOptionsAsString(
+                    sessionRetryOptionsAccessor.getRegionSwitchHint(sessionRetryOptionsWithRemoteRegionPreferred),
+                    sessionRetryOptionsAccessor.getMinInRegionRetryTime(sessionRetryOptionsWithRemoteRegionPreferred),
+                    sessionRetryOptionsAccessor.getMaxInRegionRetryCount(sessionRetryOptionsWithRemoteRegionPreferred))
+            },
+            {
+                sessionRetryOptionsWithNoDefaults,
+                reconstructSessionRetryOptionsAsString(
+                    sessionRetryOptionsAccessor.getRegionSwitchHint(sessionRetryOptionsWithNoDefaults),
+                    sessionRetryOptionsAccessor.getMinInRegionRetryTime(sessionRetryOptionsWithNoDefaults),
+                    sessionRetryOptionsAccessor.getMaxInRegionRetryCount(sessionRetryOptionsWithNoDefaults))
+            }
+        };
+    }
+
     @Test(groups = { "unit" })
     public void bareMinimum() throws Exception {
         DiagnosticsClientContext clientContext = Mockito.mock(DiagnosticsClientContext.class);
@@ -238,11 +286,43 @@ public void full(
 
         String expectedProactiveInitConfigString = reconstructProactiveInitConfigString(cosmosContainerIdentities, aggressiveWarmupDuration, proactiveConnectionRegionCount);
 
-        assertThat(objectNode.get("proactiveInit").asText()).isEqualTo(expectedProactiveInitConfigString);
+        assertThat(objectNode.get("proactiveInitCfg").asText()).isEqualTo(expectedProactiveInitConfigString);
 
         System.clearProperty("COSMOS.REPLICA_ADDRESS_VALIDATION_ENABLED");
     }
 
+    @Test(groups = {"unit"}, dataProvider = "sessionRetryOptionsConfigProvider")
+    public void sessionRetryOptionsInDiagnostics(SessionRetryOptions sessionRetryOptions, String expectedSessionRetryOptionsAsString) throws Exception {
+        DiagnosticsClientContext clientContext = Mockito.mock(DiagnosticsClientContext.class);
+
+        DiagnosticsClientContext.DiagnosticsClientConfig diagnosticsClientConfig = new DiagnosticsClientContext.DiagnosticsClientConfig();
+        String machineId = "vmId:" + UUID.randomUUID();
+        diagnosticsClientConfig.withMachineId(machineId);
+        diagnosticsClientConfig.withClientId(1);
+        diagnosticsClientConfig.withConnectionMode(ConnectionMode.DIRECT);
+        diagnosticsClientConfig.withActiveClientCounter(new AtomicInteger(2));
+        diagnosticsClientConfig.withClientMap(new HashMap<>());
+        diagnosticsClientConfig.withSessionRetryOptions(sessionRetryOptions);
+
+        Mockito.doReturn(diagnosticsClientConfig).when(clientContext).getConfig();
+
+        StringWriter jsonWriter = new StringWriter();
+        JsonGenerator jsonGenerator = new JsonFactory().createGenerator(jsonWriter);
+        SerializerProvider serializerProvider = objectMapper.getSerializerProvider();
+        DiagnosticsClientContext.DiagnosticsClientConfigSerializer.INSTANCE.serialize(clientContext.getConfig(), jsonGenerator, serializerProvider);
+        jsonGenerator.flush();
+        ObjectNode objectNode = (ObjectNode) objectMapper.readTree(jsonWriter.toString());
+
+        assertThat(objectNode.get("id").asInt()).isEqualTo(1);
+        assertThat(objectNode.get("machineId").asText()).isEqualTo(machineId);
+        assertThat(objectNode.get("numberOfClients").asInt()).isEqualTo(2);
+        assertThat(objectNode.get("consistencyCfg").asText()).isEqualTo("(consistency: null, mm: false, prgns: [null])");
+        assertThat(objectNode.get("connCfg").get("rntbd").asText()).isEqualTo("null");
+        assertThat(objectNode.get("connCfg").get("gw").asText()).isEqualTo("null");
+        assertThat(objectNode.get("connCfg").get("other").asText()).isEqualTo("(ed: false, cs: false, rv: true)");
+        assertThat(objectNode.get("sessionRetryCfg").asText()).isEqualTo(expectedSessionRetryOptionsAsString);
+    }
+
     private static String reconstructProactiveInitConfigString(
         List<CosmosContainerIdentity> containerIdentities,
         Duration aggressiveWarmupDuration,
@@ -259,4 +339,13 @@ private static String reconstructProactiveInitConfigString(
             proactiveConnectionRegionCount,
             aggressiveWarmupDuration);
     }
+
+    private static String reconstructSessionRetryOptionsAsString(CosmosRegionSwitchHint regionSwitchHint, Duration minInRegionRetryTime, int maxInRegionRetryCount) {
+        return String.format(
+            "(rsh:%s, minrrt:%s, maxrrc:%s)",
+            regionSwitchHint == CosmosRegionSwitchHint.REMOTE_REGION_PREFERRED ? "REMOTE_REGION_PREFERRED" : "LOCAL_REGION_PREFERRED",
+            minInRegionRetryTime.toString(),
+            maxInRegionRetryCount
+        );
+    }
 }
diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md
@@ -11,6 +11,7 @@
 * Fixed an issue where PartitionKeyRange request will not do cross region retry. - See [PR 37403](https://github.com/Azure/azure-sdk-for-java/pull/37403)
 
 #### Other Changes
+* Added `sessionRetryCfg` to the diagnostic string and modified `proactiveInit` key name to `proactiveInitCfg` in the diagnostic string. - See [PR 36711](https://github.com/Azure/azure-sdk-for-java/pull/36711)
 
 ### 4.52.0 (2023-10-24)
 #### Features Added
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/SessionRetryOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/SessionRetryOptions.java
@@ -3,7 +3,6 @@
 
 package com.azure.cosmos;
 
-import com.azure.cosmos.implementation.Configs;
 import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
 
 import java.time.Duration;
@@ -17,8 +16,8 @@ public final class SessionRetryOptions {
 
     private final CosmosRegionSwitchHint regionSwitchHint;
     private final Duration minInRegionRetryTime;
-
     private final int maxInRegionRetryCount;
+    private final String sessionRetryOptionsAsString;
 
     /**
      * Instantiates {@link SessionRetryOptions}
@@ -29,6 +28,42 @@ public final class SessionRetryOptions {
         this.regionSwitchHint = regionSwitchHint;
         this.minInRegionRetryTime = minInRegionRetryTime ;
         this.maxInRegionRetryCount = maxInRegionRetryCount;
+        this.sessionRetryOptionsAsString = sessionRetryOptionsAsString(this.regionSwitchHint, this.minInRegionRetryTime, this.maxInRegionRetryCount);
+    }
+
+    @Override
+    public String toString() {
+        return this.sessionRetryOptionsAsString;
+    }
+
+    private static String regionSwitchHintToString(CosmosRegionSwitchHint regionSwitchHint) {
+
+        String regionSwitchHintAsString = "";
+
+        if (regionSwitchHint == null) {
+            return regionSwitchHintAsString;
+        }
+
+        if (regionSwitchHint == CosmosRegionSwitchHint.REMOTE_REGION_PREFERRED) {
+            regionSwitchHintAsString = "REMOTE_REGION_PREFERRED";
+        }
+
+        if (regionSwitchHint == CosmosRegionSwitchHint.LOCAL_REGION_PREFERRED) {
+            regionSwitchHintAsString = "LOCAL_REGION_PREFERRED";
+        }
+
+        return regionSwitchHintAsString;
+    }
+
+    private static String sessionRetryOptionsAsString(CosmosRegionSwitchHint regionSwitchHint,
+                                                      Duration minInRegionRetryTime,
+                                                      int maxInRegionRetryCount) {
+        return String.format(
+            "(rsh:%s, minrrt:%s, maxrrc:%s)",
+            regionSwitchHintToString(regionSwitchHint),
+            minInRegionRetryTime.toString(),
+            maxInRegionRetryCount
+        );
     }
 
     static void initialize() {
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DiagnosticsClientContext.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DiagnosticsClientContext.java
@@ -8,6 +8,7 @@
 import com.azure.cosmos.CosmosContainerProactiveInitConfig;
 import com.azure.cosmos.CosmosDiagnostics;
 import com.azure.cosmos.CosmosEndToEndOperationLatencyPolicyConfig;
+import com.azure.cosmos.SessionRetryOptions;
 import com.azure.cosmos.implementation.clienttelemetry.ClientTelemetry;
 import com.azure.cosmos.implementation.guava27.Strings;
 import com.fasterxml.jackson.core.JsonGenerator;
@@ -78,8 +79,9 @@ public void serialize(DiagnosticsClientConfig clientConfig, JsonGenerator genera
                 }
                 generator.writeEndObject();
                 generator.writeStringField("consistencyCfg", clientConfig.consistencyRelatedConfig());
-                generator.writeStringField("proactiveInit", clientConfig.proactivelyInitializedContainersAsString);
+                generator.writeStringField("proactiveInitCfg", clientConfig.proactivelyInitializedContainersAsString);
                 generator.writeStringField("e2ePolicyCfg", clientConfig.endToEndOperationLatencyPolicyConfigAsString);
+                generator.writeStringField("sessionRetryCfg", clientConfig.sessionRetryOptionsAsString);
             } catch (Exception e) {
                 logger.debug("unexpected failure", e);
             }
@@ -110,6 +112,7 @@ class DiagnosticsClientConfig {
         private String machineId;
         private boolean replicaValidationEnabled = Configs.isReplicaAddressValidationEnabled();
         private ConnectionPolicy connectionPolicy;
+        private String sessionRetryOptionsAsString;
 
         public DiagnosticsClientConfig withMachineId(String machineId) {
             this.machineId = machineId;
@@ -207,6 +210,16 @@ public DiagnosticsClientConfig withConnectionMode(ConnectionMode connectionMode)
             return this;
         }
 
+        public DiagnosticsClientConfig withSessionRetryOptions(SessionRetryOptions sessionRetryOptions) {
+            if (sessionRetryOptions == null) {
+                this.sessionRetryOptionsAsString = "";
+            } else {
+                this.sessionRetryOptionsAsString = sessionRetryOptions.toString();
+            }
+
+            return this;
+        }
+
         public ConnectionMode getConnectionMode() {
             return connectionMode;
         }
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
@@ -484,6 +484,7 @@ private RxDocumentClientImpl(URI serviceEndpoint,
             this.diagnosticsClientConfig.withPreferredRegions(this.connectionPolicy.getPreferredRegions());
             this.diagnosticsClientConfig.withMachineId(tempMachineId);
             this.diagnosticsClientConfig.withProactiveContainerInitConfig(containerProactiveInitConfig);
+            this.diagnosticsClientConfig.withSessionRetryOptions(sessionRetryOptions);
 
             boolean disableSessionCapturing = (ConsistencyLevel.SESSION != consistencyLevel && !sessionCapturingOverrideEnabled);
 
