Merge pull request Kong#952 from Mashape/plugins/ratelimiting-async

Supporting async increment for rate-limiting plugin

kong/plugins/rate-limiting/handler.lua

@@ -9,7 +9,6 @@ local RateLimitingHandler = BasePlugin:extend()
 
 RateLimitingHandler.PRIORITY = 900
 
-
 local function get_identifier()
   local identifier
 
@@ -31,6 +30,15 @@ local function increment(api_id, identifier, current_timestamp, value)
   end
 end
 
+local function increment_async(premature, api_id, identifier, current_timestamp, value)
+  if premature then return end
+
+  local _, stmt_err = dao.ratelimiting_metrics:increment(api_id, identifier, current_timestamp, value)
+  if stmt_err then
+    ngx.log(ngx.ERR, "failed to increment: ", tostring(stmt_err))
+  end
+end
+
 local function get_usage(api_id, identifier, current_timestamp, limits)
   local usage = {}
   local stop
@@ -59,7 +67,6 @@ local function get_usage(api_id, identifier, current_timestamp, limits)
   return usage, stop
 end
 
-
 function RateLimitingHandler:new()
   RateLimitingHandler.super.new(self, "rate-limiting")
 end
@@ -70,10 +77,11 @@ function RateLimitingHandler:access(conf)
 
   -- Consumer is identified by ip address or authenticated_credential id
   local identifier = get_identifier()
-
   local api_id = ngx.ctx.api.id
+  local is_async = conf.async
 
   -- Load current metric for configured period
+  conf.async = nil
   local usage, stop = get_usage(api_id, identifier, current_timestamp, conf)
 
   -- Adding headers
@@ -88,7 +96,14 @@ function RateLimitingHandler:access(conf)
   end
 
   -- Increment metrics for all periods if the request goes through
-  increment(api_id, identifier, current_timestamp, 1)
+  if is_async then
+    local ok, err = ngx.timer.at(0, increment_async, api_id, identifier, current_timestamp, 1)
+    if not ok then
+      ngx.log(ngx.ERR, "failed to create timer: ", err)
+    end
+  else
+    increment(api_id, identifier, current_timestamp, 1)
+  end
 end
 
 return RateLimitingHandler

kong/plugins/rate-limiting/schema.lua

@@ -8,7 +8,8 @@ return {
     hour = { type = "number" },
     day = { type = "number" },
     month = { type = "number" },
-    year = { type = "number" }
+    year = { type = "number" },
+    async = { type = "boolean", default = false }
   },
   self_check = function(schema, plugin_t, dao, is_update)
     local ordered_periods = { "second", "minute", "hour", "day", "month", "year"}

spec/plugins/rate-limiting/access_spec.lua

@@ -24,7 +24,8 @@ describe("RateLimiting Plugin", function()
         { name = "tests-rate-limiting1", request_host = "test3.com", upstream_url = "http://mockbin.com" },
         { name = "tests-rate-limiting2", request_host = "test4.com", upstream_url = "http://mockbin.com" },
         { name = "tests-rate-limiting3", request_host = "test5.com", upstream_url = "http://mockbin.com" },
-        { name = "tests-rate-limiting4", request_host = "test6.com", upstream_url = "http://mockbin.com" }
+        { name = "tests-rate-limiting4", request_host = "test6.com", upstream_url = "http://mockbin.com" },
+        { name = "tests-rate-limiting5", request_host = "test7.com", upstream_url = "http://mockbin.com" }
       },
       consumer = {
         { custom_id = "provider_123" },
@@ -36,7 +37,8 @@ describe("RateLimiting Plugin", function()
         { name = "rate-limiting", config = { minute = 8 }, __api = 1, __consumer = 1 },
         { name = "rate-limiting", config = { minute = 6 }, __api = 2 },
         { name = "rate-limiting", config = { minute = 3, hour = 5 }, __api = 3 },
-        { name = "rate-limiting", config = { minute = 33 }, __api = 4 }
+        { name = "rate-limiting", config = { minute = 33 }, __api = 4 },
+        { name = "rate-limiting", config = { minute = 6, async = true }, __api = 5 }
       },
       keyauth_credential = {
         { key = "apikey122", __consumer = 1 },
@@ -144,4 +146,27 @@ describe("RateLimiting Plugin", function()
 
     end)
   end)
+
+  describe("Async increment", function()
+
+    it("should increment asynchronously", function()
+      -- Default rate-limiting plugin for this API says 6/minute
+        local limit = 6
+
+        for i = 1, limit do
+          local _, status, headers = http_client.get(STUB_GET_URL, {}, {host = "test7.com"})
+          assert.are.equal(200, status)
+          assert.are.same(tostring(limit), headers["x-ratelimit-limit-minute"])
+          assert.are.same(tostring(limit - i), headers["x-ratelimit-remaining-minute"])
+        end
+
+        os.execute("sleep 2") -- Wait for timers to increment
+
+        local response, status = http_client.get(STUB_GET_URL, {}, {host = "test7.com"})
+        local body = cjson.decode(response)
+        assert.are.equal(429, status)
+        assert.are.equal("API rate limit exceeded", body.message)
+    end)
+
+  end)
 end)