CosmosOperationCanceledException: Add CosmosDiagnostic info

Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs

@@ -119,58 +119,65 @@ public virtual async Task<ResponseMessage> SendAsync(
                 overallScope = diagnosticsContext.GetOverallScope();
             }
 
-            using (overallScope)
-            { 
-                HttpMethod method = RequestInvokerHandler.GetHttpMethod(operationType);
-                RequestMessage request = new RequestMessage(
-                        method,
-                        resourceUri,
-                        diagnosticsContext)
-                {
-                    OperationType = operationType,
-                    ResourceType = resourceType,
-                    RequestOptions = requestOptions,
-                    Content = streamPayload,
-                };
-
-                if (partitionKey.HasValue)
+            try
+            {
+                using (overallScope)
                 {
-                    if (cosmosContainerCore == null && object.ReferenceEquals(partitionKey, Cosmos.PartitionKey.None))
+                    HttpMethod method = RequestInvokerHandler.GetHttpMethod(operationType);
+                    RequestMessage request = new RequestMessage(
+                            method,
+                            resourceUri,
+                            diagnosticsContext)
                     {
-                        throw new ArgumentException($"{nameof(cosmosContainerCore)} can not be null with partition key as PartitionKey.None");
-                    }
-                    else if (partitionKey.Value.IsNone)
+                        OperationType = operationType,
+                        ResourceType = resourceType,
+                        RequestOptions = requestOptions,
+                        Content = streamPayload,
+                    };
+
+                    if (partitionKey.HasValue)
                     {
-                        using (diagnosticsContext.CreateScope("GetNonePkValue"))
+                        if (cosmosContainerCore == null && object.ReferenceEquals(partitionKey, Cosmos.PartitionKey.None))
                         {
-                            try
-                            {
-                                PartitionKeyInternal partitionKeyInternal = await cosmosContainerCore.GetNonePartitionKeyValueAsync(cancellationToken);
-                                request.Headers.PartitionKey = partitionKeyInternal.ToJsonString();
-                            }
-                            catch (DocumentClientException dce)
-                            {
-                                return dce.ToCosmosResponseMessage(request);
-                            }
-                            catch (CosmosException ce)
+                            throw new ArgumentException($"{nameof(cosmosContainerCore)} can not be null with partition key as PartitionKey.None");
+                        }
+                        else if (partitionKey.Value.IsNone)
+                        {
+                            using (diagnosticsContext.CreateScope("GetNonePkValue"))
                             {
-                                return ce.ToCosmosResponseMessage(request);
+                                try
+                                {
+                                    PartitionKeyInternal partitionKeyInternal = await cosmosContainerCore.GetNonePartitionKeyValueAsync(cancellationToken);
+                                    request.Headers.PartitionKey = partitionKeyInternal.ToJsonString();
+                                }
+                                catch (DocumentClientException dce)
+                                {
+                                    return dce.ToCosmosResponseMessage(request);
+                                }
+                                catch (CosmosException ce)
+                                {
+                                    return ce.ToCosmosResponseMessage(request);
+                                }
                             }
                         }
+                        else
+                        {
+                            request.Headers.PartitionKey = partitionKey.Value.ToJsonString();
+                        }
                     }
-                    else
+
+                    if (operationType == OperationType.Upsert)
                     {
-                        request.Headers.PartitionKey = partitionKey.Value.ToJsonString();
+                        request.Headers.IsUpsert = bool.TrueString;
                     }
-                }
 
-                if (operationType == OperationType.Upsert)
-                {
-                    request.Headers.IsUpsert = bool.TrueString;
+                    requestEnricher?.Invoke(request);
+                    return await this.SendAsync(request, cancellationToken);
                 }
-
-                requestEnricher?.Invoke(request);
-                return await this.SendAsync(request, cancellationToken);
+            }
+            catch (OperationCanceledException oe)
+            {
+                throw new CosmosOperationCanceledException(oe, diagnosticsContext);
             }
         }
 

Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs

