com.opower.zookeeper.test.MiniZooKeeperCluster Maven / Gradle / Ivy
The newest version!
package com.opower.zookeeper.test;
import java.io.File;
import java.io.IOException;
import java.net.BindException;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import com.google.common.base.Optional;
import com.google.common.base.Preconditions;
import org.apache.commons.io.FileUtils;
import org.apache.zookeeper.server.NIOServerCnxnFactory;
import org.apache.zookeeper.server.ServerCnxnFactory;
import org.apache.zookeeper.server.ZooKeeperServer;
/**
* Allows for configuration and start/stop of a mini ZooKeeper quorum (running in the same JVM as calling code). Lifecycle is:
*
* - Construct an instance using the {@link MiniZooKeeperCluster.Builder} or one of the {@code newDefaultInstance*}
* convenience factory methods
* - Call {@link #start()}
* - Run your test against the ZK connection string provided by {@link #getZkConnectionString()}
* - Call {@link #shutdown()}
*
*
* @author eric.chang
*/
public class MiniZooKeeperCluster {
private static final int MAX_PORT_ASSIGNMENT_ATTEMPTS = 5;
private final List miniZooKeepers;
private final List factories;
private final File snapshotRootDir;
private final File logRootDir;
private String connectionString;
private MiniZooKeeperCluster(List miniZooKeepers) {
Preconditions.checkNotNull(miniZooKeepers, "miniZooKeepers cannot be null");
Preconditions.checkState(miniZooKeepers.size() > 0, "must specify at least one MiniZooKeeper");
this.miniZooKeepers = miniZooKeepers;
this.factories = new ArrayList<>(miniZooKeepers.size());
try {
this.snapshotRootDir = Files.createTempDirectory("snapshot").toFile();
this.logRootDir = Files.createTempDirectory("log").toFile();
}
catch (IOException ioe) {
throw new RuntimeException("Could not create ZooKeeper backing directories", ioe);
}
}
/**
* @return the zookeeper connection string (host:port)
*/
public String getZkConnectionString() {
return this.connectionString;
}
/**
* Start the cluster, using temporary directories for logs and snapshots
*/
public void start() throws IOException, InterruptedException {
StringBuilder connectionBuilder = new StringBuilder();
Set assignedPorts = new HashSet<>();
int serverNum = 0;
for (MiniZooKeeper miniZooKeeper : this.miniZooKeepers) {
int port = getAvailablePort(miniZooKeeper.getPort());
Preconditions.checkState(!assignedPorts.contains(port),
String.format("requested port %d is already assigned", port));
assignedPorts.add(port);
connectionBuilder.append(String.format("%s:%d", miniZooKeeper.getLocalhost(), port)).append(",");
ServerCnxnFactory factory = NIOServerCnxnFactory.createFactory(
new InetSocketAddress(miniZooKeeper.getLocalhost(), port), miniZooKeeper.getMaxClientConnections());
File snapshotDir = new File(this.snapshotRootDir, String.format("server-%d", serverNum));
File logDir = new File(this.logRootDir, String.format("server-%d", serverNum));
factory.startup(new ZooKeeperServer(snapshotDir, logDir, miniZooKeeper.getTickTime()));
this.factories.add(factory);
serverNum++;
}
// trim trailing comma
connectionBuilder.setLength(connectionBuilder.length() - 1);
this.connectionString = connectionBuilder.toString();
}
/**
* Shutdown the cluster
*/
public void shutdown() {
if (this.factories != null) {
for (ServerCnxnFactory factory : this.factories) {
factory.shutdown();
}
}
FileUtils.deleteQuietly(this.snapshotRootDir);
FileUtils.deleteQuietly(this.logRootDir);
}
/**
* Convenience factory method that creates a mini cluster with a single server running on a dynamically assigned port.
* See constants in {@link com.opower.zookeeper.test.MiniZooKeeper.Builder} for default values set on returned instance.
*/
public static MiniZooKeeperCluster newDefaultInstance() {
return newBuilder().withMiniZooKeeper(MiniZooKeeper.newDefaultInstance()).build();
}
/**
* Convenience factory method that creates a mini cluster with a single server running on specified port.
* See constants in {@link com.opower.zookeeper.test.MiniZooKeeper.Builder} for default values set on returned instance.
*/
public static MiniZooKeeperCluster newDefaultInstanceWithPort(int port) {
return newBuilder().withMiniZooKeeper(MiniZooKeeper.newDefaultInstanceWithPort(port)).build();
}
/**
* @return a new Builder to be used to create a ZooKeeperMiniCluster
*/
public static Builder newBuilder() {
return new Builder();
}
/**
* Builder for MiniZooKeeper instances
*/
public static class Builder {
// keep track of specified ports to detect possible port contention
private Set ports = new HashSet<>();
private List miniZooKeepers = new ArrayList<>();
private Builder() {
}
/**
* Add a new MiniZooKeeper configuration to the cluster. Checks that the provided configuration's port has
* has not already been assigned.
*
* @param miniZooKeeper configuration to add to this mini cluster
*/
public Builder withMiniZooKeeper(MiniZooKeeper miniZooKeeper) {
Preconditions.checkNotNull(miniZooKeeper, "miniZooKeeper cannot be null");
if (miniZooKeeper.getPort().isPresent()) {
int port = miniZooKeeper.getPort().get();
Preconditions.checkState(!this.ports.contains(port),
String.format("Port %d has already been assigned to another server", port));
this.ports.add(port);
}
this.miniZooKeepers.add(miniZooKeeper);
return this;
}
/**
* @return a MiniZooKeeperCluster. Sorts MiniZooKeepers by port numbers ascending with empty values last to
* ensure that dynamic port creation doesn't stomp on explicitly specified ports.
*/
public MiniZooKeeperCluster build() {
Preconditions.checkState(!this.miniZooKeepers.isEmpty(),
"Must specify at least one MiniZooKeeper server by invoking withMiniZooKeeper()");
Collections.sort(this.miniZooKeepers, new Comparator() {
@Override
public int compare(MiniZooKeeper z1, MiniZooKeeper z2) {
if (z1.getPort().isPresent() && z2.getPort().isPresent()) {
return z1.getPort().get().compareTo(z2.getPort().get());
}
else if (z1.getPort().isPresent()) {
return -1;
}
else {
return 1;
}
}
});
return new MiniZooKeeperCluster(this.miniZooKeepers);
}
}
/**
* Assigns a port dynamically if {@code requestedPort} is absent and updates the {@code assignedPorts} set.
* Attempts up to {@link #MAX_PORT_ASSIGNMENT_ATTEMPTS} to calculate an assigned port. Visible for testing.
*/
static int getAvailablePort(Optional requestedPort) {
if (requestedPort.isPresent()) {
int port = requestedPort.get();
try (ServerSocket socket = new ServerSocket(port)) {
return socket.getLocalPort();
}
catch (BindException be) {
throw new RuntimeException(String.format("port %d is already bound", port));
}
catch (IOException ioe) {
throw new RuntimeException(String.format("Could not determine whether port %d is bound", port), ioe);
}
}
else {
int attemptsRemaining = MAX_PORT_ASSIGNMENT_ATTEMPTS;
while (attemptsRemaining > 0) {
attemptsRemaining--;
try (ServerSocket socket = new ServerSocket(0)) {
if (socket.isBound()) {
return socket.getLocalPort();
}
}
catch (BindException be) {
// port is already bound, try again
continue;
}
catch (IOException ioe) {
throw new RuntimeException(String.format("Error while attempting to retrieve an available port"), ioe);
}
}
throw new RuntimeException(String.format("Could not find an available port after %d attempts",
MAX_PORT_ASSIGNMENT_ATTEMPTS));
}
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy