update stream-lua-nginx

nginx/app/lib/ngx/balancer.lua

@@ -1,10 +1,11 @@
 -- Copyright (C) Yichun Zhang (agentzh)
 
 
-local ffi = require "ffi"
 local base = require "resty.core.base"
+base.allows_subsystem('http', 'stream')
 
 
+local ffi = require "ffi"
 local C = ffi.C
 local ffi_str = ffi.string
 local errmsg = base.get_errmsg_ptr()
@@ -15,18 +16,81 @@ local getfenv = getfenv
 local error = error
 local type = type
 local tonumber = tonumber
+local max = math.max
+local subsystem = ngx.config.subsystem
+local ngx_lua_ffi_balancer_set_current_peer
+local ngx_lua_ffi_balancer_set_more_tries
+local ngx_lua_ffi_balancer_get_last_failure
+local ngx_lua_ffi_balancer_set_timeouts -- used by both stream and http
+
+
+if subsystem == 'http' then
+    ffi.cdef[[
+    int ngx_http_lua_ffi_balancer_set_current_peer(ngx_http_request_t *r,
+        const unsigned char *addr, size_t addr_len, int port, char **err);
+
+    int ngx_http_lua_ffi_balancer_set_more_tries(ngx_http_request_t *r,
+        int count, char **err);
+
+    int ngx_http_lua_ffi_balancer_get_last_failure(ngx_http_request_t *r,
+        int *status, char **err);
+
+    int ngx_http_lua_ffi_balancer_set_timeouts(ngx_http_request_t *r,
+        long connect_timeout, long send_timeout,
+        long read_timeout, char **err);
+    ]]
+
+    ngx_lua_ffi_balancer_set_current_peer =
+        C.ngx_http_lua_ffi_balancer_set_current_peer
+
+    ngx_lua_ffi_balancer_set_more_tries =
+        C.ngx_http_lua_ffi_balancer_set_more_tries
+
+    ngx_lua_ffi_balancer_get_last_failure =
+        C.ngx_http_lua_ffi_balancer_get_last_failure
 
+    ngx_lua_ffi_balancer_set_timeouts =
+        C.ngx_http_lua_ffi_balancer_set_timeouts
 
-ffi.cdef[[
-int ngx_http_lua_ffi_balancer_set_current_peer(ngx_http_request_t *r,
-    const unsigned char *addr, size_t addr_len, int port, char **err);
+elseif subsystem == 'stream' then
+    ffi.cdef[[
+    int ngx_stream_lua_ffi_balancer_set_current_peer(
+        ngx_stream_lua_request_t *r,
+        const unsigned char *addr, size_t addr_len, int port, char **err);
 
-int ngx_http_lua_ffi_balancer_set_more_tries(ngx_http_request_t *r,
-    int count, char **err);
+    int ngx_stream_lua_ffi_balancer_set_more_tries(ngx_stream_lua_request_t *r,
+        int count, char **err);
 
-int ngx_http_lua_ffi_balancer_get_last_failure(ngx_http_request_t *r,
-    int *status, char **err);
-]]
+    int ngx_stream_lua_ffi_balancer_get_last_failure(
+        ngx_stream_lua_request_t *r, int *status, char **err);
+
+    int ngx_stream_lua_ffi_balancer_set_timeouts(ngx_stream_lua_request_t *r,
+        long connect_timeout, long timeout, char **err);
+    ]]
+
+    ngx_lua_ffi_balancer_set_current_peer =
+        C.ngx_stream_lua_ffi_balancer_set_current_peer
+
+    ngx_lua_ffi_balancer_set_more_tries =
+        C.ngx_stream_lua_ffi_balancer_set_more_tries
+
+    ngx_lua_ffi_balancer_get_last_failure =
+        C.ngx_stream_lua_ffi_balancer_get_last_failure
+
+    local ngx_stream_lua_ffi_balancer_set_timeouts =
+        C.ngx_stream_lua_ffi_balancer_set_timeouts
+
+    ngx_lua_ffi_balancer_set_timeouts =
+        function(r, connect_timeout, send_timeout, read_timeout, err)
+            local timeout = max(send_timeout, read_timeout)
+
+            return ngx_stream_lua_ffi_balancer_set_timeouts(r, connect_timeout,
+                                                            timeout, err)
+        end
+
+else
+    error("unknown subsystem: " .. subsystem)
+end
 
 
 local peer_state_names = {
@@ -51,8 +115,8 @@ function _M.set_current_peer(addr, port)
         port = tonumber(port)
     end
 
-    local rc = C.ngx_http_lua_ffi_balancer_set_current_peer(r, addr, #addr,
-                                                            port, errmsg)
+    local rc = ngx_lua_ffi_balancer_set_current_peer(r, addr, #addr,
+                                                     port, errmsg)
     if rc == FFI_OK then
         return true
     end
@@ -67,7 +131,7 @@ function _M.set_more_tries(count)
         return error("no request found")
     end
 
-    local rc = C.ngx_http_lua_ffi_balancer_set_more_tries(r, count, errmsg)
+    local rc = ngx_lua_ffi_balancer_set_more_tries(r, count, errmsg)
     if rc == FFI_OK then
         if errmsg[0] == nil then
             return true
@@ -85,9 +149,7 @@ function _M.get_last_failure()
         return error("no request found")
     end
 
-    local state = C.ngx_http_lua_ffi_balancer_get_last_failure(r,
-                                                               int_out,
-                                                               errmsg)
+    local state = ngx_lua_ffi_balancer_get_last_failure(r, int_out, errmsg)
 
     if state == 0 then
         return nil
@@ -101,4 +163,48 @@ function _M.get_last_failure()
 end
 
 
-return _M
+function _M.set_timeouts(connect_timeout, send_timeout, read_timeout)
+    local r = getfenv(0).__ngx_req
+    if not r then
+        return error("no request found")
+    end
+
+    if not connect_timeout then
+        connect_timeout = 0
+    elseif type(connect_timeout) ~= "number" or connect_timeout <= 0 then
+        return error("bad connect timeout")
+    else
+        connect_timeout = connect_timeout * 1000
+    end
+
+    if not send_timeout then
+        send_timeout = 0
+    elseif type(send_timeout) ~= "number" or send_timeout <= 0 then
+        return error("bad send timeout")
+    else
+        send_timeout = send_timeout * 1000
+    end
+
+    if not read_timeout then
+        read_timeout = 0
+    elseif type(read_timeout) ~= "number" or read_timeout <= 0 then
+        return error("bad read timeout")
+    else
+        read_timeout = read_timeout * 1000
+    end
+
+    local rc
+
+    rc = ngx_lua_ffi_balancer_set_timeouts(r, connect_timeout,
+                                           send_timeout, read_timeout,
+                                           errmsg)
+
+    if rc == FFI_OK then
+        return true
+    end
+
+    return false, ffi_str(errmsg[0])
+end
+
+
+return _M