From 5ca56089c04c9f762d486b5d27ff1bfe835cc4a5 Mon Sep 17 00:00:00 2001
From: Hiroaki Nakamura <hnakamur@gmail.com>
Date: Fri, 26 Jul 2019 12:08:01 +0900
Subject: [PATCH 1/5] Support lock for shared dict in limit_req

Signed-off-by: Hiroaki Nakamura <hnakamur@gmail.com>
---
 lib/resty/limit/req.lua |  49 ++++++-
 t/req_with_lock.t       | 297 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 344 insertions(+), 2 deletions(-)
 create mode 100644 t/req_with_lock.t

diff --git a/lib/resty/limit/req.lua b/lib/resty/limit/req.lua
index 8b804fc..4f9ca12 100644
--- a/lib/resty/limit/req.lua
+++ b/lib/resty/limit/req.lua
@@ -6,6 +6,7 @@
 
 local ffi = require "ffi"
 local math = require "math"
+local resty_lock = require "resty.lock"
 
 
 local ngx_shared = ngx.shared
@@ -47,18 +48,27 @@ local mt = {
 }
 
 
-function _M.new(dict_name, rate, burst)
+function _M.new(dict_name, rate, burst, lock_dict_name)
     local dict = ngx_shared[dict_name]
     if not dict then
         return nil, "shared dict not found"
     end
 
+    if lock_dict_name then
+        local lock_dict = ngx_shared[lock_dict_name]
+        if not lock_dict then
+            return nil, "lock shared dict not found"
+        end
+    end
+
     assert(rate > 0 and burst >= 0)
 
     local self = {
         dict = dict,
         rate = rate * 1000,
         burst = burst * 1000,
+        dict_name = dict_name,
+        lock_dict_name = lock_dict_name,
     }
 
     return setmetatable(self, mt)
@@ -67,21 +77,43 @@ end
 
 -- sees an new incoming event
 -- the "commit" argument controls whether should we record the event in shm.
--- FIXME we have a (small) race-condition window between dict:get() and
+-- NOTE: if lock_dict_name is not set in limit_req.new(),
+-- we have a (small) race-condition window between dict:get() and
 -- dict:set() across multiple nginx worker processes. The size of the
 -- window is proportional to the number of workers.
 function _M.incoming(self, key, commit)
     local dict = self.dict
     local rate = self.rate
+    local lock_dict_name = self.lock_dict_name
     local now = ngx_now() * 1000
 
     local excess
 
+    local lock
+    if lock_dict_name then
+        local err
+        lock, err = resty_lock:new(lock_dict_name)
+        if not lock then
+            return nil, err
+        end
+
+        local elapsed, err = lock:lock(self.dict_name)
+        if not elapsed then
+            return nil, err
+        end
+    end
+
     -- it's important to anchor the string value for the read-only pointer
     -- cdata:
     local v = dict:get(key)
     if v then
         if type(v) ~= "string" or #v ~= rec_size then
+            if lock then
+                local ok, err = lock:unlock()
+                if not ok then
+                    return nil, err
+                end
+            end
             return nil, "shdict abused by other users"
         end
         local rec = ffi_cast(const_rec_ptr_type, v)
@@ -98,6 +130,12 @@ function _M.incoming(self, key, commit)
         -- print("excess: ", excess)
 
         if excess > self.burst then
+            if lock then
+                local ok, err = lock:unlock()
+                if not ok then
+                    return nil, err
+                end
+            end
             return nil, "rejected"
         end
 
@@ -111,6 +149,13 @@ function _M.incoming(self, key, commit)
         dict:set(key, ffi_str(rec_cdata, rec_size))
     end
 
+    if lock then
+        local ok, err = lock:unlock()
+        if not ok then
+            return nil, err
+        end
+    end
+
     -- return the delay in seconds, as well as excess
     return excess / rate, excess / 1000
 end
