|  | 
|  | 1 | +diff --git lib/ngx/balancer.lua lib/ngx/balancer.lua | 
|  | 2 | +index 7d64d63..781cbd1 100644 | 
|  | 3 | +--- lib/ngx/balancer.lua | 
|  | 4 | ++++ lib/ngx/balancer.lua | 
|  | 5 | +@@ -3,6 +3,7 @@ | 
|  | 6 | +  | 
|  | 7 | + local base = require "resty.core.base" | 
|  | 8 | + base.allows_subsystem('http', 'stream') | 
|  | 9 | ++require "resty.core.hash" | 
|  | 10 | +  | 
|  | 11 | +  | 
|  | 12 | + local ffi = require "ffi" | 
|  | 13 | +@@ -17,8 +18,10 @@ local error = error | 
|  | 14 | + local type = type | 
|  | 15 | + local tonumber = tonumber | 
|  | 16 | + local max = math.max | 
|  | 17 | ++local ngx_crc32_long = ngx.crc32_long | 
|  | 18 | + local subsystem = ngx.config.subsystem | 
|  | 19 | + local ngx_lua_ffi_balancer_set_current_peer | 
|  | 20 | ++local ngx_lua_ffi_balancer_enable_keepalive | 
|  | 21 | + local ngx_lua_ffi_balancer_set_more_tries | 
|  | 22 | + local ngx_lua_ffi_balancer_get_last_failure | 
|  | 23 | + local ngx_lua_ffi_balancer_set_timeouts -- used by both stream and http | 
|  | 24 | +@@ -27,7 +30,11 @@ local ngx_lua_ffi_balancer_set_timeouts -- used by both stream and http | 
|  | 25 | + if subsystem == 'http' then | 
|  | 26 | +     ffi.cdef[[ | 
|  | 27 | +     int ngx_http_lua_ffi_balancer_set_current_peer(ngx_http_request_t *r, | 
|  | 28 | +-        const unsigned char *addr, size_t addr_len, int port, char **err); | 
|  | 29 | ++        const unsigned char *addr, size_t addr_len, int port, | 
|  | 30 | ++        unsigned int cpool_crc32, unsigned int cpool_size, char **err); | 
|  | 31 | ++ | 
|  | 32 | ++    int ngx_http_lua_ffi_balancer_enable_keepalive(ngx_http_request_t *r, | 
|  | 33 | ++        unsigned long timeout, unsigned int max_requests, char **err); | 
|  | 34 | +  | 
|  | 35 | +     int ngx_http_lua_ffi_balancer_set_more_tries(ngx_http_request_t *r, | 
|  | 36 | +         int count, char **err); | 
|  | 37 | +@@ -46,6 +53,9 @@ if subsystem == 'http' then | 
|  | 38 | +     ngx_lua_ffi_balancer_set_current_peer = | 
|  | 39 | +         C.ngx_http_lua_ffi_balancer_set_current_peer | 
|  | 40 | +  | 
|  | 41 | ++    ngx_lua_ffi_balancer_enable_keepalive = | 
|  | 42 | ++        C.ngx_http_lua_ffi_balancer_enable_keepalive | 
|  | 43 | ++ | 
|  | 44 | +     ngx_lua_ffi_balancer_set_more_tries = | 
|  | 45 | +         C.ngx_http_lua_ffi_balancer_set_more_tries | 
|  | 46 | +  | 
|  | 47 | +@@ -96,6 +106,11 @@ else | 
|  | 48 | + end | 
|  | 49 | +  | 
|  | 50 | +  | 
|  | 51 | ++local DEFAULT_KEEPALIVE_POOL_SIZE = 30 | 
|  | 52 | ++local DEFAULT_KEEPALIVE_IDLE_TIMEOUT = 60000 | 
|  | 53 | ++local DEFAULT_KEEPALIVE_MAX_REQUESTS = 100 | 
|  | 54 | ++ | 
|  | 55 | ++ | 
|  | 56 | + local peer_state_names = { | 
|  | 57 | +     [1] = "keepalive", | 
|  | 58 | +     [2] = "next", | 
|  | 59 | +@@ -106,25 +121,147 @@ local peer_state_names = { | 
|  | 60 | + local _M = { version = base.version } | 
|  | 61 | +  | 
|  | 62 | +  | 
|  | 63 | +-function _M.set_current_peer(addr, port) | 
|  | 64 | +-    local r = get_request() | 
|  | 65 | +-    if not r then | 
|  | 66 | +-        error("no request found") | 
|  | 67 | ++if subsystem == "http" then | 
|  | 68 | ++    function _M.set_current_peer(addr, port, opts) | 
|  | 69 | ++        local r = get_request() | 
|  | 70 | ++        if not r then | 
|  | 71 | ++            error("no request found") | 
|  | 72 | ++        end | 
|  | 73 | ++ | 
|  | 74 | ++        local pool_crc32 | 
|  | 75 | ++        local pool_size | 
|  | 76 | ++ | 
|  | 77 | ++        if opts then | 
|  | 78 | ++            if type(opts) ~= "table" then | 
|  | 79 | ++                error("bad argument #3 to 'set_current_peer' " .. | 
|  | 80 | ++                      "(table expected, got " .. type(opts) .. ")", 2) | 
|  | 81 | ++            end | 
|  | 82 | ++ | 
|  | 83 | ++            local pool = opts.pool | 
|  | 84 | ++            pool_size = opts.pool_size | 
|  | 85 | ++ | 
|  | 86 | ++            if pool then | 
|  | 87 | ++                if type(pool) ~= "string" then | 
|  | 88 | ++                    error("bad option 'pool' to 'set_current_peer' " .. | 
|  | 89 | ++                          "(string expected, got " .. type(pool) .. ")", 2) | 
|  | 90 | ++                end | 
|  | 91 | ++ | 
|  | 92 | ++                pool_crc32 = ngx_crc32_long(pool) | 
|  | 93 | ++            end | 
|  | 94 | ++ | 
|  | 95 | ++            if pool_size then | 
|  | 96 | ++                if type(pool_size) ~= "number" then | 
|  | 97 | ++                    error("bad option 'pool_size' to 'set_current_peer' " .. | 
|  | 98 | ++                          "(number expected, got " .. type(pool_size) .. ")", 2) | 
|  | 99 | ++ | 
|  | 100 | ++                elseif pool_size < 1 then | 
|  | 101 | ++                    error("bad option 'pool_size' to 'set_current_peer' " .. | 
|  | 102 | ++                          "(expected > 0)", 2) | 
|  | 103 | ++                end | 
|  | 104 | ++            end | 
|  | 105 | ++        end | 
|  | 106 | ++ | 
|  | 107 | ++        if not port then | 
|  | 108 | ++            port = 0 | 
|  | 109 | ++ | 
|  | 110 | ++        elseif type(port) ~= "number" then | 
|  | 111 | ++            port = tonumber(port) | 
|  | 112 | ++        end | 
|  | 113 | ++ | 
|  | 114 | ++        if not pool_crc32 then | 
|  | 115 | ++            pool_crc32 = 0 | 
|  | 116 | ++        end | 
|  | 117 | ++ | 
|  | 118 | ++        if not pool_size then | 
|  | 119 | ++            pool_size = DEFAULT_KEEPALIVE_POOL_SIZE | 
|  | 120 | ++        end | 
|  | 121 | ++ | 
|  | 122 | ++        local rc = ngx_lua_ffi_balancer_set_current_peer(r, addr, #addr, port, | 
|  | 123 | ++                                                         pool_crc32, pool_size, | 
|  | 124 | ++                                                         errmsg) | 
|  | 125 | ++        if rc == FFI_OK then | 
|  | 126 | ++            return true | 
|  | 127 | ++        end | 
|  | 128 | ++ | 
|  | 129 | ++        return nil, ffi_str(errmsg[0]) | 
|  | 130 | +     end | 
|  | 131 | +  | 
|  | 132 | +-    if not port then | 
|  | 133 | +-        port = 0 | 
|  | 134 | +-    elseif type(port) ~= "number" then | 
|  | 135 | +-        port = tonumber(port) | 
|  | 136 | ++else | 
|  | 137 | ++    function _M.set_current_peer(addr, port, opts) | 
|  | 138 | ++        local r = get_request() | 
|  | 139 | ++        if not r then | 
|  | 140 | ++            error("no request found") | 
|  | 141 | ++        end | 
|  | 142 | ++ | 
|  | 143 | ++        if opts then | 
|  | 144 | ++            error("bad argument #3 to 'set_current_peer' ('opts' not yet " .. | 
|  | 145 | ++                  "implemented in " .. subsystem .. " subsystem)", 2) | 
|  | 146 | ++        end | 
|  | 147 | ++ | 
|  | 148 | ++        if not port then | 
|  | 149 | ++            port = 0 | 
|  | 150 | ++ | 
|  | 151 | ++        elseif type(port) ~= "number" then | 
|  | 152 | ++            port = tonumber(port) | 
|  | 153 | ++        end | 
|  | 154 | ++ | 
|  | 155 | ++        local rc = ngx_lua_ffi_balancer_set_current_peer(r, addr, #addr, | 
|  | 156 | ++                                                         port, errmsg) | 
|  | 157 | ++        if rc == FFI_OK then | 
|  | 158 | ++            return true | 
|  | 159 | ++        end | 
|  | 160 | ++ | 
|  | 161 | ++        return nil, ffi_str(errmsg[0]) | 
|  | 162 | +     end | 
|  | 163 | ++end | 
|  | 164 | +  | 
|  | 165 | +-    local rc = ngx_lua_ffi_balancer_set_current_peer(r, addr, #addr, | 
|  | 166 | +-                                                     port, errmsg) | 
|  | 167 | +-    if rc == FFI_OK then | 
|  | 168 | +-        return true | 
|  | 169 | ++ | 
|  | 170 | ++if subsystem == "http" then | 
|  | 171 | ++    function _M.enable_keepalive(idle_timeout, max_requests) | 
|  | 172 | ++        local r = get_request() | 
|  | 173 | ++        if not r then | 
|  | 174 | ++            error("no request found") | 
|  | 175 | ++        end | 
|  | 176 | ++ | 
|  | 177 | ++        if not idle_timeout then | 
|  | 178 | ++            idle_timeout = DEFAULT_KEEPALIVE_IDLE_TIMEOUT | 
|  | 179 | ++ | 
|  | 180 | ++        elseif type(idle_timeout) ~= "number" then | 
|  | 181 | ++            error("bad argument #1 to 'enable_keepalive' " .. | 
|  | 182 | ++                  "(number expected, got " .. type(idle_timeout) .. ")", 2) | 
|  | 183 | ++ | 
|  | 184 | ++        elseif idle_timeout < 0 then | 
|  | 185 | ++            error("bad argument #1 to 'enable_keepalive' (expected >= 0)", 2) | 
|  | 186 | ++ | 
|  | 187 | ++        else | 
|  | 188 | ++            idle_timeout = idle_timeout * 1000 | 
|  | 189 | ++        end | 
|  | 190 | ++ | 
|  | 191 | ++        if not max_requests then | 
|  | 192 | ++            max_requests = DEFAULT_KEEPALIVE_MAX_REQUESTS | 
|  | 193 | ++ | 
|  | 194 | ++        elseif type(max_requests) ~= "number" then | 
|  | 195 | ++            error("bad argument #2 to 'enable_keepalive' " .. | 
|  | 196 | ++                  "(number expected, got " .. type(max_requests) .. ")", 2) | 
|  | 197 | ++ | 
|  | 198 | ++        elseif max_requests < 0 then | 
|  | 199 | ++            error("bad argument #2 to 'enable_keepalive' (expected >= 0)", 2) | 
|  | 200 | ++        end | 
|  | 201 | ++ | 
|  | 202 | ++        local rc = ngx_lua_ffi_balancer_enable_keepalive(r, idle_timeout, | 
|  | 203 | ++                                                         max_requests, errmsg) | 
|  | 204 | ++        if rc == FFI_OK then | 
|  | 205 | ++            return true | 
|  | 206 | ++        end | 
|  | 207 | ++ | 
|  | 208 | ++        return nil, ffi_str(errmsg[0]) | 
|  | 209 | +     end | 
|  | 210 | +  | 
|  | 211 | +-    return nil, ffi_str(errmsg[0]) | 
|  | 212 | ++else | 
|  | 213 | ++    function _M.enable_keepalive() | 
|  | 214 | ++        error("'enable_keepalive' not yet implemented in " .. subsystem .. | 
|  | 215 | ++              " subsystem", 2) | 
|  | 216 | ++    end | 
|  | 217 | + end | 
|  | 218 | +  | 
|  | 219 | +  | 
0 commit comments