Barrier strong e2 e follow

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxDocumentClientUnderTest.java

@@ -19,6 +19,8 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.function.BiFunction;
+import java.util.function.Function;
 
 import static org.mockito.Mockito.doAnswer;
 
@@ -64,7 +66,7 @@ public RxDocumentClientUnderTest(URI serviceEndpoint,
                 null,
                 false
         );
-        init(null, null);
+        init(null, null, null, null);
     }
 
     RxGatewayStoreModel createRxGatewayProxy(
@@ -93,6 +95,7 @@ RxGatewayStoreModel createRxGatewayProxy(
                 userAgentContainer,
                 globalEndpointManager,
                 spyHttpClient,
-                apiType);
+                apiType,
+                null);
     }
 }

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxGatewayStoreModelTest.java

@@ -99,7 +99,8 @@ public void readTimeout() throws Exception {
                 userAgentContainer,
                 globalEndpointManager,
                 httpClient,
-            null);
+                null,
+                null);
         storeModel.setGatewayServiceConfigurationReader(gatewayServiceConfigurationReader);
 
         RxDocumentServiceRequest dsr = RxDocumentServiceRequest.createFromName(clientContext,
@@ -143,6 +144,7 @@ public void serviceUnavailable() throws Exception {
             userAgentContainer,
             globalEndpointManager,
             httpClient,
+            null,
             null);
         storeModel.setGatewayServiceConfigurationReader(gatewayServiceConfigurationReader);
 
@@ -202,7 +204,8 @@ public void applySessionToken(
             new UserAgentContainer(),
             globalEndpointManager,
             httpClient,
-            apiType);
+            apiType,
+            null);
         storeModel.setGatewayServiceConfigurationReader(gatewayServiceConfigurationReader);
 
         RxDocumentServiceRequest dsr = RxDocumentServiceRequest.createFromName(
@@ -274,7 +277,8 @@ public void validateApiType() throws Exception {
             new UserAgentContainer(),
             globalEndpointManager,
             httpClient,
-            apiType);
+            apiType,
+            null);
 
         RxDocumentServiceRequest dsr = RxDocumentServiceRequest.createFromName(
             clientContext,

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/SpyClientUnderTestFactory.java

@@ -27,6 +27,8 @@
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Future;
+import java.util.function.BiFunction;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import static org.mockito.Mockito.doAnswer;
@@ -110,7 +112,7 @@ public static class ClientWithGatewaySpy extends SpyBaseClass<RxDocumentServiceR
                     credential,
                     contentResponseOnWriteEnabled,
                     clientTelemetryConfig);
-            init(null, null);
+            init(null, null, null, null);
             updateOrigRxGatewayStoreModel();
         }
 
@@ -126,15 +128,17 @@ RxGatewayStoreModel createRxGatewayProxy(ISessionContainer sessionContainer,
                                                  UserAgentContainer userAgentContainer,
                                                  GlobalEndpointManager globalEndpointManager,
                                                  HttpClient rxClient,
-                                                 ApiType apiType) {
+                                                 ApiType apiType,
+                                                 Function<RxDocumentServiceRequest, RxDocumentServiceResponse> httpRequestInterceptor) {
             this.origRxGatewayStoreModel = super.createRxGatewayProxy(
                 sessionContainer,
                 consistencyLevel,
                 queryCompatibilityMode,
                 userAgentContainer,
                 globalEndpointManager,
                 rxClient,
-                apiType);
+                apiType,
+                null);
             this.requests = Collections.synchronizedList(new ArrayList<>());
             this.spyRxGatewayStoreModel = Mockito.spy(this.origRxGatewayStoreModel);
             this.initRequestCapture();
@@ -201,7 +205,7 @@ public static class ClientUnderTest extends SpyBaseClass<HttpRequest> {
                     credential,
                     contentResponseOnWriteEnabled,
                     clientTelemetryConfig);
-            init(null, this::initHttpRequestCapture);
+            init(null, this::initHttpRequestCapture, null, null);
         }
 
         private Mono<HttpResponse> captureHttpRequest(InvocationOnMock invocationOnMock) {
@@ -287,7 +291,7 @@ public static class DirectHttpsClientUnderTest extends SpyBaseClass<HttpRequest>
                     contentResponseOnWriteEnabled,
                     clientTelemetryConfig);
             assert connectionPolicy.getConnectionMode() == ConnectionMode.DIRECT;
