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. Because that heap can hold nodes a picker was
not built with, a picker only ever hands out the nodes of its own conf, which
is what the rest of balancer.lua assumes of the server it gets back. The state
is held weakly, so it lives exactly as long as a picker that can still release
a connection into it.

pick_server skips the balancer entirely when the upstream has a single node.
least_conn cannot afford that: the requests routed while the upstream was alone
are the ones a later scale out needs to know about, and a single node is where
a k8s deployment or a discovery service starts from. Route them through the
balancer so they are counted.

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

Now that the state 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. Make after_balance own
that: releasing a server clears ctx.balancer_server, so the field always names
the server the request currently holds and a caller that retries simply gets a
fresh one from the next pick. Three release sites relied on the caller to
remember and did not - pick_server on entering a retry and on skipping an
unhealthy node, and ai-proxy-multi's retry_on_error, which releases before it
knows whether it will pick again and returns as-is when the fallback strategy
does not cover the status code. Each of them left the log phase free to release
the same server a second time. Reaching the first one needs nothing unusual:
the retry count is derived from all nodes while only the healthy ones are
picked from, so a single failing health check is enough.

ctx.server_picker was also published only at the very end of pick_server, so a
request that bailed out after picking a server never released it at all.

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.

Note that an upstream declared inline in a traffic-split rule has no stable
resource key, so it keeps the previous behavior.

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
@@ -331,6 +331,9 @@ local function pick_server(route, ctx)
         return nil, "failed to find valid upstream server, all upstream servers tried"
     end
     ctx.balancer_server = server
+    -- from here on the request holds a server, so the log phase must be able to
+    -- release it even if we bail out below
+    ctx.server_picker = server_picker
 
     local domain = server_picker.addr_to_domain[server]
     local res, err = lrucache_addr(server, nil, parse_addr, server)
@@ -347,7 +350,6 @@ local function pick_server(route, ctx)
     if is_http and ctx.var then
         ctx.var.upstream_unresolved_host = ctx.upstream_unresolved_host
     end
-    ctx.server_picker = server_picker
     res.upstream_host = parse_server_for_upstream_host(res, ctx.upstream_scheme)
 
     return res

apisix/balancer/chash.lua

@@ -123,6 +123,11 @@ function _M.new(up_nodes, upstream)
             return servers[id]
         end,
         after_balance = function (ctx, before_retry)
+            -- the release is what makes the request stop holding the server, so drop
+            -- the reference here instead of leaving every caller to remember
+            local server = ctx.balancer_server
+            ctx.balancer_server = nil
+
             if not before_retry then
                 if ctx.balancer_tried_servers then
                     core.tablepool.release("balancer_tried_servers", ctx.balancer_tried_servers)
@@ -132,11 +137,15 @@ function _M.new(up_nodes, upstream)
                 return nil
             end
 
+            if not server then
+                return nil
+            end
+
             if not ctx.balancer_tried_servers then
                 ctx.balancer_tried_servers = core.tablepool.fetch("balancer_tried_servers", 0, 2)
             end
 
-            ctx.balancer_tried_servers[ctx.balancer_server] = true
+            ctx.balancer_tried_servers[server] = true
             ctx.balancer_tried_servers_count = (ctx.balancer_tried_servers_count or 0) + 1
         end,
         before_retry_next_priority = function (ctx)

apisix/balancer/ewma.lua

@@ -187,12 +187,21 @@ local function _ewma_find(ctx, up_nodes)
 end
 
 local function _ewma_after_balance(ctx, before_retry)
+    -- the release is what makes the request stop holding the server, so drop the
+    -- reference here instead of leaving every caller to remember
+    local server = ctx.balancer_server
+    ctx.balancer_server = nil
+
     if before_retry then
+        if not server then
+            return nil
+        end
+
         if not ctx.balancer_tried_servers then
             ctx.balancer_tried_servers = core.tablepool.fetch("balancer_tried_servers", 0, 2)
         end
 
-        ctx.balancer_tried_servers[ctx.balancer_server] = true
+        ctx.balancer_tried_servers[server] = true
         ctx.balancer_tried_servers_count = (ctx.balancer_tried_servers_count or 0) + 1
 
         return nil

apisix/balancer/least_conn.lua

@@ -24,66 +24,181 @@ local pairs = pairs
 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). Note that
+-- the picker cache counts upstreams while this one counts priority levels, so an
+-- upstream with several of them consumes several entries.
+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 identifies the upstream and is stable across node scaling, unlike
+    -- the picker version which changes whenever the nodes change. Do not fall back to
+    -- resource_id: it is a bare id, so a route and an upstream sharing one would land
+    -- on the same heap and evict each other's nodes
+    local up_key = upstream.resource_key
+    local state
+    if up_key and priority then
+        -- each priority level owns a disjoint node set, so it needs its own heap.
+        -- Do not default the priority: a caller that does not name one has a node
+        -- set we cannot place, and folding it into level 0 would let sync_nodes
+        -- evict that level's nodes from the heap it shares
+        state = states(up_key .. "#" .. 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)
+            local tried = ctx.balancer_tried_servers
             local server, info, err
