Simplify by having just one map and synchronize access

2 files changed, 23 insertions, 36 deletions

diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/host/HostRegistry.java b/configserver/src/main/java/com/yahoo/vespa/config/server/host/HostRegistry.java
index bc1d195b080..77572856ff5 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/host/HostRegistry.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/host/HostRegistry.java
@@ -4,16 +4,14 @@ package com.yahoo.vespa.config.server.host;
 import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.logging.Logger;
+import java.util.stream.Collectors;
 
-import com.google.common.base.Predicate;
 import com.google.common.collect.Collections2;
 import com.yahoo.log.LogLevel;
 
 /**
- * A host registry that create mappings between some type T and a list of hosts, represented as
- * strings.
+ * A host registry with a mapping between hosts (hostname as a String) and some type T
  * TODO: Maybe we should have a Host type, but using String for now.
- * TODO: Is there a generalized version of this pattern? Need some sort mix of Bimap and Multimap
  *
  * @author Ulf Lilleengen
  */
@@ -21,23 +19,24 @@ public class HostRegistry<T> implements HostValidator<T> {
 
     private static final Logger log = Logger.getLogger(HostRegistry.class.getName());
 
-    private final Map<T, Collection<String>> key2HostsMap = new ConcurrentHashMap<>();
     private final Map<String, T> host2KeyMap = new ConcurrentHashMap<>();
 
     public T getKeyForHost(String hostName) {
         return host2KeyMap.get(hostName);
     }
 
-    public void update(T key, Collection<String> newHosts) {
+    public synchronized void update(T key, Collection<String> newHosts) {
         verifyHosts(key, newHosts);
-        log.log(LogLevel.DEBUG, "Setting hosts for key(" + key + "), newHosts(" + newHosts + "), " +
-                                "currentHosts(" + getCurrentHosts(key) + ")");
-        Collection<String> removedHosts = getRemovedHosts(newHosts, getCurrentHosts(key));
+        Collection<String> currentHosts = getHostsForKey(key);
+        log.log(LogLevel.DEBUG, () -> "Setting hosts for key '" + key + "', " +
+                "newHosts: " + newHosts + ", " +
+                "currentHosts: " + currentHosts);
+        Collection<String> removedHosts = getRemovedHosts(newHosts, currentHosts);
         removeHosts(removedHosts);
         addHosts(key, newHosts);
     }
 
-    public void verifyHosts(T key, Collection<String> newHosts) {
+    public synchronized void verifyHosts(T key, Collection<String> newHosts) {
         for (String host : newHosts) {
             if (hostAlreadyTaken(host, key)) {
                 throw new IllegalArgumentException("'" + key + "' tried to allocate host '" + host + 
@@ -46,51 +45,41 @@ public class HostRegistry<T> implements HostValidator<T> {
         }
     }
 
-    public void removeHostsForKey(T key) {
-        for (Iterator<Map.Entry<T, Collection<String>>> it = key2HostsMap.entrySet().iterator(); it.hasNext(); ) {
-            Map.Entry<T, Collection<String>> entry = it.next();
-            if (entry.getKey().equals(key)) {
-                Collection<String> hosts = entry.getValue();
-                it.remove();
-                removeHosts(hosts);
-            }
-        }
+    public synchronized void removeHostsForKey(T key) {
+        host2KeyMap.entrySet().removeIf(entry -> entry.getValue().equals(key));
     }
 
-    public Collection<String> getAllHosts() {
+    public synchronized Collection<String> getAllHosts() {
         return Collections.unmodifiableCollection(new ArrayList<>(host2KeyMap.keySet()));
     }
 
-    Collection<String> getCurrentHosts(T key) {
-        return key2HostsMap.containsKey(key) ? new ArrayList<>(key2HostsMap.get(key)) : new ArrayList<String>();
+    synchronized Collection<String> getHostsForKey(T key) {
+        return host2KeyMap.entrySet().stream()
+                .filter(entry -> entry.getValue().equals(key))
+                .map(Map.Entry::getKey)
+                .collect(Collectors.toSet());
     }
 
     private boolean hostAlreadyTaken(String host, T key) {
         return host2KeyMap.containsKey(host) && !key.equals(host2KeyMap.get(host));
     }
 
-    private static Collection<String> getRemovedHosts(final Collection<String> newHosts, Collection<String> previousHosts) {
-        return Collections2.filter(previousHosts, new Predicate<String>() {
-            @Override
-            public boolean apply(String host) {
-                return !newHosts.contains(host);
-            }
-        });
+    private static Collection<String> getRemovedHosts(Collection<String> newHosts, Collection<String> previousHosts) {
+        return Collections2.filter(previousHosts, host -> !newHosts.contains(host));
     }
 
     private void removeHosts(Collection<String> removedHosts) {
         for (String host : removedHosts) {
-            log.log(LogLevel.DEBUG, "Removing " + host);
+            log.log(LogLevel.DEBUG, () -> "Removing " + host);
             host2KeyMap.remove(host);
         }
     }
 
     private void addHosts(T key, Collection<String> newHosts) {
         for (String host : newHosts) {
-            log.log(LogLevel.DEBUG, "Adding " + host);
+            log.log(LogLevel.DEBUG, () -> "Adding " + host);
             host2KeyMap.put(host, key);
         }
-        key2HostsMap.put(key, new ArrayList<>(newHosts));
     }
 
 }
diff --git a/configserver/src/test/java/com/yahoo/vespa/config/server/host/HostRegistryTest.java b/configserver/src/test/java/com/yahoo/vespa/config/server/host/HostRegistryTest.java
index f9cc13ee277..b2335165105 100644
--- a/configserver/src/test/java/com/yahoo/vespa/config/server/host/HostRegistryTest.java
+++ b/configserver/src/test/java/com/yahoo/vespa/config/server/host/HostRegistryTest.java
@@ -6,7 +6,6 @@ import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
 
-import com.yahoo.vespa.config.server.host.HostRegistry;
 import org.junit.Test;
 
 import static org.hamcrest.collection.IsIterableContainingInOrder.contains;
@@ -15,7 +14,6 @@ import static org.junit.Assert.*;
 
 /**
  * @author Ulf Lilleengen
- * @since 5.3
  */
 public class HostRegistryTest {
     @Test
@@ -71,9 +69,9 @@ public class HostRegistryTest {
         HostRegistry<String> reg = new HostRegistry<>();
         List<String> hosts = new ArrayList<>(Arrays.asList("foo.com", "bar.com", "baz.com"));
         reg.update("fookey", hosts);
-        assertThat(reg.getCurrentHosts("fookey").size(), is(3));
+        assertThat(reg.getHostsForKey("fookey").size(), is(3));
         hosts.remove(2);
-        assertThat(reg.getCurrentHosts("fookey").size(), is(3));
+        assertThat(reg.getHostsForKey("fookey").size(), is(3));
     }
 
     @Test