com.linkedin.parseq.zk.server.ZKServer.java Source code

Introduction

Here is the source code for com.linkedin.parseq.zk.server.ZKServer.java

Source

/*
   Copyright (c) 2016 LinkedIn Corp.
    
   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 com.linkedin.parseq.zk.server;

import java.io.File;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.util.concurrent.CountDownLatch;
import org.apache.commons.io.FileUtils;
import org.apache.zookeeper.server.NIOServerCnxnFactory;
import org.apache.zookeeper.server.ZooKeeperServer;

/**
 * Very simple wrapper around ZooKeeper-3.4.6 server, intended only for TEST use.
 *
 * @author Steven Ihde
 * @author Ang Xu
 */

public class ZKServer {
    private volatile ZooKeeperServer _zk;
    private volatile NIOServerCnxnFactory _factory;
    private final File _dataDir;
    private final File _logDir;
    private final int _port;
    private final boolean _erase;

    private final CountDownLatch _latch = new CountDownLatch(1);

    /**
     * Create a ZK server with automatically determined port, log and data dirs, which will be
     * erased upon shutdown.
     * @throws IOException
     */
    public ZKServer() throws IOException {
        this(0);
    }

    /**
     * Create a ZK server with automatically determined log and data dirs, which will be
     * erased upon shutdown.
     * @param port The port to listen on
     * @throws IOException
     */
    public ZKServer(int port) throws IOException {
        this(createTempDir("data"), createTempDir("log"), port, true);
    }

    /**
     * Create a ZK server with specified data and log dirs.
     * @param dataDir
     * @param logDir
     * @param port
     * @param erase if true, dataDir and logDir will be erased when shutdown() is called
     * @throws IOException
     */
    public ZKServer(File dataDir, File logDir, int port, boolean erase) throws IOException {
        _dataDir = dataDir;
        _logDir = logDir;
        _port = port;
        _zk = new ZooKeeperServer(dataDir, logDir, 5000);
        _factory = new NIOServerCnxnFactory();
        _factory.configure(new InetSocketAddress(port), 60);
        _erase = erase;

    }

    public int getPort() {
        return _factory.getLocalPort();
    }

    public void startup() throws IOException, InterruptedException {
        ensureDir(_dataDir);
        ensureDir(_logDir);
        _factory.startup(_zk);
    }

    public void shutdown() throws IOException {
        shutdown(_erase);
    }

    public void shutdown(boolean erase) throws IOException {
        _factory.shutdown();
        try {
            _factory.join();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        _zk.shutdown();
        if (_zk.getZKDatabase() != null) {
            _zk.getZKDatabase().close();
        }
        if (erase) {
            FileUtils.deleteDirectory(_dataDir);
            FileUtils.deleteDirectory(_logDir);
        }
    }

    public void restart() throws IOException, InterruptedException {
        shutdown(false);

        _zk = new ZooKeeperServer(_dataDir, _logDir, 5000);
        _factory = new NIOServerCnxnFactory();
        _factory.configure(new InetSocketAddress(_port), 60);
        startup();
    }

    private static void ensureDir(File dir) throws IOException {
        if (dir.exists() && !dir.isDirectory()) {
            throw new IOException("Not a directory: " + dir.getAbsolutePath());
        }
    }

    private static File createTempDir(String suffix) {
        File parent = new File(System.getProperty("java.io.tmpdir"));
        String baseName = System.currentTimeMillis() + ".";

        for (int counter = 0; counter < 10; ++counter) {
            File tempDir = new File(parent, baseName + counter + "." + suffix);
            if (tempDir.mkdir()) {
                return tempDir;
            }
        }

        throw new IllegalStateException(
                "Failed to create directory within 10 attempts (tried " + baseName + "0 to " + baseName + 9 + ')');
    }

}