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

Microsoft.Azure.Cosmos/src/ExceptionHandlingUtility.cs

@@ -0,0 +1,44 @@
+// ------------------------------------------------------------
+// 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>
+        /// Creates a shallow copy of specific exception types (e.g., TaskCanceledException, TimeoutException, OperationCanceledException) 
+        /// to prevent excessive stack trace growth and rethrows them. All other exceptions are not processed.
+        /// </summary>
+        public static void CloneAndRethrowException(Exception e)
+        {
+            if (e is ICloneable ex)
+            {
+                throw (Exception)ex.Clone();
+            }
+
+            if (e is TaskCanceledException)
+            {
+                TaskCanceledException taskCanceledEx = new TaskCanceledException(e.Message, e.InnerException);
+                taskCanceledEx.Data["Message"] = e.Data["Message"];
+                throw taskCanceledEx;
+            }
+
+            if (e is TimeoutException)
+            {
+                TimeoutException timeoutEx = new TimeoutException(e.Message, e.InnerException);
+                timeoutEx.Data["Message"] = e.Data["Message"];
+                throw timeoutEx;
+            }
+        }
+    }
+}

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,8 +23,8 @@ internal sealed class AsyncCache<TKey, TValue>
         private readonly IEqualityComparer<TValue> valueEqualityComparer;
         private readonly IEqualityComparer<TKey> keyEqualityComparer;
 
-        private ConcurrentDictionary<TKey, AsyncLazy<TValue>> values;
-
+        private ConcurrentDictionary<TKey, AsyncLazy<TValue>> values;
+        private static readonly bool isStackTraceOptimizationEnabled = false;
         public AsyncCache(IEqualityComparer<TValue> valueEqualityComparer, IEqualityComparer<TKey> keyEqualityComparer = null)
         {
             this.keyEqualityComparer = keyEqualityComparer ?? EqualityComparer<TKey>.Default;
@@ -148,11 +148,18 @@ 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);
-                throw;
+                this.TryRemoveValue(key, actualValue);
+
+                if (isStackTraceOptimizationEnabled)
+                {
+                    // Creates a shallow copy of specific exception types to prevent stack trace proliferation 
+                    // and rethrows them, doesn't process other exceptions.
+                    ExceptionHandlingUtility.CloneAndRethrowException(ex);
+                }
+                throw ex;
             }
         }
 

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

@@ -25,7 +25,9 @@ internal sealed class AsyncCacheNonBlocking<TKey, TValue> : IDisposable
         private readonly Func<Exception, bool> removeFromCacheOnBackgroundRefreshException;
 
         private readonly IEqualityComparer<TKey> keyEqualityComparer;
-        private bool isDisposed;
+        private bool isDisposed;
+
+        private static readonly bool isStackTraceOptimizationEnabled = false;
 
         public AsyncCacheNonBlocking(
             Func<Exception, bool> removeFromCacheOnBackgroundRefreshException = null,
@@ -114,9 +116,15 @@ public async Task<TValue> GetAsync(
                             key,
                             removed,
                             e);
-                    }
-
-                    throw;
+                    }
+
+                    if (isStackTraceOptimizationEnabled)
+                    {
+                        // Creates a shallow copy of specific exception types to prevent stack trace proliferation 
+                        // and rethrows them, doesn't process other exceptions.
+                        ExceptionHandlingUtility.CloneAndRethrowException(e);
+                    }
+                    throw e;
                 }
 
                 return await this.UpdateCacheAndGetValueFromBackgroundTaskAsync(
@@ -158,7 +166,7 @@ public async Task<TValue> GetAsync(
                 this.values.TryRemove(key, out _);
                 throw;
             }
-        }
+        }
 
         public void Set(
            TKey key,
@@ -279,7 +287,7 @@ public AsyncLazyWithRefreshTask(
 
             public bool IsValueCreated => this.value != null;
 
-            public Task<T> GetValueAsync(
+            public Task<T> GetValueAsync(
                 Func<T, Task<T>> createValueFunc)
             {
                 // The task was already created so just return it.