Implement GRCA (redis/redis#14826)

Author: mgravell

src/RESPite/Messages/RespReader.cs

@@ -1852,30 +1852,54 @@ public readonly decimal ReadDecimal()
     }
 
     /// <summary>
-    /// Read the current element as a <see cref="bool"/> value.
+    /// Try to read the current element as a <see cref="bool"/> value.
     /// </summary>
-    public readonly bool ReadBoolean()
+    /// <param name="value">The parsed boolean value if successful.</param>
+    /// <returns>True if the value was successfully parsed; false otherwise.</returns>
+    public readonly bool TryReadBoolean(out bool value)
     {
         var span = Buffer(stackalloc byte[2]);
         switch (span.Length)
         {
             case 1:
                 switch (span[0])
                 {
-                    case (byte)'0' when Prefix == RespPrefix.Integer: return false;
-                    case (byte)'1' when Prefix == RespPrefix.Integer: return true;
-                    case (byte)'f' when Prefix == RespPrefix.Boolean: return false;
-                    case (byte)'t' when Prefix == RespPrefix.Boolean: return true;
+                    case (byte)'0' when Prefix == RespPrefix.Integer:
+                        value = false;
+                        return true;
+                    case (byte)'1' when Prefix == RespPrefix.Integer:
+                        value = true;
+                        return true;
+                    case (byte)'f' when Prefix == RespPrefix.Boolean:
+                        value = false;
+                        return true;
+                    case (byte)'t' when Prefix == RespPrefix.Boolean:
+                        value = true;
+                        return true;
                 }
 
                 break;
-            case 2 when Prefix == RespPrefix.SimpleString && IsOK(): return true;
+            case 2 when Prefix == RespPrefix.SimpleString && IsOK():
+                value = true;
+                return true;
         }
 
-        ThrowFormatException();
+        value = false;
         return false;
     }
 
+    /// <summary>
+    /// Read the current element as a <see cref="bool"/> value.
+    /// </summary>
+    public readonly bool ReadBoolean()
+    {
+        if (!TryReadBoolean(out var value))
+        {
+            ThrowFormatException();
+        }
+        return value;
+    }
+
     /// <summary>
     /// Parse a scalar value as an enum of type <typeparamref name="T"/>.
     /// </summary>

src/RESPite/PublicAPI/PublicAPI.Unshipped.txt

@@ -155,6 +155,7 @@
 [SER004]RESPite.Messages.RespReader.ProtocolBytesRemaining.get -> long
 [SER004]RESPite.Messages.RespReader.ReadArray<TResult>(RESPite.Messages.RespReader.Projection<TResult>! projection, bool scalar = false) -> TResult[]?
 [SER004]RESPite.Messages.RespReader.ReadBoolean() -> bool
+[SER004]RESPite.Messages.RespReader.TryReadBoolean(out bool value) -> bool
 [SER004]RESPite.Messages.RespReader.ReadByteArray() -> byte[]?
 [SER004]RESPite.Messages.RespReader.ReadDecimal() -> decimal
 [SER004]RESPite.Messages.RespReader.ReadDouble() -> double

src/StackExchange.Redis/Enums/RedisCommand.cs

@@ -60,6 +60,7 @@ internal enum RedisCommand
     GEOSEARCH,
     GEOSEARCHSTORE,
 
+    GCRA,
     GET,
     GETBIT,
     GETDEL,
@@ -323,6 +324,7 @@ internal static bool IsPrimaryOnly(this RedisCommand command)
             case RedisCommand.EXPIREAT:
             case RedisCommand.FLUSHALL:
             case RedisCommand.FLUSHDB:
+            case RedisCommand.GCRA:
             case RedisCommand.GEOSEARCHSTORE:
             case RedisCommand.GETDEL:
             case RedisCommand.GETEX:

src/StackExchange.Redis/ExtensionMethods.cs

@@ -7,6 +7,7 @@
 using System.Security.Authentication;
 using System.Security.Cryptography.X509Certificates;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using Pipelines.Sockets.Unofficial.Arenas;
 
