fix: improve least_conn balancer implementation

Author: coder2z

apisix/balancer/least_conn.lua

@@ -24,26 +24,16 @@ local ngx = ngx
 local ngx_shared = ngx.shared
 local tostring = tostring
 
-
 local _M = {}
 
 -- Shared dictionary to store connection counts across balancer recreations
 local CONN_COUNT_DICT_NAME = "balancer-least-conn"
-local conn_count_dict = ngx_shared[CONN_COUNT_DICT_NAME]
-
--- Check that the shared dictionary exists
-if ngx.shared and conn_count_dict == nil then
-    core.log.error("shared dict '", CONN_COUNT_DICT_NAME, "' not found, ",
-            "please add 'lua_shared_dict ", CONN_COUNT_DICT_NAME, " 10m;' ",
-            "to your nginx configuration")
-    error("shared dict '" .. CONN_COUNT_DICT_NAME .. "' not found")
-end
+local conn_count_dict
 
 local function least_score(a, b)
     return a.score < b.score
 end
 
-
 -- Get the connection count key for a specific upstream and server
 local function get_conn_count_key(upstream, server)
     local upstream_id = upstream.id
@@ -129,10 +119,14 @@ local function cleanup_stale_conn_counts(upstream, current_servers)
     end
 end
 
-
 function _M.new(up_nodes, upstream)
     if not conn_count_dict then
-        return nil, "dictionary not find"
+        conn_count_dict = ngx_shared[CONN_COUNT_DICT_NAME]
+    end
+
+    if not conn_count_dict then
+        core.log.error("shared dict '", CONN_COUNT_DICT_NAME, "' not found")
+        return nil, "shared dict not found"
     end
 
     local servers_heap = binaryHeap.minUnique(least_score)
@@ -162,7 +156,7 @@ function _M.new(up_nodes, upstream)
 
     return {
         upstream = upstream,
-        get = function (ctx)
+        get = function(ctx)
             local server, info, err
             if ctx.balancer_tried_servers then
                 local tried_server_list = {}
@@ -201,7 +195,7 @@ function _M.new(up_nodes, upstream)
             incr_server_conn_count(upstream, server, 1)
             return server
         end,
-        after_balance = function (ctx, before_retry)
+        after_balance = function(ctx, before_retry)
             local server = ctx.balancer_server
             local info = servers_heap:valueByPayload(server)
             if not info then
@@ -233,7 +227,7 @@ function _M.new(up_nodes, upstream)
 
             ctx.balancer_tried_servers[server] = true
         end,
-        before_retry_next_priority = function (ctx)
+        before_retry_next_priority = function(ctx)
             if ctx.balancer_tried_servers then
                 core.tablepool.release("balancer_tried_servers", ctx.balancer_tried_servers)
                 ctx.balancer_tried_servers = nil
@@ -242,6 +236,5 @@ function _M.new(up_nodes, upstream)
     }
 end
 
-
 return _M