/**
* Copyright 2016 LinkedIn Corp. All rights reserved.
*
* 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.
*/
package com.github.ambry.network;
import java.util.HashMap;
import java.util.LinkedList;
/**
* The ConnectionTracker keeps track of current connections to datanodes, and provides methods to check out and
* check in connections.
*
* This class is not thread safe.
*/
class ConnectionTracker {
private final HashMap<String, HostPortPoolManager> hostPortToPoolManager;
private final HashMap<String, HostPortPoolManager> connectionIdToPoolManager;
private final int maxConnectionsPerPortPlainText;
private final int maxConnectionsPerPortSsl;
private int totalManagedConnectionsCount;
/**
* Instantiates a ConnectionTracker
* @param maxConnectionsPerPortPlainText the connection pool limit for plain text connections to a (host, port)
* @param maxConnectionsPerPortSsl the connection pool limit for ssl connections to a (host, port)
*/
ConnectionTracker(int maxConnectionsPerPortPlainText, int maxConnectionsPerPortSsl) {
hostPortToPoolManager = new HashMap<String, HostPortPoolManager>();
connectionIdToPoolManager = new HashMap<String, HostPortPoolManager>();
totalManagedConnectionsCount = 0;
this.maxConnectionsPerPortPlainText = maxConnectionsPerPortPlainText;
this.maxConnectionsPerPortSsl = maxConnectionsPerPortSsl;
}
/**
* Returns true if a new connection may be created for the given hostPort, that is if the number of connections for
* the given hostPort has not reached the pool limit.
* @param host the host associated with this check.
* @param port the port associated with this check.
* @return true if a new connection may be created, false otherwise.
*/
boolean mayCreateNewConnection(String host, Port port) {
return !getHostPortPoolManager(host, port).hasReachedPoolLimit();
}
/**
* Start tracking a new connection id associated with the given host and port. Note that this connection will not
* be made available for checking out until a {@link #checkInConnection(String)} is called on it.
* @param host the host to which this connection belongs.
* @param port the port on the host to which this connection belongs.
* @param connId the connection id of the connection.
*/
void startTrackingInitiatedConnection(String host, Port port, String connId) {
HostPortPoolManager hostPortPoolManager = getHostPortPoolManager(host, port);
hostPortPoolManager.incrementPoolCount();
connectionIdToPoolManager.put(connId, hostPortPoolManager);
totalManagedConnectionsCount++;
}
/**
* Attempts to check out an existing connection to the hostPort provided, or returns null if none available.
* @param host The host to connect to.
* @param port The port on the host to connect to.
* @return connectionId, if there is one available to use, null otherwise.
*/
String checkOutConnection(String host, Port port) {
return getHostPortPoolManager(host, port).checkOutConnection();
}
/**
* Add connection to available pool.
* @param connectionId the id of the newly established or previously checked out connection.
* @throws {@link IllegalArgumentException} if the passed in connection id is invalid.
*/
void checkInConnection(String connectionId) {
HostPortPoolManager hostPortPoolManager = connectionIdToPoolManager.get(connectionId);
if (hostPortPoolManager == null) {
throw new IllegalArgumentException("Invalid connection id passed in");
}
hostPortPoolManager.checkInConnection(connectionId);
}
/**
* Remove and stop tracking the given connection id.
* @param connectionId connection to remove.
* @throws {@link IllegalArgumentException} if the passed in connection id is invalid.
*/
void removeConnection(String connectionId) {
HostPortPoolManager hostPortPoolManager = connectionIdToPoolManager.remove(connectionId);
if (hostPortPoolManager == null) {
throw new IllegalArgumentException("Invalid connection id passed in");
}
hostPortPoolManager.removeConnection(connectionId);
totalManagedConnectionsCount--;
}
/**
* Return the total number of connections that are managed by this connection tracker.
* @return the total number of initiated and/or established connections.
*/
int getTotalConnectionsCount() {
return totalManagedConnectionsCount;
}
/**
* Return the total available connections across all hostPortPoolManagers.
* @return total established and available connections.
*/
int getAvailableConnectionsCount() {
int count = 0;
for (HostPortPoolManager hostPortPoolManager : hostPortToPoolManager.values()) {
count += hostPortPoolManager.getAvailableConnectionsCount();
}
return count;
}
/**
* Returns the {@link HostPortPoolManager} associated with the (host, port) pair. Creates one if not available
* already.
* @param host The hostname
* @param port The port
* @return the HostPortPoolManager for the associated (host, port) pair.
*/
private HostPortPoolManager getHostPortPoolManager(String host, Port port) {
String lookupStr = host + ":" + Integer.toString(port.getPort());
HostPortPoolManager poolManager = hostPortToPoolManager.get(lookupStr);
if (poolManager == null) {
poolManager = new HostPortPoolManager(
port.getPortType() == PortType.SSL ? maxConnectionsPerPortSsl : maxConnectionsPerPortPlainText);
hostPortToPoolManager.put(lookupStr, poolManager);
}
return poolManager;
}
/**
* HostPortPoolManager manages all the connections to a specific (host,
* port) pair. The {@link ConnectionTracker} creates one for every (host, port) pair it knows of.
*/
private class HostPortPoolManager {
private final int maxConnectionsToHostPort;
private final LinkedList<String> availableConnections;
private int poolCount;
/**
* Instantiate a HostPortPoolManager
* @param poolLimit the max connections allowed for this hostPort.
*/
HostPortPoolManager(int poolLimit) {
poolCount = 0;
maxConnectionsToHostPort = poolLimit;
availableConnections = new LinkedList<String>();
}
/**
* Return true if this manager has reached the pool limit.
* @return true if this manager has reached the pool limit
*/
boolean hasReachedPoolLimit() {
return poolCount == maxConnectionsToHostPort;
}
/**
* Increment the pool count.
*/
void incrementPoolCount() {
poolCount++;
}
/**
* Attempts to check out a connection to the (host, port) associated with this manager.
* @return returns a connection id, if there is one; null otherwise.
*/
String checkOutConnection() {
return availableConnections.poll();
}
/**
* Add connection to available pool.
* @param connectionId the connection id of the connection.
*/
void checkInConnection(String connectionId) {
availableConnections.add(connectionId);
}
/**
* Remove a connection managed by this manager. This connection id could be either a checked out connection or a
* connection that was previously available to be checked out.
* @param connectionId the connection id of the connection.
*/
void removeConnection(String connectionId) {
availableConnections.remove(connectionId);
poolCount--;
}
/**
* Return the number of available connections to this hostPort
* @return number of available connections
*/
int getAvailableConnectionsCount() {
return availableConnections.size();
}
}
}