/*
* Copyright (C) 2016 The Calrissian Authors
*
* 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.calrissian.accumulorecipes.test;
import com.google.common.io.Files;
import org.apache.accumulo.core.client.*;
import org.apache.accumulo.core.client.security.tokens.PasswordToken;
import org.apache.accumulo.core.security.Authorizations;
import org.apache.accumulo.minicluster.MiniAccumuloCluster;
import org.apache.accumulo.minicluster.MiniAccumuloConfig;
import org.apache.commons.io.FileUtils;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooKeeper;
import org.junit.rules.ExternalResource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.Closeable;
import java.io.File;
import java.io.IOException;
import java.net.ServerSocket;
import java.util.concurrent.*;
import static com.google.common.base.Objects.firstNonNull;
public class AccumuloMiniClusterDriver extends ExternalResource implements Closeable {
public static final Logger log = LoggerFactory.getLogger(AccumuloMiniClusterDriver.class);
public static String INSTANCE_NAME = firstNonNull(System.getProperty("accumulo.minicluster.instance.name"),"test-instance");
public static String ROOT_PASSWORD = firstNonNull(System.getProperty("accumulo.minicluster.root.password"),"secret");
public static int ZOOKEEPER_PORT = Integer.parseInt(firstNonNull(System.getProperty("accumulo.minicluster.zookeeper.port"),findFreePort()+""));
private MiniAccumuloCluster miniAccumuloCluster;
private MiniAccumuloConfig miniAccumuloConfig;
private File tempDir;
private ClientConfiguration clientConfiguration;
private Instance instance;
private Connector connector;
public AccumuloMiniClusterDriver() {}
@Override
protected void before() throws Throwable {
start();
}
@Override
protected void after() {
try {
close();
} catch (IOException e) {
//noop
}
}
public void deleteAllTables() throws AccumuloSecurityException, AccumuloException, TableNotFoundException {
for (String table : connector.tableOperations().list()) {
if (table.startsWith("accumulo.")) {
continue;
}
connector.tableOperations().delete(table);
}
}
public void start() throws IOException {
try {
initInstanceAndConnector(300);
} catch (Throwable e) {
try {
initConfig();
miniAccumuloCluster = new MiniAccumuloCluster(miniAccumuloConfig);
miniAccumuloCluster.start();
log.info("started minicluster");
initInstanceAndConnector(null);
} catch (Throwable e1) {
throw new IOException(e1);
}
}
}
private Boolean isZookeeperRunning(String host, int timeout) {
final CountDownLatch connectedSignal = new CountDownLatch(1);
try {
new ZooKeeper(host, timeout,
new Watcher() {
public void process(WatchedEvent event) {
if (event.getState() == Event.KeeperState.SyncConnected) {
connectedSignal.countDown();
}
}
});
return connectedSignal.await(timeout,TimeUnit.MILLISECONDS);
} catch (Throwable e) {
return Boolean.FALSE;
}
}
private void initConfig() {
tempDir = Files.createTempDir();
miniAccumuloConfig = new MiniAccumuloConfig(tempDir,ROOT_PASSWORD);
miniAccumuloConfig.setZooKeeperPort(ZOOKEEPER_PORT);
miniAccumuloConfig.setInstanceName(INSTANCE_NAME);
log.info("attempting to start minicluster in " + tempDir.getAbsolutePath());
}
private void initInstanceAndConnector(final Integer zkTimeout) throws AccumuloException, AccumuloSecurityException {
clientConfiguration = new ClientConfiguration()
.withInstance(INSTANCE_NAME)
.withZkHosts("localhost:" + ZOOKEEPER_PORT);
if (zkTimeout!=null && zkTimeout!=-1) {
if (isZookeeperRunning("localhost:"+ZOOKEEPER_PORT,zkTimeout)) {
instance = new ZooKeeperInstance(clientConfiguration);
}
} else {
instance = new ZooKeeperInstance(clientConfiguration);
}
connector = instance.getConnector("root",new PasswordToken(ROOT_PASSWORD));
}
public ClientConfiguration getClientConfiguration() {
return clientConfiguration;
}
public Instance getInstance() {
return instance;
}
public Connector getConnector() {
return connector;
}
private static int findFreePort() {
ServerSocket socket = null;
try {
socket = new ServerSocket(0);
socket.setReuseAddress(true);
int port = socket.getLocalPort();
try {
socket.close();
} catch (IOException e) {
// Ignore IOException on close()
}
return port;
} catch (IOException e) {
} finally {
if (socket != null) {
try {
socket.close();
} catch (IOException e) {
}
}
}
throw new IllegalStateException("Could not find a free TCP/IP port for zookeeper");
}
@Override
public void close() throws IOException {
log.info("attempting to close");
if (miniAccumuloCluster != null) {
try {
log.info("miniAccumuloCluster found attempting to stop it");
miniAccumuloCluster.stop();
log.info("stopped miniAccumuloCluster");
} catch (InterruptedException e) {
//noop
}
}
if (tempDir!=null) {
log.info("deleting temp dir: "+ tempDir.getAbsolutePath());
FileUtils.deleteQuietly(tempDir);
}
}
public String getRootPassword() {
return ROOT_PASSWORD;
}
public void setRootAuths(Authorizations auths) throws AccumuloSecurityException, AccumuloException {
connector.securityOperations().changeUserAuthorizations("root",auths);
}
public String getZooKeepers() {
if (miniAccumuloCluster==null) {
return "localhost:"+ZOOKEEPER_PORT;
}
return miniAccumuloCluster.getZooKeepers();
}
public String getInstanceName() {
if (miniAccumuloCluster==null) {
return INSTANCE_NAME;
}
return miniAccumuloCluster.getInstanceName();
}
}