com.netflix.suro.connection.ConnectionPool.java Source code

Java tutorial

Introduction

Here is the source code for com.netflix.suro.connection.ConnectionPool.java

Source

/*
 * Copyright 2013 Netflix, Inc.
 *
 *    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.netflix.suro.connection;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Sets;
import com.google.inject.Inject;
import com.netflix.governator.guice.lazy.LazySingleton;
import com.netflix.loadbalancer.ILoadBalancer;
import com.netflix.loadbalancer.Server;
import com.netflix.servo.annotations.DataSourceType;
import com.netflix.servo.annotations.Monitor;
import com.netflix.servo.monitor.Monitors;
import com.netflix.suro.ClientConfig;
import com.netflix.suro.thrift.Result;
import com.netflix.suro.thrift.ServiceStatus;
import com.netflix.suro.thrift.SuroServer;
import com.netflix.suro.thrift.TMessageSet;
import org.apache.thrift.TException;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.TFramedTransport;
import org.apache.thrift.transport.TSocket;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.TTransportException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.PreDestroy;
import java.util.*;
import java.util.concurrent.*;

/**
 * Pooling for thrift connection to suro-server
 * After creating all connections to suro-server discovered by {@link ILoadBalancer}, a {@code ConnectionPool} returns
 * a connection when the client requests to get one. When there's no connection available, {@code ConnectionPool} will
 * create a new connection immediately. This is called OutPool connection.
 *
 * @author jbae
 */
@LazySingleton
public class ConnectionPool {
    private static final Logger logger = LoggerFactory.getLogger(ConnectionPool.class);

    private Map<Server, SuroConnection> connectionPool = new ConcurrentHashMap<Server, SuroConnection>();
    private Set<Server> serverSet = Collections.newSetFromMap(new ConcurrentHashMap<Server, Boolean>());
    private List<SuroConnection> connectionList = Collections.synchronizedList(new LinkedList<SuroConnection>());

    private final ClientConfig config;
    private final ILoadBalancer lb;

    private ScheduledExecutorService connectionSweeper;
    private ExecutorService newConnectionBuilder;
    private BlockingQueue<SuroConnection> connectionQueue = new LinkedBlockingQueue<SuroConnection>();
    private CountDownLatch populationLatch;

    /**
     *
     * @param config Client configuration
     * @param lb LoadBalancer implementation
     */
    @Inject
    public ConnectionPool(ClientConfig config, ILoadBalancer lb) {
        this.config = config;
        this.lb = lb;

        connectionSweeper = Executors.newScheduledThreadPool(1);
        newConnectionBuilder = Executors.newFixedThreadPool(1);

        Monitors.registerObject(this);

        populationLatch = new CountDownLatch(
                Math.min(lb.getServerList(true).size(), config.getAsyncSenderThreads()));
        Executors.newSingleThreadExecutor().submit(new Runnable() {
            @Override
            public void run() {
                populateClients();
            }
        });
        try {
            populationLatch.await(populationLatch.getCount() * config.getConnectionTimeout(),
                    TimeUnit.MILLISECONDS);
        } catch (InterruptedException e) {
            logger.error("Exception on CountDownLatch awaiting: " + e.getMessage(), e);
        }
        logger.info("ConnectionPool population finished with the size: " + getPoolSize() + ", will continue up to: "
                + lb.getServerList(true).size());
    }

    @PreDestroy
    public void shutdown() {
        serverSet.clear();
        connectionPool.clear();
        connectionQueue.clear();
        for (SuroConnection conn : connectionList) {
            conn.disconnect();
        }

        connectionSweeper.shutdownNow();
        newConnectionBuilder.shutdownNow();
    }

    /**
     * @return number of connections in the pool
     */
    @Monitor(name = "PoolSize", type = DataSourceType.GAUGE)
    public int getPoolSize() {
        return connectionList.size();
    }

    @Monitor(name = "OutPoolSize", type = DataSourceType.GAUGE)
    private int outPoolSize = 0;

