fix weight cannot work properly on multi-registry scenario (#9787)

also see #9786
apache · Mar 16, 2022 · fc2d4d7 · fc2d4d7
1 parent a5b1d17
commit fc2d4d7
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/...o-cluster/src/main/java/org/apache/dubbo/rpc/cluster/loadbalance/AbstractLoadBalance.java b/...o-cluster/src/main/java/org/apache/dubbo/rpc/cluster/loadbalance/AbstractLoadBalance.java
@@ -20,12 +20,12 @@
 import org.apache.dubbo.common.utils.CollectionUtils;
 import org.apache.dubbo.rpc.Invocation;
 import org.apache.dubbo.rpc.Invoker;
+import org.apache.dubbo.rpc.cluster.ClusterInvoker;
 import org.apache.dubbo.rpc.cluster.LoadBalance;
 
 import java.util.List;
 
 import static org.apache.dubbo.common.constants.CommonConstants.TIMESTAMP_KEY;
-import static org.apache.dubbo.common.constants.RegistryConstants.REGISTRY_KEY;
 import static org.apache.dubbo.common.constants.RegistryConstants.REGISTRY_SERVICE_REFERENCE_PATH;
 import static org.apache.dubbo.rpc.cluster.Constants.DEFAULT_WARMUP;
 import static org.apache.dubbo.rpc.cluster.Constants.DEFAULT_WEIGHT;
@@ -75,9 +75,12 @@ public <T> Invoker<T> select(List<Invoker<T>> invokers, URL url, Invocation invo
     int getWeight(Invoker<?> invoker, Invocation invocation) {
         int weight;
         URL url = invoker.getUrl();
+        if (invoker instanceof ClusterInvoker) {
+            url = ((ClusterInvoker<?>) invoker).getRegistryUrl();
+        }
         // Multiple registry scenario, load balance among multiple registries.
         if (REGISTRY_SERVICE_REFERENCE_PATH.equals(url.getServiceInterface())) {
-            weight = url.getParameter(REGISTRY_KEY + "." + WEIGHT_KEY, DEFAULT_WEIGHT);
+            weight = url.getParameter(WEIGHT_KEY, DEFAULT_WEIGHT);
         } else {
             weight = url.getMethodParameter(invocation.getMethodName(), WEIGHT_KEY, DEFAULT_WEIGHT);
             if (weight > 0) {

diff --git a/...uster/src/test/java/org/apache/dubbo/rpc/cluster/loadbalance/AbstractLoadBalanceTest.java b/...uster/src/test/java/org/apache/dubbo/rpc/cluster/loadbalance/AbstractLoadBalanceTest.java
@@ -20,6 +20,7 @@
 import org.apache.dubbo.rpc.Invocation;
 import org.apache.dubbo.rpc.Invoker;
 import org.apache.dubbo.rpc.RpcInvocation;
+import org.apache.dubbo.rpc.cluster.ClusterInvoker;
 
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
@@ -71,12 +72,16 @@ public void testGetRegistryWeight() {
         url1 = url1.addParameter(REGISTRY_KEY + "." + WEIGHT_KEY, 10);
         given(invoker1.getUrl()).willReturn(url1);
 
-        Invoker invoker2 = mock(Invoker.class, Mockito.withSettings().stubOnly());
+
+        ClusterInvoker invoker2 = mock(ClusterInvoker.class, Mockito.withSettings().stubOnly());
         URL url2 = new URL("", "", 0, "org.apache.dubbo.registry.RegistryService", new HashMap<>());
-        url2 = url2.addParameter(REGISTRY_KEY + "." + WEIGHT_KEY, 20);
+        url2 = url2.addParameter(WEIGHT_KEY, 20);
+        URL registryUrl2 = new URL("", "", 0, "org.apache.dubbo.registry.RegistryService", new HashMap<>());
+        registryUrl2 = registryUrl2.addParameter(WEIGHT_KEY, 30);
         given(invoker2.getUrl()).willReturn(url2);
+        given(invoker2.getRegistryUrl()).willReturn(registryUrl2);
 
         Assertions.assertEquals(100, balance.getWeight(invoker1, invocation));
-        Assertions.assertEquals(20, balance.getWeight(invoker2, invocation));
+        Assertions.assertEquals(30, balance.getWeight(invoker2, invocation));
     }
 }