AsyncCache: Adds support for stack trace optimization during exceptions for AsyncCache and AsyncCacheNonblocking

Microsoft.Azure.Cosmos/src/ExceptionHandlingUtility.cs

@@ -0,0 +1,45 @@
+// ------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+// ------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Text;
+    using System.Threading.Tasks;
+    using Microsoft.Azure.Documents;
+
+    /// <summary>
+    /// Utility for post-processing of exceptions.
+    /// </summary>
+    internal static class ExceptionHandlingUtility
+    {
+        /// <summary>
+        /// Tries to create a shallow copy of specific exception types (e.g., TaskCanceledException, TimeoutException, OperationCanceledException)
+        /// to prevent excessive stack trace growth. Returns true if the exception was cloned, otherwise false.
+        /// </summary>
+        public static bool TryCloneException(Exception e, out Exception clonedException)
+        {
+            clonedException = e switch
+            {
+                ICloneable cloneableEx => (Exception)cloneableEx.Clone(),
+                TaskCanceledException taskCanceledEx => AddMessageData(new TaskCanceledException(taskCanceledEx.Message, taskCanceledEx), e),
+                TimeoutException timeoutEx => AddMessageData(new TimeoutException(timeoutEx.Message, timeoutEx), e),
+                _ => null
+            };
+
+            return clonedException is not null;
+        }
+
+        private static Exception AddMessageData(Exception target, Exception source)
+        {
+            if (source.Data.Contains("Message"))
+            {
+                target.Data["Message"] = source.Data["Message"];
+            }
+
+            return target;
+        }
+    }
+}

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

@@ -17,7 +17,7 @@ namespace Microsoft.Azure.Cosmos
     /// <summary>
     /// The Cosmos Client exception
     /// </summary>
-    public class CosmosException : Exception
+    public class CosmosException : Exception, ICloneable
     {
         private readonly string stackTrace;
         private readonly Lazy<string> lazyMessage;
@@ -196,7 +196,7 @@ internal ResponseMessage ToCosmosResponseMessage(RequestMessage request)
                  trace: this.Trace);
 
             return responseMessage;
-        }
+        }
 
         private static string GetMessageHelper(
             HttpStatusCode statusCode,
@@ -301,6 +301,17 @@ internal static void RecordOtelAttributes(CosmosException exception, DiagnosticS
             {
                 CosmosDbEventSource.RecordDiagnosticsForExceptions(exception.Diagnostics);
             }
-        }
+        }
+
+        /// <summary>
+        /// Creates a shallow copy of the current exception instance.
+        /// This ensures that the cloned exception retains the same properties but does not
+        /// excessively proliferate stack traces or deep-copy unnecessary objects.
+        /// </summary>
+        /// <returns>A shallow copy of the current <see cref="CosmosException"/>.</returns>
+        public object Clone()
+        {
+            return this.MemberwiseClone();
+        }
     }
 }

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

@@ -18,7 +18,7 @@ namespace Microsoft.Azure.Cosmos
     ///  diagnostics of the operation that was canceled.
     /// </summary> 
     [Serializable]
-    public class CosmosOperationCanceledException : OperationCanceledException
+    public class CosmosOperationCanceledException : OperationCanceledException, ICloneable
     {
         private readonly OperationCanceledException originalException;
         private readonly Lazy<string> lazyMessage;
@@ -155,6 +155,17 @@ public override void GetObjectData(SerializationInfo info, StreamingContext cont
             info.AddValue("tokenCancellationRequested", this.tokenCancellationRequested);
             info.AddValue("lazyMessage", this.lazyMessage.Value);
             info.AddValue("toStringMessage", this.toStringMessage.Value);
-        }
+        }
+
+        /// <summary>
+        /// Creates a shallow copy of the current exception instance.
+        /// This ensures that the cloned exception retains the same properties but does not
+        /// excessively proliferate stack traces or deep-copy unnecessary objects.
+        /// </summary>
+        /// <returns>A shallow copy of the current <see cref="CosmosOperationCanceledException"/>.</returns>
+        public object Clone()
+        {
+            return this.MemberwiseClone();
+        }
     }
 }

Microsoft.Azure.Cosmos/src/Routing/AsyncCache.cs

@@ -23,7 +23,6 @@ internal sealed class AsyncCache<TKey, TValue>
         private readonly bool enableAsyncCacheExceptionNoSharing;
         private readonly IEqualityComparer<TValue> valueEqualityComparer;
         private readonly IEqualityComparer<TKey> keyEqualityComparer;