    /**
     * @return number of connections created out of the pool
     */
    public int getOutPoolSize() {
        return outPoolSize;
    }

    public void populateClients() {
        for (Server server : lb.getServerList(true)) {
            SuroConnection connection = new SuroConnection(server, config, true);
            try {
                connection.connect();
                addConnection(server, connection, true);
                logger.info(connection + " is added to SuroClientPool");
            } catch (Exception e) {
                logger.error("Error in connecting to " + connection + " message: " + e.getMessage(), e);
                lb.markServerDown(server);
            }
        }

        connectionSweeper.scheduleAtFixedRate(new Runnable() {
            @Override
            public void run() {
                removeConnection(Sets.difference(serverSet, new HashSet<Server>(lb.getServerList(true))));
            }
        }, config.getConnectionSweepInterval(), config.getConnectionSweepInterval(), TimeUnit.SECONDS);
    }

    @VisibleForTesting
    protected void addConnection(Server server, SuroConnection connection, boolean inPool) {
        if (inPool) {
            connectionPool.put(server, connection);
            if (populationLatch.getCount() > 0) {
                populationLatch.countDown();
            }
        }
        serverSet.add(server);
        connectionList.add(connection);
    }

    private synchronized void removeConnection(Set<Server> removedServers) {
        for (Server s : removedServers) {
            serverSet.remove(s);
            connectionPool.remove(s);
        }

        Iterator<SuroConnection> i = connectionQueue.iterator();
        while (i.hasNext()) {
            if (!serverSet.contains(i.next().getServer())) {
                i.remove();
                logger.info("connection was removed from the queue");
            }
        }

        i = connectionList.iterator();
        while (i.hasNext()) {
            SuroConnection c = i.next();
            if (!serverSet.contains(c.getServer())) {
                c.disconnect();
                i.remove();
            }
        }
    }

    /**
     * When the client calls this method, it will return the connection.
     * @return connection
     */
    public SuroConnection chooseConnection() {
        SuroConnection connection = connectionQueue.poll();
        if (connection == null) {
            connection = chooseFromPool();
        }

        if (config.getEnableOutPool()) {
            synchronized (this) {
                for (int i = 0; i < config.getRetryCount() && connection == null; ++i) {
                    Server server = lb.chooseServer(null);
                    if (server != null) {
                        connection = new SuroConnection(server, config, false);
                        try {
                            connection.connect();
                            ++outPoolSize;
                            logger.info(connection + " is created out of the pool");
                            break;
                        } catch (Exception e) {
                            logger.error("Error in connecting to " + connection + " message: " + e.getMessage(), e);
                            lb.markServerDown(server);
                        }
                    }
                }
            }
        }

        if (connection == null) {
            logger.error("No valid connection exists after " + config.getRetryCount() + " retries");
        }

        return connection;
    }

    private SuroConnection chooseFromPool() {
        SuroConnection connection = null;
        int count = 0;

        while (connection == null) {
            Server server = lb.chooseServer(null);
            if (server != null) {
                if (!serverSet.contains(server)) {
                    newConnectionBuilder.execute(createNewConnection(server, true));
                } else {
                    connection = connectionPool.remove(server);
                }
            } else {
                break;
            }

            ++count;
            if (count >= 10) {
                logger.error("no connection available selected in 10 retries");
                break;
            }
        }

        return connection;
    }

    private Runnable createNewConnection(final Server server, final boolean inPool) {
        return new Runnable() {
            @Override
            public void run() {
                if (connectionPool.get(server) == null) {
                    SuroConnection connection = new SuroConnection(server, config, inPool);
                    try {
                        connection.connect();
                        addConnection(server, connection, inPool);
                        logger.info(connection + " is added to ConnectionPool");
                    } catch (Exception e) {
                        logger.error("Error in connecting to " + connection + " message: " + e.getMessage(), e);
                        lb.markServerDown(server);
                    }
                }
            }
        };
    }

