Rate limiting feature

Author: HaikAsatryan

Readme.md

@@ -3,12 +3,13 @@
 Pandatech.DistributedCache is a .NET library providing an efficient and performant abstraction layer over
 `StackExchange.Redis`, specifically designed for .NET applications. This library builds on top of
 `StackExchange.Redis.Extensions.AspNetCore` and `StackExchange.Redis.Extensions.MsgPack` to offer a robust, easy-to-use
-caching solution with advanced features such as typed cache services, distributed locking, and health checks.
+caching solution with advanced features such as typed cache services, distributed locking, business logic rate limiting.
 
 ## Features
 
 - **Typed Cache Service:** Supports strongly-typed caching with MessagePack serialization.
 - **Distributed Locking:** Ensures data consistency with distributed locks.
+- **Distributed Rate Limiting:** Prevents cache abuse with rate limiting based on business logic.
 - **Key Isolation:** Modular monolith support by prefixing keys with assembly names.
 - **Stampede Protection:** Protects against cache stampede in the `GetOrCreateAsync` method.
 - **No Serializer Override:** Enforces MessagePack serialization for performance and readability.
@@ -192,6 +193,53 @@ public interface ICacheService<T> where T : class
 }
 ```
 
+### 5. Rate Limiting
+
+Implement rate limiting using `IRateLimitService` and `RateLimitConfiguration`.
+
+**Define Rate Limiting Configuration**
+
+```csharp
+public enum ActionType //your business logic actions
+{
+    SmsForTfa = 1,
+    EmailForTfa = 2
+}
+
+public static class RateLimitingConfigurations //your shared rate limiting configuration
+{
+    public static RateLimitConfiguration GetSmsConfig()
+    {
+        return new RateLimitConfiguration
+        {
+            ActionType = (int)ActionType.SmsForTfa,
+            MaxAttempts = 2,
+            TimeToLive = TimeSpan.FromSeconds(10)
+        };
+    }
+}
+```
+
+**Implement Rate Limiting in the service**
+
+```csharp
+using DistributedCache.Dtos;
+using DistributedCache.Services.Interfaces;
+
+public class SendSmsService(IRateLimitService rateLimitService)
+{
+    public async Task<RateLimitState> SendSms(CancellationToken cancellationToken = default)
+    {
+        var phoneNumber = "1234567890";
+        var rateLimitConfiguration = RateLimitingConfigurations.GetSmsConfig().SetIdentifiers(phoneNumber);
+
+        return await rateLimitService.RateLimitAsync(rateLimitConfiguration, cancellationToken);
+    }
+}
+```
+
+Based on rate limit state you can throw exception/return 427 or proceed with the business logic.
+
 ## Enforced MessagePack Serialization
 
 `Pandatech.DistributedCache` enforces the use of MessagePack serialization for several compelling reasons:

src/DistributedCache/DistributedCache.csproj

@@ -8,13 +8,13 @@
         <PackageReadmeFile>Readme.md</PackageReadmeFile>
         <Authors>Pandatech</Authors>
         <Copyright>MIT</Copyright>
-        <Version>1.0.0</Version>
+        <Version>1.1.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>InitialCommit</PackageReleaseNotes>
+        <PackageReleaseNotes>Rate limiting feature</PackageReleaseNotes>
     </PropertyGroup>
 
     <ItemGroup>

src/DistributedCache/Dtos/RateLimitCache.cs

@@ -0,0 +1,37 @@
+using DistributedCache.Services.Interfaces;
+using MessagePack;
+
+namespace DistributedCache.Dtos;
+
+[MessagePackObject]
+public class RateLimitCache : ICacheEntity
+{
+    [Key(0)] public int Attempts { get;  set; } = 1;
+    [Key(1)] public int MaxAttempts { get;  init; } 
+    [Key(2)] public DateTime Expiration { get;  init; }
+
+    public static RateLimitCache CreateRateLimitCache(RateLimitConfiguration configuration)
+    {
+        return new RateLimitCache
+        {
+            MaxAttempts = configuration.MaxAttempts,
+            Expiration = DateTime.UtcNow + configuration.TimeToLive
+        };
+    }
+
+    internal bool TryUpdateAttempts()
+    {
+        if (Attempts >= MaxAttempts)
+        {
+            return false;
+        }
+
+        Attempts++;
+        return true;
+    }
+
+    internal TimeSpan GetNewExpiration()
+    {
+        return Expiration - DateTime.UtcNow;
+    }
+}

src/DistributedCache/Dtos/RateLimitConfiguration.cs

@@ -0,0 +1,30 @@
+namespace DistributedCache.Dtos;
+
+public class RateLimitConfiguration : RateLimitKey
+{
+    private readonly int _maxAttempts;
+    private readonly TimeSpan _timeToLive;
+
+    public int MaxAttempts
+    {
+        get => _maxAttempts;
+        init => _maxAttempts = value > 0
+            ? value
+            : throw new ArgumentOutOfRangeException(nameof(MaxAttempts), "Must be greater than zero.");
+    }
+
+    public TimeSpan TimeToLive
+    {
+        get => _timeToLive;
+        init => _timeToLive = value > TimeSpan.Zero
+            ? value
+            : throw new ArgumentOutOfRangeException(nameof(TimeToLive), "Must be a positive time span.");
+    }
+    
+    public override RateLimitConfiguration SetIdentifiers(string primaryIdentifier, string? secondaryIdentifier = null)
+    {
+        PrimaryIdentifier = primaryIdentifier;
+        SecondaryIdentifier = secondaryIdentifier;
+        return this;
+    }
+}

src/DistributedCache/Dtos/RateLimitKey.cs

@@ -0,0 +1,17 @@
+namespace DistributedCache.Dtos;
+
+public abstract class RateLimitKey
+{
+    public required int ActionType { get; init; }
+    protected string PrimaryIdentifier { get; set; } = null!;
+    protected string? SecondaryIdentifier { get; set; }
+    
+    internal string GetKey()
+    {
+        return !string.IsNullOrWhiteSpace(SecondaryIdentifier)
+            ? $"{ActionType}:{PrimaryIdentifier}:{SecondaryIdentifier}:limit"
+            : $"{ActionType}:{PrimaryIdentifier}:limit";
+    }
+
+    public abstract RateLimitConfiguration SetIdentifiers(string primaryIdentifier, string? secondaryIdentifier = null);
+}

src/DistributedCache/Dtos/RateLimitState.cs

@@ -0,0 +1,5 @@
+using DistributedCache.Enums;
+
+namespace DistributedCache.Dtos;
+
+public record RateLimitState(RateLimitStatus Status, TimeSpan TimeToReset, int RemainingAttempts);

src/DistributedCache/Enums/RateLimitStatus.cs

@@ -0,0 +1,7 @@
+namespace DistributedCache.Enums;
+
+public enum RateLimitStatus
+{
+    Exceeded = 1,
+    NotExceeded = 2
+}

src/DistributedCache/Extensions/WebApplicationBuilderExtension.cs

@@ -33,8 +33,14 @@ public static WebApplicationBuilder AddDistributedCache(this WebApplicationBuild
         builder.Services.AddSingleton<IConnectionMultiplexer>(ConnectionMultiplexer.Connect(redisOptions));
 
         builder.Services.AddSingleton(typeof(ICacheService<>), typeof(RedisCacheService<>));
+        builder.Services.AddSingleton<RedisLockService>();
+        builder.Services.AddScoped<IRateLimitService, RedisRateLimitService>();
 
-        var redisConfiguration = new RedisConfiguration { ConnectionString = configurations.RedisConnectionString };
+        var redisConfiguration = new RedisConfiguration
+        {
+            ConnectionString = configurations.RedisConnectionString,
+            Name = "DistributedCacheConfiguration"
+        };
 
         builder.Services.AddStackExchangeRedisExtensions<RedisMsgPackObjectSerializer>(redisConfiguration);
         //builder.Services.AddHostedService<RedisHealthCheckService>(); //Discontinued feature

src/DistributedCache/Services/Implementations/RedisCacheService.cs

@@ -1,4 +1,5 @@
-using DistributedCache.Helpers;
+using DistributedCache.Dtos;
+using DistributedCache.Helpers;
 using DistributedCache.Options;
 using DistributedCache.Services.Interfaces;
 using Microsoft.Extensions.Options;
@@ -7,15 +8,16 @@
 
 namespace DistributedCache.Services.Implementations;
 
-internal class RedisCacheService<T>(IRedisClient redisClient, IOptions<CacheConfigurationOptions> options)
+internal class RedisCacheService<T>(
+    IRedisClient redisClient,
+    IOptions<CacheConfigurationOptions> options,
+    RedisLockService lockService)
     : ICacheService<T>
     where T : class, ICacheEntity
 {
     private readonly IRedisDatabase _redisDatabase = redisClient.GetDefaultDatabase();
     private readonly CacheConfigurationOptions _config = options.Value;
     private readonly string _moduleName = typeof(T).Assembly.GetName().Name!;
-    private readonly TimeSpan _lockExpiry = options.Value.DistributedLockDuration;
-    private readonly TimeSpan _lockRetryDelay = TimeSpan.FromMilliseconds(10);
 
     public async ValueTask<T> GetOrCreateAsync(string key, Func<CancellationToken, ValueTask<T>> factory,
         TimeSpan? expiration = null, IReadOnlyCollection<string>? tags = null, CancellationToken token = default)
@@ -24,18 +26,18 @@ public async ValueTask<T> GetOrCreateAsync(string key, Func<CancellationToken, V
             ? KeyFormatHelper.GetPrefixedKey(key)
             : KeyFormatHelper.GetPrefixedKey(key, _moduleName);
 
-        var lockKey = KeyFormatHelper.GetLockKey(prefixedKey);
+
         var lockValue = Guid.NewGuid().ToString();
 
         while (true)
         {
             token.ThrowIfCancellationRequested();
 
-            var isLocked = await _redisDatabase.Database.KeyExistsAsync(lockKey);
+            var isLocked = await lockService.CheckForLockAsync(prefixedKey);
 
             if (isLocked)
             {
-                await WaitForLockReleaseAsync(lockKey, token);
+                await lockService.WaitForLockReleaseAsync(prefixedKey, token);
                 continue;
             }
 
@@ -45,11 +47,11 @@ public async ValueTask<T> GetOrCreateAsync(string key, Func<CancellationToken, V
                 return cachedValue;
             }
 
-            var lockAcquired = await AcquireLockAsync(lockKey, lockValue);
+            var lockAcquired = await lockService.AcquireLockAsync(prefixedKey, lockValue);
 
             if (!lockAcquired)
             {
-                await WaitForLockReleaseAsync(lockKey, token);
+                await lockService.WaitForLockReleaseAsync(prefixedKey, token);
                 continue;
             }
 
@@ -64,7 +66,7 @@ public async ValueTask<T> GetOrCreateAsync(string key, Func<CancellationToken, V
         }
         finally
         {
-            await ReleaseLockAsync(lockKey, lockValue);
+            await lockService.ReleaseLockAsync(prefixedKey, lockValue);
         }
     }
 
@@ -139,30 +141,4 @@ public async ValueTask RemoveByTagsAsync(IEnumerable<string> tags, CancellationT
         var tasks = tags.Select(tag => RemoveByTagAsync(tag, token).AsTask());
         await Task.WhenAll(tasks);
     }
-
-    private async Task<bool> AcquireLockAsync(string lockKey, string lockValue)
-    {
-        return await _redisDatabase.Database.StringSetAsync(lockKey, lockValue, _lockExpiry, When.NotExists);
-    }
-
-    private async Task WaitForLockReleaseAsync(string lockKey, CancellationToken token)
-    {
-        while (await _redisDatabase.Database.KeyExistsAsync(lockKey) && !token.IsCancellationRequested)
-        {
-            await Task.Delay(_lockRetryDelay, token);
-        }
-    }
-
-    private async Task ReleaseLockAsync(string lockKey, string lockValue)
-    {
-        // Check if the current instance owns the lock before releasing it
-        const string script = @"
-                if redis.call('GET', KEYS[1]) == ARGV[1] then
-                    return redis.call('DEL', KEYS[1])
-                else
-                    return 0
-                end";
-
-        await _redisDatabase.Database.ScriptEvaluateAsync(script, [lockKey], [lockValue]);
-    }
 }

src/DistributedCache/Services/Implementations/RedisLockService.cs

@@ -0,0 +1,50 @@
+using DistributedCache.Helpers;
+using DistributedCache.Options;
+using Microsoft.Extensions.Options;
+using StackExchange.Redis;
+using StackExchange.Redis.Extensions.Core.Abstractions;
+
+namespace DistributedCache.Services.Implementations;
+
+public class RedisLockService(IRedisClient redisClient, IOptions<CacheConfigurationOptions> options)
+{
+    private readonly IRedisDatabase _redisDatabase = redisClient.GetDefaultDatabase();
+    private readonly TimeSpan _lockExpiry = options.Value.DistributedLockDuration;
+    private readonly TimeSpan _lockRetryDelay = TimeSpan.FromMilliseconds(10);
+    
+    public async Task<bool> AcquireLockAsync(string key, string lockValue)
+    {
+        var lockKey = KeyFormatHelper.GetLockKey(key);
+        return await _redisDatabase.Database.StringSetAsync(lockKey, lockValue, _lockExpiry, When.NotExists);
+    }
+    
+    public async Task<bool> CheckForLockAsync(string key)
+    {
+        var lockKey = KeyFormatHelper.GetLockKey(key);
+        return await _redisDatabase.Database.KeyExistsAsync(lockKey);
+    }
+
+    public async Task WaitForLockReleaseAsync(string key, CancellationToken token)
+    {
+        var lockKey = KeyFormatHelper.GetLockKey(key);
+        while (await _redisDatabase.Database.KeyExistsAsync(lockKey) && !token.IsCancellationRequested)
+        {
+            await Task.Delay(_lockRetryDelay, token);
+        }
+    }
+
+    public async Task ReleaseLockAsync(string key, string lockValue)
+    {
+        var lockKey = KeyFormatHelper.GetLockKey(key);
+        
+        const string script = @"
+                if redis.call('GET', KEYS[1]) == ARGV[1] then
+                    return redis.call('DEL', KEYS[1])
+                else
+                    return 0
+                end";
+
+        await _redisDatabase.Database.ScriptEvaluateAsync(script, [lockKey], [lockValue]);
+    }
+    
+}