package org.aion.db.utils; import static org.junit.Assert.fail; import com.spotify.docker.client.DefaultDockerClient; import com.spotify.docker.client.DockerClient; import com.spotify.docker.client.LogStream; import com.spotify.docker.client.messages.ContainerConfig; import com.spotify.docker.client.messages.ContainerCreation; import com.spotify.docker.client.messages.ExecCreation; import com.spotify.docker.client.messages.HostConfig; import com.spotify.docker.client.messages.PortBinding; import java.util.Arrays; import java.util.List; import java.util.Map; import org.aion.db.impl.DatabaseTestUtils; /** Helper class for spinning up a MongoDB instance to be used for unit tests. */ public class MongoTestRunner implements AutoCloseable { private int port; private DockerClient dockerClient; private String runningDockerContainerId; private static final String MONGO_IMAGE = "library/mongo:3.6.9"; private static class Holder { static final MongoTestRunner INSTANCE = new MongoTestRunner(); } public static MongoTestRunner inst() { return Holder.INSTANCE; } private MongoTestRunner() { try { // Start by getting a connection to the docker service running on the machine DefaultDockerClient.Builder clientBuilder = DefaultDockerClient.fromEnv(); System.out.println("Connecting to docker daemon at " + clientBuilder.uri().toString()); dockerClient = clientBuilder.build(); // Pull the docker image, this will be very quick if it already exists on the machine dockerClient.pull( MONGO_IMAGE, message -> System.out.println("Docker pull: " + message.status())); // Bind container port 27017 to an automatically allocated available host port. this.port = DatabaseTestUtils.findOpenPort(); final Map<String, List<PortBinding>> portBindings = Map.of( "27017", Arrays.asList(PortBinding.of("0.0.0.0", Integer.toString(this.port)))); final HostConfig hostConfig = HostConfig.builder().portBindings(portBindings).build(); // Configure how we want the image to run ContainerConfig containerConfig = ContainerConfig.builder() .attachStderr(true) .hostConfig(hostConfig) .exposedPorts("27017") .image(MONGO_IMAGE) .cmd("--replSet", "rs0", "--noauth", "--nojournal", "--quiet") .build(); // Actually start the container ContainerCreation creation = dockerClient.createContainer(containerConfig); dockerClient.startContainer(creation.id()); this.runningDockerContainerId = creation.id(); // Next we run a command to initialize the mongo server's replicas set and admin // accounts String[] initializationCommands = {"mongo", "--eval", "rs.initiate()"}; tryInitializeDb(initializationCommands, 30, 100); // Finally, add a shutdown hook to kill the Mongo server when the process dies Runtime.getRuntime() .addShutdownHook( new Thread( () -> { try { close(); } catch (Exception e) { e.printStackTrace(); fail("Failed to close MongoDB connection"); } })); } catch (Exception e) { e.printStackTrace(); fail( "Error encountered when initializing mongo docker image. Make sure docker service is running"); } } /** * Helper method to run some initialization command on Mongo with some retry logic if the * command fails. Since it's not determinate how long starting the database will take, we need * this retry logic. * * @param initializationCommands The command to actually run * @param retriesRemaining How many more times to retry the command if it fails * @param pauseTimeMillis How long to pause between retries * @throws InterruptedException Thrown when the thread gets interrupted trying to sleep. */ private void tryInitializeDb( String[] initializationCommands, int retriesRemaining, long pauseTimeMillis) throws InterruptedException { Exception exception = null; String execOutput = ""; try { final ExecCreation execCreation = dockerClient.execCreate( this.runningDockerContainerId, initializationCommands, DockerClient.ExecCreateParam.attachStdout(), DockerClient.ExecCreateParam.attachStderr(), DockerClient.ExecCreateParam.detach(false)); final LogStream output = dockerClient.execStart(execCreation.id()); execOutput = output.readFully(); } catch (Exception e) { exception = e; } // We can't get the exit code, but look for an expected message in the output to determine // success if (exception != null || !execOutput.contains("Using a default configuration for the set")) { // This is the case that the command didn't work if (retriesRemaining == 0) { // We're out of retries, we should fail if (exception != null) { exception.printStackTrace(); } fail( "Failed to initialize MongoDB, no retries remaining. Output was: " + execOutput); } else { Thread.sleep(pauseTimeMillis); tryInitializeDb(initializationCommands, retriesRemaining - 1, pauseTimeMillis); } } } /** * Returns the connection string to be used to connect to the started Mongo instance * * @return The connection string. */ public String getConnectionString() { return String.format("mongodb://localhost:%d", this.port); } @Override public void close() throws Exception { if (this.dockerClient != null && this.runningDockerContainerId != null) { System.out.println("Killing mongo docker container"); this.dockerClient.killContainer(this.runningDockerContainerId); this.dockerClient.removeContainer(this.runningDockerContainerId); this.dockerClient.close(); this.dockerClient = null; this.runningDockerContainerId = null; } } }