fix(balancer): keep least_conn load state across upstream scaling

When least_conn proxies long-lived connections (WebSocket) and the upstream
is scaled, the load stays skewed on the original nodes: the newly added ones
are not preferred and least_conn degrades to round-robin.

The picker is cached by the upstream version, so it is rebuilt whenever the
upstream changes - and scaling changes it. The binary heap holding the
per-server scores lives inside the picker, so every score is reset to the base
weight on rebuild and the connections already established are forgotten. The
requests that are still in flight keep the picker they were routed with and
release their server on it in the log phase, so their releases land on a heap
nobody reads anymore, while the rebuilt heap never learns about them.

Move both the heap and the in-flight connection counts out of the picker into
a per-worker state keyed by the upstream resource key, which is stable across
scaling (and across health status flips) unlike the picker version. The picker
now reconciles that heap with the current node set instead of rebuilding it:
surviving nodes keep their load, a freshly added node starts empty and is
preferred right away, and a node that leaves keeps its count so its score is
restored if it comes back. Every generation of pickers shares one view of the
load, so a connection established before a rebuild is released against the
heap that is actually in use.

The score is now derived from the connection count instead of being
accumulated with +/- effect_weight, which keeps it exact over time.

Because that state now outlives the picker, releasing a server the request no
longer holds is no longer self-healing: it used to be washed away by the next
rebuild, now it is written into shared state for good. pick_server released
the current server when it entered a retry but left ctx.balancer_server set,
so a request that then ran out of servers to try released it a second time in
the log phase. It only takes an active health check marking a node unhealthy
to reach that path, since the retry count is derived from all nodes while only
the healthy ones are picked from. Clear ctx.balancer_server on release, so it
always names the server the request currently holds, and skip the release when
it names none.

The priority is part of the state key, since node sets of different priorities
are disjoint and must not share a heap. A node that moves between priorities
leaves its count behind in the old level, where it drains normally.

Also call after_balance in the stream log phase, which it never did: for L4
the count was only ever incremented, so least_conn could not balance TCP long
connections at all and ctx.balancer_tried_servers was leaked.

Fixes #12217

Author: AlinsRan

apisix/balancer.lua

@@ -174,9 +174,9 @@ local function create_server_picker(upstream, checker)
             return server_picker
         end
 
-        core.log.info("upstream nodes: ",
-                      core.json.delay_encode(up_nodes[up_nodes._priority_index[1]]))
-        local server_picker = picker.new(up_nodes[up_nodes._priority_index[1]], upstream)
+        local priority = up_nodes._priority_index[1]
+        core.log.info("upstream nodes: ", core.json.delay_encode(up_nodes[priority]))
+        local server_picker = picker.new(up_nodes[priority], upstream, priority)
         server_picker.addr_to_domain = addr_to_domain
         return server_picker
     end
@@ -269,6 +269,10 @@ local function pick_server(route, ctx)
     if ctx.balancer_try_count > 1 then
         if ctx.server_picker and ctx.server_picker.after_balance then
             ctx.server_picker.after_balance(ctx, true)
+            -- the server is released, the request holds none until the next pick.
+            -- Leaving it set would let the log phase release it a second time when
+            -- no further server can be picked below.
+            ctx.balancer_server = nil
         end
 
         if checker then
@@ -324,6 +328,7 @@ local function pick_server(route, ctx)
         end
 
         server_picker.after_balance(ctx, true)
+        ctx.balancer_server = nil
         server = nil
     end
 

apisix/balancer/least_conn.lua

@@ -19,28 +19,118 @@ local core = require("apisix.core")
 local binaryHeap = require("binaryheap")
 local ipairs = ipairs
 local pairs = pairs
+local tostring = tostring
 
 
 local _M = {}
 
 
+-- Per-worker balancing state, shared by every picker built for the same upstream.
+--
+-- A picker is cached by the upstream version, so it is thrown away whenever the
+-- upstream changes: scaling, a config update, or a health status flip. The
+-- requests that are still in flight keep the picker they were routed with, and
+-- release their server on that picker in the log phase. If the heap lived inside
+-- the picker, the rebuilt one would never see those releases: the connections
+-- established before the rebuild would be forgotten on creation and then, once
+-- they closed, decremented on a heap nobody reads anymore. Long-lived
+-- connections (WebSocket) would keep the load skewed on the original nodes and
+-- least_conn would degrade to round-robin. See #12217.
+--
+-- Keeping the heap and the in-flight counts here, keyed by something stable
+-- across scaling, gives every generation of pickers a single view of the load.
+--
+-- The ttl does not drop the state: the version is constant, so an expired entry
+-- is reused and refreshed instead of rebuilt. The count is what bounds it, and
+-- it is the price of not leaking a state per upstream that is ever deleted: a
+-- worker balancing more than `count` (upstream, priority) pairs evicts the least
+-- recently rebuilt one, which degrades that upstream to the pre-fix behavior
+-- (its in-flight connections are forgotten on the next picker rebuild). Keep it
+-- comfortably above the picker cache above, which is what drives the recency.
+local STATE_VER = "least_conn"
+local states = core.lrucache.new({ttl = 300, count = 512})
+
+
 local function least_score(a, b)
     return a.score < b.score
 end
 
 
