Allow localhost zk affinity

This allows the config server's Curator component to use a connectSpec that
only specifies the localhost zk server.

2 files changed, 36 insertions, 10 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 15257e11cbe..4c932969460 100644
--- a/zkfacade/src/main/java/com/yahoo/vespa/curator/Curator.java
+++ b/zkfacade/src/main/java/com/yahoo/vespa/curator/Curator.java
@@ -3,6 +3,7 @@ package com.yahoo.vespa.curator;
 
 import com.google.inject.Inject;
 import com.yahoo.cloud.config.ConfigserverConfig;
+import com.yahoo.net.HostName;
 import com.yahoo.path.Path;
 import com.yahoo.vespa.curator.recipes.CuratorCounter;
 import com.yahoo.vespa.zookeeper.ZooKeeperServer;
@@ -21,7 +22,6 @@ import org.apache.curator.framework.state.ConnectionState;
 import org.apache.curator.framework.state.ConnectionStateListener;
 import org.apache.curator.retry.ExponentialBackoffRetry;
 
-import java.io.Closeable;
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collections;
@@ -68,16 +68,26 @@ public class Curator implements AutoCloseable {
         this(createConnectionSpec(configserverConfig));
     }
     
-    private static String createConnectionSpec(ConfigserverConfig config) {
+    static String createConnectionSpec(ConfigserverConfig config) {
+        String thisServer = HostName.getLocalhost();
+
         StringBuilder sb = new StringBuilder();
         for (int i = 0; i < config.zookeeperserver().size(); i++) {
             ConfigserverConfig.Zookeeperserver server = config.zookeeperserver(i);
-            sb.append(server.hostname());
-            sb.append(":");
-            sb.append(server.port());
-            if (i < config.zookeeperserver().size() - 1) {
-                sb.append(",");
+
+            String spec = String.format("%s:%d", server.hostname(), server.port());
+
+            if (config.zookeeperLocalhostAffinity() && server.hostname().equals(thisServer)) {
+                // Only connect to localhost server if possible, to save network traffic
+                // and balance load.
+                return spec;
             }
+
+            if (sb.length() > 0) {
+                sb.append(',');
+            }
+
+            sb.append(spec);
         }
         return sb.toString();
     }
diff --git a/zkfacade/src/test/java/com/yahoo/vespa/zookeeper/CuratorTest.java b/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorTest.java
index 36205bdaca3..1899dcfe7cd 100644
--- a/zkfacade/src/test/java/com/yahoo/vespa/zookeeper/CuratorTest.java
+++ b/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorTest.java
@@ -1,8 +1,8 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package com.yahoo.vespa.zookeeper;
+package com.yahoo.vespa.curator;
 
 import com.yahoo.cloud.config.ConfigserverConfig;
-import com.yahoo.vespa.curator.Curator;
+import com.yahoo.net.HostName;
 import org.apache.curator.test.TestingServer;
 import org.junit.After;
 import org.junit.Before;
@@ -11,7 +11,6 @@ import org.junit.Test;
 import java.io.IOException;
 
 import static org.hamcrest.core.Is.is;
-import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertThat;
 
 /**
@@ -74,6 +73,23 @@ public class CuratorTest {
         }
     }
 
+    @Test
+    public void localhost_affinity() {
+        String localhostHostName = "myhost";
+        int localhostPort = 123;
+        String localhostSpec = localhostHostName + ":" + localhostPort;
+
+        ConfigserverConfig.Builder builder = new ConfigserverConfig.Builder();
+        builder.zookeeperserver(createZKBuilder(localhostHostName, localhostPort));
+        builder.zookeeperserver(createZKBuilder("otherhost", 345));
+        builder.zookeeperLocalhostAffinity(true);
+        ConfigserverConfig config = new ConfigserverConfig(builder);
+
+        HostName.setHostNameForTestingOnly(localhostHostName);
+
+        assertThat(Curator.createConnectionSpec(config), is(localhostSpec));
+    }
+
     private ConfigserverConfig createTestConfig() {
         ConfigserverConfig.Builder builder = new ConfigserverConfig.Builder();
         builder.zookeeperserver(createZKBuilder("localhost", port1));