add set_peer_weight api

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)

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]