@@ -337,5 +338,69 @@ internal static int VectorSafeIndexOfCRLF(this ReadOnlySpan<byte> span)
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         internal static TTo[]? ToArray<TTo, TState>(in this RawResult result, Projection<RawResult, TState, TTo> selector, in TState state)
             => result.IsNull ? null : result.GetItems().ToArray(selector, in state);
+
+        /// <summary>
+        /// Attempts to acquire a GCRA rate limit token, retrying with delays if rate limited.
+        /// </summary>
+        /// <param name="database">The database instance.</param>
+        /// <param name="key">The key for the rate limiter.</param>
+        /// <param name="maxBurst">The maximum burst size.</param>
+        /// <param name="requestsPerPeriod">The number of requests allowed per period.</param>
+        /// <param name="allow">The maximum time to wait for a successful acquisition.</param>
+        /// <param name="periodSeconds">The period in seconds (default: 1.0).</param>
+        /// <param name="count">The number of tokens to acquire (default: 1).</param>
+        /// <param name="flags">The command flags to use.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns>True if the token was acquired within the allowed time; false otherwise.</returns>
+        public static async ValueTask<bool> TryAcquireGcraAsync(
+            this IDatabaseAsync database,
+            RedisKey key,
+            int maxBurst,
+            int requestsPerPeriod,
+            TimeSpan allow,
+            double periodSeconds = 1.0,
+            int count = 1,
+            CommandFlags flags = CommandFlags.None,
+            CancellationToken cancellationToken = default)
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+
+            var startTime = DateTime.UtcNow;
+            var allowMilliseconds = allow.TotalMilliseconds;
+
+            while (true)
+            {
+                var result = await database.StringGcraRateLimitAsync(key, maxBurst, requestsPerPeriod, periodSeconds, count, flags).ConfigureAwait(false);
+
+                if (!result.Limited)
+                {
+                    return true;
+                }
+
+                var elapsed = (DateTime.UtcNow - startTime).TotalMilliseconds;
+                var remaining = allowMilliseconds - elapsed;
+
+                if (remaining <= 0)
+                {
+                    return false;
+                }
+
+                var delaySeconds = result.RetryAfterSeconds;
+                if (delaySeconds <= 0)
+                {
+                    // Shouldn't happen when Limited is true, but handle defensively
+                    return false;
+                }
+
+                var delayMilliseconds = delaySeconds * 1000.0;
+                if (delayMilliseconds > remaining)
+                {
+                    // Not enough time left to wait for retry
+                    return false;
+                }
+
+                await Task.Delay(TimeSpan.FromSeconds(delaySeconds), cancellationToken).ConfigureAwait(false);
+            }
+        }
     }
 }

src/StackExchange.Redis/Gcra.GcraMessage.cs

@@ -0,0 +1,40 @@
+namespace StackExchange.Redis;
+
+internal partial class RedisDatabase
+{
+    internal sealed class GcraMessage(
+        int database,
+        CommandFlags flags,
+        RedisKey key,
+        int maxBurst,
+        int requestsPerPeriod,
+        double periodSeconds,
+        int count) : Message(database, flags, RedisCommand.GCRA)
+    {
+        protected override void WriteImpl(PhysicalConnection connection)
+        {
+            // GCRA key max_burst requests_per_period period [NUM_REQUESTS count]
+            connection.WriteHeader(Command, ArgCount);
+            connection.WriteBulkString(key);
+            connection.WriteBulkString(maxBurst);
+            connection.WriteBulkString(requestsPerPeriod);
+            connection.WriteBulkString(periodSeconds);
+
+            if (count != 1)
+            {
+                connection.WriteBulkString("NUM_REQUESTS"u8);
+                connection.WriteBulkString(count);
+            }
+        }
+
+        public override int ArgCount
+        {
+            get
+            {
+                int argCount = 4; // key, max_burst, requests_per_period, period
+                if (count != 1) argCount += 2; // NUM_REQUESTS, count
+                return argCount;
+            }
+        }
+    }
+}

src/StackExchange.Redis/Gcra.GcraRateLimitResult.cs

