From 57418849edabce3c9547dc8f6854c80ac138bed9 Mon Sep 17 00:00:00 2001
From: lepdou <ledouzhang@tencent.com>
Date: Mon, 7 Nov 2022 14:18:36 +0800
Subject: [PATCH] fix weight lb bug (#704)

* fix weight lb bug

* add change log
---
 CHANGELOG.md                                   |  1 +
 .../PolarisLoadBalancerCompositeRule.java      |  2 +-
 .../src/main/resources/bootstrap.yml           |  6 ++++++
 .../loadbalancer/PolarisWeightedRule.java      | 18 ++++++++++++++++--
 4 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index be798dea3..1ff48cf6e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,3 +3,4 @@
 
 - [fix router concurrent bug&fix spring-retry circuit breaker not work bug](https://github.com/Tencent/spring-cloud-tencent/pull/679)
 - [fix router unit test](https://github.com/Tencent/spring-cloud-tencent/pull/684)
+- [fix weight load balancer bug](https://github.com/Tencent/spring-cloud-tencent/pull/704)
diff --git a/spring-cloud-starter-tencent-polaris-router/src/main/java/com/tencent/cloud/polaris/router/PolarisLoadBalancerCompositeRule.java b/spring-cloud-starter-tencent-polaris-router/src/main/java/com/tencent/cloud/polaris/router/PolarisLoadBalancerCompositeRule.java
index 947444de4..f06a8f323 100644
--- a/spring-cloud-starter-tencent-polaris-router/src/main/java/com/tencent/cloud/polaris/router/PolarisLoadBalancerCompositeRule.java
+++ b/spring-cloud-starter-tencent-polaris-router/src/main/java/com/tencent/cloud/polaris/router/PolarisLoadBalancerCompositeRule.java
@@ -122,7 +122,7 @@ public class PolarisLoadBalancerCompositeRule extends AbstractLoadBalancerRule {
 		}
 
 		// set load balancer to delegate rule
-		if (!initializedLoadBalancerToDelegateRule.compareAndSet(false, true)) {
+		if (initializedLoadBalancerToDelegateRule.compareAndSet(false, true)) {
 			delegateRule.setLoadBalancer(loadBalancer);
 		}
 
diff --git a/spring-cloud-tencent-examples/polaris-router-example/router-caller-service/src/main/resources/bootstrap.yml b/spring-cloud-tencent-examples/polaris-router-example/router-caller-service/src/main/resources/bootstrap.yml
index 7a53c2f07..47d4b1d10 100644
--- a/spring-cloud-tencent-examples/polaris-router-example/router-caller-service/src/main/resources/bootstrap.yml
+++ b/spring-cloud-tencent-examples/polaris-router-example/router-caller-service/src/main/resources/bootstrap.yml
@@ -14,6 +14,12 @@ spring:
       enabled: true
       loadbalancer:
         enabled: true
+    loadbalancer:
+      retry:
+        retry-on-all-operations: true
+        max-retries-on-same-service-instance: 0
+        max-retries-on-next-service-instance: 3
+#        strategy: polarisWeighted
 management:
   endpoints:
     web:
diff --git a/spring-cloud-tencent-polaris-loadbalancer/src/main/java/com/tencent/cloud/polaris/loadbalancer/PolarisWeightedRule.java b/spring-cloud-tencent-polaris-loadbalancer/src/main/java/com/tencent/cloud/polaris/loadbalancer/PolarisWeightedRule.java
index 31278f558..8c38d3a56 100644
--- a/spring-cloud-tencent-polaris-loadbalancer/src/main/java/com/tencent/cloud/polaris/loadbalancer/PolarisWeightedRule.java
+++ b/spring-cloud-tencent-polaris-loadbalancer/src/main/java/com/tencent/cloud/polaris/loadbalancer/PolarisWeightedRule.java
@@ -21,7 +21,10 @@ package com.tencent.cloud.polaris.loadbalancer;
 import java.util.List;
 
 import com.netflix.client.config.IClientConfig;
-import com.netflix.loadbalancer.AbstractLoadBalancerRule;
+import com.netflix.loadbalancer.AbstractServerPredicate;
+import com.netflix.loadbalancer.AvailabilityPredicate;
+import com.netflix.loadbalancer.CompositePredicate;
+import com.netflix.loadbalancer.PredicateBasedRule;
 import com.netflix.loadbalancer.Server;
 import com.tencent.cloud.common.pojo.PolarisServer;
 import com.tencent.polaris.api.config.consumer.LoadBalanceConfig;
@@ -38,9 +41,10 @@ import org.springframework.util.CollectionUtils;
  *
  * @author lepdou 2022-05-17
  */
-public class PolarisWeightedRule extends AbstractLoadBalancerRule {
+public class PolarisWeightedRule extends PredicateBasedRule {
 
 	private final RouterAPI routerAPI;
+	private CompositePredicate compositePredicate;
 
 	public PolarisWeightedRule(RouterAPI routerAPI) {
 		this.routerAPI = routerAPI;
@@ -48,6 +52,13 @@ public class PolarisWeightedRule extends AbstractLoadBalancerRule {
 
 	@Override
 	public void initWithNiwsConfig(IClientConfig clientConfig) {
+		AvailabilityPredicate availabilityPredicate = new AvailabilityPredicate(this, clientConfig);
+		compositePredicate = CompositePredicate.withPredicates(availabilityPredicate).build();
+	}
+
+	@Override
+	public AbstractServerPredicate getPredicate() {
+		return compositePredicate;
 	}
 
 	@Override
@@ -57,6 +68,9 @@ public class PolarisWeightedRule extends AbstractLoadBalancerRule {
 			return null;
 		}
 
+		// filter circuit breaker servers by ribbon
+		servers = compositePredicate.getEligibleServers(servers);
+
 		ServiceInstances serviceInstances = LoadBalancerUtils.transferServersToServiceInstances(servers);
 
 		ProcessLoadBalanceRequest request = new ProcessLoadBalanceRequest();