add redis initialization #667

3scale · May 14, 2018 · b04fbf2 · b04fbf2
1 parent 356be50
commit b04fbf2
Show file tree

Hide file tree

Showing 3 changed files with 137 additions and 9 deletions.
diff --git a/gateway/src/apicast/policy/rate_limit/rate_limit.lua b/gateway/src/apicast/policy/rate_limit/rate_limit.lua
@@ -72,6 +72,9 @@ local function redis_shdict(url)
         return nil
       end
       return val
+    end,
+    del = function(_, key)
+      return redis:del(key)
     end
   }
 end
@@ -108,6 +111,40 @@ local function init_error_settings(limits_exceeded_error, configuration_error)
   return error_settings
 end
 
+local function initialize_redis_records(type, redis, seed, limiters)
+  for _, limiter in ipairs(limiters) do
+    local key
+    if limiter.key.scope == "service" then
+      key = limiter.key.service_name.."_"..type.."_"..limiter.key.name
+    else
+      key = type.."_"..limiter.key.name
+    end
+
+    local seed_key = key.."_seed"
+    local seed_value = redis:get(seed_key)
+
+    if not seed_value or tonumber(seed_value) ~= seed then
+      redis:set(seed_key, seed)
+
+      if type == "connections" or type == "leaky_bucket" then
+        redis:del(key)
+      else
+        local count_key = key.."_count"
+        local count = redis:get(count_key)
+
+        local window_key = key.."_window"
+        local window = redis:get(window_key)
+
+        if not count or not window or tonumber(count) ~= limiter.count or tonumber(window) ~= limiter.window then
+          redis:del(key)
+          redis:set(count_key, limiter.count)
+          redis:set(window_key, limiter.window)
+        end
+      end
+    end
+  end
+end
+
 local function build_limiters_and_keys(type, limiters, redis, error_settings)
   local res_limiters = {}
   local res_keys = {}
@@ -146,6 +183,7 @@ function _M.new(config)
   self.redis_url = config.redis_url
   self.error_settings = init_error_settings(
     config.limits_exceeded_error, config.configuration_error)
+  self.seed = ngx.time()
 
   return self
 end
@@ -160,6 +198,10 @@ function _M:access()
       error(self.error_settings, "configuration_issue")
       return
     end
+
+    initialize_redis_records('connections', red, self.seed, self.connection_limiters)
+    initialize_redis_records('leaky_bucket', red, self.seed, self.leaky_bucket_limiters)
+    initialize_redis_records('fixed_window', red, self.seed, self.fixed_window_limiters)
   end
 
   local conn_limiters, conn_keys = build_limiters_and_keys(

diff --git a/spec/policy/rate_limit/rate_limit_spec.lua b/spec/policy/rate_limit/rate_limit_spec.lua
@@ -50,7 +50,10 @@ describe('Rate limit policy', function()
     local redis = require('resty.redis'):new()
     redis:connect(redis_host, redis_port)
     redis:select(1)
-    redis:del('connections_test1', 'leaky_bucket_test2', 'fixed_window_test3', 'bank_A_leaky_bucket_test4')
+    redis:del('connections_test1', 'leaky_bucket_test2', 'bank_A_leaky_bucket_test4')
+    redis:del('fixed_window_test3', 'fixed_window_test3_count', 'fixed_window_test3_window')
+    redis:del('connections_test1_seed', 'leaky_bucket_test2_seed')
+    redis:del('fixed_window_test3_seed', 'bank_A_leaky_bucket_test4_seed')
     init_val()
   end)
 
@@ -181,6 +184,88 @@ describe('Rate limit policy', function()
       rate_limit_policy:access()
       assert.spy(ngx_sleep_spy).was_called_with(match.is_gt(0.001))
     end)
+
+    it('initialize redis records', function()
+      local config = {
+        connection_limiters = {
+          {key = {name = 'test1'}, conn = 20, burst = 10, delay = 0.5}
+        },
+        fixed_window_limiters = {
+          {key = {name = 'test3'}, count = 10, window = 10}
+        },
+        redis_url = 'redis://'..redis_host..':'..redis_port..'/1'
+      }
+      local config2 = {
+        connection_limiters = {
+          {key = {name = 'test1'}, conn = 20, burst = 10, delay = 0.5}
+        },
+        fixed_window_limiters = {
+          {key = {name = 'test3'}, count = 15, window = 10}
+        },
+        redis_url = 'redis://'..redis_host..':'..redis_port..'/1'
+      }
+
+      local rate_limit_policy = RateLimitPolicy.new(config)
+      rate_limit_policy:access()
+
+      local redis = require('resty.redis'):new()
+      redis:connect(redis_host, redis_port)
+      redis:select(1)
+      local conn = redis:get('connections_test1')
+      local fixed_window = redis:get('fixed_window_test3')
+      local count = redis:get('fixed_window_test3_count')
+      local window = redis:get('fixed_window_test3_window')
+      assert.equal('1', conn)
+      assert.equal('9', fixed_window)
+      assert.equal('10', count)
+      assert.equal('10', window)
+
+      ngx.sleep(1)
+      rate_limit_policy = RateLimitPolicy.new(config2)
+      rate_limit_policy:access()
+
+      conn = redis:get('connections_test1')
+      fixed_window = redis:get('fixed_window_test3')
+      count = redis:get('fixed_window_test3_count')
+      window = redis:get('fixed_window_test3_window')
+      assert.equal('1', conn)
+      assert.equal('14', fixed_window)
+      assert.equal('15', count)
+      assert.equal('10', window)
+    end)
+
+    it('do not initialize redis records', function()
+      local config = {
+        fixed_window_limiters = {
+          {key = {name = 'test3'}, count = 10, window = 10}
+        },
+        redis_url = 'redis://'..redis_host..':'..redis_port..'/1'
+      }
+
+      local rate_limit_policy = RateLimitPolicy.new(config)
+      rate_limit_policy:access()
+
+      local redis = require('resty.redis'):new()
+      redis:connect(redis_host, redis_port)
+      redis:select(1)
+      local fixed_window = redis:get('fixed_window_test3')
+      local count = redis:get('fixed_window_test3_count')
+      local window = redis:get('fixed_window_test3_window')
+      assert.equal('9', fixed_window)
+      assert.equal('10', count)
+      assert.equal('10', window)
+
+      ngx.sleep(1)
+      rate_limit_policy = RateLimitPolicy.new(config)
+      rate_limit_policy:access()
+
+      fixed_window = redis:get('fixed_window_test3')
+      count = redis:get('fixed_window_test3_count')
+      window = redis:get('fixed_window_test3_window')
+      assert.equal('8', fixed_window)
+      assert.equal('10', count)
+      assert.equal('10', window)
+    end)
   end)
 
   describe('.log', function()

diff --git a/t/apicast-policy-rate-limit.t b/t/apicast-policy-rate-limit.t
@@ -77,7 +77,7 @@ Return 200 code.
       local redis = require('resty.redis'):new()
       redis:connect(redis_host, redis_port)
       redis:select(1)
-      redis:del("service_C_connections_test1")
+      redis:del("service_C_connections_test1", "service_C_connections_test1_seed")
     }
   }
 