@@ -0,0 +1,45 @@
+namespace StackExchange.Redis;
+
+/// <summary>
+/// Represents the result of a GCRA (Generic Cell Rate Algorithm) rate limit check.
+/// </summary>
+public readonly partial struct GcraRateLimitResult
+{
+    /// <summary>
+    /// Indicates whether the request was rate limited (true) or allowed (false).
+    /// </summary>
+    public bool Limited { get; }
+
+    /// <summary>
+    /// The maximum number of requests allowed. Always equal to max_burst + 1.
+    /// </summary>
+    public int MaxRequests { get; }
+
+    /// <summary>
+    /// The number of requests available immediately without being rate limited.
+    /// </summary>
+    public int AvailableRequests { get; }
+
+    /// <summary>
+    /// The number of seconds after which the caller should retry.
+    /// Returns -1 if the request is not limited.
+    /// </summary>
+    public int RetryAfterSeconds { get; }
+
+    /// <summary>
+    /// The number of seconds after which a full burst will be allowed.
+    /// </summary>
+    public int FullBurstAfterSeconds { get; }
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="GcraRateLimitResult"/> struct.
+    /// </summary>
+    public GcraRateLimitResult(bool limited, int maxRequests, int availableRequests, int retryAfterSeconds, int fullBurstAfterSeconds)
+    {
+        Limited = limited;
+        MaxRequests = maxRequests;
+        AvailableRequests = availableRequests;
+        RetryAfterSeconds = retryAfterSeconds;
+        FullBurstAfterSeconds = fullBurstAfterSeconds;
+    }
+}

src/StackExchange.Redis/Gcra.ResultProcessor.cs

@@ -0,0 +1,70 @@
+namespace StackExchange.Redis;
+
+public readonly partial struct GcraRateLimitResult
+{
+    internal static readonly ResultProcessor<GcraRateLimitResult> Processor = new GcraRateLimitResultProcessor();
+
+    private sealed class GcraRateLimitResultProcessor : ResultProcessor<GcraRateLimitResult>
+    {
+        protected override bool SetResultCore(PhysicalConnection connection, Message message, in RawResult result)
+        {
+            // GCRA returns an array with 5 elements:
+            // 1) <limited> # 0 or 1
+            // 2) <max-req-num> # max number of request. Always equal to max_burst+1
+            // 3) <num-avail-req> # number of requests available immediately
+            // 4) <reply-after> # number of seconds after which caller should retry. Always returns -1 if request isn't limited.
+            // 5) <full-burst-after> # number of seconds after which a full burst will be allowed
+            if (result.Resp2TypeArray == ResultType.Array && result.ItemsCount >= 5)
+            {
+                var items = result.GetItems();
+                bool limited = items[0].GetBoolean();
+                if (items[1].TryGetInt64(out long maxRequests)
+                    && items[2].TryGetInt64(out long availableRequests)
+                    && items[3].TryGetInt64(out long retryAfterSeconds)
+                    && items[4].TryGetInt64(out long fullBurstAfterSeconds))
+                {
+                    var grca = new GcraRateLimitResult(
+                        limited: limited,
+                        maxRequests: (int)maxRequests,
+                        availableRequests: (int)availableRequests,
+                        retryAfterSeconds: (int)retryAfterSeconds,
+                        fullBurstAfterSeconds: (int)fullBurstAfterSeconds);
+                    SetResult(message, grca);
+                    return true;
+                }
+            }
+
+            return false;
+        }
+
+        /* for v3, already done (due to branch choice)
+        protected override bool SetResultCore(PhysicalConnection connection, Message message, ref RespReader reader)
+        {
+            // GCRA returns an array with 5 elements:
+            // 1) <limited> # 0 or 1
+            // 2) <max-req-num> # max number of request. Always equal to max_burst+1
+            // 3) <num-avail-req> # number of requests available immediately
+            // 4) <reply-after> # number of seconds after which caller should retry. Always returns -1 if request isn't limited.
+            // 5) <full-burst-after> # number of seconds after which a full burst will be allowed
+            if (reader.IsAggregate
+                && reader.TryMoveNext() && reader.IsScalar && reader.TryReadBoolean(out bool limited)
+                && reader.TryMoveNext() && reader.IsScalar && reader.TryReadInt64(out long maxRequests)
+                && reader.TryMoveNext() && reader.IsScalar && reader.TryReadInt64(out long availableRequests)
+                && reader.TryMoveNext() && reader.IsScalar && reader.TryReadInt64(out long retryAfterSeconds)
+                && reader.TryMoveNext() && reader.IsScalar && reader.TryReadInt64(out long fullBurstAfterSeconds))
+            {
+                var result = new GcraRateLimitResult(
+                    limited: limited,
+                    maxRequests: (int)maxRequests,
+                    availableRequests: (int)availableRequests,
+                    retryAfterSeconds: (int)retryAfterSeconds,
+                    fullBurstAfterSeconds: (int)fullBurstAfterSeconds);
+                SetResult(message, result);
+                return true;
+            }
+
+            return false;
+        }
+        */
+    }
+}

