fix(plugin/rate-limiting): use redis:eval() to improve rate-limiting accuracy

CHANGELOG.md

@@ -68,6 +68,8 @@
   [#10836](https://github.com/Kong/kong/pull/10836)
 - **ACME**: Fixed sanity test can't work with "kong" storage in Hybrid mode
   [#10852](https://github.com/Kong/kong/pull/10852)
+- **rate-limiting**: Fixed an issue that impact the accuracy with the `redis` policy.
+  [#10559](https://github.com/Kong/kong/pull/10559)
 
 #### PDK
 

kong/plugins/rate-limiting/policies/init.lua

@@ -179,52 +179,9 @@ return {
   },
   ["redis"] = {
     increment = function(conf, limits, identifier, current_timestamp, value)
-      local red, err = get_redis_connection(conf)
-      if not red then
-        return nil, err
-      end
-
-      local keys = {}
-      local expiration = {}
-      local idx = 0
-      local periods = timestamp.get_timestamps(current_timestamp)
-      for period, period_date in pairs(periods) do
-        if limits[period] then
-          local cache_key = get_local_key(conf, identifier, period, period_date)
-          local exists, err = red:exists(cache_key)
-          if err then
-            kong.log.err("failed to query Redis: ", err)
-            return nil, err
-          end
-
-          idx = idx + 1
-          keys[idx] = cache_key
-          if not exists or exists == 0 then
-            expiration[idx] = EXPIRATION[period]
-          end
-
-          red:init_pipeline()
-          for i = 1, idx do
-            red:incrby(keys[i], value)
-            if expiration[i] then
-              red:expire(keys[i], expiration[i])
-            end
-          end
-        end
-      end
-
-      local _, err = red:commit_pipeline()
-      if err then
-        kong.log.err("failed to commit increment pipeline in Redis: ", err)
-        return nil, err
-      end
-
-      local ok, err = red:set_keepalive(10000, 100)
-      if not ok then
-        kong.log.err("failed to set Redis keepalive: ", err)
-        return nil, err
-      end
+            -- the usage function already incremented the value of redis key to avoid race condition in concurrent calls
 
+      -- because of that we don't need to increment here
       return true
     end,
     usage = function(conf, identifier, period, current_timestamp)
@@ -238,7 +195,17 @@ return {
       local periods = timestamp.get_timestamps(current_timestamp)
       local cache_key = get_local_key(conf, identifier, period, periods[period])
 
-      local current_metric, err = red:get(cache_key)
+      -- the usage of redis command incr instead of get is to avoid race conditions in concurrent calls
+      local current_metric, err = red:eval([[
+        local cache_key, expiration = KEYS[1], ARGV[1]
+        local result_incr = redis.call("incr", cache_key)
+        if result_incr == 1 then
+          redis.call("expire", cache_key, expiration)
+        end
+
+        return result_incr - 1
+      ]], 1, cache_key, EXPIRATION[period])
+
       if err then
         return nil, err
       end
@@ -255,4 +222,4 @@ return {
       return current_metric or 0
     end
   }
-}
+}

spec/03-plugins/23-rate-limiting/02-policies_spec.lua

@@ -2,6 +2,10 @@ local uuid      = require("kong.tools.utils").uuid
 local helpers   = require "spec.helpers"
 local timestamp = require "kong.tools.timestamp"
 
+--[[
+  basically a copy of `get_local_key()` 
+  in `kong/plugins/rate-limiting/policies/init.lua`
+--]]
 local get_local_key = function(conf, identifier, period, period_date)
   return string.format("ratelimit:%s:%s:%s:%s:%s",
     conf.route_id, conf.service_id, identifier, period_date, period)
@@ -142,4 +146,51 @@ describe("Plugin: rate-limiting (policies)", function()
     end)
   end
 
+  describe("redis", function()
+    local identifier = uuid()
+    local conf       = {
+      route_id = uuid(),
+      service_id = uuid(),
+      redis_host = helpers.redis_host,
+      redis_port = helpers.redis_port,
+      redis_database = 0,
+    }
+
+    before_each(function()
+      local red = require "resty.redis"
+      local redis = assert(red:new())
+      redis:set_timeout(1000)
+      assert(redis:connect(conf.redis_host, conf.redis_port))
+      redis:flushall()
+      redis:close()
+    end)
+
+    it("increase & usage", function()
+      --[[
+        Just a simple test:
+        - increase 1
+        - check usage == 1
+        - increase 1
+        - check usage == 2
+        - increase 1 (beyond the limit)
+        - check usage == 3
+      --]]
+
+      local current_timestamp = 1424217600
+      local periods = timestamp.get_timestamps(current_timestamp)
+
+      for period in pairs(periods) do
+
+        local metric = assert(policies.redis.usage(conf, identifier, period, current_timestamp))
+        assert.equal(0, metric)
+
+        for i = 1, 3 do
+          assert(policies.redis.increment(conf, { [period] = 2 }, identifier, current_timestamp, 1))
+          metric = assert(policies.redis.usage(conf, identifier, period, current_timestamp))
+          assert.equal(i, metric)
+        end
+      end
+    end)
+  end)
+
 end)

spec/03-plugins/23-rate-limiting/05-integration_spec.lua

@@ -32,7 +32,7 @@ local function flush_redis(red, db)
 end
 
 local function add_redis_user(red)
-  assert(red:acl("setuser", REDIS_USER_VALID, "on", "allkeys", "+incrby", "+select", "+info", "+expire", "+get", "+exists", ">" .. REDIS_PASSWORD))
+  assert(red:acl("setuser", REDIS_USER_VALID, "on", "allkeys", "allcommands", ">" .. REDIS_PASSWORD))
   assert(red:acl("setuser", REDIS_USER_INVALID, "on", "allkeys", "+get", ">" .. REDIS_PASSWORD))
 end
 
@@ -92,7 +92,7 @@ describe("Plugin: rate-limiting (integration)", function()
     describe("config.policy = redis #" .. strategy, function()
       -- Regression test for the following issue:
       -- https://github.com/Kong/kong/issues/3292
-  
+
       lazy_setup(function()
         flush_redis(red, REDIS_DB_1)
         flush_redis(red, REDIS_DB_2)
@@ -108,7 +108,7 @@ describe("Plugin: rate-limiting (integration)", function()
         }, {
           "rate-limiting"
         })
-  
+
         local route1 = assert(bp.routes:insert {
           hosts        = { "redistest1.com" },
         })
@@ -128,7 +128,7 @@ describe("Plugin: rate-limiting (integration)", function()
             redis_timeout     = 10000,
           },
         })
-  
+
         local route2 = assert(bp.routes:insert {
           hosts        = { "redistest2.com" },
         })
@@ -148,7 +148,7 @@ describe("Plugin: rate-limiting (integration)", function()
             redis_timeout     = 10000,
           },
         })
