feat(ai-rate-limiting): support redis policy for shared counters

apisix/plugins/ai-rate-limiting.lua

@@ -25,6 +25,7 @@ local math_floor = math.floor
 local math_huge = math.huge
 local core = require("apisix.core")
 local limit_count = require("apisix.plugins.limit-count.init")
+local policy_to_additional_properties = limit_count.policy_to_additional_properties
 
 local plugin_name = "ai-rate-limiting"
 
@@ -90,6 +91,12 @@ local schema = {
         rejected_msg = {
             type = "string", minLength = 1
         },
+        policy = {
+            type = "string",
+            enum = {"local", "redis", "redis-cluster", "redis-sentinel"},
+            default = "local",
+        },
+        allow_degradation = {type = "boolean", default = false},
         rules = {
             type = "array",
             items = {
@@ -135,7 +142,31 @@ local schema = {
         {
             required = {"rules"},
         }
-    }
+    },
+    ["if"] = {
+        properties = {
+            policy = {
+                enum = {"redis"},
+            },
+        },
+    },
+    ["then"] = policy_to_additional_properties.redis,
+    ["else"] = {
+        ["if"] = {
+            properties = {
+                policy = {
+                    enum = {"redis-cluster"},
+                },
+            },
+        },
+        ["then"] = policy_to_additional_properties["redis-cluster"],
+        ["else"] = {
+            ["if"] = {
+                properties = { policy = { enum = { "redis-sentinel" } } },
+            },
+            ["then"] = policy_to_additional_properties["redis-sentinel"],
+        },
+    },
 }
 
 local _M = {
@@ -191,43 +222,69 @@ end
 
 local function transform_limit_conf(plugin_conf, instance_conf, instance_name)
     local limit_conf = {
+        _meta = plugin_conf._meta,
         rejected_code = plugin_conf.rejected_code,
         rejected_msg = plugin_conf.rejected_msg,
         show_limit_quota_header = plugin_conf.show_limit_quota_header,
 
-        -- we may expose those fields to ai-rate-limiting later
-        policy = "local",
+        -- counters can be shared across nodes via redis policies
+        policy = plugin_conf.policy or "local",
         key_type = "constant",
-        allow_degradation = false,
+        allow_degradation = plugin_conf.allow_degradation,
         sync_interval = -1,
         limit_header = "X-AI-RateLimit-Limit",
         remaining_header = "X-AI-RateLimit-Remaining",
         reset_header = "X-AI-RateLimit-Reset",
     }
+    local name = instance_name or ""
     if plugin_conf.rules and #plugin_conf.rules > 0 then
         limit_conf.rules = plugin_conf.rules
-        limit_conf._meta = plugin_conf._meta
-        return limit_conf
+    else
+        local key = plugin_name .. "#global"
+        local limit = plugin_conf.limit
+        local time_window = plugin_conf.time_window
+        if instance_conf then
+            name = instance_conf.name
+            key = instance_conf.name
+            limit = instance_conf.limit
+            time_window = instance_conf.time_window
+        end
+        limit_conf._vid = key
+        limit_conf.key = key
+        limit_conf.count = limit
+        limit_conf.time_window = time_window
+        limit_conf.limit_header = "X-AI-RateLimit-Limit-" .. name
+        limit_conf.remaining_header = "X-AI-RateLimit-Remaining-" .. name
+        limit_conf.reset_header = "X-AI-RateLimit-Reset-" .. name
     end
 
-    local key = plugin_name .. "#global"
-    local limit = plugin_conf.limit
-    local time_window = plugin_conf.time_window
-    local name = instance_name or ""
-    if instance_conf then
-        name = instance_conf.name
-        key = instance_conf.name
-        limit = instance_conf.limit
-        time_window = instance_conf.time_window
+    if plugin_conf.policy == "redis" then
+        limit_conf.redis_host = plugin_conf.redis_host
+        limit_conf.redis_port = plugin_conf.redis_port
+        limit_conf.redis_username = plugin_conf.redis_username
+        limit_conf.redis_password = plugin_conf.redis_password
+        limit_conf.redis_database = plugin_conf.redis_database
+        limit_conf.redis_timeout = plugin_conf.redis_timeout
+        limit_conf.redis_ssl = plugin_conf.redis_ssl
+        limit_conf.redis_ssl_verify = plugin_conf.redis_ssl_verify
+    elseif plugin_conf.policy == "redis-cluster" then
+        limit_conf.redis_cluster_nodes = plugin_conf.redis_cluster_nodes
+        limit_conf.redis_cluster_name = plugin_conf.redis_cluster_name
+        limit_conf.redis_password = plugin_conf.redis_password
+        limit_conf.redis_timeout = plugin_conf.redis_timeout
+        limit_conf.redis_cluster_ssl = plugin_conf.redis_cluster_ssl
+        limit_conf.redis_cluster_ssl_verify = plugin_conf.redis_cluster_ssl_verify
+    elseif plugin_conf.policy == "redis-sentinel" then
+        limit_conf.redis_sentinels = plugin_conf.redis_sentinels
+        limit_conf.redis_master_name = plugin_conf.redis_master_name
+        limit_conf.sentinel_username = plugin_conf.sentinel_username
+        limit_conf.sentinel_password = plugin_conf.sentinel_password
+        limit_conf.redis_role = plugin_conf.redis_role
+        limit_conf.redis_connect_timeout = plugin_conf.redis_connect_timeout
+        limit_conf.redis_read_timeout = plugin_conf.redis_read_timeout
+        limit_conf.redis_keepalive_timeout = plugin_conf.redis_keepalive_timeout
+        limit_conf.redis_database = plugin_conf.redis_database
     end
-    limit_conf._vid = key
-    limit_conf.key = key
-    limit_conf._meta = plugin_conf._meta
-    limit_conf.count = limit
-    limit_conf.time_window = time_window
-    limit_conf.limit_header = "X-AI-RateLimit-Limit-" .. name
-    limit_conf.remaining_header = "X-AI-RateLimit-Remaining-" .. name
-    limit_conf.reset_header = "X-AI-RateLimit-Reset-" .. name
     return limit_conf
 end
 

apisix/plugins/limit-count/init.lua

@@ -227,6 +227,7 @@ local schema = {
 local schema_copy = core.table.deepcopy(schema)
 
 local _M = {
+    policy_to_additional_properties = policy_to_additional_properties,
     schema = schema,
     metadata_schema = metadata_schema,
 }

docs/en/latest/plugins/ai-rate-limiting.md

@@ -55,6 +55,28 @@ The `ai-rate-limiting` Plugin enforces token-based rate limiting for requests se
 | instances.time_window | integer | True | | >0 | The time interval corresponding to the rate limiting `limit` in seconds for an instance. |
 | rejected_code | integer | False | 503 | [200, 599] | The HTTP status code returned when a request exceeding the quota is rejected. |
 | rejected_msg | string | False | | | The response body returned when a request exceeding the quota is rejected. |
+| policy | string | False | local | [`local`, `redis`, `redis-cluster`, `redis-sentinel`] | The policy for the rate limiting counter. If it is `local`, the counter is stored in local memory. If it is `redis`, the counter is stored on a Redis instance. If it is `redis-cluster`, the counter is stored in a Redis cluster. If it is `redis-sentinel`, the counter is stored on the Redis master discovered through Sentinel. Use a Redis-based policy to share token counters across gateway nodes in a distributed deployment. |
+| allow_degradation | boolean | False | false | | If true, allow APISIX to continue serving requests without the Plugin when the Redis backend is unavailable. |
+| redis_host | string | False | | | The address of the Redis node. Required when `policy` is `redis`. |
+| redis_port | integer | False | 6379 | [1,...] | The port of the Redis node when `policy` is `redis`. |
+| redis_username | string | False | | | The username for Redis if Redis ACL is used. If you use the legacy authentication method `requirepass`, configure only the `redis_password`. Used when `policy` is `redis` or `redis-sentinel`. |
+| redis_password | string | False | | | The password of the Redis node when `policy` is `redis`, `redis-cluster`, or `redis-sentinel`. |
+| redis_database | integer | False | 0 | >= 0 | The database number in Redis when `policy` is `redis` or `redis-sentinel`. |
+| redis_timeout | integer | False | 1000 | [1,...] | The Redis timeout value in milliseconds when `policy` is `redis` or `redis-cluster`. |
+| redis_ssl | boolean | False | false | | If true, use SSL to connect to Redis when `policy` is `redis`. |
+| redis_ssl_verify | boolean | False | false | | If true, verify the server SSL certificate when `policy` is `redis`. |
+| redis_cluster_nodes | array[string] | False | | | The list of Redis cluster nodes with at least one address. Required when `policy` is `redis-cluster`. |
+| redis_cluster_name | string | False | | | The name of the Redis cluster. Required when `policy` is `redis-cluster`. |
+| redis_cluster_ssl | boolean | False | false | | If true, use SSL to connect to Redis when `policy` is `redis-cluster`. |
+| redis_cluster_ssl_verify | boolean | False | false | | If true, verify the server SSL certificate when `policy` is `redis-cluster`. |
+| redis_sentinels | array[object] | False | | | The list of Sentinel nodes. Required when `policy` is `redis-sentinel`. Each item must contain `host` and `port`. |
+| redis_master_name | string | False | | | The Redis master name monitored by Sentinel. Required when `policy` is `redis-sentinel`. |
+| sentinel_username | string | False | | | The username for Sentinel when `policy` is `redis-sentinel`. |
+| sentinel_password | string | False | | | The password for Sentinel when `policy` is `redis-sentinel`. |
+| redis_role | string | False | master | [`master`, `slave`] | The Redis role selected through Sentinel when `policy` is `redis-sentinel`. |
+| redis_connect_timeout | integer | False | 1000 | [1,...] | The Redis connection timeout in milliseconds when `policy` is `redis-sentinel`. |
+| redis_read_timeout | integer | False | 1000 | [1,...] | The Redis read timeout in milliseconds when `policy` is `redis-sentinel`. |
+| redis_keepalive_timeout | integer | False | 60000 | [1,...] | The Redis keepalive timeout in milliseconds when `policy` is `redis-sentinel`. |
 | rules | array[object] | False | | | An array of rate-limiting rules that are applied sequentially. If configured, this takes precedence over `limit` and `time_window`. |
 | rules.count | integer or string | True | | >0 or variable expression | The maximum number of tokens allowed within a given time interval. Can be a static integer or a variable expression like `$http_custom_limit`. |
 | rules.time_window | integer or string | True | | >0 or variable expression | The time interval corresponding to the rate limiting `count` in seconds. Can be a static integer or a variable expression. |
@@ -286,6 +308,63 @@ You should receive a response similar to the following:
 
 If the rate limiting quota of 300 prompt tokens has been consumed in a 30-second window, all additional requests will be rejected.
 
+### Share Rate Limiting Counters Across Nodes with Redis
+
+By default, the `local` policy keeps token counters in each node's shared memory, so counters are not shared across a multi-node deployment and the effective quota scales with the number of nodes. Set `policy` to `redis`, `redis-cluster`, or `redis-sentinel` to centralize counters and enforce a single quota across all nodes.
+
+The following example demonstrates how you can use the `redis` policy to share the rate limiting counter across gateway nodes.
+
+Create a Route as such and update with your LLM providers, models, API keys, endpoints, and Redis address, if applicable:
+
+```shell
+curl "http://127.0.0.1:9180/apisix/admin/routes" -X PUT \
+  -H "X-API-KEY: ${admin_key}" \
+  -d '{
+    "id": "ai-rate-limiting-route",
+    "uri": "/anything",
+    "methods": ["POST"],
+    "plugins": {
+      "ai-proxy": {
+        "provider": "openai",
+        "auth": {
+          "header": {
+            "Authorization": "Bearer '"$OPENAI_API_KEY"'"
+          }
+        },
+        "options": {
+          "model": "gpt-35-turbo-instruct",
+          "max_tokens": 512,
+          "temperature": 1.0
+        }
+      },
+      "ai-rate-limiting": {
+        "limit": 300,
+        "time_window": 30,
+        "limit_strategy": "total_tokens",
+        "policy": "redis",
+        "redis_host": "127.0.0.1",
+        "redis_port": 6379,
+        "redis_database": 1
+      }
+    }
+  }'
+```
+
+Send a POST request to the Route with a sample question in the request body from any node:
+
+```shell
+curl "http://127.0.0.1:9080/anything" -X POST \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [
+      { "role": "system", "content": "You are a mathematician" },
+      { "role": "user", "content": "What is 1+1?" }
+    ]
+  }'
+```
+
+The consumed tokens are counted against the shared Redis counter. Once the quota of 300 tokens is consumed in a 30-second window, additional requests are rejected on every node.
+
 ### Rate Limit One Instance Among Multiple
 
 The following example demonstrates how you can use `ai-proxy-multi` to configure two models for load balancing, forwarding 80% of the traffic to one instance and 20% to the other. Additionally, use `ai-rate-limiting` to configure token-based rate limiting on the instance that receives 80% of the traffic, such that when the configured quota is fully consumed, the additional traffic will be forwarded to the other instance.

t/plugin/ai-rate-limiting.t

@@ -1538,3 +1538,133 @@ X-AI-Fixture: openai/chat-model-echo.json
 --- error_code: 200
 --- no_error_log
 [error]
+
+
+
+=== TEST 35: schema check, redis policy requires redis_host
+--- config
+    location /t {
+        content_by_lua_block {
+            local plugin = require("apisix.plugins.ai-rate-limiting")
+            local ok, err = plugin.check_schema({
+                limit = 30,
+                time_window = 60,
+                policy = "redis",
+            })
+            if not ok then
+                ngx.say(err)
+            else
+                ngx.say("passed")
+            end
+        }
+    }
+--- response_body
+then clause did not match
+
+
+
+=== TEST 36: flush redis and set route with redis policy
+--- config
+    location /t {
+        content_by_lua_block {
+            local redis = require("resty.redis")
+            local red = redis:new()
+            red:set_timeout(1000)
+            local ok, err = red:connect("127.0.0.1", 6379)
+            if not ok then
+                ngx.say("failed to connect: ", err)
+                return
+            end
+            red:flushall()
+
+            local t = require("lib.test_admin").test
+            local code, body = t('/apisix/admin/routes/1',
+                 ngx.HTTP_PUT,
+                 [[{
+                    "uri": "/ai",
+                    "plugins": {
+                        "ai-proxy": {
+                            "provider": "openai",
+                            "auth": {
+                                "header": {
+                                    "Authorization": "Bearer token"
+                                }
+                            },
+                            "options": {
+                                "model": "gpt-35-turbo-instruct",
+                                "max_tokens": 512,
+                                "temperature": 1.0
+                            },
+                            "override": {
+                                "endpoint": "http://127.0.0.1:1980"
+                            },
+                            "ssl_verify": false
+                        },
+                        "ai-rate-limiting": {
+                            "limit": 30,
+                            "time_window": 60,
+                            "policy": "redis",
+                            "redis_host": "127.0.0.1",
+                            "redis_port": 6379,
+                            "redis_database": 1
+                        }
+                    },
+                    "upstream": {
+                        "type": "roundrobin",
+                        "nodes": {
+                            "canbeanything.com": 1
+                        }
+                    }
+                }]]
+            )
+
+            if code >= 300 then
+                ngx.status = code
+            end
+            ngx.say(body)
+        }
+    }
+--- response_body
+passed
+
+
+
+=== TEST 37: redis policy shares counter and rejects the 4th request
+--- pipelined_requests eval
+[
+    "POST /ai\n" . "{ \"messages\": [ { \"role\": \"system\", \"content\": \"You are a mathematician\" }, { \"role\": \"user\", \"content\": \"What is 1+1?\"} ] }",
+    "POST /ai\n" . "{ \"messages\": [ { \"role\": \"system\", \"content\": \"You are a mathematician\" }, { \"role\": \"user\", \"content\": \"What is 1+1?\"} ] }",
+    "POST /ai\n" . "{ \"messages\": [ { \"role\": \"system\", \"content\": \"You are a mathematician\" }, { \"role\": \"user\", \"content\": \"What is 1+1?\"} ] }",
+    "POST /ai\n" . "{ \"messages\": [ { \"role\": \"system\", \"content\": \"You are a mathematician\" }, { \"role\": \"user\", \"content\": \"What is 1+1?\"} ] }",
+]
+--- more_headers
+Authorization: Bearer token
+X-AI-Fixture: openai/chat-model-echo.json
+--- error_code eval
+[200, 200, 200, 503]
+
+
+
+=== TEST 38: counter is stored in redis, not local memory
+--- config
+    location /t {
+        content_by_lua_block {
+            local redis = require("resty.redis")
+            local red = redis:new()
+            red:set_timeout(1000)
+            local ok, err = red:connect("127.0.0.1", 6379)
+            if not ok then
+                ngx.say("failed to connect: ", err)
+                return
+            end
+            red:select(1)
+            local keys, err = red:keys("*ai-rate-limiting*")
+            if not keys then
+                ngx.say("failed to get keys: ", err)
+                return
+            end
+            ngx.say(#keys > 0 and "counter in redis" or "no counter")
+        }
+    }
+--- response_body
+counter in redis