-function _M.new(up_nodes, upstream)
-    local servers_heap = binaryHeap.minUnique(least_score)
+local function new_state()
+    return {
+        heap = binaryHeap.minUnique(least_score),
+        -- server -> in-flight connections, only holds positive counts
+        conns = {},
+        -- server -> true, mirrors the payloads currently in the heap
+        members = {},
+    }
+end
+
+
+local function update_score(state, server)
+    local info = state.heap:valueByPayload(server)
+    -- the server may have left the upstream while it still held connections
+    if not info then
+        return
+    end
+
+    info.score = (1 + (state.conns[server] or 0)) * info.effect_weight
+    state.heap:update(server, info)
+end
+
+
+-- Align the long-lived heap with the current node set, keeping the in-flight
+-- counts of the nodes that survive. A node that is added back later (scaled in
+-- again, or reported healthy again) gets its score restored from `conns`.
+local function sync_nodes(state, up_nodes)
+    local heap = state.heap
+
+    for server in pairs(state.members) do
+        if not up_nodes[server] then
+            heap:remove(server)
+            state.members[server] = nil
+        end
+    end
+
     for server, weight in pairs(up_nodes) do
-        local score = 1 / weight
-        -- Note: the argument order of insert is different from others
-        servers_heap:insert({
-            server = server,
-            effect_weight = 1 / weight,
-            score = score,
-        }, server)
+        local effect_weight = 1 / weight
+        local score = (1 + (state.conns[server] or 0)) * effect_weight
+        local info = heap:valueByPayload(server)
+        if info then
+            info.effect_weight = effect_weight
+            info.score = score
+            heap:update(server, info)
+        else
+            -- Note: the argument order of insert is different from others
+            heap:insert({
+                server = server,
+                effect_weight = effect_weight,
+                score = score,
+            }, server)
+            state.members[server] = true
+        end
+    end
+end
+
+
+function _M.new(up_nodes, upstream, priority)
+    -- 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 state
+    if up_key then
+        -- each priority level owns a disjoint node set, so it needs its own heap
+        state = states(up_key .. "#" .. tostring(priority), STATE_VER, new_state)
+    else
+        -- no stable identity, fall back to a state private to this picker
+        state = new_state()
     end
 
+    sync_nodes(state, up_nodes)
+
+    local servers_heap = state.heap
+    local conns = state.conns
+
     return {
         upstream = upstream,
         get = function (ctx)
@@ -68,22 +158,29 @@ function _M.new(up_nodes, upstream)
                     servers_heap:insert(info, info.server)
                 end
             else
-                server, info = servers_heap:peek()
+                server = servers_heap:peek()
             end
 
             if not server then
                 return nil, err
             end
 
-            info.score = info.score + info.effect_weight
-            servers_heap:update(server, info)
+            conns[server] = (conns[server] or 0) + 1
+            update_score(state, server)
             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)
+            -- the request may hold no server: balancer.lua releases the current one
+            -- when it enters a retry, and the log phase still runs when no further
+            -- server could be picked
+            if server then
+                local count = (conns[server] or 0) - 1
+                -- a count is never negative, and keeping it >= 0 also keeps the
+                -- score away from `0 * inf` (a node may be configured weight 0)
+                conns[server] = count > 0 and count or nil
+                update_score(state, server)
+            end
 
             if not before_retry then
                 if ctx.balancer_tried_servers then

apisix/balancer/priority.lua

@@ -33,7 +33,9 @@ function _M.new(up_nodes, upstream, picker_mod)
 
     local pickers = core.table.new(#priority_index, 0)
     for i, priority in ipairs(priority_index) do
-        local picker, err = picker_mod.new(up_nodes[priority], upstream)
+        -- the priority is part of the picker's identity: node sets of different
+        -- priorities are disjoint and must not share balancing state
+        local picker, err = picker_mod.new(up_nodes[priority], upstream, priority)
         if not picker then
             return nil, "failed to create picker with priority " .. priority .. ": " .. err
         end

apisix/init.lua

@@ -1541,6 +1541,10 @@ function _M.stream_log_phase()
 
     healthcheck_passive(api_ctx)
 
+    if api_ctx.server_picker and api_ctx.server_picker.after_balance then
+        api_ctx.server_picker.after_balance(api_ctx, false)
+    end
+
     core.ctx.release_vars(api_ctx)
     if api_ctx.plugins then
         core.tablepool.release("plugins", api_ctx.plugins)

t/node/least_conn.t

@@ -149,3 +149,20 @@ qr/proxy request to \S+ while connecting to upstream/
 --- grep_error_log_out
 proxy request to 127.0.0.1:1999 while connecting to upstream
 proxy request to 0.0.0.0:1999 while connecting to upstream
+
+
+
+=== TEST 5: more retries than nodes, the request ends up holding no server
+--- apisix_yaml
+upstreams:
+  - id: 1
+    type: least_conn
+    retries: 3
+    nodes:
+        "127.0.0.1:1999": 2
+        "0.0.0.0:1999": 1
+--- error_code: 502
+--- error_log
+failed to find valid upstream server, all upstream servers tried
+--- no_error_log
+table index is nil