Balancer by lua stream

lib/ngx/balancer_stream.lua

@@ -0,0 +1,82 @@
+-- Copyright (C) Yichun Zhang (agentzh)
+
+
+local ffi = require "ffi"
+local base = require "resty.core.base"
+
+
+local C = ffi.C
+local ffi_str = ffi.string
+local errmsg = base.get_errmsg_ptr()
+local FFI_OK = base.FFI_OK
+local FFI_ERROR = base.FFI_ERROR
+local int_out = ffi.new("int[1]")
+local getfenv = getfenv
+local error = error
+local type = type
+local tonumber = tonumber
+
+
+ffi.cdef[[
+int ngx_stream_lua_ffi_balancer_set_current_peer(ngx_stream_session_t *r,
+    const unsigned char *addr, size_t addr_len, int port, char **err);
+
+int ngx_stream_lua_ffi_balancer_set_more_tries(ngx_stream_session_t *r,
+    int count, char **err);
+
+int ngx_stream_lua_ffi_balancer_get_last_failure(ngx_stream_session_t *r,
+    int *status, char **err);
+]]
+
+
+local peer_state_names = {
+    [1] = "keepalive",
+    [2] = "next",
+    [4] = "failed",
+}
+
+
+local _M = { version = base.version }
+
+
+function _M.set_current_peer(addr, port)
+    local r = getfenv(0).__ngx_sess
+    if not r then
+        return error("no request found")
+    end
+
+    if not port then
+        port = 0
+    elseif type(port) ~= "number" then
+        port = tonumber(port)
+    end
+
+    local rc = C.ngx_stream_lua_ffi_balancer_set_current_peer(r, addr, #addr,
+                                                            port, errmsg)
+    if rc == FFI_OK then
+        return true
+    end
+
+    return nil, ffi_str(errmsg[0])
+end
+
+
+function _M.set_more_tries(count)
+    local r = getfenv(0).__ngx_sess
+    if not r then
+        return error("no request found")
+    end
+
+    local rc = C.ngx_stream_lua_ffi_balancer_set_more_tries(r, count, errmsg)
+    if rc == FFI_OK then
+        if errmsg[0] == nil then
+            return true
+        end
+        return true, ffi_str(errmsg[0])  -- return the warning
+    end
+
+    return nil, ffi_str(errmsg[0])
+end
+
+
+return _M

lib/resty/core/base.lua

@@ -81,6 +81,14 @@ if not pcall(ffi.typeof, "ngx_http_request_t") then
 end
 
 
+if not pcall(ffi.typeof, "ngx_stream_session_t") then
+    ffi.cdef[[
+        struct ngx_stream_session_s;
+        typedef struct ngx_stream_session_s  ngx_stream_session_t;
+    ]]
+end
+
+
 if not pcall(ffi.typeof, "ngx_http_lua_ffi_str_t") then
     ffi.cdef[[
         typedef struct {

lib/resty/core/shdict.lua

@@ -26,7 +26,7 @@ ffi.cdef[[
         int get_stale, int *is_stale);
 
     int ngx_http_lua_ffi_shdict_incr(void *zone, const unsigned char *key,
-        size_t key_len, double *value, char **err);
+        size_t key_len, double *value, int exptime, char **err);
 
     int ngx_http_lua_ffi_shdict_store(void *zone, int op,
         const unsigned char *key, size_t key_len, int value_type,
@@ -313,7 +313,7 @@ local function shdict_get_stale(zone, key)
 end
 
 
-local function shdict_incr(zone, key, value)
+local function shdict_incr(zone, key, value, exptime)
     zone = check_zone(zone)
 
     if key == nil then
@@ -336,9 +336,11 @@ local function shdict_incr(zone, key, value)
         value = tonumber(value)
     end
     num_value[0] = value
+
+    exptime = exptime or -1
 
     local rc = C.ngx_http_lua_ffi_shdict_incr(zone, key, key_len, num_value,
-                                             errmsg)
+                                             exptime, errmsg)
     if rc ~= 0 then  -- ~= NGX_OK
         return nil, ffi_str(errmsg[0])
     end

t/shdict.t

@@ -915,3 +915,27 @@ qr/\[TRACE   \d+ content_by_lua\(nginx\.conf:\d+\):7 loop\]/
  -- NYI:
 stitch
 
+
+
+
+=== TEST 27: incr key expire
+--- http_config eval: $::HttpConfig
+--- config
+    location = /t {
+        content_by_lua '
+            local val, flags
+            local dogs = ngx.shared.dogs
+            local value, err = dogs:incr(nil, 32, 10)
+            if not ok then
+                ngx.say("failed to incr: ", err)
+            end
+        ';
+    }
+--- request
+GET /t
+--- response_body
+failed to incr: nil key
+--- no_error_log
+[error]
+[alert]
+[crit]