diff --git a/t/req_with_lock.t b/t/req_with_lock.t
new file mode 100644
index 0000000..06d232b
--- /dev/null
+++ b/t/req_with_lock.t
@@ -0,0 +1,297 @@
+# vim:set ft= ts=4 sw=4 et fdm=marker:
+
+use Test::Nginx::Socket::Lua;
+use Cwd qw(cwd);
+
+repeat_each(2);
+
+plan tests => repeat_each() * (blocks() * 4);
+
+#no_diff();
+#no_long_string();
+
+my $pwd = cwd();
+
+our $HttpConfig = <<_EOC_;
+    lua_package_path "$pwd/lib/?.lua;;";
+_EOC_
+
+no_long_string();
+run_tests();
+
+__DATA__
+
+=== TEST 1: a single key (always commit)
+--- http_config eval
+qq{
+$::HttpConfig
+
+    init_by_lua_block {
+        local v = require "jit.v"
+        -- v.on("/tmp/a.dump")
+    }
+    lua_shared_dict store 1m;
+    lua_shared_dict lock 12k;
+}
+--- config
+    location = /t {
+        content_by_lua '
+            local limit_req = require "resty.limit.req"
+            ngx.shared.store:flush_all()
+            local lim = limit_req.new("store", 40, 40, "lock")
+            local begin = ngx.now()
+            local uri = ngx.var.uri
+            for i = 1, 80 do
+                local delay, err = lim:incoming(uri, true)
+                if not delay then
+                    ngx.say("failed to limit request: ", err)
+                    return
+                end
+                ngx.sleep(delay)
+            end
+            ngx.say("elapsed: ", ngx.now() - begin, " sec.")
+        ';
+    }
+--- request
+    GET /t
+--- response_body_like eval
+qr/^elapsed: 1\.9[6-9]\d* sec\.$/
+--- no_error_log
+[error]
+[lua]
+--- timeout: 10
+
+
+
+=== TEST 2: multiple keys
+--- http_config eval
+"
+$::HttpConfig
+
+    lua_shared_dict store 1m;
+    lua_shared_dict lock 12k;
+"
+--- config
+    location = /t {
+        content_by_lua '
+            local limit_req = require "resty.limit.req"
+            ngx.shared.store:flush_all()
+            local lim = limit_req.new("store", 2, 10, "lock")
+            local delay1, excess1 = lim:incoming("foo", true)
+            local delay2, excess2 = lim:incoming("foo", true)
+            local delay3 = lim:incoming("bar", true)
+            local delay4 = lim:incoming("bar", true)
+            ngx.say("delay1: ", delay1)
+            ngx.say("excess1: ", excess1)
+            ngx.say("delay2: ", delay2)
+            ngx.say("excess2: ", excess2)
+            ngx.say("delay3: ", delay3)
+            ngx.say("delay4: ", delay4)
+        ';
+    }
+--- request
+    GET /t
+--- response_body
+delay1: 0
+excess1: 0
+delay2: 0.5
+excess2: 1
+delay3: 0
+delay4: 0.5
+--- no_error_log
+[error]
+[lua]
+
+
+
+=== TEST 3: burst
+--- http_config eval
+"
+$::HttpConfig
+
+    lua_shared_dict store 1m;
+    lua_shared_dict lock 12k;
+"
+--- config
+    location = /t {
+        content_by_lua '
+            local limit_req = require "resty.limit.req"
+            local lim = limit_req.new("store", 2, 0, "lock")
+
+            for burst = 0, 2 do
+                ngx.shared.store:flush_all()
+                if burst > 0 then
+                    lim:set_burst(burst)
+                end
+
+                for i = 1, 10 do
+                    local delay, err = lim:incoming("foo", true)
+                    if not delay then
+                        ngx.say(i, ": error: ", err)
+                        break
+                    end
+                end
+            end
+        ';
+    }
+--- request
+    GET /t
+--- response_body
+2: error: rejected
+3: error: rejected
+4: error: rejected
+--- no_error_log
+[error]
+[lua]
+
+
+
+=== TEST 4: a single key (do not commit since the 3rd time)
+--- http_config eval
+"
+$::HttpConfig
+
+    lua_shared_dict store 1m;
+    lua_shared_dict lock 12k;
+"
+--- config
+    location = /t {
+        content_by_lua '
+            local limit_req = require "resty.limit.req"
+            ngx.shared.store:flush_all()
+            local lim = limit_req.new("store", 2, 10, "lock")
+            local key = "bar"
+            for i = 1, 4 do
+                local delay, err = lim:incoming(key, i < 3 and true or false)
+                if not delay then
+                    ngx.say("failed to limit request: ", err)
+                else
+                    ngx.say("delay: ", delay)
+                end
+            end
+        ';
+    }
+--- request
+    GET /t
+--- response_body
+delay: 0
+delay: 0.5
+delay: 1
+delay: 1
+--- no_error_log
+[error]
+[lua]
+
+
+
+=== TEST 5: bad value in shdict (integer type)
+--- http_config eval
+"
+$::HttpConfig
+
+    lua_shared_dict store 1m;
+    lua_shared_dict lock 12k;
+"
+--- config
+    location = /t {
+        content_by_lua '
+            local limit_req = require "resty.limit.req"
+            ngx.shared.store:flush_all()
+            local key = "bar"
+            ngx.shared.store:set("bar", 32)
+            local lim = limit_req.new("store", 2, 10, "lock")
+            local delay, err = lim:incoming(key, true)
+            if not delay then
+                ngx.say("failed to limit request: ", err)
+            else
+                ngx.say("delay: ", delay)
+            end
+        ';
+    }
+--- request
+    GET /t
+--- response_body
+failed to limit request: shdict abused by other users
+--- no_error_log
+[error]
+[lua]
+
+
+
+=== TEST 6: bad value in shdict (string type, and wrong size)
+--- http_config eval
+"
+$::HttpConfig
+
+    lua_shared_dict store 1m;
+    lua_shared_dict lock 12k;
+"
+--- config
+    location = /t {
+        content_by_lua '
+            local limit_req = require "resty.limit.req"
+            ngx.shared.store:flush_all()
+            local key = "bar"
+            ngx.shared.store:set("bar", "a")
+            local lim = limit_req.new("store", 2, 10, "lock")
+            local delay, err = lim:incoming(key, true)
+            if not delay then
+                ngx.say("failed to limit request: ", err)
+            else
+                ngx.say("delay: ", delay)
+            end
+        ';
+    }
+--- request
+    GET /t
+--- response_body
+failed to limit request: shdict abused by other users
+--- no_error_log
+[error]
+[lua]
+
+
+
+=== TEST 7: a single key (commit & uncommit)
+--- http_config eval
+"
+$::HttpConfig
+
+    lua_shared_dict store 1m;
+    lua_shared_dict lock 12k;
+"
+--- config
+    location = /t {
+        content_by_lua '
+            local limit_req = require "resty.limit.req"
+            ngx.shared.store:flush_all()
+            local lim = limit_req.new("store", 40, 40, "lock")
+            local begin = ngx.now()
+            local uri = ngx.var.uri
+            for i = 1, 5 do
+                local delay, err = lim:incoming(uri, true)
+                if not delay then
+                    ngx.say("failed to limit request: ", err)
+                    return
+                end
+                ngx.say(i, ": delay: ", delay)
+                -- --[[
+                local ok, err = lim:uncommit(uri)
+                if not ok then
+                    ngx.say("failed to uncommit: ", err)
+                end
+                -- ]]
+            end
+        ';
+    }
+--- request
+    GET /t
+--- response_body
+1: delay: 0
+2: delay: 0.025
+3: delay: 0.025
+4: delay: 0.025
+5: delay: 0.025
+--- no_error_log
+[error]
+[lua]