-            if ctx.balancer_tried_servers then
-                local tried_server_list = {}
-                while true do
-                    server, info = servers_heap:peek()
-                    -- we need to let the retry > #nodes so this branch can be hit and
-                    -- the request will retry next priority of nodes
-                    if server == nil then
-                        err = "all upstream servers tried"
-                        break
-                    end
-
-                    if not ctx.balancer_tried_servers[server] then
-                        break
-                    end
-
-                    servers_heap:pop()
-                    core.table.insert(tried_server_list, info)
+            local skipped
+
+            while true do
+                server, info = servers_heap:peek()
+                -- we need to let the retry > #nodes so this branch can be hit and
+                -- the request will retry next priority of nodes
+                if server == nil then
+                    err = "all upstream servers tried"
+                    break
                 end
 
-                for _, info in ipairs(tried_server_list) do
-                    servers_heap:insert(info, info.server)
+                -- the heap is shared with the pickers built for later versions of
+                -- the upstream, so it can hold nodes this request's conf does not
+                -- know about. Only hand out the ones it does
+                if up_nodes[server] and not (tried and tried[server]) then
+                    break
+                end
+
+                servers_heap:pop()
+                if not skipped then
+                    skipped = {}
+                end
+                core.table.insert(skipped, info)
+            end
+
+            if skipped then
+                for _, skipped_info in ipairs(skipped) do
+                    servers_heap:insert(skipped_info, skipped_info.server)
                 end
-            else
-                server, info = 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)
+            -- the release is what makes the request stop holding the server, so drop
+            -- the reference here instead of leaving every caller to remember. A caller
+            -- that goes on to retry gets a fresh one from the next pick
             local server = ctx.balancer_server
-            local info = servers_heap:valueByPayload(server)
-            info.score = info.score - info.effect_weight
-            servers_heap:update(server, info)
+            ctx.balancer_server = nil
+
+            if server then
+                local count = (conns[server] or 0) - 1
+                if count < 0 then
+                    -- released more than picked: the score would go below the
+                    -- baseline and, with weight 0, `0 * inf` would poison the heap
+                    -- order with a NaN. Clamp, but do not hide the accounting bug
+                    core.log.error("released a connection never picked on ", server)
+                    count = 0
+                end
+                conns[server] = count > 0 and count or nil
+                update_score(state, server)
+            end
 
             if not before_retry then
                 if ctx.balancer_tried_servers then
@@ -94,6 +209,10 @@ function _M.new(up_nodes, upstream)
                 return nil
             end
 
+            if not server then
+                return nil
+            end
+
             if not ctx.balancer_tried_servers then
                 ctx.balancer_tried_servers = core.tablepool.fetch("balancer_tried_servers", 0, 2)
             end

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/balancer/roundrobin.lua

@@ -58,6 +58,11 @@ function _M.new(up_nodes, upstream)
             return server
         end,
         after_balance = function (ctx, before_retry)
+            -- the release is what makes the request stop holding the server, so drop
+            -- the reference here instead of leaving every caller to remember
+            local server = ctx.balancer_server
+            ctx.balancer_server = nil
+
             if not before_retry then
                 if ctx.balancer_tried_servers then
                     core.tablepool.release("balancer_tried_servers", ctx.balancer_tried_servers)
@@ -67,11 +72,15 @@ function _M.new(up_nodes, upstream)
                 return nil
             end
 
+            if not server then
+                return nil
+            end
+
             if not ctx.balancer_tried_servers then
                 ctx.balancer_tried_servers = core.tablepool.fetch("balancer_tried_servers", 0, 2)
             end
 
-            ctx.balancer_tried_servers[ctx.balancer_server] = true
+            ctx.balancer_tried_servers[server] = true
             ctx.balancer_tried_servers_count = (ctx.balancer_tried_servers_count or 0) + 1
         end,
         before_retry_next_priority = function (ctx)

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/lib/server.lua

@@ -386,6 +386,29 @@ end
 _M.websocket_handshake_route = _M.websocket_handshake
 
 
+-- keep the session open until the client closes it, so that the request stays
+-- in flight in the balancer the way a real WebSocket session does
+function _M.websocket_hold()
+    local websocket = require "resty.websocket.server"
+    local wb, err = websocket:new({timeout = 30000})
+    if not wb then
+        ngx.log(ngx.ERR, "failed to new websocket: ", err)
+        return ngx.exit(400)
+    end
+
+    while true do
+        local _, typ, err = wb:recv_frame()
+        if not typ then
+            ngx.log(ngx.ERR, "failed to receive websocket frame: ", err)
+            return
+        end
+        if typ == "close" then
+            return
+        end
+    end
+end
+
+
 function _M.api_breaker()
     ngx.exit(tonumber(ngx.var.arg_code))
 end