From 9549dd0308fd889a86a65f21175be45e9fe420c3 Mon Sep 17 00:00:00 2001
From: Debdatta Kunda <dkunda@microsoft.com>
Date: Tue, 17 Jan 2023 14:07:00 -0800
Subject: [PATCH] Code changes to refactor async non-blocking cache code.

---
 .../src/Routing/AsyncCacheNonBlocking.cs      | 84 +++++++++++--------
 1 file changed, 48 insertions(+), 36 deletions(-)

diff --git a/Microsoft.Azure.Cosmos/src/Routing/AsyncCacheNonBlocking.cs b/Microsoft.Azure.Cosmos/src/Routing/AsyncCacheNonBlocking.cs
index fc6731d4f7..37e35f5761 100644
--- a/Microsoft.Azure.Cosmos/src/Routing/AsyncCacheNonBlocking.cs
+++ b/Microsoft.Azure.Cosmos/src/Routing/AsyncCacheNonBlocking.cs
@@ -120,30 +120,11 @@ public async Task<TValue> GetAsync(
                     throw;
                 }
 
-                try
-                {
-                    return await initialLazyValue.CreateAndWaitForBackgroundRefreshTaskAsync(
-                       createRefreshTask: singleValueInitFunc);
-                }
-                catch (Exception e)
-                {
-                    if (initialLazyValue.ShouldRemoveFromCacheThreadSafe())
-                    {
-                        DefaultTrace.TraceError(
-                            "AsyncCacheNonBlocking.GetAsync with ForceRefresh Failed. key: {0}, Exception: {1}",
-                            key,
-                            e);
-
-                        // In some scenarios when a background failure occurs like a 404
-                        // the initial cache value should be removed.
-                        if (this.removeFromCacheOnBackgroundRefreshException(e))
-                        {
-                            this.TryRemove(key);
-                        }
-                    }
-
-                    throw;
-                }
+                return await this.UpdateCacheAndGetValueFromBackgroundTaskAsync(
+                    key: key,
+                    initialValue: initialLazyValue,
+                    callbackDelegate: singleValueInitFunc,
+                    operationName: "GetAsync");
             }
 
             // The AsyncLazyWithRefreshTask is lazy and won't create the task until GetValue is called.
@@ -208,24 +189,55 @@ public async Task RefreshAsync(
         {
             if (this.values.TryGetValue(key, out AsyncLazyWithRefreshTask<TValue> initialLazyValue))
             {
-                try
-                {
-                    await initialLazyValue.CreateAndWaitForBackgroundRefreshTaskAsync(singleValueInitFunc);
-                }
-                catch (Exception ex)
+                await this.UpdateCacheAndGetValueFromBackgroundTaskAsync(
+                    key: key,
+                    initialValue: initialLazyValue,
+                    callbackDelegate: singleValueInitFunc,
+                    operationName: "RefreshAsync");
+            }
+        }
+
+        /// <summary>
+        /// Creates a background task to invoke the callback delegate and updates the cache with the value returned from the delegate.
+        /// </summary>
+        /// <param name="key">The requested key to be updated.</param>
+        /// <param name="initialValue">An instance of <see cref="AsyncLazyWithRefreshTask{T}"/> containing the initial cached value.</param>
+        /// <param name="callbackDelegate">A func callback delegate to be invoked at a later point of time.</param>
+        /// <param name="operationName">A string indicating the operation on the cache.</param>
+        /// <returns>A <see cref="Task{TValue}"/> containing the updated, refreshed value.</returns>
+        private async Task<TValue> UpdateCacheAndGetValueFromBackgroundTaskAsync(
+            TKey key,
+            AsyncLazyWithRefreshTask<TValue> initialValue,
+            Func<TValue, Task<TValue>> callbackDelegate,
+            string operationName)
+        {
+            try
+            {
+                return await initialValue.CreateAndWaitForBackgroundRefreshTaskAsync(
+                   createRefreshTask: callbackDelegate);
+            }
+            catch (Exception ex)
+            {
+                if (initialValue.ShouldRemoveFromCacheThreadSafe())
                 {
+                    bool removed = false;
+
+                    // In some scenarios when a background failure occurs like a 404
+                    // the initial cache value should be removed.
                     if (this.removeFromCacheOnBackgroundRefreshException(ex))
                     {
-                        bool removed = this.TryRemove(key);
-                        DefaultTrace.TraceError(
-                            "AsyncCacheNonBlocking Failed RefreshAsync. key: {0}, tryRemoved: {1}, Exception: {2}",
-                            key,
-                            removed,
-                            ex);
+                        removed = this.TryRemove(key);
                     }
 
-                    throw;
+                    DefaultTrace.TraceError(
+                        "AsyncCacheNonBlocking Failed. key: {0}, operation: {1}, tryRemoved: {2}, Exception: {3}",
+                        key,
+                        operationName,
+                        removed,
+                        ex);
                 }
+
+                throw;
             }
         }