diff --git a/src/ngx_http_lua_upstream_module.c b/src/ngx_http_lua_upstream_module.c index b299457..d6eb0e9 100644 --- a/src/ngx_http_lua_upstream_module.c +++ b/src/ngx_http_lua_upstream_module.c @@ -34,6 +34,7 @@ static ngx_http_upstream_srv_conf_t * static ngx_http_upstream_rr_peer_t * ngx_http_lua_upstream_lookup_peer(lua_State *L); static int ngx_http_lua_upstream_set_peer_down(lua_State * L); +static int ngx_http_lua_upstream_set_peer_weight(lua_State * L); static ngx_http_module_t ngx_http_lua_upstream_ctx = { @@ -98,6 +99,9 @@ ngx_http_lua_upstream_create_module(lua_State * L) lua_pushcfunction(L, ngx_http_lua_upstream_set_peer_down); lua_setfield(L, -2, "set_peer_down"); + lua_pushcfunction(L, ngx_http_lua_upstream_set_peer_weight); + lua_setfield(L, -2, "set_peer_weight"); + return 1; } @@ -348,6 +352,50 @@ ngx_http_lua_upstream_set_peer_down(lua_State * L) return 1; } +static int +ngx_http_lua_upstream_set_peer_weight(lua_State * L) +{ + ngx_http_upstream_rr_peer_t *peer; + ngx_str_t host; + ngx_http_upstream_srv_conf_t *us; + ngx_http_upstream_rr_peers_t *peers; + + if (lua_gettop(L) != 4) { + return luaL_error(L, "exactly 4 arguments expected"); + } + + peer = ngx_http_lua_upstream_lookup_peer(L); + if (peer == NULL) { + return 2; + } + + int new_weight = (int)lua_tointeger(L, 4); + if (new_weight < 1){ + lua_pushnil(L); + lua_pushliteral(L, "ilegal weight"); + return 2; + } + + + int diff = new_weight - peer->weight; + peer->weight = new_weight; + peer->effective_weight = new_weight; + peer->current_weight += diff; + + // find upstream, in order to update weighted & total_weight + host.data = (u_char *) luaL_checklstring(L, 1, &host.len); + us = ngx_http_lua_upstream_find_upstream(L, &host); + if (us == NULL) { + lua_pushnil(L); + lua_pushliteral(L, "upstream not found"); + return 2; + } + peers = us->peer.data; + peers->total_weight += diff; + peers->weighted = (peers->total_weight == peers->number); + + return 1; +} static ngx_http_upstream_rr_peer_t * ngx_http_lua_upstream_lookup_peer(lua_State *L) diff --git a/t/sanity.t b/t/sanity.t index b6142a4..4ee445c 100644 --- a/t/sanity.t +++ b/t/sanity.t @@ -564,3 +564,37 @@ upstream 127.0.0.1:1130: --- no_error_log [error] +=== TEST 15: set peer weight +--- http_config + $TEST_NGINX_MY_INIT_CONFIG + upstream bar { + server 127.0.0.2 weight=2; + server 127.0.0.3 weight=3; + server 127.0.0.4 fail_timeout=23 weight=7 max_fails=200 backup; + } +--- config + location /t { + content_by_lua ' + local upstream = require "ngx.upstream" + local ljson = require "ljson" + local u = "bar" + local ok, err = upstream.set_peer_weight(u, false, 0, 1) + if not ok then + ngx.say("failed to set peer weight: ", err) + return + end + + local peers, err = upstream.get_primary_peers(u) + if not peers then + ngx.say("failed to get peers: ", err) + return + end + ngx.say(ljson.encode(peers)) + '; + } +--- request + GET /t +--- response_body +[{"current_weight":-1,"effective_weight":1,"fail_timeout":10,"fails":0,"id":0,"max_fails":1,"name":"127.0.0.2:80","weight":1},{"current_weight":0,"effective_weight":3,"fail_timeout":10,"fails":0,"id":1,"max_fails":1,"name":"127.0.0.3:80","weight":3}] +--- no_error_log +[error]