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

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
@@ -248,7 +248,13 @@ local function pick_server(route, ctx)
     local up_conf = ctx.upstream_conf
 
     local nodes_count = #up_conf.nodes
-    if nodes_count == 1 then
+    -- least_conn counts the in-flight connections of every request it routes, so it
+    -- has to see them even while the upstream has a single node: those connections
+    -- are what tells a later scale out that the node is not empty. Skipping the
+    -- balancer here would leave it blind to everything routed before the second
+    -- node showed up, which is the state a k8s deployment or a discovery service
+    -- starts from. See #12217
+    if nodes_count == 1 and up_conf.type ~= "least_conn" then
         local node = up_conf.nodes[1]
         ctx.balancer_ip = node.host
         ctx.balancer_port = node.port
@@ -268,7 +274,10 @@ local function pick_server(route, ctx)
     ctx.balancer_try_count = (ctx.balancer_try_count or 0) + 1
     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)
+            -- remembering the server as tried is what keeps the next pick off it, so
+            -- only do it when there is another one to move to. With a single node the
+            -- retry has to land on it again, the way the fast path below always did
+            ctx.server_picker.after_balance(ctx, nodes_count > 1)
         end
 
         if checker then
@@ -331,6 +340,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 +359,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,184 @@ 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 values are weak, which is exactly the lifetime this state needs. A picker
+-- holds its state, an in-flight request holds the picker it was routed with, and
+-- the picker cache holds the current one - so a state survives for as long as
+-- anything can still release a connection into it. Once nothing references it
+-- there is no connection left to count, and dropping it costs nothing. That also
+-- means there is no size to tune and no eviction that could quietly forget a busy
+-- upstream, which is what an LRU would do here: it would rank states by how
+-- recently they were rebuilt, and a stable upstream holding many long-lived
+-- connections is precisely the one that is never rebuilt.
+local states = setmetatable({}, {__mode = "v"})
+
+
 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 info = heap:valueByPayload(server)
+        if info then
+            info.effect_weight = effect_weight
+        else
+            -- Note: the argument order of insert is different from others
+            heap:insert({
+                server = server,
+                effect_weight = effect_weight,
+                score = effect_weight,
+            }, server)
+            state.members[server] = true
+        end
+        -- one place decides what a score is worth
+        update_score(state, server)
     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
+        local state_key = up_key .. "#" .. priority
+        state = states[state_key]
+        if not state then
+            state = new_state()
+            states[state_key] = state
+        end
+    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
+                    -- a release with no matching pick. The store below floors the
+                    -- count either way, so the score cannot be corrupted - but the
+                    -- accounting is wrong and it should not pass in silence
+                    core.log.error("released a connection never picked on ", server)
+                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 +212,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)