fix(balancer): keep least_conn in-flight count across upstream scaling

The least_conn picker stores per-server connection scores inside a binary
heap that lives in the picker instance. The picker is cached by upstream
version and rebuilt whenever the upstream changes, e.g. when nodes are
scaled. On rebuild every score is reset to the base weight, so connections
already established are forgotten. For long-lived connections such as
WebSocket this leaves the load skewed on the original nodes and the newly
added nodes are not preferred, degrading least_conn to round-robin.

Keep a per-worker in-flight connection count for each (upstream, server)
outside the picker, keyed by the upstream resource key which is stable
across scaling. The rebuilt heap seeds each score from this count, so
surviving nodes keep their load while freshly added nodes start empty and
are preferred right after scaling. Drained servers that leave the upstream
are pruned to bound memory.

Fixes #12217

Author: AlinsRan

apisix/balancer/least_conn.lua

@@ -24,20 +24,64 @@ local pairs = pairs
 local _M = {}
 
 
+-- Per-worker in-flight connection count for each (upstream, server), kept outside
+-- the picker so it survives balancer recreation. The picker (and its heap) is
+-- rebuilt whenever the upstream changes, e.g. on scaling. Without this, every
+-- score is reset to the base weight and connections already established are
+-- forgotten, so the newly added nodes are not preferred and long-lived
+-- connections (e.g. WebSocket) stay skewed on the original nodes. See #12217.
+-- structure: conn_count[upstream_key][server] = in-flight count
+local conn_count = {}
+
+
 local function least_score(a, b)
     return a.score < b.score
 end
 
 
 function _M.new(up_nodes, upstream)
+    -- resource_key/resource_id identifies the upstream and is stable across node
+    -- scaling, unlike the picker version which changes whenever the nodes change
+    local up_key = upstream.resource_key or upstream.resource_id
+    local counts
+    if up_key then
+        counts = conn_count[up_key]
+        if not counts then
+            counts = {}
+            conn_count[up_key] = counts
+        end
+        -- drop drained servers no longer in the upstream to bound memory on churn
+        for server, count in pairs(counts) do
+            if count <= 0 and not up_nodes[server] then
+                counts[server] = nil
+            end
+        end
+    end
+
+    local function update_conn_count(server, delta)
+        if not counts then
+            return
+        end
+        local count = (counts[server] or 0) + delta
+        if count <= 0 and not up_nodes[server] then
+            counts[server] = nil
+        else
+            counts[server] = count
+        end
+    end
+
     local servers_heap = binaryHeap.minUnique(least_score)
     for server, weight in pairs(up_nodes) do
-        local score = 1 / weight
+        local effect_weight = 1 / weight
+        -- seed the score with the connections this worker already holds so that
+        -- surviving nodes keep their load while freshly added nodes start empty
+        -- and are preferred right after scaling
+        local held = counts and counts[server] or 0
         -- Note: the argument order of insert is different from others
         servers_heap:insert({
             server = server,
-            effect_weight = 1 / weight,
-            score = score,
+            effect_weight = effect_weight,
+            score = (held + 1) * effect_weight,
         }, server)
     end
 
@@ -77,13 +121,15 @@ function _M.new(up_nodes, upstream)
 
             info.score = info.score + info.effect_weight
             servers_heap:update(server, info)
+            update_conn_count(server, 1)
             return server
         end,
         after_balance = function (ctx, before_retry)
             local server = ctx.balancer_server
             local info = servers_heap:valueByPayload(server)
             info.score = info.score - info.effect_weight
             servers_heap:update(server, info)
+            update_conn_count(server, -1)
 
             if not before_retry then
                 if ctx.balancer_tried_servers then

t/node/least_conn3.t

@@ -0,0 +1,106 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+use t::APISIX 'no_plan';
+
+repeat_each(2);
+log_level('info');
+no_root_location();
+no_shuffle();
+
+run_tests();
+
+__DATA__
+
+=== TEST 1: keep in-flight conn count across balancer recreation on scaling
+--- config
+    location /t {
+        content_by_lua_block {
+            local least_conn = require("apisix.balancer.least_conn")
+            -- resource_key is stable across scaling, so both pickers share the
+            -- same connection count table
+            local up = {resource_key = "/upstreams/lc-scale"}
+
+            -- 2 nodes serving long-lived connections
+            local nodes = {["127.0.0.1:1980"] = 1, ["127.0.0.1:1981"] = 1}
+            local p1 = least_conn.new(nodes, up)
+
+            -- establish 4 in-flight connections (get without after_balance)
+            local ctx = {}
+            local held = {}
+            for _ = 1, 4 do
+                held[#held + 1] = p1.get(ctx)
+            end
+
+            -- scale out: add a third node, the picker is recreated
+            local scaled = {["127.0.0.1:1980"] = 1, ["127.0.0.1:1981"] = 1,
+                            ["127.0.0.1:1982"] = 1}
+            local p2 = least_conn.new(scaled, up)
+
+            -- the freshly added node has no connection, so it must be picked first
+            for _ = 1, 2 do
+                local s = p2.get(ctx)
+                held[#held + 1] = s
+                ngx.say(s)
+            end
+
+            -- release everything so repeated runs start from a clean state
+            for _, s in ipairs(held) do
+                ctx.balancer_server = s
+                p2.after_balance(ctx, false)
+            end
+        }
+    }
+--- request
+GET /t
+--- response_body
+127.0.0.1:1982
+127.0.0.1:1982
+
+
+
+=== TEST 2: scale down drops the removed node, remaining nodes balance
+--- config
+    location /t {
+        content_by_lua_block {
+            local least_conn = require("apisix.balancer.least_conn")
+            local up = {resource_key = "/upstreams/lc-scale-down"}
+
+            local nodes = {["127.0.0.1:1980"] = 1, ["127.0.0.1:1981"] = 1}
+            local p1 = least_conn.new(nodes, up)
+
+            local ctx = {}
+            -- fully complete two requests, one per node
+            for _ = 1, 2 do
+                local s = p1.get(ctx)
+                ctx.balancer_server = s
+                p1.after_balance(ctx, false)
+            end
+
+            -- scale down to a single remaining node, picker recreated
+            local scaled = {["127.0.0.1:1981"] = 1}
+            local p2 = least_conn.new(scaled, up)
+
+            local s = p2.get(ctx)
+            ctx.balancer_server = s
+            p2.after_balance(ctx, false)
+            ngx.say(s)
+        }
+    }
+--- request
+GET /t
+--- response_body
+127.0.0.1:1981