From 9a3ccff4c84b4c39719c3d4c58ee416836a104f6 Mon Sep 17 00:00:00 2001
From: rinfx <yucheng.lxr@alibaba-inc.com>
Date: Tue, 23 Sep 2025 19:24:33 +0800
Subject: [PATCH] opt(ai-load-balancer): update global least request lua script
 for ai-load-balancer (#2945)

---
 .../global_least_request/lb_policy.go         | 34 ++++++-------------
 .../prefix_cache/lb_policy.go                 | 13 +++----
 2 files changed, 16 insertions(+), 31 deletions(-)

diff --git a/plugins/wasm-go/extensions/ai-load-balancer/global_least_request/lb_policy.go b/plugins/wasm-go/extensions/ai-load-balancer/global_least_request/lb_policy.go
index 23db26d2a..46fcb819c 100644
--- a/plugins/wasm-go/extensions/ai-load-balancer/global_least_request/lb_policy.go
+++ b/plugins/wasm-go/extensions/ai-load-balancer/global_least_request/lb_policy.go
@@ -28,31 +28,19 @@ local function randomBool()
     return math.random() >= 0.5
 end
 
-local function is_healthy(addr)
-    for i = 4, #KEYS do
-        if addr == KEYS[i] then
-            return true
-        end
-    end
-    return false
-end
-
 if redis.call('HEXISTS', hset_key, current_target) == 1 then
     current_count = redis.call('HGET', hset_key, current_target)
-    local hash = redis.call('HGETALL', hset_key)
-    for i = 1, #hash, 2 do
-		local addr = hash[i]
-		local count = hash[i+1]
-        if is_healthy(addr) then
-            if tonumber(count) < tonumber(current_count) then
-                current_target = addr
-                current_count = count
-            elseif count == current_count and randomBool() then
-                current_target = addr
-                current_count = count
-            end
-        end
-    end
+	for i = 4, #KEYS do
+		if redis.call('HEXISTS', hset_key, KEYS[i]) == 1 then
+			local count = redis.call('HGET', hset_key, KEYS[i])
+			if tonumber(count) < tonumber(current_count) then
+				current_target = KEYS[i]
+				current_count = count
+			elseif count == current_count and randomBool() then
+				current_target = KEYS[i]
+			end
+		end
+	end
 end
 
 redis.call("HINCRBY", hset_key, current_target, 1)
diff --git a/plugins/wasm-go/extensions/ai-load-balancer/prefix_cache/lb_policy.go b/plugins/wasm-go/extensions/ai-load-balancer/prefix_cache/lb_policy.go
index 5663e759c..c71254997 100644
--- a/plugins/wasm-go/extensions/ai-load-balancer/prefix_cache/lb_policy.go
+++ b/plugins/wasm-go/extensions/ai-load-balancer/prefix_cache/lb_policy.go
@@ -121,17 +121,14 @@ if target == "" then
 	target = default_target
 	if redis.call('HEXISTS', hset_key, target) == 1 then
 		current_count = redis.call('HGET', hset_key, target)
-		local hash = redis.call('HGETALL', hset_key)
-		for i = 1, #hash, 2 do
-			local addr = hash[i]
-			local count = hash[i+1]
-			if is_healthy(addr) then
+		for i = 4, #KEYS do
+			if redis.call('HEXISTS', hset_key, KEYS[i]) == 1 then
+				local count = redis.call('HGET', hset_key, KEYS[i])
 				if tonumber(count) < tonumber(current_count) then
-					target = addr
+					target = KEYS[i]
 					current_count = count
 				elseif count == current_count and randomBool() then
-					target = addr
-					current_count = count
+					target = KEYS[i]
 				end
 			end
 		end