-
         private ConcurrentDictionary<TKey, AsyncLazy<TValue>> values;
 
         public AsyncCache(
@@ -154,12 +153,38 @@ public async Task<TValue> GetAsync(
             {
                 return await generator;
             }
-            catch (Exception) when (object.ReferenceEquals(actualValue, newLazyValue))
+            catch (Exception ex) when (object.ReferenceEquals(actualValue, newLazyValue))
             {
                 // If the lambda this thread added to values triggered an exception remove it from the cache.
                 this.TryRemoveValue(key, actualValue);
+
+                if (this.enableAsyncCacheExceptionNoSharing)
+                {
+                    // Creates a shallow copy of specific exception types to prevent stack trace proliferation 
+                    // and rethrows them, doesn't process other exceptions.
+                    if (ExceptionHandlingUtility.TryCloneException(ex, out Exception clone))
+                    {
+                        throw clone;
+                    }
+                }
+
                 throw;
-            }
+            }
+            catch (Exception ex)
+            {
+                if (this.enableAsyncCacheExceptionNoSharing)
+                {
+                    // Creates a shallow copy of specific exception types to prevent stack trace proliferation
+                    // and rethrows them, doesn't process other exceptions.
+                    if (ExceptionHandlingUtility.TryCloneException(ex, out Exception clone))
+                    {
+                        throw clone;
+                    }
+                }
+
+                throw;
+            }
+
         }
 
         public void Remove(TKey key)

Microsoft.Azure.Cosmos/src/Routing/AsyncCacheNonBlocking.cs

@@ -26,8 +26,8 @@ internal sealed class AsyncCacheNonBlocking<TKey, TValue> : IDisposable
         private readonly Func<Exception, bool> removeFromCacheOnBackgroundRefreshException;
 
         private readonly IEqualityComparer<TKey> keyEqualityComparer;
-        private bool isDisposed;
-
+        private bool isDisposed;
+
         public AsyncCacheNonBlocking(
             Func<Exception, bool> removeFromCacheOnBackgroundRefreshException = null,
             IEqualityComparer<TKey> keyEqualityComparer = null,
@@ -119,8 +119,17 @@ public async Task<TValue> GetAsync(
                             key,
                             removed,
                             e);
-                    }
-
+                    }
+
+                    if (this.enableAsyncCacheExceptionNoSharing)
+                    {
+                        // Creates a shallow copy of specific exception types to prevent stack trace proliferation 
+                        // and rethrows them, doesn't process other exceptions.
+                        if (ExceptionHandlingUtility.TryCloneException(e, out Exception clone))
+                        {
+                            throw clone;
+                        }
+                    }
                     throw;
                 }
 
@@ -160,10 +169,20 @@ public async Task<TValue> GetAsync(
                             e.ToString());
 
                 // Remove the failed task from the dictionary so future requests can send other calls..
-                this.values.TryRemove(key, out _);
+                this.values.TryRemove(key, out _);
+
+                if (this.enableAsyncCacheExceptionNoSharing)
+                {
+                    // Creates a shallow copy of specific exception types to prevent stack trace proliferation
+                    // and rethrows them, doesn't process other exceptions.
+                    if (ExceptionHandlingUtility.TryCloneException(e, out Exception clone))
+                    {
+                        throw clone;
+                    }
+                }
                 throw;
             }
-        }
+        }
 
         public void Set(
            TKey key,

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

@@ -937,6 +937,85 @@ public async Task MultiRegionAccountTest()
             Assert.IsNotNull(cosmosClient);
             AccountProperties properties = await cosmosClient.ReadAccountAsync();
             Assert.IsNotNull(properties);
+        }
+
+        [DataTestMethod]
+        [DataRow(true)]
+        [DataRow(false)]
+        [Owner("amudumba")]
+        public async Task ValidateAsyncExceptionNoSharing(bool asyncCacheExceptionNoSharing)
+        {
+            TimeoutException exception = new TimeoutException("HTTP Timeout exception", new TimeoutException("Inner exception message"));
+            int expectedHandlers = 3;
+            int enteredHandlers = 0;
+            TaskCompletionSource<object> blockSendingHandlers = new(TaskCreationOptions.RunContinuationsAsynchronously);
+
+            CosmosClientOptions cosmosClientOptions = new CosmosClientOptions()
+            {
+                ConsistencyLevel = Cosmos.ConsistencyLevel.Session,
+                SendingRequestEventArgs = async (sender, e) =>
+                {
+                    if (e.IsHttpRequest())
+                    {
+                        string endWith = "partitionKeyRangeIds=0";
+                        if (e.HttpRequest.Method == HttpMethod.Get &&
+                            e.HttpRequest.RequestUri.OriginalString.EndsWith(endWith))
+                        {
+                            // Wait until all expected threads are in the handler
+                            if (Interlocked.Increment(ref enteredHandlers) == expectedHandlers)
+                            {
+                                blockSendingHandlers.SetResult(null);
+                            }
+
+                            await blockSendingHandlers.Task; // block here until all enter
+
+                            throw exception;
+                        }
+                    }
+                },
+                EnableAsyncCacheExceptionNoSharing = asyncCacheExceptionNoSharing
+            };
+
+            Cosmos.Database db = null;
+            try
+            {
+                CosmosClient cosmosClient = TestCommon.CreateCosmosClient(clientOptions: cosmosClientOptions);
+
+                db = await cosmosClient.CreateDatabaseIfNotExistsAsync("TimeoutFaultTest");
+                Container container = await db.CreateContainerIfNotExistsAsync("TimeoutFaultContainer", "/pk");
+
+                int iterations = expectedHandlers;
+                List<Task> createTasks = new();
+
+                for (int i = 0; i < iterations; i++)
+                {
+                    createTasks.Add(Task.Run(async () =>
+                    {
+                        ToDoActivity testItem = ToDoActivity.CreateRandomToDoActivity();
+                        await container.CreateItemAsync(testItem);
+                    }));
+                }
+
+                // Wait for all tasks to complete (they should all fail)
+                await Task.WhenAll(createTasks);
+            }
+            catch (TimeoutException tex)
+            {
+                if (asyncCacheExceptionNoSharing)
+                {
+                    //asyncCacheExceptionNoSharing feature is enabled. Shallow copies of the exception will be thrown.
+                    Assert.IsFalse(Object.ReferenceEquals(tex, exception), "Exception should not be the same");
+                }
+                else
+                {
+                    //asyncCacheExceptionNoSharing feature is disabled. Exceptions will be thrown as is.
+                    Assert.IsTrue(Object.ReferenceEquals(tex, exception), "Exception should be the same");
+                }
+            }
+            finally
+            {
+                if (db != null) await db.DeleteAsync();
+            }
         }
 
         [TestMethod]

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/AsyncCacheTest.cs

