From e344fa942f9ae3909b809619e9cdfa63b430c254 Mon Sep 17 00:00:00 2001
From: Leonid Tsarev <leo@leotsarev.ru>
Date: Tue, 28 May 2024 14:12:31 +0300
Subject: [PATCH] Fix race condition in CacheEntryKeyLock.WaitAsync

---
 src/CacheTower/CacheStack.cs                 | 23 ++++---------
 src/CacheTower/Internal/CacheEntryKeyLock.cs | 36 +++++++++++++-------
 2 files changed, 31 insertions(+), 28 deletions(-)

diff --git a/src/CacheTower/CacheStack.cs b/src/CacheTower/CacheStack.cs
index 62a68e02..c19dec63 100644
--- a/src/CacheTower/CacheStack.cs
+++ b/src/CacheTower/CacheStack.cs
@@ -329,8 +329,14 @@ private async ValueTask BackPopulateCacheAsync<T>(int fromIndexExclusive, string
 
 	private async ValueTask<CacheEntry<T>?> RefreshValueAsync<T>(string cacheKey, Func<T, Task<T>> asyncValueFactory, CacheSettings settings, CacheEntryStatus entryStatus)
 	{
-		if (KeyLock.AcquireLock(cacheKey))
+		if (KeyLock.TryAcquireLockOrWaitAsync(cacheKey) is Task<ICacheEntry> task)
 		{
+			//We lost race to refresh
+			return await task as CacheEntry<T>;
+		}
+		else
+		{ 
+			// we got lock
 			try
 			{
 				var previousEntry = await GetAsync<T>(cacheKey);
@@ -374,21 +380,6 @@ private async ValueTask BackPopulateCacheAsync<T>(int fromIndexExclusive, string
 				throw;
 			}
 		}
-		else if (entryStatus != CacheEntryStatus.Stale)
-		{
-			//Last minute check to confirm whether waiting is required
-			var currentEntry = await GetAsync<T>(cacheKey);
-			if (currentEntry != null && currentEntry.GetStaleDate(settings) > DateTimeProvider.Now)
-			{
-				KeyLock.ReleaseLock(cacheKey, currentEntry);
-				return currentEntry;
-			}
-
-			//Lock until we are notified to be unlocked
-			return await KeyLock.WaitAsync(cacheKey) as CacheEntry<T>;
-		}
-
-		return default;
 	}
 
 	/// <summary>
diff --git a/src/CacheTower/Internal/CacheEntryKeyLock.cs b/src/CacheTower/Internal/CacheEntryKeyLock.cs
index 62287389..0d9dca67 100644
--- a/src/CacheTower/Internal/CacheEntryKeyLock.cs
+++ b/src/CacheTower/Internal/CacheEntryKeyLock.cs
@@ -15,33 +15,45 @@ public bool AcquireLock(string cacheKey)
 	{
 		lock (keyLocks)
 		{
+			return TryAddImpl(cacheKey, null);
+		}
+	}
+
+	private bool TryAddImpl (string cacheKey, TaskCompletionSource<ICacheEntry>? value)
+	{
 #if NETSTANDARD2_0
-			var hasLock = !keyLocks.ContainsKey(cacheKey);
-			if (hasLock)
+			var needToAdd = !keyLocks.ContainsKey(cacheKey);
+			if (needToAdd)
 			{
-				keyLocks[cacheKey] = null;
+				keyLocks[cacheKey] = value;
 			}
-			return hasLock;
+			return needToAdd;
 #elif NETSTANDARD2_1
-			return keyLocks.TryAdd(cacheKey, null);
+		return keyLocks.TryAdd(cacheKey, value);
 #endif
-		}
 	}
 
-	public Task<ICacheEntry> WaitAsync(string cacheKey)
+	/// <summary>
+	/// Returns null if acquire lock succeeded or returns task that wait until another thread will release lock and publish value.
+	/// 
+	/// </summary>
+	public Task<ICacheEntry>? TryAcquireLockOrWaitAsync(string cacheKey)
 	{
-		TaskCompletionSource<ICacheEntry>? completionSource;
-
 		lock (keyLocks)
 		{
-			if (!keyLocks.TryGetValue(cacheKey, out completionSource) || completionSource == null)
+			if (TryAddImpl(cacheKey, null))
+			{
+				// Lock acquired
+				return null;
+			}
+			var completionSource = keyLocks[cacheKey]; // Always exists here
+			if (completionSource == null)
 			{
 				completionSource = new TaskCompletionSource<ICacheEntry>();
 				keyLocks[cacheKey] = completionSource;
 			}
+			return completionSource.Task;
 		}
-
-		return completionSource.Task;
 	}
 
 	[MethodImpl(MethodImplOptions.AggressiveInlining)]