Merge pull request #4633 from vespa-engine/hakonhall/some-curator-clients-require-ensemble-connect-string

Some Curator clients require ensemble connect string

13 files changed, 131 insertions, 90 deletions

diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/session/FileDistributionFactory.java b/configserver/src/main/java/com/yahoo/vespa/config/server/session/FileDistributionFactory.java
index 243c47ba3d7..02c854f4130 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/session/FileDistributionFactory.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/session/FileDistributionFactory.java
@@ -27,7 +27,7 @@ public class FileDistributionFactory {
 
     @Inject
     public FileDistributionFactory(Curator curator) {
-        this(curator, curator.connectionSpec());
+        this(curator, curator.zooKeeperEnsembleConnectionSpec());
     }
 
     public FileDistributionFactory(Curator curator, String zkSpec) {
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClient.java b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClient.java
index 6a90eaacb02..f7d6b8f05fe 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClient.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClient.java
@@ -117,7 +117,7 @@ public class SessionZooKeeperClient {
 
     /** Returns the number of node members needed in a barrier */
     private int getNumberOfMembers() {
-        return (curator.serverCount() / 2) + 1; // majority
+        return (curator.zooKeeperEnsembleCount() / 2) + 1; // majority
     }
 
     private Curator.CompletionWaiter createCompletionWaiter(String waiterNode) {
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
index 9b05478eb07..19894050ff4 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
@@ -630,7 +630,7 @@ public class NodeRepository extends AbstractComponent {
     // Public for testing
     public List<Node> getConfigNodes() {
         // TODO: Revisit this when config servers are added to the repository
-        return Arrays.stream(curator.connectionSpec().split(","))
+        return Arrays.stream(curator.zooKeeperEnsembleConnectionSpec().split(","))
                 .map(hostPort -> hostPort.split(":")[0])
                 .map(host -> createNode(host, host, Optional.empty(),
                         flavors.getFlavorOrThrow("v-4-8-100"), // Must be a flavor that exists in Hosted Vespa
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/ZooKeeperAccessMaintainer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/ZooKeeperAccessMaintainer.java
index 18987ccc97e..b392d670a77 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/ZooKeeperAccessMaintainer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/ZooKeeperAccessMaintainer.java
@@ -39,7 +39,7 @@ public class ZooKeeperAccessMaintainer extends Maintainer {
             hosts.add(node.hostname());
 
         if ( ! hosts.isEmpty()) { // no nodes -> not a hosted instance: Pass an empty list to deactivate restriction
-            for (String hostPort : curator.connectionSpec().split(","))
+            for (String hostPort : curator.zooKeeperEnsembleConnectionSpec().split(","))
                 hosts.add(hostPort.split(":")[0]);
         }
 
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNodeRepository.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNodeRepository.java
index 928d8082c6b..4596b71300a 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNodeRepository.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNodeRepository.java
@@ -51,7 +51,7 @@ public class MockNodeRepository extends NodeRepository {
               new DockerImage("docker-registry.domain.tld:8080/dist/vespa"));
         this.flavors = flavors;
 
-        curator.setConnectionSpec("cfg1:1234,cfg2:1234,cfg3:1234");
+        curator.setZooKeeperEnsembleConnectionSpec("cfg1:1234,cfg2:1234,cfg3:1234");
         populate();
     }
 
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTester.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTester.java
index 3d01bde4291..583d5658ef4 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTester.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTester.java
@@ -32,7 +32,7 @@ public class NodeRepositoryTester {
         nodeFlavors = new NodeFlavors(createConfig());
         clock = new ManualClock();
         curator = new MockCurator();
-        curator.setConnectionSpec("server1:1234,server2:5678");
+        curator.setZooKeeperEnsembleConnectionSpec("server1:1234,server2:5678");
         nodeRepository = new NodeRepository(nodeFlavors, curator, clock, Zone.defaultZone(),
                                             new MockNameResolver().mockAnyLookup(),
                                             new DockerImage("docker-registry.domain.tld:8080/dist/vespa"));
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/ZooKeeperAccessMaintainerTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/ZooKeeperAccessMaintainerTest.java
index bba5aa2db8d..93cf19f5450 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/ZooKeeperAccessMaintainerTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/ZooKeeperAccessMaintainerTest.java
@@ -23,7 +23,7 @@ public class ZooKeeperAccessMaintainerTest {
     @Test
     public void test() {
         NodeRepositoryTester tester = new NodeRepositoryTester();
-        tester.curator().setConnectionSpec("server1:1234,server2:5678");
+        tester.curator().setZooKeeperEnsembleConnectionSpec("server1:1234,server2:5678");
         ZooKeeperAccessMaintainer maintainer = new ZooKeeperAccessMaintainer(tester.nodeRepository(), 
                                                                              tester.curator(), Duration.ofHours(1), new JobControl(tester.nodeRepository().database()));
         assertTrue(ZooKeeperServer.getAllowedClientHostnames().isEmpty());
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/AclProvisioningTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/AclProvisioningTest.java
index c38d6f6bd40..492bcaa5462 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/AclProvisioningTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/AclProvisioningTest.java
@@ -18,7 +18,6 @@ import org.junit.Test;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.Comparator;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
@@ -230,7 +229,7 @@ public class AclProvisioningTest {
     }
 
     private List<Node> setConfigServers(String connectionSpec) {
-        curator.setConnectionSpec(connectionSpec);
+        curator.setZooKeeperEnsembleConnectionSpec(connectionSpec);
         return tester.nodeRepository().getConfigNodes();
     }
 
diff --git a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/status/ZookeeperStatusServiceTest.java b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/status/ZookeeperStatusServiceTest.java
index 1e742b5940a..2e914718e20 100644
--- a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/status/ZookeeperStatusServiceTest.java
+++ b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/status/ZookeeperStatusServiceTest.java
@@ -8,8 +8,6 @@ import com.yahoo.vespa.orchestrator.TestIds;
 import org.apache.commons.lang.exception.ExceptionUtils;
 import org.apache.curator.SessionFailRetryLoop.SessionFailedException;
 import org.apache.curator.framework.CuratorFramework;
-import org.apache.curator.framework.CuratorFrameworkFactory;
-import org.apache.curator.retry.ExponentialBackoffRetry;
 import org.apache.curator.test.KillSession;
 import org.apache.curator.test.TestingServer;
 import org.hamcrest.Description;
@@ -56,7 +54,7 @@ public class ZookeeperStatusServiceTest {
     }
 
     private static Curator createConnectedCurator(TestingServer server) throws InterruptedException {
-        Curator curator = new Curator(server.getConnectString());
+        Curator curator = Curator.create(server.getConnectString());
         curator.framework().blockUntilConnected(1, TimeUnit.MINUTES);
         return curator;
     }
diff --git a/zkfacade/pom.xml b/zkfacade/pom.xml
index f9cb8c42688..0aa689d896d 100644
--- a/zkfacade/pom.xml
+++ b/zkfacade/pom.xml
@@ -77,7 +77,6 @@
             <arg>-Xlint:all</arg>
             <arg>-Xlint:-serial</arg>
             <arg>-Xlint:-try</arg>
-            <arg>-Werror</arg>
           </compilerArgs>
         </configuration>
       </plugin>
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 4c932969460..5f7fa4a4e51 100644
--- a/zkfacade/src/main/java/com/yahoo/vespa/curator/Curator.java
+++ b/zkfacade/src/main/java/com/yahoo/vespa/curator/Curator.java
@@ -26,10 +26,10 @@ import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
-import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
 
 /**
  * Curator interface for Vespa.
@@ -47,77 +47,99 @@ public class Curator implements AutoCloseable {
     private static final int ZK_SESSION_TIMEOUT = 30000;
     private static final int ZK_CONNECTION_TIMEOUT = 30000;
 
-    private static final int baseSleepTime = 1000; //ms
-    private static final int maxRetries = 10;
+    private static final int BASE_SLEEP_TIME = 1000; //ms
+    private static final int MAX_RETRIES = 10;
 
-    private final CuratorFramework curatorFramework;
     protected final RetryPolicy retryPolicy;
 
-    private final String connectionSpec;
-    private final int serverCount;
+    private final CuratorFramework curatorFramework;
+    private final String connectionSpec; // May be a subset of the servers in the ensemble
+
+    private final String zooKeeperEnsembleConnectionSpec;
+    private final int zooKeeperEnsembleCount;
 
     /** Creates a curator instance from a comma-separated string of ZooKeeper host:port strings */
     public static Curator create(String connectionSpec) {
-        return new Curator(connectionSpec);
+        return new Curator(connectionSpec, connectionSpec);
     }
 
     // Depend on ZooKeeperServer to make sure it is started first
     // TODO: Move zookeeperserver config out of configserverconfig (requires update of controller services.xml as well)
     @Inject
     public Curator(ConfigserverConfig configserverConfig, ZooKeeperServer server) {
-        this(createConnectionSpec(configserverConfig));
+        this(configserverConfig, createConnectionSpec(configserverConfig));
     }
-    
-    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);
-
-            String spec = String.format("%s:%d", server.hostname(), server.port());
+    private Curator(ConfigserverConfig configserverConfig, String zooKeeperEnsembleConnectionSpec) {
+        this(configserverConfig.zookeeperLocalhostAffinity() ?
+                createConnectionSpecForLocalhost(configserverConfig) : zooKeeperEnsembleConnectionSpec,
+                zooKeeperEnsembleConnectionSpec);
+    }
 
-            if (config.zookeeperLocalhostAffinity() && server.hostname().equals(thisServer)) {
-                // Only connect to localhost server if possible, to save network traffic
-                // and balance load.
-                return spec;
-            }
+    private Curator(String connectionSpec, String zooKeeperEnsembleConnectionSpec) {
+        this(connectionSpec,
+                zooKeeperEnsembleConnectionSpec,
+                (retryPolicy) -> CuratorFrameworkFactory
+                        .builder()
+                        .retryPolicy(retryPolicy)
+                        .sessionTimeoutMs(ZK_SESSION_TIMEOUT)
+                        .connectionTimeoutMs(ZK_CONNECTION_TIMEOUT)
+                        .connectString(connectionSpec)
+                        .zookeeperFactory(new DNSResolvingFixerZooKeeperFactory(UNKNOWN_HOST_TIMEOUT_MILLIS))
+                        .build());
+    }
 
-            if (sb.length() > 0) {
-                sb.append(',');
-            }
+    protected Curator(String connectionSpec,
+                      String zooKeeperEnsembleConnectionSpec,
+                      Function<RetryPolicy, CuratorFramework> curatorFactory) {
+        this(connectionSpec, zooKeeperEnsembleConnectionSpec, curatorFactory,
+                new ExponentialBackoffRetry(BASE_SLEEP_TIME, MAX_RETRIES));
+    }
 
-            sb.append(spec);
+    private Curator(String connectionSpec,
+                    String zooKeeperEnsembleConnectionSpec,
+                    Function<RetryPolicy, CuratorFramework> curatorFactory,
+                    RetryPolicy retryPolicy) {
+        this.connectionSpec = connectionSpec;
+        this.retryPolicy = retryPolicy;
+        this.curatorFramework = curatorFactory.apply(retryPolicy);
+        if (this.curatorFramework != null) {
+            validateConnectionSpec(connectionSpec);
+            validateConnectionSpec(zooKeeperEnsembleConnectionSpec);
+            addFakeListener();
+            curatorFramework.start();
         }
-        return sb.toString();
+
+        this.zooKeeperEnsembleConnectionSpec = zooKeeperEnsembleConnectionSpec;
+        this.zooKeeperEnsembleCount = zooKeeperEnsembleConnectionSpec.split(",").length;
     }
 
-    /** 
-     * Create a curator instance which connects to the zookeeper servers given by a connection spec
-     * on the format "hostname1:port,hostname2:port" ...
-     */
-    public Curator(String connectionSpec) {
-        Objects.requireNonNull(connectionSpec, "The curator connection spec cannot be null");
-        this.connectionSpec = connectionSpec;
-        this.serverCount = connectionSpec.split(",").length;
-        validateConnectionSpec(connectionSpec);
-        retryPolicy = new ExponentialBackoffRetry(baseSleepTime, maxRetries);
-        curatorFramework = CuratorFrameworkFactory.builder()
-                .retryPolicy(retryPolicy)
-                .sessionTimeoutMs(ZK_SESSION_TIMEOUT)
-                .connectionTimeoutMs(ZK_CONNECTION_TIMEOUT)
-                .connectString(connectionSpec)
-                .zookeeperFactory(new DNSResolvingFixerZooKeeperFactory(UNKNOWN_HOST_TIMEOUT_MILLIS))
-                .build();
-        addFakeListener();
-        curatorFramework.start();
+    static String createConnectionSpec(ConfigserverConfig config) {
+        StringBuilder connectionSpec = new StringBuilder();
+        for (int i = 0; i < config.zookeeperserver().size(); i++) {
+            if (connectionSpec.length() > 0) {
+                connectionSpec.append(',');
+            }
+            ConfigserverConfig.Zookeeperserver server = config.zookeeperserver(i);
+            connectionSpec.append(server.hostname());
+            connectionSpec.append(':');
+            connectionSpec.append(server.port());
+        }
+        return connectionSpec.toString();
     }
 
-    protected Curator() {
-        this.connectionSpec = "";
-        this.serverCount = 0;
-        retryPolicy = new ExponentialBackoffRetry(baseSleepTime, maxRetries);
-        curatorFramework = null;
+    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) {
@@ -125,18 +147,19 @@ public class Curator implements AutoCloseable {
             throw new IllegalArgumentException(String.format("Connections spec '%s' is not valid", connectionSpec));
     }
 
-    /** Returns the number of zooKeeper servers in this cluster */
-    public int serverCount() { return serverCount; }
-
-    /** 
-     * Returns the servers in this cluster as a comma-separated list of host:port strings. 
+    /**
+     * Returns the ZooKeeper "connect string" used by curator: a comma-separated list of
+     * host:port of ZooKeeper endpoints to connect to. This may be a subset of
+     * zooKeeperEnsembleConnectionSpec() if there's some affinity, e.g. for
+     * performance reasons.
+     *
      * This may be empty but never null 
      */
     public String connectionSpec() { return connectionSpec; }
 
     /** For internal use; prefer creating a {@link CuratorCounter} */
     public DistributedAtomicLong createAtomicCounter(String path) {
-        return new DistributedAtomicLong(curatorFramework, path, new ExponentialBackoffRetry(baseSleepTime, maxRetries));
+        return new DistributedAtomicLong(curatorFramework, path, new ExponentialBackoffRetry(BASE_SLEEP_TIME, MAX_RETRIES));
     }
 
     /** For internal use; prefer creating a {@link com.yahoo.vespa.curator.recipes.CuratorLock} */
@@ -339,4 +362,19 @@ public class Curator implements AutoCloseable {
 
     }
 
+    /**
+     * @return The non-null connect string containing all ZooKeeper servers in the ensemble.
+     * WARNING: This may be different from the servers this Curator may connect to.
+     * TODO: Move method out of this class.
+     */
+    public String zooKeeperEnsembleConnectionSpec() {
+        return zooKeeperEnsembleConnectionSpec;
+    }
+
+    /**
+     * Returns the number of zooKeeper servers in this ensemble.
+     * WARNING: This may be different from the number of servers this Curator may connect to.
+     * TODO: Move method out of this class.
+     */
+    public int zooKeeperEnsembleCount() { return zooKeeperEnsembleCount; }
 }
diff --git a/zkfacade/src/main/java/com/yahoo/vespa/curator/mock/MockCurator.java b/zkfacade/src/main/java/com/yahoo/vespa/curator/mock/MockCurator.java
index ccd02e5c6d6..4013cf1d649 100644
--- a/zkfacade/src/main/java/com/yahoo/vespa/curator/mock/MockCurator.java
+++ b/zkfacade/src/main/java/com/yahoo/vespa/curator/mock/MockCurator.java
@@ -7,8 +7,6 @@ import com.yahoo.collections.Pair;
 import com.yahoo.concurrent.Lock;
 import com.yahoo.concurrent.Locks;
 import com.yahoo.path.Path;
-import static com.yahoo.vespa.curator.mock.MemoryFileSystem.Node;
-
 import com.yahoo.vespa.curator.CompletionTimeoutException;
 import com.yahoo.vespa.curator.Curator;
 import com.yahoo.vespa.curator.recipes.CuratorLockException;
@@ -86,6 +84,8 @@ import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 
+import static com.yahoo.vespa.curator.mock.MemoryFileSystem.Node;
+
 /**
  * <p>A <b>non thread safe</b> mock of the curator API.
  * The methods are implemented lazily, due to laziness.
@@ -106,7 +106,7 @@ public class MockCurator extends Curator {
     private boolean shouldTimeoutOnEnter = false;
     private int monotonicallyIncreasingNumber = 0;
     private final boolean stableOrdering;
-    private String connectionSpec = "";
+    private String zooKeeperEnsembleConnectionSpec = "";
     private final Locks<String> locks = new Locks<>(Long.MAX_VALUE, TimeUnit.DAYS);
 
     /** The file system used by this mock to store zookeeper files and directories */
@@ -133,6 +133,7 @@ public class MockCurator extends Curator {
      *                       This is not what ZooKeeper does.
      */
     public MockCurator(boolean stableOrdering) {
+        super("", "", (retryPolicy) -> null);
         this.stableOrdering = stableOrdering;
         curatorFramework = new MockCuratorFramework();
         curatorFramework.start();
@@ -152,11 +153,18 @@ public class MockCurator extends Curator {
         return Optional.ofNullable(atomicCounters.get(path));
     }
     
-    /** Assigns the connection string, which must be on the form host1:port,host2:port ... */
-    public void setConnectionSpec(String connectionSpec) { this.connectionSpec = connectionSpec; }
+    /**
+     * Sets the ZooKeeper ensemble connection spec, which must be on the form
+     * host1:port,host2:port ...
+     */
+    public void setZooKeeperEnsembleConnectionSpec(String ensembleSpec) {
+        this.zooKeeperEnsembleConnectionSpec = ensembleSpec;
+    }
 
     @Override
-    public String connectionSpec() { return connectionSpec; }
+    public String zooKeeperEnsembleConnectionSpec() {
+        return zooKeeperEnsembleConnectionSpec;
+    }
 
     // ----- Start of adaptor methods from Curator to the mock file system -----
 
@@ -368,7 +376,7 @@ public class MockCurator extends Curator {
         public void notify(Path path, PathChildrenCacheEvent event) {
             try {
                 // Snapshot directoryListeners in case notification leads to new directoryListeners added
-                Set<Map.Entry<Path, PathChildrenCacheListener>>directoryLlistenerSnapshot = new HashSet<>(directoryListeners.entrySet());
+                Set<Map.Entry<Path, PathChildrenCacheListener>> directoryLlistenerSnapshot = new HashSet<>(directoryListeners.entrySet());
                 for (Map.Entry<Path, PathChildrenCacheListener> listener : directoryLlistenerSnapshot) {
                     if (path.isChildOf(listener.getKey()))
                         listener.getValue().childEvent(curatorFramework, event);
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 1899dcfe7cd..c5227fcbaa5 100644
--- a/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorTest.java
+++ b/zkfacade/src/test/java/com/yahoo/vespa/curator/CuratorTest.java
@@ -60,7 +60,7 @@ public class CuratorTest {
     public void require_that_curator_can_produce_spec() {
         try (Curator curator = createCurator(createTestConfig())) {
             assertThat(curator.connectionSpec(), is(spec1 + "," + spec2));
-            assertThat(curator.serverCount(), is(2));
+            assertThat(curator.zooKeeperEnsembleCount(), is(2));
         }
     }
 
@@ -69,26 +69,25 @@ public class CuratorTest {
         ConfigserverConfig.Builder builder = new ConfigserverConfig.Builder();
         builder.zookeeperserver(createZKBuilder("localhost", port1));
         try (Curator curator = createCurator(new ConfigserverConfig(builder))) {
-            assertThat(curator.serverCount(), is(1));
+            assertThat(curator.zooKeeperEnsembleCount(), is(1));
         }
     }
 
     @Test
-    public void localhost_affinity() {
-        String localhostHostName = "myhost";
-        int localhostPort = 123;
-        String localhostSpec = localhostHostName + ":" + localhostPort;
+     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));
-        builder.zookeeperLocalhostAffinity(true);
-        ConfigserverConfig config = new ConfigserverConfig(builder);
+         ConfigserverConfig.Builder builder = new ConfigserverConfig.Builder();
+         builder.zookeeperserver(createZKBuilder(localhostHostName, localhostPort));
+         builder.zookeeperserver(createZKBuilder("otherhost", 345));
+         ConfigserverConfig config = new ConfigserverConfig(builder);
 
-        HostName.setHostNameForTestingOnly(localhostHostName);
+         HostName.setHostNameForTestingOnly(localhostHostName);
 
-        assertThat(Curator.createConnectionSpec(config), is(localhostSpec));
-    }
+        String localhostSpec = localhostHostName + ":" + localhostPort;
+         assertThat(Curator.createConnectionSpecForLocalhost(config), is(localhostSpec));
+     }
 
     private ConfigserverConfig createTestConfig() {
         ConfigserverConfig.Builder builder = new ConfigserverConfig.Builder();