Build proper connect string instead of just hostname

2 files changed, 32 insertions, 1 deletions

diff --git a/zkfacade/src/main/java/com/yahoo/vespa/curator/Curator.java b/zkfacade/src/main/java/com/yahoo/vespa/curator/Curator.java
index 5ba16232221..5f7fa4a4e51 100644
--- a/zkfacade/src/main/java/com/yahoo/vespa/curator/Curator.java
+++ b/zkfacade/src/main/java/com/yahoo/vespa/curator/Curator.java
@@ -72,7 +72,7 @@ public class Curator implements AutoCloseable {
 
     private Curator(ConfigserverConfig configserverConfig, String zooKeeperEnsembleConnectionSpec) {
         this(configserverConfig.zookeeperLocalhostAffinity() ?
-                HostName.getLocalhost() : zooKeeperEnsembleConnectionSpec,
+                createConnectionSpecForLocalhost(configserverConfig) : zooKeeperEnsembleConnectionSpec,
                 zooKeeperEnsembleConnectionSpec);
     }
 
@@ -128,6 +128,20 @@ public class Curator implements AutoCloseable {
         return connectionSpec.toString();
     }
 
+    static String createConnectionSpecForLocalhost(ConfigserverConfig config) {
+        String thisServer = HostName.getLocalhost();
+
+        for (int i = 0; i < config.zookeeperserver().size(); i++) {
+            ConfigserverConfig.Zookeeperserver server = config.zookeeperserver(i);
+            if (thisServer.equals(server.hostname())) {
+                return String.format("%s:%d", server.hostname(), server.port());
+            }
+        }
+
+        throw new IllegalArgumentException("Unable to create connect string to localhost: " +
+                "There is no localhost servers specified in config: " + config);
+    }
+
     private static void validateConnectionSpec(String connectionSpec) {
         if (connectionSpec == null || connectionSpec.isEmpty())
             throw new IllegalArgumentException(String.format("Connections spec '%s' is not valid", connectionSpec));
diff --git a/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorTest.java b/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorTest.java
index 2fc4c2a7fc4..c5227fcbaa5 100644
--- a/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorTest.java
+++ b/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorTest.java
@@ -2,6 +2,7 @@
 package com.yahoo.vespa.curator;
 
 import com.yahoo.cloud.config.ConfigserverConfig;
+import com.yahoo.net.HostName;
 import org.apache.curator.test.TestingServer;
 import org.junit.After;
 import org.junit.Before;
@@ -72,6 +73,22 @@ public class CuratorTest {
         }
     }
 
+    @Test
+     public void localhost_affinity() {
+         String localhostHostName = "myhost";
+         int localhostPort = 123;
+
+         ConfigserverConfig.Builder builder = new ConfigserverConfig.Builder();
+         builder.zookeeperserver(createZKBuilder(localhostHostName, localhostPort));
+         builder.zookeeperserver(createZKBuilder("otherhost", 345));
+         ConfigserverConfig config = new ConfigserverConfig(builder);
+
+         HostName.setHostNameForTestingOnly(localhostHostName);
+
+        String localhostSpec = localhostHostName + ":" + localhostPort;
+         assertThat(Curator.createConnectionSpecForLocalhost(config), is(localhostSpec));
+     }
+
     private ConfigserverConfig createTestConfig() {
         ConfigserverConfig.Builder builder = new ConfigserverConfig.Builder();
         builder.zookeeperserver(createZKBuilder("localhost", port1));