Adding retry context on query diagnostics by redesigning retry context logic (Azure#20746)

* Adding request level diagnostic information for queries

* Removing unused variable

* Fixing spot bugs
Refactoring spotbug suppressions as the class ClientSideRequestStatistics is moved to implementation

* test fixes

* PR comments and minor improvements

* Fixing tests

* Retry context refactor and adding on query diagnostics

* fixing null pointer

* incremental changes

* adding test cases

* fixing warning

* cleaning code for PR

* fixing test case

* resolving comments and fixing test cases

* resolving comments

* resolving comments

* adding session failure test

* fixing session test

* adding null check again in BridgeInternal.java for non document resources

* adding null check for readablity uniformity

Co-authored-by: Bhaskar Mallapragada <bhaskar.mallapragada@microsoft.com>

Author: simplynaveen20

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java

@@ -21,6 +21,7 @@
 import com.azure.cosmos.implementation.RequestTimeline;
 import com.azure.cosmos.implementation.Resource;
 import com.azure.cosmos.implementation.ResourceResponse;
+import com.azure.cosmos.implementation.RetryContext;
 import com.azure.cosmos.implementation.RxDocumentServiceRequest;
 import com.azure.cosmos.implementation.RxDocumentServiceResponse;
 import com.azure.cosmos.implementation.SerializationDiagnosticsContext;
@@ -159,7 +160,6 @@ public static <T> FeedResponse<T> createFeedResponseWithQueryMetrics(
         boolean useEtagAsContinuation,
         boolean isNoChangesResponse,
         CosmosDiagnostics cosmosDiagnostics) {
-
         FeedResponse<T> feedResponseWithQueryMetrics = ModelBridgeInternal.createFeedResponseWithQueryMetrics(
             results,
             headers,
@@ -551,9 +551,8 @@ public static void recordResponse(CosmosDiagnostics cosmosDiagnostics,
     }
 
     @Warning(value = INTERNAL_USE_ONLY_WARNING)
-    public static void recordRetryContext(CosmosDiagnostics cosmosDiagnostics,
-                                          RxDocumentServiceRequest request) {
-        cosmosDiagnostics.clientSideRequestStatistics().recordRetryContext(request);
+    public static void recordRetryContextEndTime(CosmosDiagnostics cosmosDiagnostics) {
+        cosmosDiagnostics.clientSideRequestStatistics().recordRetryContextEndTime();
     }
 
     @Warning(value = INTERNAL_USE_ONLY_WARNING)
@@ -809,4 +808,13 @@ public static SqlQuerySpec getOfferQuerySpecFromResourceId(CosmosAsyncContainer
     public static CosmosAsyncContainer getControlContainerFromThroughputGlobalControlConfig(GlobalThroughputControlConfig globalControlConfig) {
         return globalControlConfig.getControlContainer();
     }
+
+    @Warning(value = INTERNAL_USE_ONLY_WARNING)
+    public static RetryContext getRetryContext(CosmosDiagnostics cosmosDiagnostics) {
+        if(cosmosDiagnostics != null && cosmosDiagnostics.clientSideRequestStatistics() != null) {
+            return cosmosDiagnostics.clientSideRequestStatistics().getRetryContext();
+        } else {
+            return null;
+        }
+    }
 }

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

@@ -37,6 +37,11 @@ public void onBeforeSendRequest(RxDocumentServiceRequest request) {
         this.retryPolicy.onBeforeSendRequest(request);
     }
 
+    @Override
+    public RetryContext getRetryContext() {
+        return this.retryPolicy.getRetryContext();
+    }
+
     @Override
     public Mono<ShouldRetryResult> shouldRetry(Exception e) {
 

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

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 package com.azure.cosmos.implementation;
 
+import com.azure.cosmos.BridgeInternal;
 import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList;
 import com.azure.cosmos.implementation.directconnectivity.WebExceptionUtility;
 import com.azure.cosmos.CosmosException;
@@ -52,15 +53,16 @@ public ClientRetryPolicy(DiagnosticsClientContext diagnosticsClientContext,
                              boolean enableEndpointDiscovery,
                              ThrottlingRetryOptions throttlingRetryOptions) {
 
-        this.throttlingRetry = new ResourceThrottleRetryPolicy(
-                throttlingRetryOptions.getMaxRetryAttemptsOnThrottledRequests(),
-                throttlingRetryOptions.getMaxRetryWaitTime());
         this.globalEndpointManager = globalEndpointManager;
         this.failoverRetryCount = 0;
         this.enableEndpointDiscovery = enableEndpointDiscovery;
         this.sessionTokenRetryCount = 0;
         this.canUseMultipleWriteLocations = false;
         this.cosmosDiagnostics = diagnosticsClientContext.createDiagnostics();
+        this.throttlingRetry = new ResourceThrottleRetryPolicy(
+            throttlingRetryOptions.getMaxRetryAttemptsOnThrottledRequests(),
+            throttlingRetryOptions.getMaxRetryWaitTime(),
+            BridgeInternal.getRetryContext(this.getCosmosDiagnostics()));
     }
 
     @Override
@@ -289,6 +291,16 @@ public void onBeforeSendRequest(RxDocumentServiceRequest request) {
             request.requestContext.routeToLocation(this.locationEndpoint);
         }
     }
+
+    @Override
+    public com.azure.cosmos.implementation.RetryContext getRetryContext() {
+        return BridgeInternal.getRetryContext(this.getCosmosDiagnostics());
+    }
+
+    CosmosDiagnostics getCosmosDiagnostics() {
+        return cosmosDiagnostics;
+    }
+
     private static class RetryContext {
 
         public int retryCount;

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

@@ -63,6 +63,7 @@ public ClientSideRequestStatistics(DiagnosticsClientContext diagnosticsClientCon
         this.connectionMode = ConnectionMode.DIRECT;
         this.metadataDiagnosticsContext = new MetadataDiagnosticsContext();
         this.serializationDiagnosticsContext = new SerializationDiagnosticsContext();
+        this.retryContext = new RetryContext();
     }
 
     public Duration getDuration() {
@@ -82,7 +83,6 @@ public void recordResponse(RxDocumentServiceRequest request, StoreResult storeRe
 
         URI locationEndPoint = null;
         if (request.requestContext != null) {
-            this.retryContext = new RetryContext(request.requestContext.retryContext);
             if (request.requestContext.locationEndpointToRoute != null) {
                 locationEndPoint = request.requestContext.locationEndpointToRoute;
             }
@@ -121,11 +121,8 @@ public void recordGatewayResponse(
             URI locationEndPoint = null;
             if (rxDocumentServiceRequest != null && rxDocumentServiceRequest.requestContext != null) {
                 locationEndPoint = rxDocumentServiceRequest.requestContext.locationEndpointToRoute;
-                if (rxDocumentServiceRequest.requestContext.retryContext != null) {
-                    rxDocumentServiceRequest.requestContext.retryContext.retryEndTime = Instant.now();
-                    this.retryContext = new RetryContext(rxDocumentServiceRequest.requestContext.retryContext);
-                }
             }
+            this.recordRetryContextEndTime();
 
             if (locationEndPoint != null) {
                 this.regionsContacted.add(locationEndPoint);
@@ -224,11 +221,12 @@ public SerializationDiagnosticsContext getSerializationDiagnosticsContext(){
         return this.serializationDiagnosticsContext;
     }
 
-    public void recordRetryContext(RxDocumentServiceRequest request) {
-        if(request.requestContext.retryContext != null) {
-            request.requestContext.retryContext.retryEndTime =  Instant.now();
-            this.retryContext = new RetryContext(request.requestContext.retryContext);
-        }
+    public void recordRetryContextEndTime() {
+        this.retryContext.updateEndTime();
+    }
+
+    public RetryContext getRetryContext() {
+        return retryContext;
     }
 
     public static class StoreResponseStatistics {

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

@@ -6,7 +6,7 @@
  * While this class is public, but it is not part of our published public APIs.
  * This is meant to be internally used only by our sdk.
  */
-public abstract class DocumentClientRetryPolicy extends RetryPolicyWithDiagnostics {
+public abstract class DocumentClientRetryPolicy implements IRetryPolicy {
 
     // TODO: this is just a place holder for now. As .Net has this method.
     // I have to spend more time figure out what's the right pattern for this (if method needed)

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

@@ -38,10 +38,8 @@ public class DocumentServiceRequestContext implements Cloneable {
     public volatile PartitionKeyInternal effectivePartitionKey;
     public volatile CosmosDiagnostics cosmosDiagnostics;
     public volatile String resourcePhysicalAddress;
-    public RetryContext retryContext;
 
     public DocumentServiceRequestContext() {
-        retryContext = new RetryContext();
     }
 
     /**
@@ -78,52 +76,6 @@ public void clearRouteToLocation() {
         this.usePreferredLocations = null;
     }
 
-    public void updateRetryContext(IRetryPolicy retryPolicy, boolean isGenericRetry) {
-        if (isGenericRetry) {
-            if (this.retryContext.directRetrySpecificStatusAndSubStatusCodes != null && this.retryContext.directRetrySpecificStatusAndSubStatusCodes.size() > 0) {
-                for (int i = this.retryContext.directRetrySpecificStatusAndSubStatusCodes.size() - 1; i >= 0; i--) {
-                    retryPolicy.incrementRetry();
-                    retryPolicy.addStatusAndSubStatusCode(0, this.retryContext.directRetrySpecificStatusAndSubStatusCodes.get(i)[0],
-                        this.retryContext.directRetrySpecificStatusAndSubStatusCodes.get(i)[1]);
-                }
-                this.retryContext.directRetrySpecificStatusAndSubStatusCodes.clear();
-            }
-
-            if (retryPolicy.getStatusAndSubStatusCodes() != null) {
-                this.retryContext.genericRetrySpecificStatusAndSubStatusCodes = Collections.synchronizedList(new ArrayList<>(retryPolicy.getStatusAndSubStatusCodes()));
-            } else {
-                this.retryContext.genericRetrySpecificStatusAndSubStatusCodes = Collections.synchronizedList(new ArrayList<>());
-            }
-            this.retryContext.retryCount = retryPolicy.getRetryCount();
-            this.retryContext.statusAndSubStatusCodes = retryPolicy.getStatusAndSubStatusCodes();
-            if (this.retryContext.retryStartTime == null) {
-                this.retryContext.retryStartTime = retryPolicy.getStartTime();
-            }
-            this.retryContext.retryEndTime = retryPolicy.getEndTime();
-        } else {
-            if (this.retryContext.genericRetrySpecificStatusAndSubStatusCodes != null && this.retryContext.genericRetrySpecificStatusAndSubStatusCodes.size() > 0) {
-                for (int i = this.retryContext.genericRetrySpecificStatusAndSubStatusCodes.size() - 1; i >= 0; i--) {
-                    retryPolicy.incrementRetry();
-                    retryPolicy.addStatusAndSubStatusCode(0, this.retryContext.genericRetrySpecificStatusAndSubStatusCodes.get(i)[0],
-                        this.retryContext.genericRetrySpecificStatusAndSubStatusCodes.get(i)[1]);
-                }
-                this.retryContext.genericRetrySpecificStatusAndSubStatusCodes.clear();
-            }
-
-            if (retryPolicy.getStatusAndSubStatusCodes() != null) {
-                this.retryContext.directRetrySpecificStatusAndSubStatusCodes = Collections.synchronizedList(new ArrayList<>(retryPolicy.getStatusAndSubStatusCodes()));
-            } else {
-                this.retryContext.directRetrySpecificStatusAndSubStatusCodes = Collections.synchronizedList(new ArrayList<>());
-            }
-            this.retryContext.retryCount = retryPolicy.getRetryCount();
-            this.retryContext.statusAndSubStatusCodes = retryPolicy.getStatusAndSubStatusCodes();
-            if (this.retryContext.retryStartTime == null) {
-                this.retryContext.retryStartTime = retryPolicy.getStartTime();
-            }
-            this.retryContext.retryEndTime = retryPolicy.getEndTime();
-        }
-    }
-
     @Override
     public DocumentServiceRequestContext clone() {
         DocumentServiceRequestContext context = new DocumentServiceRequestContext();
@@ -147,7 +99,6 @@ public DocumentServiceRequestContext clone() {
         context.performedBackgroundAddressRefresh = this.performedBackgroundAddressRefresh;
         context.cosmosDiagnostics = this.cosmosDiagnostics;
         context.resourcePhysicalAddress = this.resourcePhysicalAddress;
-        context.retryContext = new RetryContext(this.retryContext);
 
         return context;
     }

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

@@ -5,16 +5,12 @@
 
 import reactor.core.publisher.Mono;
 
-import java.time.Duration;
-import java.time.Instant;
-import java.util.List;
-
 // TODO update documentation
 /**
  * While this class is public, but it is not part of our published public APIs.
  * This is meant to be internally used only by our sdk.
  */
-public interface IRetryPolicy  {
+public interface IRetryPolicy {
     // this capture all the retry logic
     // TODO: design decision should this return a single or an observable?
 
@@ -26,21 +22,5 @@ public interface IRetryPolicy  {
     /// <returns>If the retry needs to be attempted or not</returns>
     Mono<ShouldRetryResult> shouldRetry(Exception e);
 
-    int getRetryCount();
-
-    void incrementRetry();
-
-    void captureStartTimeIfNotSet();
-
-    void updateEndTime();
-
-    Duration getRetryLatency();
-
-    Instant getStartTime();
-
-    Instant getEndTime();
-
-    void addStatusAndSubStatusCode(Integer index, int statusCode, int subStatusCode);
-
-    List<int[]> getStatusAndSubStatusCodes();
+    RetryContext getRetryContext();
 }

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

@@ -2,10 +2,13 @@
 // Licensed under the MIT License.
 package com.azure.cosmos.implementation;
 
+import com.azure.cosmos.CosmosDiagnostics;
+
 /**
  * While this class is public, but it is not part of our published public APIs.
  * This is meant to be internally used only by our sdk.
  */
 public interface IRetryPolicyFactory {
     DocumentClientRetryPolicy getRequestPolicy();
+    RetryContext getRetryContext();
 }

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

@@ -45,6 +45,15 @@ public void onBeforeSendRequest(RxDocumentServiceRequest request) {
         }
     }
 
+    @Override
+    public RetryContext getRetryContext() {
+        if (this.nextPolicy != null) {
+            return this.nextPolicy.getRetryContext();
+        } else {
+            return null;
+        }
+    }
+
     @Override
     public Mono<ShouldRetryResult> shouldRetry(Exception e) {
         CosmosException clientException = Utils.as(e, CosmosException.class);

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

@@ -92,4 +92,13 @@ public void onBeforeSendRequest(RxDocumentServiceRequest request) {
         this.request = request;
         this.nextRetryPolicy.onBeforeSendRequest(request);
     }
+
+    @Override
+    public RetryContext getRetryContext() {
+        if (this.nextRetryPolicy != null) {
+            return this.nextRetryPolicy.getRetryContext();
+        } else {
+            return null;
+        }
+    }
 }