CassandraDriver.java

package cassandra;

import cassandra.protocol.CassandraMessageCodec;
import cassandra.protocol.CassandraMessageHandler;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelOption;
import io.netty.channel.EventLoopGroup;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.logging.LogLevel;
import io.netty.handler.logging.LoggingHandler;
import io.netty.handler.ssl.SslHandler;
import io.netty.util.concurrent.*;

import javax.net.ssl.SSLEngine;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.util.LinkedList;
import java.util.Queue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

public class CassandraDriver {

    public static final String CQL_VERSION = "3.1.1";
    public static final String NATIVE_PROTOCOL_VERSION = "2";
    public static final int NATIVE_PROTOCOL_VERSION_NUMBER = Integer.parseInt(NATIVE_PROTOCOL_VERSION);

    private static final int FUTURE_MAP_CLEAN_INTERVAL_SECS = 10;
    private static final AtomicReference<EventExecutorGroup> GLOBAL_EVENT_EXECUTOR;

    static {
        GLOBAL_EVENT_EXECUTOR = new AtomicReference<EventExecutorGroup>(GlobalEventExecutor.INSTANCE);
    }

    private final EventLoopGroup worker;
    private final LoggingHandler tracer;
    private final CassandraFutureMap futureMap;
    private final CassandraMessageHandler handler;
    private final ShutdownListener shutdownListener;

    public CassandraDriver() {
        this(0);
    }

    public CassandraDriver(int workers) {
        this(new NioEventLoopGroup(workers, new DefaultThreadFactory("cassandra-worker")));
    }

    public CassandraDriver(EventLoopGroup worker) {
        if (worker == null) {
            throw new NullPointerException("worker");
        }
        if (worker.isShuttingDown() || worker.isShutdown()) {
            throw new IllegalArgumentException("worker shutdown");
        }
        this.worker = worker;
        tracer = new LoggingHandler(CassandraConnection.class, LogLevel.TRACE);
        futureMap = new CassandraFutureMap();
        handler = new CassandraMessageHandler(futureMap);
        shutdownListener = new ShutdownListener();
        worker.terminationFuture().addListener(shutdownListener);
        startCleaner();
    }

    public CassandraCluster.Builder newClusterBuilder() {
        return CassandraCluster.newBuilder().setDriver(this);
    }

    public CassandraConnection newConnection(InetAddress inetAddress) {
        return newConnection(inetAddress, CassandraOptions.DEFAULT);
    }

    public CassandraConnection newConnection(InetAddress inetAddress, CassandraOptions options) {
        return newConnection(new InetSocketAddress(inetAddress, options.getPort()), options);
    }

    public CassandraConnection newConnection(InetSocketAddress socketAddress) {
        return newConnection(socketAddress, CassandraOptions.DEFAULT);
    }

    public CassandraConnection newConnection(InetSocketAddress socketAddress, CassandraOptions options) {
        if (isShutdown()) {
            throw new IllegalStateException("driver shutdown");
        }
        Channel channel = new NioSocketChannel();
        channel.pipeline().addLast(tracer);
        if (options.getSslContext() != null) {
            SSLEngine engine = options.getSslContext().createSSLEngine();
            engine.setUseClientMode(true);
            if (options.getCipherSuites() != null) {
                engine.setEnabledCipherSuites(options.getCipherSuites());
            }
            channel.pipeline().addLast(new SslHandler(engine));
        }
        channel.pipeline().addLast(new CassandraMessageCodec(options.getCompression()));
        channel.pipeline().addLast(handler);
        channel.config().setOption(ChannelOption.SO_KEEPALIVE, true);
        channel.config().setOption(ChannelOption.TCP_NODELAY, true);
        channel.config().setOption(ChannelOption.CONNECT_TIMEOUT_MILLIS, options.getConnectTimeoutMillis());
        ChannelFuture registerFuture = worker.register(channel);
        if (registerFuture.cause() != null) {
            if (channel.isRegistered()) {
                channel.close();
            } else {
                channel.unsafe().closeForcibly();
            }
            throw new IllegalStateException(registerFuture.cause().getMessage());
        }
        return new CassandraConnection(options, futureMap, channel, socketAddress);
    }

    public CassandraFutureMap futureMap() {
        return futureMap;
    }

    public boolean isShutdown() {
        return worker.isShuttingDown() || worker.isShutdown();
    }

    public void shutdown() {
        worker.shutdownGracefully();
        shutdownEventExecutor(GLOBAL_EVENT_EXECUTOR.get());
    }

    public void addShutdownHook(CassandraCluster cluster) {
        shutdownListener.clusters.add(cluster);
    }

    public void removeShutdownHook(CassandraCluster cluster) {
        shutdownListener.clusters.remove(cluster);
    }

    private void startCleaner() {
        worker.scheduleAtFixedRate(new Runnable() {
            @Override
            public void run() {
                futureMap.clearExpiredFutures();
            }
        }, FUTURE_MAP_CLEAN_INTERVAL_SECS, FUTURE_MAP_CLEAN_INTERVAL_SECS, TimeUnit.SECONDS);
    }

    public static EventExecutorGroup getGlobalEventExecutor() {
        return GLOBAL_EVENT_EXECUTOR.get();
    }

    public static boolean setGlobalEventExecutor(int nThread) {
        return setGlobalEventExecutor(nThread, "cassandra-eventexecutor");
    }

    public static boolean setGlobalEventExecutor(int nThread, String name) {
        if (nThread < 1) {
            throw new IllegalArgumentException(String.format("nThread: %d (expected: >= 1)", nThread));
        }
        if (name == null) {
            throw new NullPointerException("name");
        }
        if (name.isEmpty()) {
            throw new IllegalArgumentException("empty name");
        }
        return setGlobalEventExecutor(new DefaultEventExecutorGroup(nThread, new DefaultThreadFactory(name)));
    }

    public static boolean setGlobalEventExecutor(EventExecutorGroup globalEventExecutor) {
        if (globalEventExecutor == null) {
            throw new NullPointerException("globalEventExecutor");
        }
        if (!(globalEventExecutor instanceof GlobalEventExecutor)) {
            if (globalEventExecutor.isShuttingDown() || globalEventExecutor.isShutdown()) {
                throw new IllegalArgumentException("globalEventExecutor shutdown");
            }
        }
        EventExecutorGroup old = GLOBAL_EVENT_EXECUTOR.get();
        boolean updated = GLOBAL_EVENT_EXECUTOR.compareAndSet(old, globalEventExecutor);
        if (updated) {
            shutdownEventExecutor(old);
        }
        return updated;
    }

    private static void shutdownEventExecutor(EventExecutorGroup eventExecutor) {
        if (eventExecutor != GlobalEventExecutor.INSTANCE) {
            eventExecutor.shutdownGracefully();
        }
    }

    private class ShutdownListener implements GenericFutureListener<Future<Object>> {

        private Queue<CassandraCluster> clusters = new LinkedList<CassandraCluster>();

        @Override
        public void operationComplete(Future<Object> future) throws Exception {
            try {
                while (!clusters.isEmpty()) {
                    CassandraCluster cluster = clusters.remove();
                    if (cluster != null) {
                        cluster.close();
                    }
                }
            } catch (Exception e) {
                // ignore
            }
        }
    }
}