@@ -12,8 +12,7 @@ namespace Microsoft.Azure.Cosmos.Tests
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.Common;
     using Microsoft.Azure.Documents;
-    using Microsoft.VisualStudio.TestTools.UnitTesting;
-
+    using Microsoft.VisualStudio.TestTools.UnitTesting;
     [TestClass]
     public class AsyncCacheTest
     {
@@ -306,6 +305,45 @@ await Task.Factory.StartNew(() =>
             TaskCreationOptions.None,
             new SingleTaskScheduler()
             );
+        }
+
+        [DataTestMethod]
+        [DataRow(true)]
+        [DataRow(false)]
+        public async Task GetAsync_Assert_OnExceptionPostProcessing(bool enabled)
+        {
+            AsyncCache<string, string> cacheDefault = new AsyncCache<string, string>(enableAsyncCacheExceptionNoSharing: enabled);
+
+            // Arrange
+            Exception testException = new TimeoutException("Simulated timeout exception");
+            CancellationToken cancellationToken = CancellationToken.None;
+
+            // Simulate a failing lambda function
+            Func<Task<string>> failingLambda = () => Task.FromException<string>(testException);
+
+            // Act
+            try
+            {
+                await cacheDefault.GetAsync(
+                    "testKey",
+                    obsoleteValue: null,
+                    singleValueInitFunc: failingLambda,
+                    cancellationToken,
+                    forceRefresh: false);
+            }
+            catch (TimeoutException ex)
+            {
+                if (enabled)
+                {
+                    // Assert that the expected exception was rethrown
+                    Assert.IsFalse(Object.ReferenceEquals(testException, ex));
+                } else
+                {
+                    // Assert that no cloning and rethrowing was done on the exceptions,
+                    Assert.IsTrue(Object.ReferenceEquals(testException, ex));
+                }
+
+            }
         }
 
         private int GenerateIntFuncThatThrows()

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Contracts/DotNetSDKAPI.json

@@ -3537,6 +3537,11 @@
           "Attributes": [],
           "MethodInfo": "System.Nullable`1[System.TimeSpan] RetryAfter;CanRead:True;CanWrite:False;System.Nullable`1[System.TimeSpan] get_RetryAfter();IsAbstract:False;IsStatic:False;IsVirtual:True;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
         },
+        "System.Object Clone()": {
+          "Type": "Method",
+          "Attributes": [],
+          "MethodInfo": "System.Object Clone();IsAbstract:False;IsStatic:False;IsVirtual:True;IsGenericMethod:False;IsConstructor:False;IsFinal:True;"
+        },
         "System.String ActivityId": {
           "Type": "Property",
           "Attributes": [],
@@ -3663,6 +3668,11 @@
           "Attributes": [],
           "MethodInfo": "System.Exception GetBaseException();IsAbstract:False;IsStatic:False;IsVirtual:True;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
         },
+        "System.Object Clone()": {
+          "Type": "Method",
+          "Attributes": [],
+          "MethodInfo": "System.Object Clone();IsAbstract:False;IsStatic:False;IsVirtual:True;IsGenericMethod:False;IsConstructor:False;IsFinal:True;"
+        },
         "System.String get_HelpLink()": {
           "Type": "Method",
           "Attributes": [],