diff --git a/dubbo-cluster/src/main/java/org/apache/dubbo/rpc/cluster/loadbalance/AbstractLoadBalance.java b/dubbo-cluster/src/main/java/org/apache/dubbo/rpc/cluster/loadbalance/AbstractLoadBalance.java index de057cf3016..310520a63c1 100644 --- a/dubbo-cluster/src/main/java/org/apache/dubbo/rpc/cluster/loadbalance/AbstractLoadBalance.java +++ b/dubbo-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 Invoker select(List> 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/dubbo-cluster/src/test/java/org/apache/dubbo/rpc/cluster/loadbalance/AbstractLoadBalanceTest.java b/dubbo-cluster/src/test/java/org/apache/dubbo/rpc/cluster/loadbalance/AbstractLoadBalanceTest.java index eec4d89a327..5e0b1f79958 100644 --- a/dubbo-cluster/src/test/java/org/apache/dubbo/rpc/cluster/loadbalance/AbstractLoadBalanceTest.java +++ b/dubbo-cluster/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)); } }