From 484aa03153d07f7086527fa36173a2519041b4da Mon Sep 17 00:00:00 2001
From: GZTime <Time.GZ@outlook.com>
Date: Wed, 4 Sep 2024 11:10:38 +0800
Subject: [PATCH] feat(cache): single flight cache update

---
 src/GZCTF/Extensions/CacheExtensions.cs | 74 ++++++++++++++++++++-----
 src/GZCTF/Services/Cache/CacheHelper.cs |  2 +-
 2 files changed, 60 insertions(+), 16 deletions(-)

diff --git a/src/GZCTF/Extensions/CacheExtensions.cs b/src/GZCTF/Extensions/CacheExtensions.cs
index da8b44cee..050edd896 100644
--- a/src/GZCTF/Extensions/CacheExtensions.cs
+++ b/src/GZCTF/Extensions/CacheExtensions.cs
@@ -1,4 +1,5 @@
-using MemoryPack;
+using GZCTF.Services.Cache;
+using MemoryPack;
 using Microsoft.Extensions.Caching.Distributed;
 
 namespace GZCTF.Extensions;
@@ -19,26 +20,28 @@ public static async Task<TResult> GetOrCreateAsync<TResult, TLogger>(this IDistr
         var value = await cache.GetAsync(key, token);
         TResult? result = default;
 
-        if (value is not null)
-        {
-            try
-            {
-                result = MemoryPackSerializer.Deserialize<TResult>(value);
-            }
-            catch
-            {
-                // ignored
-            }
-
-            if (result is not null)
-                return result;
-        }
+        // most of the time, the cache is already been set
+        if (TryDeserialize(value, ref result))
+            return result!;
+
+        // wait if the cache is updating
+        value = await WaitLockAsync(cache, key, token);
 
+        if (TryDeserialize(value, ref result))
+            return result!;
+
+        var lockKey = CacheKey.UpdateLock(key);
+        await SetLockAsync(cache, lockKey, token);
+
+        // begin the update
         var cacheOptions = new DistributedCacheEntryOptions();
         result = await func(cacheOptions);
         var bytes = MemoryPackSerializer.Serialize(result);
 
         await cache.SetAsync(key, bytes, cacheOptions, token);
+        // finish the update
+
+        await ReleaseLockAsync(cache, lockKey, token);
 
         logger.SystemLog(Program.StaticLocalizer[
             nameof(Resources.Program.Cache_Updated),
@@ -47,4 +50,45 @@ public static async Task<TResult> GetOrCreateAsync<TResult, TLogger>(this IDistr
 
         return result;
     }
+
+    static bool TryDeserialize<TResult>(byte[]? value, ref TResult? result)
+    {
+        if (value is null)
+            return false;
+
+        try
+        {
+            result = MemoryPackSerializer.Deserialize<TResult>(value);
+            return result is not null;
+        }
+        catch
+        {
+            return false;
+        }
+    }
+
+    static async Task<byte[]?> WaitLockAsync(IDistributedCache cache, string key, CancellationToken token = default)
+    {
+        var lockKey = CacheKey.UpdateLock(key);
+        var lockValue = await cache.GetAsync(lockKey, token);
+
+        if (lockValue is null)
+            return null;
+
+        while (lockValue is not null)
+        {
+            await Task.Delay(100, token);
+            lockValue = await cache.GetAsync(lockKey, token);
+        }
+
+        // if we wait for the lock, we should try to get the value again
+        return await cache.GetAsync(key, token);
+    }
+
+    static Task SetLockAsync(IDistributedCache cache, string lockKey, CancellationToken token = default)
+        => cache.SetAsync(lockKey, [],
+            new DistributedCacheEntryOptions { AbsoluteExpirationRelativeToNow = TimeSpan.FromMinutes(1) }, token);
+
+    static Task ReleaseLockAsync(IDistributedCache cache, string lockKey, CancellationToken token = default) =>
+        cache.RemoveAsync(lockKey, token);
 }
diff --git a/src/GZCTF/Services/Cache/CacheHelper.cs b/src/GZCTF/Services/Cache/CacheHelper.cs
index f93159f8b..6c9b1fdb5 100644
--- a/src/GZCTF/Services/Cache/CacheHelper.cs
+++ b/src/GZCTF/Services/Cache/CacheHelper.cs
@@ -57,7 +57,7 @@ public static class CacheKey
     /// <summary>
     /// The cache update lock
     /// </summary>
-    public static string UpdateLock(string key) => $"_UpdateLock_{key}";
+    public static string UpdateLock(string key) => $"_UpdateLock{key}";
 
     /// <summary>
     /// The last update time