@@ -137,11 +137,21 @@ public override async Task<ResponseMessage> ReadNextAsync(CancellationToken canc
             CosmosDiagnosticsContext diagnostics = CosmosDiagnosticsContext.Create(this.requestOptions);
             using (diagnostics.GetOverallScope())
             {
-                // This catches exception thrown by the pipeline and converts it to QueryResponse
-                QueryResponseCore responseCore = await this.cosmosQueryExecutionContext.ExecuteNextAsync(cancellationToken);
-
-                // This swaps the diagnostics in the context. This shows all the page reads between the previous ReadNextAsync and the current ReadNextAsync
-                diagnostics.AddDiagnosticsInternal(this.cosmosQueryContext.GetAndResetDiagnostics());
+                QueryResponseCore responseCore;
+                try
+                {
+                    // This catches exception thrown by the pipeline and converts it to QueryResponse
+                    responseCore = await this.cosmosQueryExecutionContext.ExecuteNextAsync(cancellationToken);
+                }
+                catch (OperationCanceledException ex) when (!(ex is CosmosOperationCanceledException))
+                {
+                    throw new CosmosOperationCanceledException(ex, diagnostics);
+                }
+                finally
+                {
+                    // This swaps the diagnostics in the context. This shows all the page reads between the previous ReadNextAsync and the current ReadNextAsync
+                    diagnostics.AddDiagnosticsInternal(this.cosmosQueryContext.GetAndResetDiagnostics());
+                }
 
                 if (responseCore.IsSuccess)
                 {

Microsoft.Azure.Cosmos/src/Resource/CosmosExceptions/CosmosOperationCanceledException.cs

@@ -0,0 +1,93 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos
+{
+    using System;
+    using System.Collections;
+    using System.Threading;
+
+    /// <summary>
+    /// The exception that is thrown in a thread upon cancellation of an operation that
+    ///  the thread was executing. This extends the OperationCanceledException to include the
+    ///  diagnostics of the operation that was canceled.
+    /// </summary>
+    public class CosmosOperationCanceledException : OperationCanceledException
+    {
+        private readonly OperationCanceledException originalException;
+
+        internal CosmosOperationCanceledException(
+            OperationCanceledException originalException,
+            CosmosDiagnosticsContext diagnosticsContext)
+            : this(
+                originalException,
+                diagnosticsContext?.Diagnostics)
+        {
+        }
+
+        /// <summary>
+        /// Create an instance of CosmosOperationCanceledException
+        /// </summary>
+        /// <param name="originalException">The original operation canceled exception</param>
+        /// <param name="diagnostics"></param>
+        public CosmosOperationCanceledException(
+            OperationCanceledException originalException,
+            CosmosDiagnostics diagnostics)
+            : base(originalException.CancellationToken)
+        {
+            if (originalException == null)
+            {
+                throw new ArgumentNullException(nameof(originalException));
+            }
+
+            if (diagnostics == null)
+            {
+                throw new ArgumentNullException(nameof(diagnostics));
+            }
+
+            this.originalException = originalException;
+            this.Diagnostics = diagnostics;
+        }
+
+        /// <inheritdoc/>
+        public override string Source
+        {
+            get => this.originalException.Source;
+            set => this.originalException.Source = value;
+        }
+
+        /// <inheritdoc/>
+        public override string Message => this.originalException.Message;
+
+        /// <inheritdoc/>
+        public override string StackTrace => this.originalException.StackTrace;
+
+        /// <inheritdoc/>
+        public override IDictionary Data => this.originalException.Data;
+
+        /// <summary>
+        /// Gets the diagnostics for the request
+        /// </summary>
+        public CosmosDiagnostics Diagnostics { get; }
+
+        /// <inheritdoc/>
+        public override string HelpLink
+        {
+            get => this.originalException.HelpLink;
+            set => this.originalException.HelpLink = value;
+        }
+
+        /// <inheritdoc/>
+        public override Exception GetBaseException()
+        {
+            return this.originalException.GetBaseException();
+        }
+
+        /// <inheritdoc/>
+        public override string ToString()
+        {
+            return $"{this.originalException.ToString()} {Environment.NewLine}CosmosDiagnostics: {this.Diagnostics.ToString()}";
+        }
+    }
+}

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosOperationCanceledExceptionTests.cs

@@ -0,0 +1,123 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests
+{
+    using System;
+    using System.Threading;
+    using System.Threading.Tasks;
+    using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+    [TestClass]
+    public class CosmosOperationCanceledExceptionTests : BaseCosmosClientHelper
+    {
+        private ContainerInternal Container = null;
+
+        [TestInitialize]
+        public async Task TestInitialize()
+        {
+            await base.TestInit();
+            string PartitionKey = "/status";
+            ContainerResponse response = await this.database.CreateContainerAsync(
+                new ContainerProperties(id: Guid.NewGuid().ToString(), partitionKeyPath: PartitionKey),
+                cancellationToken: this.cancellationToken);
+            Assert.IsNotNull(response);
+            Assert.IsNotNull(response.Container);
+            Assert.IsNotNull(response.Resource);
+            this.Container = (ContainerInternal)response;
+        }
+
+        [TestCleanup]
+        public async Task Cleanup()
+        {
+            await base.TestCleanup();
+        }
+
+        [TestMethod]
+        public async Task CheckCancellationTokenGatewayTestAsync()
+        {
+            using (CosmosClient gatewayClient = TestCommon.CreateCosmosClient(
+                builder => builder.WithConnectionModeGateway()))
+            {
+                Container gatewayContainer = gatewayClient.GetContainer(this.database.Id, this.Container.Id);
+                CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
+                cancellationTokenSource.Cancel();
+                await this.CheckCancellationTokenTestAsync(gatewayContainer, cancellationTokenSource.Token);
+            }
+        }
+
+        [TestMethod]
+        public async Task CheckCancellationWithTransportIntercepterTestAsync()
+        {
+            CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
+            Container withCancellationToken = TransportClientHelper.GetContainerWithIntercepter(
+                 this.database.Id,
+                 this.Container.Id,
+                 (uri, resourceOperation, documentServiceRequest) =>
+                 {
+                     if (documentServiceRequest.ResourceType == Documents.ResourceType.Document)
+                     {
+                         cancellationTokenSource.Cancel();
+                         cancellationTokenSource.Token.ThrowIfCancellationRequested();
+                     }
+                 });
+
+            await this.CheckCancellationTokenTestAsync(withCancellationToken, cancellationTokenSource.Token);
+        }
+
+        [TestMethod]
+        public async Task CheckCancellationTokenDirectTestAsync()
+        {
+            CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
+            cancellationTokenSource.Cancel();
+            await this.CheckCancellationTokenTestAsync(this.Container, cancellationTokenSource.Token);
+        }
+
+
+        private async Task CheckCancellationTokenTestAsync(
+            Container container,
+            CancellationToken cancellationToken)
+        {
+            ToDoActivity toDoActivity = ToDoActivity.CreateRandomToDoActivity();
+
+            try
+            {
+                await container.CreateItemAsync<ToDoActivity>(
+                    toDoActivity,
+                    new Cosmos.PartitionKey(toDoActivity.status),
+                    cancellationToken: cancellationToken);
+
+                Assert.Fail("Should have thrown");
+            }
+            catch (CosmosOperationCanceledException ce)
+            {
+                Assert.IsNotNull(ce);
+                string message = ce.Message;
+                string diagnostics = ce.Diagnostics.ToString();
+                string toString = ce.ToString();
+                Assert.IsTrue(toString.Contains(diagnostics));
+                Assert.IsTrue(toString.Contains(message));
+            }
+
+            try
+            {
+                FeedIterator feedIterator = container.GetItemQueryStreamIterator(
+                    "select * from T");
+
+                await feedIterator.ReadNextAsync(cancellationToken: cancellationToken);
+
+                Assert.Fail("Should have thrown");
+            }
+            catch (CosmosOperationCanceledException ce)
+            {
+                Assert.IsNotNull(ce);
+                string message = ce.Message;
+                string diagnostics = ce.Diagnostics.ToString();
+                string toString = ce.ToString();
+                Assert.IsTrue(toString.Contains(diagnostics));
+                Assert.IsTrue(toString.Contains(message));
+            }
+        }
+    }
+}

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Utils/TransportClientHelper.cs

@@ -36,11 +36,17 @@ internal static Container GetContainerWithItemTransportException(
         internal static Container GetContainerWithIntercepter(
             string databaseId,
             string containerId,
-            Action<Uri, ResourceOperation, DocumentServiceRequest> interceptor)
+            Action<Uri, ResourceOperation, DocumentServiceRequest> interceptor,
+            bool useGatewayMode = false)
         {
             CosmosClient clientWithIntercepter = TestCommon.CreateCosmosClient(
                builder =>
                {
+                   if (useGatewayMode)
+                   {
+                       builder.WithConnectionModeGateway();
+                   }
+
                    builder.WithTransportClientHandlerFactory(transportClient => new TransportClientWrapper(
                        transportClient,
                        interceptor));