-            init(null, null);
+            init(null, null, null, null);
 
             this.origHttpClient = ReflectionUtils.getDirectHttpsHttpClient(this);
             this.spyHttpClient = spy(this.origHttpClient);

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/GlobalAddressResolverTest.java

@@ -110,7 +110,7 @@ public void resolveAsync() throws Exception {
 
         GlobalAddressResolver globalAddressResolver = new GlobalAddressResolver(mockDiagnosticsClientContext(), httpClient, endpointManager, Protocol.HTTPS, authorizationTokenProvider, collectionCache, routingMapProvider,
                 userAgentContainer,
-                serviceConfigReader, connectionPolicy, null);
+                serviceConfigReader, connectionPolicy, null, null);
         RxDocumentServiceRequest request;
         request = RxDocumentServiceRequest.createFromName(mockDiagnosticsClientContext(),
                 OperationType.Read,
@@ -145,6 +145,7 @@ public void submitOpenConnectionTasksAndInitCaches() {
                         userAgentContainer,
                         serviceConfigReader,
                         connectionPolicy,
+                        null,
                         null);
         GlobalAddressResolver.EndpointCache endpointCache = new GlobalAddressResolver.EndpointCache();
         GatewayAddressCache gatewayAddressCache = Mockito.mock(GatewayAddressCache.class);

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

@@ -19,6 +19,8 @@
 import com.azure.cosmos.implementation.QueryFeedOperationState;
 import com.azure.cosmos.implementation.RequestOptions;
 import com.azure.cosmos.implementation.ResourceType;
+import com.azure.cosmos.implementation.RxDocumentServiceRequest;
+import com.azure.cosmos.implementation.RxDocumentServiceResponse;
 import com.azure.cosmos.implementation.Strings;
 import com.azure.cosmos.implementation.WriteRetryPolicy;
 import com.azure.cosmos.implementation.clienttelemetry.ClientMetricsDiagnosticsHandler;
@@ -27,6 +29,7 @@
 import com.azure.cosmos.implementation.clienttelemetry.CosmosMeterOptions;
 import com.azure.cosmos.implementation.clienttelemetry.MetricCategory;
 import com.azure.cosmos.implementation.clienttelemetry.TagName;
+import com.azure.cosmos.implementation.directconnectivity.StoreResponse;
 import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdMetrics;
 import com.azure.cosmos.implementation.faultinjection.IFaultInjectorProvider;
 import com.azure.cosmos.implementation.throughputControl.sdk.config.SDKThroughputControlGroupInternal;
@@ -60,6 +63,7 @@
 import java.util.EnumSet;
 import java.util.List;
 import java.util.Objects;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -115,6 +119,8 @@ public final class CosmosAsyncClient implements Closeable {
     private final List<CosmosOperationPolicy> requestPolicies;
     private final CosmosItemSerializer defaultCustomSerializer;
     private final java.util.function.Function<CosmosAsyncContainer, CosmosAsyncContainer> containerFactory;
+    private final BiFunction<RxDocumentServiceRequest, StoreResponse, StoreResponse> storeResponseInterceptor;
+    private Function<RxDocumentServiceRequest, RxDocumentServiceResponse> httpRequestInterceptor;
 
     CosmosAsyncClient(CosmosClientBuilder builder) {
         // Async Cosmos client wrapper
@@ -134,6 +140,9 @@ public final class CosmosAsyncClient implements Closeable {
         this.nonIdempotentWriteRetryPolicy = builder.getNonIdempotentWriteRetryPolicy();
         this.requestPolicies = builder.getOperationPolicies();
         this.defaultCustomSerializer = builder.getCustomItemSerializer();
+        this.storeResponseInterceptor = builder.getStoreResponseInterceptor();
+        this.httpRequestInterceptor = builder.getHttpRequestInterceptor();
+
         if (builder.containerCreationInterceptor() != null) {
             this.containerFactory = builder.containerCreationInterceptor();
         } else {
@@ -185,6 +194,8 @@ public final class CosmosAsyncClient implements Closeable {
                                        .withDefaultSerializer(this.defaultCustomSerializer)
                                        .withRegionScopedSessionCapturingEnabled(builder.isRegionScopedSessionCapturingEnabled())
                                        .withPerPartitionAutomaticFailoverEnabled(builder.isPerPartitionAutomaticFailoverEnabled())
+                                       .withStoreResponseInterceptor(this.storeResponseInterceptor)
+                                       .withHttpRequestInterceptor(this.httpRequestInterceptor)
                                        .build();
 
         this.accountConsistencyLevel = this.asyncDocumentClient.getDefaultConsistencyLevelOfAccount();

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

@@ -13,13 +13,16 @@
 import com.azure.cosmos.implementation.ConnectionPolicy;
 import com.azure.cosmos.implementation.CosmosClientMetadataCachesSnapshot;
 import com.azure.cosmos.implementation.DiagnosticsProvider;
+import com.azure.cosmos.implementation.RxDocumentServiceRequest;
+import com.azure.cosmos.implementation.RxDocumentServiceResponse;
 import com.azure.cosmos.implementation.Strings;
 import com.azure.cosmos.implementation.WriteRetryPolicy;
 import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList;
 import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
 import com.azure.cosmos.implementation.apachecommons.lang.time.StopWatch;
+import com.azure.cosmos.implementation.directconnectivity.StoreResponse;
 import com.azure.cosmos.implementation.guava25.base.Preconditions;
-import com.azure.cosmos.implementation.perPartitionCircuitBreaker.PartitionLevelCircuitBreakerConfig;
+import com.azure.cosmos.implementation.http.HttpClient;
 import com.azure.cosmos.implementation.routing.LocationHelper;
 import com.azure.cosmos.models.CosmosAuthorizationTokenResolver;
 import com.azure.cosmos.models.CosmosClientTelemetryConfig;
@@ -39,6 +42,7 @@
 import java.util.Locale;
 import java.util.Objects;
 import java.util.Set;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
@@ -155,6 +159,8 @@ public class CosmosClientBuilder implements
     private boolean serverCertValidationDisabled = false;
 
     private Function<CosmosAsyncContainer, CosmosAsyncContainer> containerFactory = null;
+    private BiFunction<RxDocumentServiceRequest, StoreResponse, StoreResponse> storeResponseInterceptor = null;
+    private Function<RxDocumentServiceRequest, RxDocumentServiceResponse> httpRequestInterceptor = null;
 
     /**
      * Instantiates a new Cosmos client builder.
@@ -170,6 +176,25 @@ public CosmosClientBuilder() {
         this.requestPolicies = new LinkedList<>();
     }
 
+    CosmosClientBuilder httpRequestInterceptor(Function<RxDocumentServiceRequest, RxDocumentServiceResponse> httpRequestInterceptor) {
+        this.httpRequestInterceptor = httpRequestInterceptor;
+        return this;
+    }
+
+    Function<RxDocumentServiceRequest, RxDocumentServiceResponse> getHttpRequestInterceptor() {
+        return this.httpRequestInterceptor;
+    }
+
+    CosmosClientBuilder storeResponseInterceptor(
+            BiFunction<RxDocumentServiceRequest, StoreResponse, StoreResponse> storeResponseInterceptor) {
+        this.storeResponseInterceptor = storeResponseInterceptor;
+        return this;
+    }
+
+    BiFunction<RxDocumentServiceRequest, StoreResponse, StoreResponse> getStoreResponseInterceptor() {
+        return this.storeResponseInterceptor;
+    }
+
     CosmosClientBuilder metadataCaches(CosmosClientMetadataCachesSnapshot metadataCachesSnapshot) {
         this.state = metadataCachesSnapshot;
         return this;

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

@@ -12,6 +12,7 @@
 import com.azure.cosmos.implementation.Utils;
 import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
 import com.azure.cosmos.implementation.batch.BatchExecUtils;
+import com.azure.cosmos.implementation.directconnectivity.StoreResponse;
 import com.azure.cosmos.implementation.directconnectivity.Uri;
 import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdChannelAcquisitionTimeline;
 import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdChannelStatistics;
@@ -92,6 +93,8 @@ public class CosmosException extends AzureException {
      */
     private RntbdChannelStatistics rntbdChannelStatistics;
 
+    private StoreResponse interceptedStoreResponse;
+
     /**
      * LSN
      */
@@ -615,6 +618,14 @@ Map<String, Set<String>> getReplicaStatusList() {
         return this.replicaStatusList;
     }
 
+    public void setInterceptedStoreResponse(StoreResponse storeResponse) {
+        this.interceptedStoreResponse = storeResponse;
+    }
+
+    public StoreResponse getInterceptedStoreResponse() {
+        return this.interceptedStoreResponse;
+    }
+
     ///////////////////////////////////////////////////////////////////////////////////////////
     // the following helper/accessor only helps to access this class outside of this package.//
     ///////////////////////////////////////////////////////////////////////////////////////////

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

@@ -15,6 +15,8 @@
 import com.azure.cosmos.implementation.batch.ServerBatchRequest;
 import com.azure.cosmos.implementation.caches.RxClientCollectionCache;
 import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache;
+import com.azure.cosmos.implementation.directconnectivity.StoreResponse;
+import com.azure.cosmos.implementation.http.HttpClient;
 import com.azure.cosmos.implementation.perPartitionAutomaticFailover.GlobalPartitionEndpointManagerForPerPartitionAutomaticFailover;
 import com.azure.cosmos.implementation.perPartitionCircuitBreaker.GlobalPartitionEndpointManagerForPerPartitionCircuitBreaker;
 import com.azure.cosmos.implementation.clienttelemetry.ClientTelemetry;
@@ -42,6 +44,8 @@
 import java.net.URISyntaxException;
 import java.util.List;
 import java.util.Map;
+import java.util.function.BiFunction;
+import java.util.function.Function;
 
 /**
  * Provides a client-side logical representation of the Azure Cosmos DB
@@ -115,6 +119,8 @@ class Builder {
         private boolean isRegionScopedSessionCapturingEnabled;
         private boolean isPerPartitionAutomaticFailoverEnabled;
         private List<CosmosOperationPolicy> operationPolicies;
+        private BiFunction<RxDocumentServiceRequest, StoreResponse, StoreResponse> storeResponseInterceptor;
+        private Function<RxDocumentServiceRequest, RxDocumentServiceResponse> httpRequestInterceptor;
 
         public Builder withServiceEndpoint(String serviceEndpoint) {
             try {
@@ -287,6 +293,16 @@ public Builder withPerPartitionAutomaticFailoverEnabled(boolean isPerPartitionAu
             return this;
         }
 
+        public Builder withStoreResponseInterceptor(BiFunction<RxDocumentServiceRequest, StoreResponse, StoreResponse> storeResponseInterceptor) {
+            this.storeResponseInterceptor = storeResponseInterceptor;
+            return this;
+        }
+
+        public Builder withHttpRequestInterceptor(Function<RxDocumentServiceRequest, RxDocumentServiceResponse> httpRequestInterceptor) {
+            this.httpRequestInterceptor = httpRequestInterceptor;
+            return this;
+        }
+
         private void ifThrowIllegalArgException(boolean value, String error) {
             if (value) {
                 throw new IllegalArgumentException(error);
@@ -329,7 +345,7 @@ public AsyncDocumentClient build() {
                     operationPolicies,
                     isPerPartitionAutomaticFailoverEnabled);
 
-            client.init(state, null);
+            client.init(state, null, httpRequestInterceptor, storeResponseInterceptor);
             return client;
         }
 

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

@@ -3,6 +3,7 @@
 package com.azure.cosmos.implementation;
 
 import com.azure.cosmos.implementation.directconnectivity.AddressSelector;
+import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.util.retry.Retry;
@@ -16,7 +17,6 @@
  * This is meant to be internally used only by our sdk.
  */
 public class BackoffRetryUtility {
-
     // transforms a retryFunc to a function which can be used by Observable.retryWhen(.)
     // also it invokes preRetryCallback prior to doing retry.
     public static final Quadruple<Boolean, Boolean, Duration, Integer> InitialArgumentValuePolicyArg = Quadruple.with(false, false,
@@ -30,7 +30,6 @@ public class BackoffRetryUtility {
     // a helper method for invoking callback method given the retry policy
     static public <T> Mono<T> executeRetry(Callable<Mono<T>> callbackMethod,
                                            IRetryPolicy retryPolicy) {
-
         return Mono.defer(() -> {
             try {
                 return callbackMethod.call();