From 3f04331309c4c1827d75426f095db586873f953b Mon Sep 17 00:00:00 2001
From: Hiroaki Nakamura <hnakamur@gmail.com>
Date: Fri, 26 Jul 2019 03:47:58 +0000
Subject: [PATCH 2/5] Update document for lock in resty.limit.req

Signed-off-by: Hiroaki Nakamura <hnakamur@gmail.com>
---
 README.md              |  8 ++++-
 lib/resty/limit/req.md | 66 +++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 72 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index ac38265..8a9e46e 100644
--- a/README.md
+++ b/README.md
@@ -285,7 +285,13 @@ Installation
 This library is enabled by default in [OpenResty](https://openresty.org/) 1.11.2.2+.
 
 If you have to install this library manually,
-then ensure you are using at least OpenResty 1.11.2.1 or a custom nginx build including ngx_lua 0.10.6+. Also, You need to configure
+then ensure you are using at least OpenResty 1.11.2.1 or a custom nginx build including ngx_lua 0.10.6+.
+
+If you use a lock by passing non-nil value to `lock_shdict_name` parameter of
+[resty.limit.req new](lib/resty/limit/req.md#new) method, then ensure you install
+[openresty/lua-resty-lock](https://github.com/openresty/lua-resty-lock#prerequisites) too.
+
+Also, You need to configure
 the [lua_package_path](https://github.com/openresty/lua-nginx-module#lua_package_path) directive to
 add the path of your `lua-resty-limit-traffic` source tree to ngx_lua's Lua module search path, as in
 
diff --git a/lib/resty/limit/req.md b/lib/resty/limit/req.md
index 1681087..ee0da76 100644
--- a/lib/resty/limit/req.md
+++ b/lib/resty/limit/req.md
@@ -87,6 +87,65 @@ http {
 }
 ```
 
+```nginx
+# demonstrate the usage of the resty.limit.req module (alone!) with lock
+http {
+    lua_shared_dict my_limit_req_store 100m;
+    lua_shared_dict my_limit_req_lock 12k;
+
+    server {
+        location / {
+            access_by_lua_block {
+                -- well, we could put the require() and new() calls in our own Lua
+                -- modules to save overhead. here we put them below just for
+                -- convenience.
+
+                local limit_req = require "resty.limit.req"
+
+                -- limit the requests under 200 req/sec with a burst of 100 req/sec,
+                -- that is, we delay requests under 300 req/sec and above 200
+                -- req/sec, and reject any requests exceeding 300 req/sec.
+                local lim, err = limit_req.new("my_limit_req_store", 200, 100, "my_limit_req_lock")
+                if not lim then
+                    ngx.log(ngx.ERR,
+                            "failed to instantiate a resty.limit.req object: ", err)
+                    return ngx.exit(500)
+                end
+
+                -- the following call must be per-request.
+                -- here we use the remote (IP) address as the limiting key
+                local key = ngx.var.binary_remote_addr
+                local delay, err = lim:incoming(key, true)
+                if not delay then
+                    if err == "rejected" then
+                        return ngx.exit(503)
+                    end
+                    ngx.log(ngx.ERR, "failed to limit req: ", err)
+                    return ngx.exit(500)
+                end
+
+                if delay >= 0.001 then
+                    -- the 2nd return value holds  the number of excess requests
+                    -- per second for the specified key. for example, number 31
+                    -- means the current request rate is at 231 req/sec for the
+                    -- specified key.
+                    local excess = err
+
+                    -- the request exceeding the 200 req/sec but below 300 req/sec,
+                    -- so we intentionally delay it here a bit to conform to the
+                    -- 200 req/sec rate.
+                    ngx.sleep(delay)
+                end
+            }
+
+            # content handler goes here. if it is content_by_lua, then you can
+            # merge the Lua code above in access_by_lua into your content_by_lua's
+            # Lua handler to save a little bit of CPU time.
+        }
+    }
+}
+```
+
 Description
 ===========
 
@@ -108,7 +167,7 @@ Methods
 
 new
 ---
-**syntax:** `obj, err = class.new(shdict_name, rate, burst)`
+**syntax:** `obj, err = class.new(shdict_name, rate, burst, lock_shdict_name)`
 
 Instantiates an object of this class. The `class` value is returned by the call `require "resty.limit.req"`.
 
@@ -125,6 +184,11 @@ This method takes the following arguments:
     Requests exceeding this hard limit
 will get rejected immediately.
 
+* `lock_shdict_name` is the opitional argument for the name of the [lua_shared_dict](https://github.com/openresty/lua-nginx-module#lua_shared_dict) shm zone
+used to lock the shared dict specified by `shdict_name`. If `lock_shdict_name` is omitted or the value is `nil`, no lock is used and there is a small race-condition window
+between getting and setting of the shared dict specified by `shdict_name` across multiple nginx worker processes. The size of the window is proportional to the number of workers.
+
+
 On failure, this method returns `nil` and a string describing the error (like a bad `lua_shared_dict` name).
 
 [Back to TOC](#table-of-contents)

From 98ab0415b2806b26f2a254cfc921aa4b990bc203 Mon Sep 17 00:00:00 2001
From: Hiroaki Nakamura <hnakamur@gmail.com>
Date: Fri, 26 Jul 2019 04:49:57 +0000
Subject: [PATCH 3/5] Install lua-resty-lock on Travis CI

Signed-off-by: Hiroaki Nakamura <hnakamur@gmail.com>
---
 .travis.yml       | 1 +
 t/req_with_lock.t | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/.travis.yml b/.travis.yml
index 58d1f64..9099501 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -36,6 +36,7 @@ install:
   - git clone https://github.com/openresty/no-pool-nginx.git ../no-pool-nginx
   - git clone https://github.com/openresty/lua-resty-lrucache.git ../lua-resty-lrucache
   - git clone https://github.com/openresty/lua-resty-core.git ../lua-resty-core
+  - git clone https://github.com/openresty/lua-resty-lock.git ../lua-resty-lock
   - git clone -b v2.1-agentzh https://github.com/openresty/luajit2.git
 
 script:
diff --git a/t/req_with_lock.t b/t/req_with_lock.t
index 06d232b..f2d6d91 100644
--- a/t/req_with_lock.t
+++ b/t/req_with_lock.t
@@ -13,7 +13,7 @@ plan tests => repeat_each() * (blocks() * 4);
 my $pwd = cwd();
 
 our $HttpConfig = <<_EOC_;
-    lua_package_path "$pwd/lib/?.lua;;";
+    lua_package_path "$pwd/lib/?.lua;$pwd/..lua-resty-lock/lib/?.lua;;";
 _EOC_
 
 no_long_string();

From 6935ba59fe055cedf4ae2ab8a3fa0c5760e05705 Mon Sep 17 00:00:00 2001
From: Hiroaki Nakamura <hnakamur@gmail.com>
Date: Fri, 26 Jul 2019 05:01:35 +0000
Subject: [PATCH 4/5] Fix lua_package_path in t/req_with_lock.t

Signed-off-by: Hiroaki Nakamura <hnakamur@gmail.com>
---
 t/req_with_lock.t | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/t/req_with_lock.t b/t/req_with_lock.t
index f2d6d91..5aa283e 100644
--- a/t/req_with_lock.t
+++ b/t/req_with_lock.t
@@ -13,7 +13,7 @@ plan tests => repeat_each() * (blocks() * 4);
 my $pwd = cwd();
 
 our $HttpConfig = <<_EOC_;
-    lua_package_path "$pwd/lib/?.lua;$pwd/..lua-resty-lock/lib/?.lua;;";
+    lua_package_path "$pwd/lib/?.lua;$pwd/../lua-resty-lock/lib/?.lua;;";
 _EOC_
 
 no_long_string();

From 2cc38f758c3ed3566918ce93760df1658f441016 Mon Sep 17 00:00:00 2001
From: Hiroaki Nakamura <hnakamur@gmail.com>
Date: Fri, 26 Jul 2019 05:09:47 +0000
Subject: [PATCH 5/5] Require resty-lock only when needed

Signed-off-by: Hiroaki Nakamura <hnakamur@gmail.com>
---
 lib/resty/limit/req.lua | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/lib/resty/limit/req.lua b/lib/resty/limit/req.lua
index 4f9ca12..40566d6 100644
--- a/lib/resty/limit/req.lua
+++ b/lib/resty/limit/req.lua
@@ -6,7 +6,6 @@
 
 local ffi = require "ffi"
 local math = require "math"
-local resty_lock = require "resty.lock"
 
 
 local ngx_shared = ngx.shared
@@ -91,6 +90,8 @@ function _M.incoming(self, key, commit)
 
     local lock
     if lock_dict_name then
+        local resty_lock = require "resty.lock"
+
         local err
         lock, err = resty_lock:new(lock_dict_name)
         if not lock then