HandlingForCancelledExceptions

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -9,6 +9,7 @@
 #### Bugs Fixed
 
 #### Other Changes
+* Fixed issue on noisy `CancellationException` log - See [PR 31882](https://github.com/Azure/azure-sdk-for-java/pull/31882)
 
 ### 4.39.0 (2022-11-16)
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/RntbdTransportClient.java

@@ -226,33 +226,43 @@ public Mono<StoreResponse> invokeStoreAsync(final Uri addressUri, final RxDocume
         final URI address = addressUri.getURI();
 
         final RntbdRequestArgs requestArgs = new RntbdRequestArgs(request, addressUri);
+
         final RntbdEndpoint endpoint = this.endpointProvider.get(address);
         final RntbdRequestRecord record = endpoint.request(requestArgs);
 
         final Context reactorContext = Context.of(KEY_ON_ERROR_DROPPED, onErrorDropHookWithReduceLogLevel);
 
-        final Mono<StoreResponse> result = Mono.fromFuture(record.whenComplete((response, throwable) -> {
+        // Since reactor-core 3.4.23, if the Mono.fromCompletionStage is cancelled, then it will also cancel the internal future
+        // If SDK has not sent the request to server, then SDK will not send the request to server
+        // If the request has been sent to server, then SDK will discard the response once get from server
+        return Mono.fromFuture(record).map(storeResponse -> {
             record.stage(RntbdRequestRecord.Stage.COMPLETED);
 
             if (request.requestContext.cosmosDiagnostics == null) {
                 request.requestContext.cosmosDiagnostics = request.createCosmosDiagnostics();
             }
 
-            if (response != null) {
-                RequestTimeline timeline = record.takeTimelineSnapshot();
-                response.setRequestTimeline(timeline);
-                response.setEndpointStatistics(record.serviceEndpointStatistics());
-                response.setRntbdResponseLength(record.responseLength());
-                response.setRntbdRequestLength(record.requestLength());
-                response.setRequestPayloadLength(request.getContentLength());
-                response.setRntbdChannelTaskQueueSize(record.channelTaskQueueLength());
-                response.setRntbdPendingRequestSize(record.pendingRequestQueueSize());
-                if(this.channelAcquisitionContextEnabled) {
-                    response.setChannelAcquisitionTimeline(record.getChannelAcquisitionTimeline());
-                }
+            RequestTimeline timeline = record.takeTimelineSnapshot();
+            storeResponse.setRequestTimeline(timeline);
+            storeResponse.setEndpointStatistics(record.serviceEndpointStatistics());
+            storeResponse.setRntbdResponseLength(record.responseLength());
+            storeResponse.setRntbdRequestLength(record.requestLength());
+            storeResponse.setRequestPayloadLength(request.getContentLength());
+            storeResponse.setRntbdChannelTaskQueueSize(record.channelTaskQueueLength());
+            storeResponse.setRntbdPendingRequestSize(record.pendingRequestQueueSize());
+            if (this.channelAcquisitionContextEnabled) {
+                storeResponse.setChannelAcquisitionTimeline(record.getChannelAcquisitionTimeline());
             }
 
-        })).onErrorMap(throwable -> {
+            return storeResponse;
+
+        }).onErrorMap(throwable -> {
+
+            record.stage(RntbdRequestRecord.Stage.COMPLETED);
+
+            if (request.requestContext.cosmosDiagnostics == null) {
+                request.requestContext.cosmosDiagnostics = request.createCosmosDiagnostics();
+            }
 
             Throwable error = throwable instanceof CompletionException ? throwable.getCause() : throwable;
 
@@ -285,67 +295,21 @@ public Mono<StoreResponse> invokeStoreAsync(final Uri addressUri, final RxDocume
             BridgeInternal.setRntbdPendingRequestQueueSize(cosmosException, record.pendingRequestQueueSize());
             BridgeInternal.setChannelTaskQueueSize(cosmosException, record.channelTaskQueueLength());
             BridgeInternal.setSendingRequestStarted(cosmosException, record.hasSendingRequestStarted());
-            if(this.channelAcquisitionContextEnabled) {
+            if (this.channelAcquisitionContextEnabled) {
                 BridgeInternal.setChannelAcquisitionTimeline(cosmosException, record.getChannelAcquisitionTimeline());
             }
 
             return cosmosException;
-        });
-
-        return result.doFinally(signalType -> {
-
-            // This lambda ensures that a pending Direct TCP request in a reactive stream dropped by an end user or the
-            // HA layer completes without bubbling up to reactor.core.publisher.Hooks#onErrorDropped as a
-            // CompletionException error. Pending requests may be left outstanding when, for example, an end user calls
-            // CosmosAsyncClient#close or the HA layer detects that a partition split has occurred. This code guarantees
-            // that each pending Mono<StoreResponse> in the stream will run to completion with a new subscriber.
-            // Consequently the default Hooks#onErrorDropped method will not be called thus preventing distracting error
-            // messages.
-            //
-            // This lambda does not prevent requests that complete exceptionally before the call to this lambda from
-            // bubbling up to Hooks#onErrorDropped as CompletionException errors. We will still see some onErrorDropped
-            // messages due to CompletionException errors. Anecdotal evidence shows that this is more likely to be seen
-            // in low latency environments on Azure cloud. To avoid the onErrorDropped events to get logged in the
-            // default hook (which logs with level ERROR) we inject a local hook in the Reactor Context to just log it
-            // as DEBUG level for the lifecycle of this Mono (safe here because we know the onErrorDropped doesn't have
-            // any functional issues.
-            //
-            // One might be tempted to complete a pending request here, but that is ill advised. Testing and
-            // inspection of the reactor code shows that this does not prevent errors from bubbling up to
-            // reactor.core.publisher.Hooks#onErrorDropped. Worse than this it has been seen to cause failures in
-            // the HA layer:
-            //
-            // * Calling record.cancel or record.completeExceptionally causes failures in (low-latency) cloud
-            //   environments and all errors bubble up Hooks#onErrorDropped.
-            //
-            // * Calling record.complete with a null value causes failures in all environments, depending on the
-            //   operation being performed. In short: many of our tests fail.
-
+        }).doFinally(signalType -> {
             if (signalType != SignalType.CANCEL) {
                 return;
             }
 
-            result.subscribe(
-                response -> {
-                    if (logger.isDebugEnabled()) {
-                        logger.debug(
-                            "received response to cancelled request: {\"request\":{},\"response\":{\"type\":{},"
-                                + "\"value\":{}}}}",
-                            RntbdObjectMapper.toJson(record),
-                            response.getClass().getSimpleName(),
-                            RntbdObjectMapper.toJson(response));
-                    }
-                },
-                throwable -> {
-                    if (logger.isDebugEnabled()) {
-                        logger.debug(
-                            "received response to cancelled request: {\"request\":{},\"response\":{\"type\":{},"
-                                + "\"value\":{}}}",
-                            RntbdObjectMapper.toJson(record),
-                            throwable.getClass().getSimpleName(),
-                            RntbdObjectMapper.toJson(throwable));
-                    }
-                });
+            // Since reactor-core 3.4.23, if the Mono.fromCompletionStage is cancelled, then it will also cancel the internal future
+            // But the stated behavior may change in later versions (https://github.com/reactor/reactor-core/issues/3235).
+            // In order to keep consistent behavior, we internally will always cancel the future.
+            record.cancel(true);
+
         }).contextWrite(reactorContext);
     }
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdMetrics.java

@@ -241,8 +241,11 @@ public void markComplete(RntbdRequestRecord requestRecord) {
         requestRecord.stop(this.requests, requestRecord.isCompletedExceptionally()
             ? this.responseErrors
             : this.responseSuccesses);
-        this.requestSize.record(requestRecord.requestLength());
-        this.responseSize.record(requestRecord.responseLength());
+
+        if (!requestRecord.isCancelled()) {
+            this.requestSize.record(requestRecord.requestLength());
+            this.responseSize.record(requestRecord.responseLength());
+        }
     }
 
     @Override

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestManager.java

@@ -66,6 +66,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 
 import static com.azure.cosmos.implementation.HttpConstants.StatusCodes;
 import static com.azure.cosmos.implementation.HttpConstants.SubStatusCodes;
@@ -583,12 +584,14 @@ public void write(final ChannelHandlerContext context, final Object message, fin
 
             record.setSendingRequestHasStarted();
 
-            context.write(this.addPendingRequestRecord(context, record), promise).addListener(completed -> {
-                record.stage(RntbdRequestRecord.Stage.SENT);
-                if (completed.isSuccess()) {
-                    this.timestamps.channelWriteCompleted();
-                }
-            });
+            if (!record.isCancelled()) {
+                context.write(this.addPendingRequestRecord(context, record), promise).addListener(completed -> {
+                    record.stage(RntbdRequestRecord.Stage.SENT);
+                    if (completed.isSuccess()) {
+                        this.timestamps.channelWriteCompleted();
+                    }
+                });
+            }
 
             return;
         }
@@ -664,30 +667,34 @@ Timestamps snapshotTimestamps() {
 
     private RntbdRequestRecord addPendingRequestRecord(final ChannelHandlerContext context, final RntbdRequestRecord record) {
 
-        return this.pendingRequests.compute(record.transportRequestId(), (id, current) -> {
+        AtomicReference<Timeout> pendingRequestTimeout = new AtomicReference<>();
+        this.pendingRequests.compute(record.transportRequestId(), (id, current) -> {
 
             reportIssueUnless(current == null, context, "id: {}, current: {}, request: {}", record);
             record.pendingRequestQueueSize(pendingRequests.size());
 
-            final Timeout pendingRequestTimeout = record.newTimeout(timeout -> {
+            pendingRequestTimeout.set(record.newTimeout(timeout -> {
 
                 // We don't wish to complete on the timeout thread, but rather on a thread doled out by our executor
                 requestExpirationExecutor.execute(record::expire);
-            });
-
-            record.whenComplete((response, error) -> {
-                this.pendingRequests.remove(id);
-                pendingRequestTimeout.cancel();
-            });
+            }));
 
             return record;
+        });
 
+        // NOTE: please do not put the following logic inside the compute block. It may cause dead lock when a record is cancelled early
+        record.whenComplete((response, error) -> {
+            this.pendingRequests.remove(record.transportRequestId());
+            if (pendingRequestTimeout.get() != null) {
+                pendingRequestTimeout.get().cancel();
+            }
         });
+
+        return record;
     }
 
     private void completeAllPendingRequestsExceptionally(
-        final ChannelHandlerContext context, final Throwable throwable
-    ) {
+        final ChannelHandlerContext context, final Throwable throwable) {
 
         reportIssueUnless(!this.closingExceptionally, context, "", throwable);
         this.closingExceptionally = true;

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

@@ -53,6 +53,7 @@
 import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdRequestTimer;
 import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdResponse;
 import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdResponseDecoder;
+import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdServiceEndpoint;
 import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdUUID;
 import com.azure.cosmos.implementation.guava25.base.Strings;
 import com.azure.cosmos.implementation.guava25.collect.ImmutableMap;
@@ -67,9 +68,11 @@
 import io.netty.handler.ssl.SslContextBuilder;
 import io.reactivex.subscribers.TestSubscriber;
 import org.apache.commons.lang3.StringUtils;
+import org.mockito.Mockito;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Schedulers;
 
 import java.io.IOException;
 import java.net.ConnectException;
@@ -89,6 +92,7 @@
 import static com.azure.cosmos.implementation.TestUtils.mockDiagnosticsClientContext;
 import static com.azure.cosmos.implementation.guava27.Strings.lenientFormat;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertTrue;
 import static org.testng.Assert.fail;
@@ -792,6 +796,36 @@ public void sslHandshakeTimeoutTests() throws IOException {
         }
     }
 
+    @Test(groups = "unit")
+    public void cancelRequestMono() throws InterruptedException {
+        RxDocumentServiceRequest request =
+            RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document);
+        RntbdRequestArgs requestArgs = new RntbdRequestArgs(request, physicalAddress);
+        RntbdRequestTimer requestTimer = new RntbdRequestTimer(5000, 5000);
+        RntbdRequestRecord rntbdRequestRecord = new AsyncRntbdRequestRecord(requestArgs, requestTimer);
+
+        RntbdEndpoint rntbdEndpoint = Mockito.mock(RntbdServiceEndpoint.class);
+        Mockito.when(rntbdEndpoint.request(any())).thenReturn(rntbdRequestRecord);
+
+        RntbdEndpoint.Provider endpointProvider = Mockito.mock(RntbdEndpoint.Provider.class);
+        Mockito.when(endpointProvider.get(physicalAddress.getURI())).thenReturn(rntbdEndpoint);
+
+        RntbdTransportClient transportClient = new RntbdTransportClient(endpointProvider);
+        transportClient
+            .invokeStoreAsync(
+                physicalAddress,
+                RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document))
+            .cancelOn(Schedulers.boundedElastic())
+            .subscribe()
+            .dispose();
+
+        // wait for the cancel signal to propagate
+        Thread.sleep(500);
+
+        assertThat(rntbdRequestRecord.isCancelled()).isTrue();
+        assertThat(rntbdRequestRecord.isCompletedExceptionally()).isTrue();
+    }
+
     private static RntbdTransportClient getRntbdTransportClientUnderTest(
         final UserAgentContainer userAgent,
         final ConnectionPolicy connectionPolicy,

sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestRecordTests.java

@@ -8,12 +8,16 @@
 import com.azure.cosmos.implementation.RequestTimeoutException;
 import com.azure.cosmos.implementation.ResourceType;
 import com.azure.cosmos.implementation.RxDocumentServiceRequest;
+import com.azure.cosmos.implementation.directconnectivity.StoreResponse;
 import com.azure.cosmos.implementation.directconnectivity.Uri;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
+import reactor.core.publisher.Mono;
+import reactor.core.scheduler.Schedulers;
 
 import java.net.URI;
 import java.net.URISyntaxException;
+import java.time.Duration;
 import java.util.concurrent.ExecutionException;
 
 import static com.azure.cosmos.implementation.TestUtils.mockDiagnosticsClientContext;
@@ -57,4 +61,24 @@ public void expireRecord(OperationType operationType, boolean requestSent, Class
             fail("Wrong exception");
         }
     }
+
+    @Test(groups = { "unit" })
+    public void cancelRecord() throws URISyntaxException, InterruptedException {
+
+        RntbdRequestArgs requestArgs = new RntbdRequestArgs(
+            RxDocumentServiceRequest.create(mockDiagnosticsClientContext(), OperationType.Read, ResourceType.Document),
+            new Uri(new URI("http://localhost/replica-path").toString())
+        );
+
+        RntbdRequestTimer requestTimer = new RntbdRequestTimer(5000, 5000);
+        RntbdRequestRecord record = new AsyncRntbdRequestRecord(requestArgs, requestTimer);
+        Mono<StoreResponse> result = Mono.fromFuture(record)
+            .doOnNext(storeResponse -> fail("Record got cancelled should not reach here"))
+            .doOnError(throwable -> fail("Record got cancelled should not reach here"));
+
+        result.cancelOn(Schedulers.boundedElastic()).subscribe().dispose();
+
+        Thread.sleep(100);
+        assertThat(record.isCancelled()).isTrue();
+    }
 }