Fix MappingListener override in ServiceDiscoveryRegistry (#13277)

apache · Nov 1, 2023 · 06e9410 · 06e9410
1 parent 6b634b7
commit 06e9410
Showing 1 changed file with 30 additions and 5 deletions.
diff --git a/...registry-api/src/main/java/org/apache/dubbo/registry/client/ServiceDiscoveryRegistry.java b/...registry-api/src/main/java/org/apache/dubbo/registry/client/ServiceDiscoveryRegistry.java
@@ -22,6 +22,7 @@
 import org.apache.dubbo.common.logger.LoggerFactory;
 import org.apache.dubbo.common.utils.CollectionUtils;
 import org.apache.dubbo.common.utils.ConcurrentHashMapUtils;
+import org.apache.dubbo.common.utils.ConcurrentHashSet;
 import org.apache.dubbo.metadata.AbstractServiceNameMapping;
 import org.apache.dubbo.metadata.MappingChangedEvent;
 import org.apache.dubbo.metadata.MappingListener;
@@ -37,11 +38,13 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
+import java.util.stream.Collectors;
 
 import static org.apache.dubbo.common.constants.CommonConstants.INTERFACE_KEY;
 import static org.apache.dubbo.common.constants.CommonConstants.PROVIDER_SIDE;
@@ -79,7 +82,7 @@ public class ServiceDiscoveryRegistry extends FailbackRegistry {
 
     /* apps - listener */
     private final Map<String, ServiceInstancesChangedListener> serviceListeners = new ConcurrentHashMap<>();
-    private final Map<String, MappingListener> mappingListeners = new ConcurrentHashMap<>();
+    private final Map<String, Set<MappingListener>> mappingListeners = new ConcurrentHashMap<>();
     /* This lock has the same scope and lifecycle as its corresponding instance listener.
     It's used to make sure that only one interface mapping to the same app list can do subscribe or unsubscribe at the same moment.
     And the lock should be destroyed when listener destroying its corresponding instance listener.
@@ -112,7 +115,7 @@ public ServiceDiscovery getServiceDiscovery() {
      */
     protected ServiceDiscovery createServiceDiscovery(URL registryURL) {
         return getServiceDiscovery(registryURL.addParameter(INTERFACE_KEY, ServiceDiscovery.class.getName())
-                .removeParameter(REGISTRY_TYPE_KEY));
+            .removeParameter(REGISTRY_TYPE_KEY));
     }
 
     /**
@@ -206,7 +209,10 @@ public void doSubscribe(URL url, NotifyListener listener) {
                 try {
                     MappingListener mappingListener = new DefaultMappingListener(url, mappingByUrl, listener);
                     mappingByUrl = serviceNameMapping.getAndListen(this.getUrl(), url, mappingListener);
-                    mappingListeners.put(url.getProtocolServiceKey(), mappingListener);
+                    synchronized (mappingListeners) {
+                        mappingListeners.computeIfAbsent(url.getProtocolServiceKey(), (k) -> new ConcurrentHashSet<>())
+                            .add(mappingListener);
+                    }
                 } catch (Exception e) {
                     logger.warn(INTERNAL_ERROR, "", "", "Cannot find app mapping for service " + url.getServiceInterface() + ", will not migrate.", e);
                 }
@@ -248,8 +254,23 @@ public void doUnsubscribe(URL url, NotifyListener listener) {
         serviceDiscovery.unsubscribe(url, listener);
         String protocolServiceKey = url.getProtocolServiceKey();
         Set<String> serviceNames = serviceNameMapping.getMapping(url);
-        if (mappingListeners.get(protocolServiceKey) != null) {
-            serviceNameMapping.stopListen(url, mappingListeners.remove(protocolServiceKey));
+
+        synchronized (mappingListeners) {
+            Set<MappingListener> keyedListeners = mappingListeners.get(protocolServiceKey);
+            if (keyedListeners != null) {
+                List<MappingListener> matched = keyedListeners.stream()
+                    .filter(mappingListener ->
+                        mappingListener instanceof DefaultMappingListener
+                            && (Objects.equals(((DefaultMappingListener) mappingListener).getListener(), listener)))
+                    .collect(Collectors.toList());
+                for (MappingListener mappingListener : matched) {
+                    serviceNameMapping.stopListen(url, mappingListener);
+                    keyedListeners.remove(mappingListener);
+                }
+                if (keyedListeners.isEmpty()) {
+                    mappingListeners.remove(protocolServiceKey, Collections.emptySet());
+                }
+            }
         }
         if (CollectionUtils.isNotEmpty(serviceNames)) {
             String serviceNamesKey = toStringKeys(serviceNames);
@@ -430,6 +451,10 @@ public synchronized void onEvent(MappingChangedEvent event) {
             }
         }
 
+        protected NotifyListener getListener() {
+            return listener;
+        }
+
         @Override
         public void stop() {
             stopped = true;