Use CuratorConfig

3 files changed, 30 insertions, 20 deletions

diff --git a/zkfacade/abi-spec.json b/zkfacade/abi-spec.json
index e026559b283..4aa8775940e 100644
--- a/zkfacade/abi-spec.json
+++ b/zkfacade/abi-spec.json
@@ -68,6 +68,7 @@
     "methods": [
       "public static com.yahoo.vespa.curator.Curator create(java.lang.String)",
       "public static com.yahoo.vespa.curator.Curator create(java.lang.String, java.util.Optional)",
+      "public void <init>(com.yahoo.cloud.config.CuratorConfig, com.yahoo.vespa.zookeeper.VespaZooKeeperServer)",
       "public void <init>(com.yahoo.cloud.config.ConfigserverConfig, com.yahoo.vespa.zookeeper.VespaZooKeeperServer)",
       "protected void <init>(java.lang.String, java.lang.String, java.util.function.Function)",
       "public java.lang.String connectionSpec()",
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 4aaae38f939..e281cb01f74 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.cloud.config.CuratorConfig;
 import com.yahoo.io.IOUtils;
 import com.yahoo.path.Path;
 import com.yahoo.text.Utf8;
@@ -74,19 +75,23 @@ public class Curator implements AutoCloseable {
         return new Curator(ConnectionSpec.create(connectionSpec), clientConfigFile);
     }
 
-    // Depend on ZooKeeperServer to make sure it is started first
-    // TODO: This can be removed when this package is no longer public API.
     @Inject
-    public Curator(ConfigserverConfig configserverConfig, @SuppressWarnings("unused") VespaZooKeeperServer server) {
-        this(configserverConfig, Optional.of(ZK_CLIENT_CONFIG_FILE));
+    public Curator(CuratorConfig curatorConfig, @SuppressWarnings("unused") VespaZooKeeperServer server) {
+        // Depends on ZooKeeperServer to make sure it is started first
+        this(ConnectionSpec.create(curatorConfig.server(),
+                                   CuratorConfig.Server::hostname,
+                                   CuratorConfig.Server::port,
+                                   curatorConfig.zookeeperLocalhostAffinity()),
+             Optional.of(ZK_CLIENT_CONFIG_FILE));
     }
 
-    Curator(ConfigserverConfig configserverConfig, Optional<File> clientConfigFile) {
+    // TODO: This can be removed when this package is no longer public API.
+    public Curator(ConfigserverConfig configserverConfig, @SuppressWarnings("unused") VespaZooKeeperServer server) {
         this(ConnectionSpec.create(configserverConfig.zookeeperserver(),
                                    ConfigserverConfig.Zookeeperserver::hostname,
                                    ConfigserverConfig.Zookeeperserver::port,
                                    configserverConfig.zookeeperLocalhostAffinity()),
-             clientConfigFile);
+             Optional.of(ZK_CLIENT_CONFIG_FILE));
     }
 
     protected Curator(String connectionSpec,
@@ -95,7 +100,7 @@ public class Curator implements AutoCloseable {
         this(ConnectionSpec.create(connectionSpec, zooKeeperEnsembleConnectionSpec), curatorFactory, DEFAULT_RETRY_POLICY);
     }
 
-    private Curator(ConnectionSpec connectionSpec, Optional<File> clientConfigFile) {
+    Curator(ConnectionSpec connectionSpec, Optional<File> clientConfigFile) {
         this(connectionSpec,
              (retryPolicy) -> CuratorFrameworkFactory
                      .builder()
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 1c7cb3695a8..5341efaefe5 100644
--- a/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorTest.java
+++ b/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorTest.java
@@ -1,7 +1,7 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.curator;
 
-import com.yahoo.cloud.config.ConfigserverConfig;
+import com.yahoo.cloud.config.CuratorConfig;
 import com.yahoo.net.HostName;
 import org.apache.curator.test.TestingServer;
 import org.junit.After;
@@ -61,29 +61,33 @@ public class CuratorTest {
 
     @Test
     public void require_that_server_count_is_correct() {
-        ConfigserverConfig.Builder builder = new ConfigserverConfig.Builder();
-        builder.zookeeperserver(createZKBuilder(localhost, port1));
-        try (Curator curator = createCurator(new ConfigserverConfig(builder))) {
+        CuratorConfig.Builder builder = new CuratorConfig.Builder();
+        builder.server(createZKBuilder(localhost, port1));
+        try (Curator curator = createCurator(new CuratorConfig(builder))) {
             assertEquals(1, curator.zooKeeperEnsembleCount());
         }
     }
 
-    private ConfigserverConfig createTestConfig() {
-        ConfigserverConfig.Builder builder = new ConfigserverConfig.Builder();
-        builder.zookeeperserver(createZKBuilder(localhost, port1));
-        builder.zookeeperserver(createZKBuilder(localhost, port2));
-        return new ConfigserverConfig(builder);
+    private CuratorConfig createTestConfig() {
+        CuratorConfig.Builder builder = new CuratorConfig.Builder();
+        builder.server(createZKBuilder(localhost, port1));
+        builder.server(createZKBuilder(localhost, port2));
+        return new CuratorConfig(builder);
     }
 
-    private ConfigserverConfig.Zookeeperserver.Builder createZKBuilder(String hostname, int port) {
-        ConfigserverConfig.Zookeeperserver.Builder zkBuilder = new ConfigserverConfig.Zookeeperserver.Builder();
+    private CuratorConfig.Server.Builder createZKBuilder(String hostname, int port) {
+        CuratorConfig.Server.Builder zkBuilder = new CuratorConfig.Server.Builder();
         zkBuilder.hostname(hostname);
         zkBuilder.port(port);
         return zkBuilder;
     }
 
-    private Curator createCurator(ConfigserverConfig configserverConfig)  {
-        return new Curator(configserverConfig, Optional.empty());
+    private Curator createCurator(CuratorConfig curatorConfig)  {
+        return new Curator(ConnectionSpec.create(curatorConfig.server(),
+                                                 CuratorConfig.Server::hostname,
+                                                 CuratorConfig.Server::port,
+                                                 curatorConfig.zookeeperLocalhostAffinity()),
+                           Optional.empty());
     }
 
     private static class PortAllocator {