@@ -195,7 +195,7 @@ Return 200 code.
       local redis = require('resty.redis'):new()
       redis:connect(redis_host, redis_port)
       redis:select(1)
-      redis:del("connections_test4")
+      redis:del("connections_test4", "connections_test4_seed")
     }
   }
 
@@ -312,7 +312,8 @@ Return 200 code.
       local redis = require('resty.redis'):new()
       redis:connect(redis_host, redis_port)
       redis:select(1)
-      redis:del("leaky_bucket_test6_1", "connections_test6_2", "fixed_window_test6_3")
+      redis:del("leaky_bucket_test6_1", "connections_test6_2", "fixed_window_test6_3", "fixed_window_test6_3_count", "fixed_window_test6_3_window")
+      redis:del("leaky_bucket_test6_1_seed", "connections_test6_2_seed", "fixed_window_test6_3_seed")
     }
   }
 
@@ -387,7 +388,7 @@ Return 429 code.
       local redis = require('resty.redis'):new()
       redis:connect(redis_host, redis_port)
       redis:select(1)
-      redis:del("connections_test7")
+      redis:del("connections_test7", "connections_test7_seed")
     }
   }
 
@@ -447,7 +448,7 @@ Return 503 code.
       local redis = require('resty.redis'):new()
       redis:connect(redis_host, redis_port)
       redis:select(1)
-      redis:del("leaky_bucket_test8")
+      redis:del("leaky_bucket_test8", "leaky_bucket_test8_seed")
     }
   }
 
@@ -507,7 +508,7 @@ Return 429 code.
       local redis = require('resty.redis'):new()
       redis:connect(redis_host, redis_port)
       redis:select(1)
-      redis:del("fixed_window_test9")
+      redis:del("fixed_window_test9", "fixed_window_test9_count", "fixed_window_test9_window", "fixed_window_test9_seed")
     }
   }
 
@@ -585,7 +586,7 @@ Return 200 code.
       local redis = require('resty.redis'):new()
       redis:connect(redis_host, redis_port)
       redis:select(1)
-      redis:del("connections_test10")
+      redis:del("connections_test10", "connections_test10_seed")
     }
   }
 
@@ -646,7 +647,7 @@ Return 200 code.
       local redis = require('resty.redis'):new()
       redis:connect(redis_host, redis_port)
       redis:select(1)
-      redis:del("leaky_bucket_test11")
+      redis:del("leaky_bucket_test11", "leaky_bucket_test11_seed")
     }
   }