package org.xbib.elasticsearch.util;
import org.elasticsearch.ElasticsearchTimeoutException;
import org.elasticsearch.action.admin.cluster.health.ClusterHealthAction;
import org.elasticsearch.action.admin.cluster.health.ClusterHealthRequest;
import org.elasticsearch.action.admin.cluster.health.ClusterHealthResponse;
import org.elasticsearch.action.admin.cluster.node.info.NodesInfoRequest;
import org.elasticsearch.action.admin.cluster.node.info.NodesInfoResponse;
import org.elasticsearch.client.support.AbstractClient;
import org.elasticsearch.common.logging.ESLogger;
import org.elasticsearch.common.logging.ESLoggerFactory;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.transport.InetSocketTransportAddress;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.env.Environment;
import org.elasticsearch.node.MockNode;
import org.elasticsearch.node.Node;
import org.elasticsearch.plugins.Plugin;
import org.junit.After;
import org.junit.Before;
import org.xbib.elasticsearch.plugin.knapsack.KnapsackPlugin;
import java.io.IOException;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.nio.file.FileVisitResult;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.nio.file.SimpleFileVisitor;
import java.nio.file.attribute.BasicFileAttributes;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Random;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;
import static org.elasticsearch.common.settings.Settings.settingsBuilder;
public class NodeTestUtils {
protected final static ESLogger logger = ESLoggerFactory.getLogger("test");
private Map<String, Node> nodes = new HashMap<>();
private Map<String, AbstractClient> clients = new HashMap<>();
private AtomicInteger counter = new AtomicInteger();
private String cluster;
private String host;
private int port;
@Before
public void startNodes() {
try {
setClusterName();
startNode("1");
startNode("2"); // we need 2 nodes for knapsack
findNodeAddress();
try {
ClusterHealthResponse healthResponse =
client("1").execute(ClusterHealthAction.INSTANCE, new ClusterHealthRequest()
.waitForYellowStatus()
.timeout(TimeValue.timeValueSeconds(30)))
.actionGet();
if (healthResponse != null && healthResponse.isTimedOut()) {
throw new IOException("cluster state is " + healthResponse.getStatus().name()
+ ", from here on, everything will fail!");
}
} catch (ElasticsearchTimeoutException e) {
throw new IOException("timeout, cluster does not respond to health request, cowardly refusing to continue with operations");
}
logger.info("ready");
} catch (Throwable t) {
logger.error("startNodes failed", t);
}
}
@After
public void stopNodes() {
try {
closeNodes();
} catch (Exception e) {
logger.error("can not close nodes", e);
} finally {
try {
deleteFiles();
logger.info("data files wiped");
} catch (IOException e) {
logger.error(e.getMessage(), e);
}
}
}
protected void setClusterName() throws UnknownHostException {
this.cluster = "test-helper-plugin-cluster-"
+ InetAddress.getLocalHost().getHostName()
+ "-" + System.getProperty("user.name")
+ "-" + counter.incrementAndGet();
}
protected String getClusterName() {
return cluster;
}
protected Settings getSettings() {
return settingsBuilder()
.put("host", host)
.put("port", port)
.put("cluster.name", cluster)
.put("path.home", getHome())
.build();
}
protected Settings getNodeSettings() {
return settingsBuilder()
.put("cluster.name", cluster)
.put("cluster.routing.schedule", "50ms")
.put("cluster.routing.allocation.disk.threshold_enabled", false)
.put("discovery.zen.multicast.enabled", true)
.put("discovery.zen.multicast.ping_timeout", "5s")
.put("http.enabled", false)
.put("threadpool.bulk.size", Runtime.getRuntime().availableProcessors())
.put("threadpool.bulk.queue_size", 16 * Runtime.getRuntime().availableProcessors()) // default is 50, too low
.put("index.number_of_replicas", 0)
.put("path.home", getHome())
.build();
}
protected String getHome() {
return System.getProperty("path.home");
}
public static Node createNode() {
Settings nodeSettings = Settings.settingsBuilder()
.put("path.home", System.getProperty("path.home"))
.put("index.number_of_shards", 1)
.put("index.number_of_replica", 0)
.build();
// ES 2.1 renders NodeBuilder as useless
//Node node = NodeBuilder.nodeBuilder().settings(nodeSettings).local(true).build().start();
Set<Class<? extends Plugin>> plugins = new HashSet<>();
plugins.add(KnapsackPlugin.class);
Environment environment = new Environment(nodeSettings);
Node node = new MockNode(environment, plugins);
node.start();
return node;
}
public void startNode(String id) throws IOException {
buildNode(id).start();
}
public AbstractClient client(String id) {
return clients.get(id);
}
private void closeNodes() throws IOException {
logger.info("closing all clients");
for (AbstractClient client : clients.values()) {
client.close();
}
clients.clear();
logger.info("closing all nodes");
for (Node node : nodes.values()) {
if (node != null) {
node.close();
}
}
nodes.clear();
logger.info("all nodes closed");
}
protected void findNodeAddress() {
NodesInfoRequest nodesInfoRequest = new NodesInfoRequest().transport(true);
NodesInfoResponse response = client("1").admin().cluster().nodesInfo(nodesInfoRequest).actionGet();
Object obj = response.iterator().next().getTransport().getAddress()
.publishAddress();
if (obj instanceof InetSocketTransportAddress) {
InetSocketTransportAddress address = (InetSocketTransportAddress) obj;
host = address.address().getHostName();
port = address.address().getPort();
}
}
private Node buildNode(String id) throws IOException {
Settings nodeSettings = settingsBuilder()
.put(getNodeSettings())
.put("name", id)
.build();
logger.info("settings={}", nodeSettings.getAsMap());
// ES 2.1 renders NodeBuilder as useless
Set<Class<? extends Plugin>> plugins = new HashSet<>();
plugins.add(KnapsackPlugin.class);
Environment environment = new Environment(nodeSettings);
Node node = new MockNode(environment, plugins);
AbstractClient client = (AbstractClient)node.client();
nodes.put(id, node);
clients.put(id, client);
logger.info("clients={}", clients);
return node;
}
private static void deleteFiles() throws IOException {
Path directory = Paths.get(System.getProperty("path.home") + "/data");
Files.walkFileTree(directory, new SimpleFileVisitor<Path>() {
@Override
public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
Files.delete(file);
return FileVisitResult.CONTINUE;
}
@Override
public FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOException {
Files.delete(dir);
return FileVisitResult.CONTINUE;
}
});
}
private static Random random = new Random();
private static char[] numbersAndLetters = ("0123456789abcdefghijklmnopqrstuvwxyz").toCharArray();
protected String randomString(int len) {
final char[] buf = new char[len];
final int n = numbersAndLetters.length - 1;
for (int i = 0; i < buf.length; i++) {
buf[i] = numbersAndLetters[random.nextInt(n)];
}
return new String(buf);
}
}