-  
+
         if red_version >= version("6.0.0") then
           local route3 = assert(bp.routes:insert {
             hosts        = { "redistest3.com" },
@@ -171,7 +171,7 @@ describe("Plugin: rate-limiting (integration)", function()
               redis_timeout     = 10000,
             },
           })
-  
+
           local route4 = assert(bp.routes:insert {
             hosts        = { "redistest4.com" },
           })
@@ -194,37 +194,37 @@ describe("Plugin: rate-limiting (integration)", function()
             },
           })
         end
-  
-  
+
+
         assert(helpers.start_kong({
           nginx_conf = "spec/fixtures/custom_nginx.template",
           lua_ssl_trusted_certificate = config.lua_ssl_trusted_certificate,
         }))
         client = helpers.proxy_client()
       end)
-  
+
       lazy_teardown(function()
         helpers.stop_kong()
         if red_version >= version("6.0.0") then
           remove_redis_user(red)
         end
       end)
-  
+
       it("connection pool respects database setting", function()
         assert(red:select(REDIS_DB_1))
         local size_1 = assert(red:dbsize())
-  
+
         assert(red:select(REDIS_DB_2))
         local size_2 = assert(red:dbsize())
-  
+
         assert.equal(0, tonumber(size_1))
         assert.equal(0, tonumber(size_2))
         if red_version >= version("6.0.0") then
           assert(red:select(REDIS_DB_3))
           local size_3 = assert(red:dbsize())
           assert.equal(0, tonumber(size_3))
         end
-  
+
         local res = assert(client:send {
           method = "GET",
           path = "/status/200",
@@ -233,27 +233,27 @@ describe("Plugin: rate-limiting (integration)", function()
           }
         })
         assert.res_status(200, res)
-  
+
         -- Wait for async timer to increment the limit
-  
+
         ngx.sleep(SLEEP_TIME)
-  
+
         assert(red:select(REDIS_DB_1))
         size_1 = assert(red:dbsize())
-  
+
         assert(red:select(REDIS_DB_2))
         size_2 = assert(red:dbsize())
-  
+
         -- TEST: DB 1 should now have one hit, DB 2 and 3 none
-  
+
         assert.equal(1, tonumber(size_1))
         assert.equal(0, tonumber(size_2))
         if red_version >= version("6.0.0") then
           assert(red:select(REDIS_DB_3))
           local size_3 = assert(red:dbsize())
           assert.equal(0, tonumber(size_3))
         end
-  
+
         -- rate-limiting plugin will reuses the redis connection
         local res = assert(client:send {
           method = "GET",
@@ -263,27 +263,27 @@ describe("Plugin: rate-limiting (integration)", function()
           }
         })
         assert.res_status(200, res)
-  
+
         -- Wait for async timer to increment the limit
-  
+
         ngx.sleep(SLEEP_TIME)
-  
+
         assert(red:select(REDIS_DB_1))
         size_1 = assert(red:dbsize())
-  
+
         assert(red:select(REDIS_DB_2))
         size_2 = assert(red:dbsize())
-  
+
         -- TEST: DB 1 and 2 should now have one hit, DB 3 none
-  
+
         assert.equal(1, tonumber(size_1))
         assert.equal(1, tonumber(size_2))
         if red_version >= version("6.0.0") then
           assert(red:select(REDIS_DB_3))
           local size_3 = assert(red:dbsize())
           assert.equal(0, tonumber(size_3))
         end
-  
+
         if red_version >= version("6.0.0") then
           -- rate-limiting plugin will reuses the redis connection
           local res = assert(client:send {
@@ -294,30 +294,30 @@ describe("Plugin: rate-limiting (integration)", function()
             }
           })
           assert.res_status(200, res)
-  
+
           -- Wait for async timer to increment the limit
-  
+
           ngx.sleep(SLEEP_TIME)
-  
+
           assert(red:select(REDIS_DB_1))
           size_1 = assert(red:dbsize())
-  
+
           assert(red:select(REDIS_DB_2))
           size_2 = assert(red:dbsize())
-  
+
           assert(red:select(REDIS_DB_3))
           local size_3 = assert(red:dbsize())
-  
+
           -- TEST: All DBs should now have one hit, because the
           -- plugin correctly chose to select the database it is
           -- configured to hit
-  
+
           assert.is_true(tonumber(size_1) > 0)
           assert.is_true(tonumber(size_2) > 0)
           assert.is_true(tonumber(size_3) > 0)
         end
       end)
-  
+
       it("authenticates and executes with a valid redis user having proper ACLs", function()
         if red_version >= version("6.0.0") then
           local res = assert(client:send {
@@ -333,7 +333,7 @@ describe("Plugin: rate-limiting (integration)", function()
             "'authenticates and executes with a valid redis user having proper ACLs' will be skipped")
         end
       end)
-  
+
       it("fails to rate-limit for a redis user with missing ACLs", function()
         if red_version >= version("6.0.0") then
           local res = assert(client:send {
@@ -349,7 +349,7 @@ describe("Plugin: rate-limiting (integration)", function()
             "'fails to rate-limit for a redis user with missing ACLs' will be skipped")
         end
       end)
-  
+
     end)
   end