|
| 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