Handle negative e2e timeout more gracefully by throwing CosmsoException instead of IllegalArgumentException (Azure#36507)

* Handling negative e2e timeout gracefully by throwin OperationCancelledException instead of IllegalArguemntException when creating the config

* Update CosmosDiagnosticsTest.java

* Update CHANGELOG.md

* Update CosmosDiagnosticsTest.java

* Update CosmosDiagnosticsTest.java

Author: FabianMeiswinkel

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosDiagnosticsTest.java

@@ -1405,6 +1405,79 @@ public void responseStatisticRequestStartTimeUTCForDirectCall() {
         }
     }
 
+    @Test(groups = {"emulator"}, timeOut = TIMEOUT)
+    public void negativeE2ETimeoutWithPointOperation() {
+        CosmosAsyncClient client = null;
+        String databaseId = DatabaseForTest.generateId();
+
+        try {
+            client = new CosmosClientBuilder()
+                .key(TestConfigurations.MASTER_KEY)
+                .endpoint(TestConfigurations.HOST)
+                .buildAsyncClient();
+
+            createDatabase(client, databaseId);
+            CosmosAsyncContainer container = createCollection(client, databaseId, getCollectionDefinition());
+
+            TestItem testItem = TestItem.createNewItem();
+            CosmosItemRequestOptions requestOptions = new CosmosItemRequestOptions();
+            requestOptions.setCosmosEndToEndOperationLatencyPolicyConfig(
+                new CosmosEndToEndOperationLatencyPolicyConfigBuilder(Duration.ofSeconds(-1)).build()
+            );
+            container.createItem(testItem, requestOptions).block();
+            fail("This should have failed with an exception");
+        } catch(OperationCancelledException cancelledException) {
+            assertThat(cancelledException).isNotNull();
+            assertThat(cancelledException.getStatusCode()).isEqualTo(408);
+            assertThat(cancelledException.getSubStatusCode())
+                .isEqualTo(HttpConstants.SubStatusCodes.NEGATIVE_TIMEOUT_PROVIDED);
+            assertThat(cancelledException.getDiagnostics()).isNotNull();
+            logger.info("Expected request timeout: ", cancelledException);
+        }
+        finally {
+            safeClose(client);
+        }
+    }
+
+    @Test(groups = {"emulator"}, timeOut = TIMEOUT)
+    public void negativeE2ETimeoutWithQueryOperation() {
+        CosmosAsyncClient client = null;
+        String databaseId = DatabaseForTest.generateId();
+
+        try {
+            client = new CosmosClientBuilder()
+                .key(TestConfigurations.MASTER_KEY)
+                .endpoint(TestConfigurations.HOST)
+                .buildAsyncClient();
+
+            createDatabase(client, databaseId);
+            CosmosAsyncContainer container = createCollection(client, databaseId, getCollectionDefinition());
+
+            TestItem testItem = TestItem.createNewItem();
+            container.createItem(testItem).block();
+            
+            CosmosQueryRequestOptions requestOptions = new CosmosQueryRequestOptions();
+            requestOptions.setCosmosEndToEndOperationLatencyPolicyConfig(
+                new CosmosEndToEndOperationLatencyPolicyConfigBuilder(Duration.ofSeconds(-1)).build()
+            );
+            CosmosPagedFlux<ObjectNode> flux = container.readAllItems(requestOptions, ObjectNode.class);
+            List<ObjectNode> results = flux.collectList().block();
+                
+            fail("This should have failed with an exception");
+        } catch(OperationCancelledException cancelledException) {
+            assertThat(cancelledException).isNotNull();
+            assertThat(cancelledException.getStatusCode()).isEqualTo(408);
+            assertThat(cancelledException.getSubStatusCode())
+                .isEqualTo(HttpConstants.SubStatusCodes.NEGATIVE_TIMEOUT_PROVIDED);
+            // No pending requests - so, diagnostics are not guaranteed to be there
+            // assertThat(cancelledException.getDiagnostics()).isNotNull();
+            logger.info("Expected request timeout: ", cancelledException);
+        }
+        finally {
+            safeClose(client);
+        }
+    }
+
     private InternalObjectNode getInternalObjectNode() {
         InternalObjectNode internalObjectNode = new InternalObjectNode();
         String uuid = UUID.randomUUID().toString();

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -9,6 +9,7 @@
 #### Bugs Fixed
 
 #### Other Changes
+* Handling negative end-to-end timeouts provided more gracefully by throwing a `CosmsoException` (`OperationCancelledException`) instead of `IllegalArgumentException`. - See [PR 36507](https://github.com/Azure/azure-sdk-for-java/pull/36507)
 
 ### 4.49.0 (2023-08-21)
 

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

@@ -34,9 +34,7 @@ public CosmosEndToEndOperationLatencyPolicyConfig build() {
         if (endToEndOperationTimeout == null && isEnabled) {
             throw new IllegalArgumentException("endToEndOperationTimeout must be set if the policy is enabled");
         }
-        if (endToEndOperationTimeout != null && endToEndOperationTimeout.isNegative()) {
-            throw new IllegalArgumentException("endToEndOperationTimeout must be a positive Duration");
-        }
+
         return new CosmosEndToEndOperationLatencyPolicyConfig(isEnabled, endToEndOperationTimeout, availabilityStrategy);
     }
 

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

@@ -425,14 +425,18 @@ public static class SubStatusCodes {
         public static final int USER_REQUEST_RATE_TOO_LARGE = 3200;
 
         //SDK Codes(Client)
-        // IMPORTANT - whenever possible use consistency substaus codes that .Net SDK also uses
+        // IMPORTANT - whenever possible use consistency substatus codes that .Net SDK also uses
         public static final int TRANSPORT_GENERATED_410 = 20001;
         public static final int TIMEOUT_GENERATED_410 = 20002;
         // Client generated operation timeout exception
         public static final int CLIENT_OPERATION_TIMEOUT = 20008;
 
+        // IMPORTANT - below sub status codes have no corresponding .Net
+        // version, because they are only applicable in Java
+        public static final int NEGATIVE_TIMEOUT_PROVIDED = 20901; // .Net has different cancellation concept
+
         //SDK Codes (Server)
-        // IMPORTANT - whenever possible use consistency substaus codes that .Net SDK also uses
+        // IMPORTANT - whenever possible use consistency substatus codes that .Net SDK also uses
         public static final int NAME_CACHE_IS_STALE_EXCEEDED_RETRY_LIMIT = 21001;
         public static final int PARTITION_KEY_RANGE_GONE_EXCEEDED_RETRY_LIMIT = 21002;
         public static final int COMPLETING_SPLIT_EXCEEDED_RETRY_LIMIT = 21003;

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

@@ -90,6 +90,7 @@
 import java.net.URI;
 import java.net.URLEncoder;
 import java.nio.ByteBuffer;
+import java.time.Duration;
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -1001,57 +1002,86 @@ private <T> Flux<FeedResponse<T>> createQueryInternal(
         }, Queues.SMALL_BUFFER_SIZE, 1);
     }
 
+    private static void applyExceptionToMergedDiagnostics(
+        CosmosQueryRequestOptions requestOptions,
+        CosmosException exception) {
+
+        List<CosmosDiagnostics> cancelledRequestDiagnostics =
+            ImplementationBridgeHelpers
+                .CosmosQueryRequestOptionsHelper
+                .getCosmosQueryRequestOptionsAccessor()
+                .getCancelledRequestDiagnosticsTracker(requestOptions);
+
+        // if there is any cancelled requests, collect cosmos diagnostics
+        if (cancelledRequestDiagnostics != null && !cancelledRequestDiagnostics.isEmpty()) {
+            // combine all the cosmos diagnostics
+            CosmosDiagnostics aggregratedCosmosDiagnostics =
+                cancelledRequestDiagnostics
+                    .stream()
+                    .reduce((first, toBeMerged) -> {
+                        ClientSideRequestStatistics clientSideRequestStatistics =
+                            ImplementationBridgeHelpers
+                                .CosmosDiagnosticsHelper
+                                .getCosmosDiagnosticsAccessor()
+                                .getClientSideRequestStatisticsRaw(first);
+
+                        ClientSideRequestStatistics toBeMergedClientSideRequestStatistics =
+                            ImplementationBridgeHelpers
+                                .CosmosDiagnosticsHelper
+                                .getCosmosDiagnosticsAccessor()
+                                .getClientSideRequestStatisticsRaw(first);
+
+                        if (clientSideRequestStatistics == null) {
+                            return toBeMerged;
+                        } else {
+                            clientSideRequestStatistics.mergeClientSideRequestStatistics(toBeMergedClientSideRequestStatistics);
+                            return first;
+                        }
+                    })
+                    .get();
+
+            BridgeInternal.setCosmosDiagnostics(exception, aggregratedCosmosDiagnostics);
+        }
+    }
+
     private static <T> Flux<FeedResponse<T>> getFeedResponseFluxWithTimeout(
         Flux<FeedResponse<T>> feedResponseFlux,
         CosmosEndToEndOperationLatencyPolicyConfig endToEndPolicyConfig,
         CosmosQueryRequestOptions requestOptions,
         final AtomicBoolean isQueryCancelledOnTimeout) {
 
+        Duration endToEndTimeout = endToEndPolicyConfig.getEndToEndOperationTimeout();
+
+
+        Flux<FeedResponse<T>> flux;
+        if (endToEndTimeout.isNegative()) {
+            return feedResponseFlux
+                .timeout(endToEndTimeout)
+                .onErrorMap(throwable -> {
+                    if (throwable instanceof TimeoutException) {
+                        CosmosException cancellationException = getNegativeTimeoutException(null, endToEndTimeout);
+                        cancellationException.setStackTrace(throwable.getStackTrace());
+
+                        isQueryCancelledOnTimeout.set(true);
+
+                        applyExceptionToMergedDiagnostics(requestOptions, cancellationException);
+
+                        return cancellationException;
+                    }
+                    return throwable;
+                });
+        }
+
         return feedResponseFlux
-            .timeout(endToEndPolicyConfig.getEndToEndOperationTimeout())
+            .timeout(endToEndTimeout)
             .onErrorMap(throwable -> {
                 if (throwable instanceof TimeoutException) {
                     CosmosException exception = new OperationCancelledException();
                     exception.setStackTrace(throwable.getStackTrace());
 
                     isQueryCancelledOnTimeout.set(true);
 
-                    List<CosmosDiagnostics> cancelledRequestDiagnostics =
-                        ImplementationBridgeHelpers
-                            .CosmosQueryRequestOptionsHelper
-                            .getCosmosQueryRequestOptionsAccessor()
-                            .getCancelledRequestDiagnosticsTracker(requestOptions);
-
-                    // if there is any cancelled requests, collect cosmos diagnostics
-                    if (cancelledRequestDiagnostics != null && !cancelledRequestDiagnostics.isEmpty()) {
-                        // combine all the cosmos diagnostics
-                        CosmosDiagnostics aggregratedCosmosDiagnostics =
-                            cancelledRequestDiagnostics
-                                .stream()
-                                .reduce((first, toBeMerged) -> {
-                                    ClientSideRequestStatistics clientSideRequestStatistics =
-                                        ImplementationBridgeHelpers
-                                            .CosmosDiagnosticsHelper
-                                            .getCosmosDiagnosticsAccessor()
-                                            .getClientSideRequestStatisticsRaw(first);
-
-                                    ClientSideRequestStatistics toBeMergedClientSideRequestStatistics =
-                                        ImplementationBridgeHelpers
-                                            .CosmosDiagnosticsHelper
-                                            .getCosmosDiagnosticsAccessor()
-                                            .getClientSideRequestStatisticsRaw(first);
-
-                                    if (clientSideRequestStatistics == null) {
-                                        return toBeMerged;
-                                    } else {
-                                        clientSideRequestStatistics.mergeClientSideRequestStatistics(toBeMergedClientSideRequestStatistics);
-                                        return first;
-                                    }
-                                })
-                                .get();
-
-                        BridgeInternal.setCosmosDiagnostics(exception, aggregratedCosmosDiagnostics);
-                    }
+                    applyExceptionToMergedDiagnostics(requestOptions, exception);
 
                     return exception;
                 }
@@ -2022,9 +2052,15 @@ private static <T> Mono<T> getRxDocumentServiceResponseMonoWithE2ETimeout(RxDocu
                                                                                                   CosmosEndToEndOperationLatencyPolicyConfig endToEndPolicyConfig,
                                                                                                   Mono<T> rxDocumentServiceResponseMono) {
         if (endToEndPolicyConfig != null && endToEndPolicyConfig.isEnabled()) {
+
+            Duration endToEndTimeout = endToEndPolicyConfig.getEndToEndOperationTimeout();
+            if (endToEndTimeout.isNegative()) {
+                return Mono.error(getNegativeTimeoutException(request, endToEndTimeout));
+            }
+
             request.requestContext.setEndToEndOperationLatencyPolicyConfig(endToEndPolicyConfig);
             return rxDocumentServiceResponseMono
-                .timeout(endToEndPolicyConfig.getEndToEndOperationTimeout())
+                .timeout(endToEndTimeout)
                 .onErrorMap(throwable -> getCancellationException(request, throwable));
         }
         return rxDocumentServiceResponseMono;
@@ -2045,6 +2081,23 @@ private static Throwable getCancellationException(RxDocumentServiceRequest reque
         return throwable;
     }
 
+    private static CosmosException getNegativeTimeoutException(RxDocumentServiceRequest request, Duration negativeTimeout) {
+        checkNotNull(negativeTimeout, "Argument 'negativeTimeout' must not be null");
+        checkArgument(
+            negativeTimeout.isNegative(),
+            "This exception should only be used for negative timeouts");
+
+        String message = String.format("Negative timeout '%s' provided.",  negativeTimeout);
+        CosmosException exception = new OperationCancelledException(message, null);
+        BridgeInternal.setSubStatusCode(exception, HttpConstants.SubStatusCodes.NEGATIVE_TIMEOUT_PROVIDED);
+        if (request != null && request.requestContext != null) {
+            request.requestContext.setIsRequestCancelledOnTimeout(new AtomicBoolean(true));
+            return BridgeInternal.setCosmosDiagnostics(exception, request.requestContext.cosmosDiagnostics);
+        }
+
+        return exception;
+    }
+
     @Override
     public Mono<ResourceResponse<Document>> upsertDocument(String collectionLink, Object document,
                                                                  RequestOptions options, boolean disableAutomaticIdGeneration) {