    /**
     * When the client finishes communication with the client, this method
     * should be called to release the connection and return it to the pool.
     * @param connection
     */
    public void endConnection(SuroConnection connection) {
        if (connection != null && shouldChangeConnection(connection)) {
            connection.initStat();
            connectionPool.put(connection.getServer(), connection);
            connection = chooseFromPool();
        }

        if (connection != null) {
            connectionQueue.offer(connection);
        }
    }

    /**
     * Mark up the server related with the connection as down
     * When the client fails to communicate with the connection,
     * this method should be called to remove the server from the pool
     * @param connection
     */
    public void markServerDown(SuroConnection connection) {
        if (connection != null) {
            lb.markServerDown(connection.getServer());
            removeConnection(new ImmutableSet.Builder<Server>().add(connection.getServer()).build());
        }
    }

    private boolean shouldChangeConnection(SuroConnection connection) {
        if (!connection.isInPool()) {
            return false;
        }

        long now = System.currentTimeMillis();

        long minimumTimeSpan = connection.getTimeUsed() + config.getMinimumReconnectTimeInterval();
        return connectionExpired(connection, now, minimumTimeSpan);

    }

    private boolean connectionExpired(SuroConnection connection, long now, long minimumTimeSpan) {
        return minimumTimeSpan <= now && (connection.getSentCount() >= config.getReconnectInterval()
                || connection.getTimeUsed() + config.getReconnectTimeInterval() <= now);
    }

    /**
     * Thrift socket connection wrapper with configuration
     */
    public static class SuroConnection {
        private TTransport transport;
        private SuroServer.Client client;

        private final Server server;
        private final ClientConfig config;
        private final boolean inPool;

        private int sentCount = 0;
        private long timeUsed = 0;

        /**
         * @param server hostname and port information
         * @param config properties including timeout, etc
         * @param inPool whether this connection is in the pool or out of it
         */
        public SuroConnection(Server server, ClientConfig config, boolean inPool) {
            this.server = server;
            this.config = config;
            this.inPool = inPool;
        }

        public void connect() throws Exception {
            TSocket socket = new TSocket(server.getHost(), server.getPort(), config.getConnectionTimeout());
            socket.getSocket().setTcpNoDelay(true);
            socket.getSocket().setKeepAlive(true);
            socket.getSocket().setSoLinger(true, 0);
            transport = new TFramedTransport(socket);
            transport.open();

            TProtocol protocol = new TBinaryProtocol(transport);

            client = new SuroServer.Client(protocol);
            ServiceStatus status = client.getStatus();
            if (status != ServiceStatus.ALIVE) {
                transport.close();
                throw new RuntimeException(server + " IS NOT ALIVE!!!");
            }
        }

        public void disconnect() {
            try {
                transport.flush();
            } catch (TTransportException e) {
                logger.error("Exception on disconnect: " + e.getMessage(), e);
            } finally {
                transport.close();
            }
        }

        public Result send(TMessageSet messageSet) throws TException {
            ++sentCount;
            if (sentCount == 1) {
                timeUsed = System.currentTimeMillis();
            }

            return client.process(messageSet);
        }

        public Server getServer() {
            return server;
        }

        /**
         * @return How many times send() method is called
         */
        public int getSentCount() {
            return sentCount;
        }

        /**
         * For the connection retention control
         * @return how long it has been used from the client
         */
        public long getTimeUsed() {
            return timeUsed;
        }

        public boolean isInPool() {
            return inPool;
        }

        public void initStat() {
            sentCount = 0;
            timeUsed = 0;
        }

        @Override
        public String toString() {
            return server.getHostPort();
        }

        @Override
        public boolean equals(Object o) {
            if (o instanceof Server) {
                return server.equals(o);
            } else if (o instanceof SuroConnection) {
                return server.equals(((SuroConnection) o).server);
            } else {
                return false;
            }
        }

        @Override
        public int hashCode() {
            return server.hashCode();
        }
    }
}