Add unit tests for ExceptionHandlingUtility

Author: ananth7592

Microsoft.Azure.Cosmos/src/ExceptionHandlingUtility.cs

@@ -21,24 +21,28 @@ internal static class ExceptionHandlingUtility
         /// </summary>
         public static void CloneAndRethrowException(Exception e)
         {
-            if (e is ICloneable ex)
+            Exception ex = e switch
             {
-                throw (Exception)ex.Clone();
-            }
+                ICloneable cloneableEx => (Exception)cloneableEx.Clone(),
+                TaskCanceledException taskCanceledEx => AddMessageData(new TaskCanceledException(taskCanceledEx.Message, taskCanceledEx), e),
+                TimeoutException timeoutEx => AddMessageData(new TimeoutException(timeoutEx.Message, timeoutEx), e),
+                _ => null
+            };
 
-            if (e is TaskCanceledException)
+            if (ex is not null)
             {
-                TaskCanceledException taskCanceledEx = new TaskCanceledException(e.Message, e.InnerException);
-                taskCanceledEx.Data["Message"] = e.Data["Message"];
-                throw taskCanceledEx;
+                throw ex;
             }
+        }
 
-            if (e is TimeoutException)
+        private static Exception AddMessageData(Exception target, Exception source)
+        {
+            if (source.Data.Contains("Message"))
             {
-                TimeoutException timeoutEx = new TimeoutException(e.Message, e.InnerException);
-                timeoutEx.Data["Message"] = e.Data["Message"];
-                throw timeoutEx;
+                target.Data["Message"] = source.Data["Message"];
             }
+
+            return target;
         }
     }
 }

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

@@ -164,8 +164,19 @@ public async Task<TValue> GetAsync(
                     // and rethrows them, doesn't process other exceptions.
                     ExceptionHandlingUtility.CloneAndRethrowException(ex);
                 }
-                throw ex;
-            }
+                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.
+                    ExceptionHandlingUtility.CloneAndRethrowException(ex);
+                }
+                throw;
+            }
+
         }
 
         public void Remove(TKey key)

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

@@ -127,7 +127,7 @@ public async Task<TValue> GetAsync(
                         // and rethrows them, doesn't process other exceptions.
                         ExceptionHandlingUtility.CloneAndRethrowException(e);
                     }
-                    throw e;
+                    throw;
                 }
 
                 return await this.UpdateCacheAndGetValueFromBackgroundTaskAsync(
@@ -166,7 +166,14 @@ 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.
+                    ExceptionHandlingUtility.CloneAndRethrowException(e);
+                }
                 throw;
             }
         }

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

@@ -938,6 +938,86 @@ public async Task MultiRegionAccountTest()
             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"));
+            CosmosClientOptions cosmosClientOptions = new CosmosClientOptions()
+            {
+                ConsistencyLevel = Cosmos.ConsistencyLevel.Session,
+                //Set the pre-request event handler to throw a pre-defined timeout exception.
+                SendingRequestEventArgs = (sender, e) =>
+                {
+                    if (e.IsHttpRequest())
+                    {
+                        string endWith = "partitionKeyRangeIds=0";
+                        if (e.HttpRequest.RequestUri.OriginalString.EndsWith(endWith))
+                        {
+                            if (e.HttpRequest.Method == HttpMethod.Get &&
+                            e.HttpRequest.RequestUri.OriginalString.EndsWith(endWith))
+                            {
+                                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");
+
+                // Shared TaskCompletionSource ensures all tasks start together
+                TaskCompletionSource<object> tcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
+                int iterations = 3;// Simulating 3 concurrent write requests
+                List<Task> createTasks = new List<Task>();
+
+                for (int i = 0; i < iterations; i++)
+                {
+                    createTasks.Add(Task.Run(async () =>
+                    {
+                        await tcs.Task; // All tasks wait before proceeding
+
+                        ToDoActivity testItem = ToDoActivity.CreateRandomToDoActivity();
+                        await container.CreateItemAsync(testItem);
+                    }));
+                }
+
+                // Signal all tasks to start
+                tcs.SetResult(null);
+
+                // 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]
         [Owner("amudumba")]

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()