http://git-wip-us.apache.org/repos/asf/aurora/blob/06ddaadb/commons/src/main/java/org/apache/aurora/common/zookeeper/testing/ZooKeeperTestServer.java
----------------------------------------------------------------------
diff --git
a/commons/src/main/java/org/apache/aurora/common/zookeeper/testing/ZooKeeperTestServer.java
b/commons/src/main/java/org/apache/aurora/common/zookeeper/testing/ZooKeeperTestServer.java
new file mode 100644
index 0000000..db9e5ca
--- /dev/null
+++
b/commons/src/main/java/org/apache/aurora/common/zookeeper/testing/ZooKeeperTestServer.java
@@ -0,0 +1,221 @@
+/**
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.aurora.common.zookeeper.testing;
+
+import java.io.File;
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.util.Arrays;
+
+import com.google.common.base.Optional;
+import com.google.common.base.Preconditions;
+
+import org.apache.aurora.common.zookeeper.ZooKeeperClient;
+import org.apache.zookeeper.server.NIOServerCnxn;
+import org.apache.zookeeper.server.ZooKeeperServer;
+import org.apache.zookeeper.server.ZooKeeperServer.BasicDataTreeBuilder;
+import org.apache.zookeeper.server.persistence.FileTxnSnapLog;
+
+import org.apache.aurora.common.application.ShutdownRegistry;
+import org.apache.aurora.common.base.Command;
+import org.apache.aurora.common.base.ExceptionalCommand;
+import org.apache.aurora.common.io.FileUtils;
+import org.apache.aurora.common.quantity.Amount;
+import org.apache.aurora.common.quantity.Time;
+
+/**
+ * A helper class for starting in-process ZooKeeper server and clients.
+ *
+ * <p>This is ONLY meant to be used for testing.
+ */
+public class ZooKeeperTestServer {
+
+ /**
+ * The default session timeout for clients created by servers constructed
with
+ * {@link #ZooKeeperTestServer(int, ShutdownRegistry)}.
+ */
+ public static final Amount<Integer, Time> DEFAULT_SESSION_TIMEOUT =
+ Amount.of(100, Time.MILLISECONDS);
+
+ protected final ZooKeeperServer zooKeeperServer;
+ private final ShutdownRegistry shutdownRegistry;
+ private NIOServerCnxn.Factory connectionFactory;
+ private int port;
+ private final Amount<Integer, Time> defaultSessionTimeout;
+
+ /**
+ * @param port the port to start the zoo keeper server on - {@code 0} picks
an ephemeral port
+ * @param shutdownRegistry a registry that will be used to register client
and server shutdown
+ * commands. It is up to the caller to execute the registered actions
at an appropriate time.
+ * @throws IOException if there was aproblem creating the server's database
+ */
+ public ZooKeeperTestServer(int port, ShutdownRegistry shutdownRegistry)
throws IOException {
+ this(port, shutdownRegistry, DEFAULT_SESSION_TIMEOUT);
+ }
+
+ /**
+ * @param port the port to start the zoo keeper server on - {@code 0} picks
an ephemeral port
+ * @param shutdownRegistry a registry that will be used to register client
and server shutdown
+ * commands. It is up to the caller to execute the registered actions
at an appropriate time.
+ * @param defaultSessionTimeout the default session timeout for clients
created with
+ * {@link #createClient()}.
+ * @throws IOException if there was aproblem creating the server's database
+ */
+ public ZooKeeperTestServer(int port, ShutdownRegistry shutdownRegistry,
+ Amount<Integer, Time> defaultSessionTimeout) throws IOException {
+ Preconditions.checkArgument(0 <= port && port <= 0xFFFF);
+ this.port = port;
+ this.shutdownRegistry = Preconditions.checkNotNull(shutdownRegistry);
+ this.defaultSessionTimeout =
Preconditions.checkNotNull(defaultSessionTimeout);
+
+ zooKeeperServer =
+ new ZooKeeperServer(
+ new FileTxnSnapLog(createTempDir(), createTempDir()),
+ new BasicDataTreeBuilder()) {
+
+ // TODO(John Sirois): Introduce a builder to configure the
in-process server if and when
+ // some folks need JMX for in-process tests.
+ @Override protected void registerJMX() {
+ // noop
+ }
+ };
+ }
+
+ /**
+ * Starts zookeeper up on the configured port. If the configured port is
the ephemeral port
+ * (@{code 0}), then the actual chosen port is returned.
+ */
+ public final int startNetwork() throws IOException, InterruptedException {
+ connectionFactory = new NIOServerCnxn.Factory(new InetSocketAddress(port));
+ connectionFactory.startup(zooKeeperServer);
+ shutdownRegistry.addAction(new Command() {
+ @Override public void execute() {
+ shutdownNetwork();
+ }
+ });
+ port = zooKeeperServer.getClientPort();
+ return port;
+ }
+
+ /**
+ * Starts zookeeper back up on the last used port.
+ */
+ public final void restartNetwork() throws IOException, InterruptedException {
+ checkEphemeralPortAssigned();
+ Preconditions.checkState(!connectionFactory.isAlive());
+ startNetwork();
+ }
+
+ /**
+ * Shuts down the in-process zookeeper network server.
+ */
+ public final void shutdownNetwork() {
+ if (connectionFactory != null && connectionFactory.isAlive()) {
+ connectionFactory.shutdown();
+ }
+ }
+
+ /**
+ * Expires the active session for the given client. The client should be
one returned from
+ * {@link #createClient}.
+ *
+ * @param zkClient the client to expire
+ * @throws ZooKeeperClient.ZooKeeperConnectionException if a problem is
encountered connecting to
+ * the local zk server while trying to expire the session
+ * @throws InterruptedException if interrupted while requesting expiration
+ */
+ public final void expireClientSession(ZooKeeperClient zkClient)
+ throws ZooKeeperClient.ZooKeeperConnectionException,
InterruptedException {
+ zooKeeperServer.closeSession(zkClient.get().getSessionId());
+ }
+
+ /**
+ * Returns the current port to connect to the in-process zookeeper instance.
+ */
+ public final int getPort() {
+ checkEphemeralPortAssigned();
+ return port;
+ }
+
+ /**
+ * Returns a new unauthenticated zookeeper client connected to the
in-process zookeeper server
+ * with the default session timeout.
+ */
+ public final ZooKeeperClient createClient() {
+ return createClient(defaultSessionTimeout);
+ }
+
+ /**
+ * Returns a new unauthenticated zookeeper client connected to the
in-process zookeeper server
+ * with the default session timeout and a custom {@code chrootPath}.
+ */
+ public final ZooKeeperClient createClient(String chrootPath) {
+ return createClient(defaultSessionTimeout,
ZooKeeperClient.Credentials.NONE, Optional.of(chrootPath));
+ }
+
+ /**
+ * Returns a new authenticated zookeeper client connected to the in-process
zookeeper server with
+ * the default session timeout.
+ */
+ public final ZooKeeperClient createClient(ZooKeeperClient.Credentials
credentials) {
+ return createClient(defaultSessionTimeout, credentials,
Optional.<String>absent());
+ }
+
+ /**
+ * Returns a new unauthenticated zookeeper client connected to the
in-process zookeeper server
+ * with a custom {@code sessionTimeout}.
+ */
+ public final ZooKeeperClient createClient(Amount<Integer, Time>
sessionTimeout) {
+ return createClient(sessionTimeout, ZooKeeperClient.Credentials.NONE,
Optional.<String>absent());
+ }
+
+ /**
+ * Returns a new authenticated zookeeper client connected to the in-process
zookeeper server with
+ * a custom {@code sessionTimeout}.
+ */
+ public final ZooKeeperClient createClient(Amount<Integer, Time>
sessionTimeout,
+ ZooKeeperClient.Credentials credentials) {
+ return createClient(sessionTimeout, credentials,
Optional.<String>absent());
+ }
+
+ /**
+ * Returns a new authenticated zookeeper client connected to the in-process
zookeeper server with
+ * a custom {@code sessionTimeout} and a custom {@code chrootPath}.
+ */
+ public final ZooKeeperClient createClient(Amount<Integer, Time>
sessionTimeout,
+ ZooKeeperClient.Credentials credentials, Optional<String> chrootPath) {
+ final ZooKeeperClient client = new ZooKeeperClient(sessionTimeout,
credentials,
+ chrootPath,
Arrays.asList(InetSocketAddress.createUnresolved("127.0.0.1", port)));
+ shutdownRegistry.addAction(new ExceptionalCommand<InterruptedException>() {
+ @Override public void execute() {
+ client.close();
+ }
+ });
+ return client;
+ }
+
+ private void checkEphemeralPortAssigned() {
+ Preconditions.checkState(port > 0, "startNetwork must be called first");
+ }
+
+ private File createTempDir() {
+ final File tempDir = FileUtils.createTempDir();
+ shutdownRegistry.addAction(new ExceptionalCommand<IOException>() {
+ @Override public void execute() throws IOException {
+ org.apache.commons.io.FileUtils.deleteDirectory(tempDir);
+ }
+ });
+ return tempDir;
+ }
+}
