Rate limiting feature implementation redesigned for handling locks and using module prefix

Author: HaikAsatryan

src/DistributedCache/DistributedCache.csproj

@@ -8,13 +8,13 @@
         <PackageReadmeFile>Readme.md</PackageReadmeFile>
         <Authors>Pandatech</Authors>
         <Copyright>MIT</Copyright>
-        <Version>1.1.0</Version>
+        <Version>1.2.0</Version>
         <PackageId>Pandatech.DistributedCache</PackageId>
         <Title>Pandatech Distributed Cache</Title>
         <PackageTags>Pandatech, library, redis, distributed locks, cache</PackageTags>
         <Description>Pandatech.DistributedCache is a comprehensive caching library designed for .NET applications, leveraging the power of Redis. It provides easy-to-use and highly configurable caching mechanisms, including support for tagged cache entries, customizable expiration policies, and robust health check services. The library also features built-in distributed lock mechanisms to ensure data consistency and prevent cache stampedes. This ensures high performance, scalability, and reliability, making it an ideal choice for enterprise-level distributed caching needs.</Description>
         <RepositoryUrl>https://github.com/PandaTechAM/be-lib-distributed-cache</RepositoryUrl>
-        <PackageReleaseNotes>Rate limiting feature</PackageReleaseNotes>
+        <PackageReleaseNotes>Rate limiting feature implementation redesigned for handling locks and using module prefix</PackageReleaseNotes>
     </PropertyGroup>
 
     <ItemGroup>

src/DistributedCache/Dtos/RateLimitCache.cs

@@ -1,10 +1,9 @@
-using DistributedCache.Services.Interfaces;
-using MessagePack;
+using MessagePack;
 
 namespace DistributedCache.Dtos;
 
 [MessagePackObject]
-public class RateLimitCache : ICacheEntity
+public class RateLimitCache
 {
     [Key(0)] public int Attempts { get;  set; } = 1;
     [Key(1)] public int MaxAttempts { get;  init; } 

src/DistributedCache/Services/Implementations/RedisCacheService.cs

@@ -1,9 +1,7 @@
-using DistributedCache.Dtos;
-using DistributedCache.Helpers;
+using DistributedCache.Helpers;
 using DistributedCache.Options;
 using DistributedCache.Services.Interfaces;
 using Microsoft.Extensions.Options;
-using StackExchange.Redis;
 using StackExchange.Redis.Extensions.Core.Abstractions;
 
 namespace DistributedCache.Services.Implementations;

src/DistributedCache/Services/Implementations/RedisRateLimitService.cs

@@ -1,50 +1,97 @@
-using DistributedCache.Dtos;
+using System.Reflection;
+using DistributedCache.Dtos;
 using DistributedCache.Enums;
+using DistributedCache.Helpers;
+using DistributedCache.Options;
 using DistributedCache.Services.Interfaces;
-using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.Options;
+using StackExchange.Redis.Extensions.Core.Abstractions;
 
 namespace DistributedCache.Services.Implementations;
 
-public class RedisRateLimitService(ICacheService<RateLimitCache> cacheService, RedisLockService lockService)
-    : IRateLimitService
+public class RedisRateLimitService(
+   IRedisClient redisClient,
+   IOptions<CacheConfigurationOptions> options,
+   RedisLockService lockService) : IRateLimitService
 {
-    public async ValueTask<RateLimitState> RateLimitAsync(RateLimitConfiguration rateLimitConfiguration,
-        CancellationToken cancellationToken = default)
-    {
-        var key = rateLimitConfiguration.GetKey();
-      
-        var cache = await cacheService.GetAsync(key, cancellationToken);
-
-        if (cache == null)
-        {
+   private readonly IRedisDatabase _redisDatabase = redisClient.GetDefaultDatabase();
+   private readonly CacheConfigurationOptions _config = options.Value;
+
+   private readonly string _moduleName = Assembly.GetExecutingAssembly()
+                                                 .FullName!;
+
+   public async ValueTask<RateLimitState> RateLimitAsync(RateLimitConfiguration rateLimitConfiguration,
+      CancellationToken cancellationToken = default)
+   {
+      var key = _config.KeyPrefixForIsolation == KeyPrefix.None
+         ? KeyFormatHelper.GetPrefixedKey(rateLimitConfiguration.GetKey())
+         : KeyFormatHelper.GetPrefixedKey(rateLimitConfiguration.GetKey(), _moduleName);
+
+
+      var lockValue = Guid.NewGuid()
+                          .ToString();
+
+      while (true)
+      {
+         cancellationToken.ThrowIfCancellationRequested();
+
+         var isLocked = await lockService.CheckForLockAsync(key);
+
+         if (isLocked)
+         {
+            await lockService.WaitForLockReleaseAsync(key, cancellationToken);
+            continue;
+         }
+
+         var lockAcquired = await lockService.AcquireLockAsync(key, lockValue);
+         if (!lockAcquired)
+         {
+            await lockService.WaitForLockReleaseAsync(key, cancellationToken);
+            continue;
+         }
+
+         break;
+      }
+
+      try
+      {
+         var cache = await _redisDatabase.GetAsync<RateLimitCache>(key);
+
+         if (cache is null)
+         {
             var newCache = RateLimitCache.CreateRateLimitCache(rateLimitConfiguration);
-            await cacheService.SetAsync(key, newCache, rateLimitConfiguration.TimeToLive,null ,cancellationToken);
-            
+            await _redisDatabase.AddAsync(key, newCache, rateLimitConfiguration.TimeToLive);
+
             return new RateLimitState(
-                RateLimitStatus.NotExceeded, 
-                rateLimitConfiguration.TimeToLive, 
-                rateLimitConfiguration.MaxAttempts - 1
+               RateLimitStatus.NotExceeded,
+               rateLimitConfiguration.TimeToLive,
+               rateLimitConfiguration.MaxAttempts - 1
             );
-        }
+         }
 
-        var isUpdated = cache.TryUpdateAttempts();
-        var newExpiration = cache.GetNewExpiration();
+         var isUpdated = cache.TryUpdateAttempts();
+         var newExpiration = cache.GetNewExpiration();
 
-        if (!isUpdated)
-        {
+         if (!isUpdated)
+         {
             return new RateLimitState(
-                RateLimitStatus.Exceeded, 
-                newExpiration, 
-                0
+               RateLimitStatus.Exceeded,
+               newExpiration,
+               0
             );
-        }
+         }
 
-        await cacheService.SetAsync(key, cache, newExpiration, null, cancellationToken);
+         await _redisDatabase.AddAsync(key, cache, newExpiration);
 
-        return new RateLimitState(
-            RateLimitStatus.NotExceeded, 
-            newExpiration, 
+         return new RateLimitState(
+            RateLimitStatus.NotExceeded,
+            newExpiration,
             cache.MaxAttempts - cache.Attempts
-        );
-    }
+         );
+      }
+      finally
+      {
+         await lockService.ReleaseLockAsync(key, lockValue);
+      }
+   }
 }