From 2210dccf9370bfe2a6bbdf3f8030ec21ed951c74 Mon Sep 17 00:00:00 2001
From: Michal Cichra <michal@cichra.cz>
Date: Wed, 24 Jan 2018 19:28:16 +0100
Subject: [PATCH] proof of concept to use redis for distributed rate limiting

---
 src/conn-limiter.lua | 84 +++++++++++++++++++++++++++++++++++---------
 1 file changed, 67 insertions(+), 17 deletions(-)

diff --git a/src/conn-limiter.lua b/src/conn-limiter.lua
index ec5f557..6a3f113 100644
--- a/src/conn-limiter.lua
+++ b/src/conn-limiter.lua
@@ -11,6 +11,26 @@ local _M = {
 
 }
 
+local ngx_semaphore = require "ngx.semaphore"
+
+local resty_redis = require('resty.redis')
+
+local assert = assert
+
+
+local function redis_shdict(host, port)
+  local redis = assert(resty_redis:new())
+
+  if redis:connect(host or '127.0.0.1', port or 6379) then
+    return {
+      incr = function(_, key, increment, _)
+        return redis:incrby(key, increment)
+      end,
+      close = function() return redis:set_keepalive() end
+    }
+  end
+end
+
 
 function _M:access()
   -- limit the requests under 200 concurrent requests (normally just
@@ -28,6 +48,8 @@ function _M:access()
     return ngx.exit(500)
   end
 
+  lim.dict = redis_shdict()
+
   -- the following call must be per-request.
   -- here we use the remote (IP) address as the limiting key
   local key = ngx.var.host
@@ -48,8 +70,11 @@ function _M:access()
     ctx.limit_conn_key = key
     ctx.limit_conn_delay = delay
     ctx.metric_labels = labels
+    ctx.redis = redis
   end
 
+  lim.dict:close()
+
   -- the 2nd return value holds the current concurrency level
   -- for the specified key.
   local conn = err
@@ -79,28 +104,53 @@ function _M:metrics()
   return prometheus:collect()
 end
 
-function _M:log()
+local function checkin(_, ctx, time, semaphore)
+  local lim = ctx.limit_conn
+
+  lim.dict = redis_shdict()
+  -- if you are using an upstream module in the content phase,
+  -- then you probably want to use $upstream_response_time
+  -- instead of ($request_time - ctx.limit_conn_delay) below.
+  local latency = tonumber(time) - ctx.limit_conn_delay
+  local key = ctx.limit_conn_key
+  assert(key)
+  local conn, err = lim:leaving(key, latency)
+
+  if not conn then
+    ngx.log(ngx.ERR,
+      "failed to record the connection leaving ",
+      "request: ", err)
+    return
+  end
+
+  opened_connections:set(conn,  ctx.metric_labels)
+
+  lim.dict:close()
+
+  if semaphore then
+    semaphore:post(1)
+  end
+end
+
+function _M:post_action()
   local ctx = ngx.ctx
   local lim = ctx.limit_conn
-  if lim then
-    -- if you are using an upstream module in the content phase,
-    -- then you probably want to use $upstream_response_time
-    -- instead of ($request_time - ctx.limit_conn_delay) below.
-    local latency = tonumber(ngx.var.request_time) - ctx.limit_conn_delay
-    local key = ctx.limit_conn_key
-    assert(key)
-    local conn, err = lim:leaving(key, latency)
-
-    if not conn then
-      ngx.log(ngx.ERR,
-        "failed to record the connection leaving ",
-        "request: ", err)
-      return
-    end
 
-    opened_connections:set(conn,  ctx.metric_labels)
+  if lim then
+    -- checkin(false, ctx, ngx.var.request_time)
   end
+end
+
+function _M:log()
+  local lim = ngx.ctx.limit_conn
 
+  if lim then
+    local semaphore = ngx_semaphore.new()
+
+    ngx.timer.at(0, checkin, ngx.ctx, ngx.var.request_time, semaphore)
+
+    semaphore:wait(10)
+  end
 
   metric_requests:inc(1, {ngx.var.status})
 end