src/StackExchange.Redis/Interfaces/IDatabase.cs

@@ -3254,6 +3254,19 @@ IEnumerable<SortedSetEntry> SortedSetScan(
         [Experimental(Experiments.Server_8_4, UrlFormat = Experiments.UrlFormat)]
         ValueCondition? StringDigest(RedisKey key, CommandFlags flags = CommandFlags.None);
 
+        /// <summary>
+        /// Performs a GCRA (Generic Cell Rate Algorithm) rate limit check on the specified key.
+        /// </summary>
+        /// <param name="key">The key to rate limit.</param>
+        /// <param name="maxBurst">The maximum burst size.</param>
+        /// <param name="requestsPerPeriod">The number of requests allowed per period.</param>
+        /// <param name="periodSeconds">The period duration in seconds. Default is 1.0.</param>
+        /// <param name="count">The number of requests to consume. Default is 1.</param>
+        /// <param name="flags">The flags to use for this operation.</param>
+        /// <returns>A <see cref="GcraRateLimitResult"/> containing the rate limit decision and metadata.</returns>
+        /// <remarks><seealso href="https://redis.io/commands/gcra"/></remarks>
+        GcraRateLimitResult StringGcraRateLimit(RedisKey key, int maxBurst, int requestsPerPeriod, double periodSeconds = 1.0, int count = 1, CommandFlags flags = CommandFlags.None);
+
         /// <summary>
         /// Get the value of key. If the key does not exist the special value <see cref="RedisValue.Null"/> is returned.
         /// An error is returned if the value stored at key is not a string, because GET only handles string values.

src/StackExchange.Redis/Interfaces/IDatabaseAsync.cs

@@ -796,6 +796,9 @@ IAsyncEnumerable<SortedSetEntry> SortedSetScanAsync(
         [Experimental(Experiments.Server_8_4, UrlFormat = Experiments.UrlFormat)]
         Task<ValueCondition?> StringDigestAsync(RedisKey key, CommandFlags flags = CommandFlags.None);
 
+        /// <inheritdoc cref="IDatabase.StringGcraRateLimit(RedisKey, int, int, double, int, CommandFlags)"/>
+        Task<GcraRateLimitResult> StringGcraRateLimitAsync(RedisKey key, int maxBurst, int requestsPerPeriod, double periodSeconds = 1.0, int count = 1, CommandFlags flags = CommandFlags.None);
+
         /// <inheritdoc cref="IDatabase.StringGet(RedisKey, CommandFlags)"/>
         Task<RedisValue> StringGetAsync(RedisKey key, CommandFlags flags = CommandFlags.None);
 

src/StackExchange.Redis/KeyspaceIsolation/KeyPrefixed.cs

@@ -792,6 +792,9 @@ public Task<long> StringIncrementAsync(RedisKey key, long value = 1, CommandFlag
         public Task<long> StringLengthAsync(RedisKey key, CommandFlags flags = CommandFlags.None) =>
             Inner.StringLengthAsync(ToInner(key), flags);
 
+        public Task<GcraRateLimitResult> StringGcraRateLimitAsync(RedisKey key, int maxBurst, int requestsPerPeriod, double period = 1.0, int count = 1, CommandFlags flags = CommandFlags.None) =>
+            Inner.StringGcraRateLimitAsync(ToInner(key), maxBurst, requestsPerPeriod, period, count, flags);
+
         public Task<bool> StringSetAsync(RedisKey key, RedisValue value, Expiration expiry, ValueCondition when, CommandFlags flags = CommandFlags.None)
             => Inner.StringSetAsync(ToInner(key), value, expiry, when, flags);