/solr/core/src/test/org/apache/solr/cloud/ZkTestServer.java
Java | 345 lines | 256 code | 48 blank | 41 comment | 28 complexity | ec37806145479f39cc75befa7611004d MD5 | raw file
Possible License(s): LGPL-2.1, CPL-1.0, MPL-2.0-no-copyleft-exception, JSON, Apache-2.0, AGPL-1.0, GPL-2.0, GPL-3.0, MIT, BSD-3-Clause
- package org.apache.solr.cloud;
- /**
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with this
- * work for additional information regarding copyright ownership. The ASF
- * licenses this file to You 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.
- */
- import java.io.BufferedReader;
- import java.io.File;
- import java.io.IOException;
- import java.io.InputStreamReader;
- import java.io.OutputStream;
- import java.net.InetAddress;
- import java.net.InetSocketAddress;
- import java.net.Socket;
- import java.net.UnknownHostException;
- import java.util.ArrayList;
- import java.util.List;
- import javax.management.JMException;
- import org.apache.zookeeper.jmx.ManagedUtil;
- import org.apache.zookeeper.server.NIOServerCnxn;
- import org.apache.zookeeper.server.ServerConfig;
- import org.apache.zookeeper.server.SessionTracker.Session;
- import org.apache.zookeeper.server.ZKDatabase;
- import org.apache.zookeeper.server.ZooKeeperServer;
- import org.apache.zookeeper.server.persistence.FileTxnSnapLog;
- import org.apache.zookeeper.server.quorum.QuorumPeerConfig.ConfigException;
- import org.slf4j.Logger;
- import org.slf4j.LoggerFactory;
- public class ZkTestServer {
- public static final int TICK_TIME = 3000;
- private static Logger log = LoggerFactory.getLogger(ZkTestServer.class);
-
- protected final ZKServerMain zkServer = new ZKServerMain();
- private String zkDir;
- private int clientPort;
- private Thread zooThread;
-
- private int theTickTime = TICK_TIME;
- class ZKServerMain {
- private NIOServerCnxn.Factory cnxnFactory;
- private ZooKeeperServer zooKeeperServer;
-
- protected void initializeAndRun(String[] args) throws ConfigException,
- IOException {
- try {
- ManagedUtil.registerLog4jMBeans();
- } catch (JMException e) {
- }
- ServerConfig config = new ServerConfig();
- if (args.length == 1) {
- config.parse(args[0]);
- } else {
- config.parse(args);
- }
- runFromConfig(config);
- }
- /**
- * Run from a ServerConfig.
- * @param config ServerConfig to use.
- * @throws IOException
- */
- public void runFromConfig(ServerConfig config) throws IOException {
- try {
- // Note that this thread isn't going to be doing anything else,
- // so rather than spawning another thread, we will just call
- // run() in this thread.
- // create a file logger url from the command line args
- zooKeeperServer = new ZooKeeperServer();
- FileTxnSnapLog ftxn = new FileTxnSnapLog(new File(config
- .getDataLogDir()), new File(config.getDataDir()));
- zooKeeperServer.setTxnLogFactory(ftxn);
- zooKeeperServer.setTickTime(config.getTickTime());
- cnxnFactory = new NIOServerCnxn.Factory(config.getClientPortAddress(), config
- .getMaxClientCnxns());
- cnxnFactory.startup(zooKeeperServer);
- cnxnFactory.join();
- if (zooKeeperServer.isRunning()) {
- zooKeeperServer.shutdown();
- }
- } catch (InterruptedException e) {
- }
- }
- /**
- * Shutdown the serving instance
- * @throws IOException
- */
- protected void shutdown() throws IOException {
- zooKeeperServer.shutdown();
- ZKDatabase zkDb = zooKeeperServer.getZKDatabase();
- if (zkDb != null) {
- zkDb.close();
- }
- if (cnxnFactory != null && cnxnFactory.getLocalPort() != 0) {
- waitForServerDown(getZkHost() + ":" + getPort(), 5000);
- }
- if (cnxnFactory != null) {
- cnxnFactory.shutdown();
- }
- }
- public int getLocalPort() {
- if (cnxnFactory == null) {
- throw new IllegalStateException("A port has not yet been selected");
- }
- int port;
- try {
- port = cnxnFactory.getLocalPort();
- } catch (NullPointerException e) {
- throw new IllegalStateException("A port has not yet been selected");
- }
- if (port == 0) {
- throw new IllegalStateException("A port has not yet been selected");
- }
- return port;
- }
- }
- public ZkTestServer(String zkDir) {
- this.zkDir = zkDir;
- }
- public ZkTestServer(String zkDir, int port) {
- this.zkDir = zkDir;
- this.clientPort = port;
- }
- public String getZkHost() {
- return "127.0.0.1:" + zkServer.getLocalPort();
- }
- public String getZkAddress() {
- return "127.0.0.1:" + zkServer.getLocalPort() + "/solr";
- }
- public int getPort() {
- return zkServer.getLocalPort();
- }
-
- public void expire(final long sessionId) {
- zkServer.zooKeeperServer.expire(new Session() {
- @Override
- public long getSessionId() {
- return sessionId;
- }
- @Override
- public int getTimeout() {
- return 4000;
- }
- @Override
- public boolean isClosing() {
- return false;
- }});
- }
- public void run() throws InterruptedException {
- log.info("STARTING ZK TEST SERVER");
- // we don't call super.setUp
- zooThread = new Thread() {
-
- @Override
- public void run() {
- ServerConfig config = new ServerConfig() {
- {
- setClientPort(ZkTestServer.this.clientPort);
- this.dataDir = zkDir;
- this.dataLogDir = zkDir;
- this.tickTime = theTickTime;
- }
-
- public void setClientPort(int clientPort) {
- if (clientPortAddress != null) {
- try {
- this.clientPortAddress = new InetSocketAddress(
- InetAddress.getByName(clientPortAddress.getHostName()), clientPort);
- } catch (UnknownHostException e) {
- throw new RuntimeException(e);
- }
- } else {
- this.clientPortAddress = new InetSocketAddress(clientPort);
- }
- }
- };
- try {
- zkServer.runFromConfig(config);
- } catch (Throwable e) {
- throw new RuntimeException(e);
- }
- }
- };
- zooThread.setDaemon(true);
- zooThread.start();
- int cnt = 0;
- int port = -1;
- try {
- port = getPort();
- } catch(IllegalStateException e) {
-
- }
- while (port < 1) {
- Thread.sleep(100);
- try {
- port = getPort();
- } catch(IllegalStateException e) {
-
- }
- if (cnt == 100) {
- throw new RuntimeException("Could not get the port for ZooKeeper server");
- }
- cnt++;
- }
- }
- @SuppressWarnings("deprecation")
- public void shutdown() throws IOException {
- // TODO: this can log an exception while trying to unregister a JMX MBean
- zkServer.shutdown();
- }
-
-
- public static boolean waitForServerDown(String hp, long timeout) {
- long start = System.currentTimeMillis();
- while (true) {
- try {
- HostPort hpobj = parseHostPortList(hp).get(0);
- send4LetterWord(hpobj.host, hpobj.port, "stat");
- } catch (IOException e) {
- return true;
- }
-
- if (System.currentTimeMillis() > start + timeout) {
- break;
- }
- try {
- Thread.sleep(250);
- } catch (InterruptedException e) {
- // ignore
- }
- }
- return false;
- }
-
- public static class HostPort {
- String host;
- int port;
-
- HostPort(String host, int port) {
- this.host = host;
- this.port = port;
- }
- }
-
- /**
- * Send the 4letterword
- * @param host the destination host
- * @param port the destination port
- * @param cmd the 4letterword
- * @return
- * @throws IOException
- */
- public static String send4LetterWord(String host, int port, String cmd)
- throws IOException
- {
- Socket sock = new Socket(host, port);
- BufferedReader reader = null;
- try {
- OutputStream outstream = sock.getOutputStream();
- outstream.write(cmd.getBytes("US-ASCII"));
- outstream.flush();
- // this replicates NC - close the output stream before reading
- sock.shutdownOutput();
- reader =
- new BufferedReader(
- new InputStreamReader(sock.getInputStream()));
- StringBuilder sb = new StringBuilder();
- String line;
- while((line = reader.readLine()) != null) {
- sb.append(line + "\n");
- }
- return sb.toString();
- } finally {
- sock.close();
- if (reader != null) {
- reader.close();
- }
- }
- }
-
- public static List<HostPort> parseHostPortList(String hplist) {
- ArrayList<HostPort> alist = new ArrayList<HostPort>();
- for (String hp : hplist.split(",")) {
- int idx = hp.lastIndexOf(':');
- String host = hp.substring(0, idx);
- int port;
- try {
- port = Integer.parseInt(hp.substring(idx + 1));
- } catch (RuntimeException e) {
- throw new RuntimeException("Problem parsing " + hp + e.toString());
- }
- alist.add(new HostPort(host, port));
- }
- return alist;
- }
- public int getTheTickTime() {
- return theTickTime;
- }
- public void setTheTickTime(int theTickTime) {
- this.theTickTime = theTickTime;
- }
- }