package org.littleshoot.proxy.impl;
import org.littleshoot.proxy.HttpProxyServer;
import org.littleshoot.proxy.TransportProtocol;
import org.littleshoot.proxy.UnknownTransportProtocolException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.nio.channels.spi.SelectorProvider;
import java.util.ArrayList;
import java.util.EnumMap;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import io.netty.channel.EventLoopGroup;
/**
* Manages thread pools for one or more proxy server instances. When servers are created, they must register with the
* ServerGroup using {@link #registerProxyServer(HttpProxyServer)}, and when they shut down, must unregister with the
* ServerGroup using {@link #unregisterProxyServer(HttpProxyServer, boolean)}.
*/
public class ServerGroup {
private static final Logger log = LoggerFactory.getLogger(ServerGroup.class);
/**
* The default number of threads to accept incoming requests from clients. (Requests are serviced by worker threads,
* not acceptor threads.)
*/
public static final int DEFAULT_INCOMING_ACCEPTOR_THREADS = 2;
/**
* The default number of threads to service incoming requests from clients.
*/
public static final int DEFAULT_INCOMING_WORKER_THREADS = 8;
/**
* The default number of threads to service outgoing requests to servers.
*/
public static final int DEFAULT_OUTGOING_WORKER_THREADS = 8;
/**
* Global counter for the {@link #serverGroupId}.
*/
private static final AtomicInteger serverGroupCount = new AtomicInteger(0);
/**
* A name for this ServerGroup to use in naming threads.
*/
private final String name;
/**
* The ID of this server group. Forms part of the name of each thread created for this server group. Useful for
* differentiating threads when multiple proxy instances are running.
*/
private final int serverGroupId;
private final int incomingAcceptorThreads;
private final int incomingWorkerThreads;
private final int outgoingWorkerThreads;
/**
* List of all servers registered to use this ServerGroup. Any access to this list should be synchronized using the
* {@link #SERVER_REGISTRATION_LOCK}.
*/
public final List<HttpProxyServer> registeredServers = new ArrayList<HttpProxyServer>(1);
/**
* A mapping of {@link TransportProtocol}s to their initialized {@link ProxyThreadPools}. Each transport uses a
* different thread pool, since the initialization parameters are different.
*/
private final EnumMap<TransportProtocol, ProxyThreadPools> protocolThreadPools = new EnumMap<TransportProtocol, ProxyThreadPools>(TransportProtocol.class);
/**
* A mapping of selector providers to transport protocols. Avoids special-casing each transport protocol during
* transport protocol initialization.
*/
private static final EnumMap<TransportProtocol, SelectorProvider> TRANSPORT_PROTOCOL_SELECTOR_PROVIDERS = new EnumMap<TransportProtocol, SelectorProvider>(TransportProtocol.class);
static {
TRANSPORT_PROTOCOL_SELECTOR_PROVIDERS.put(TransportProtocol.TCP, SelectorProvider.provider());
// allow the proxy to operate without UDT support. this allows clients that do not use UDT to exclude the barchart
// dependency completely.
// if (ProxyUtils.isUdtAvailable()) {
// TRANSPORT_PROTOCOL_SELECTOR_PROVIDERS.put(TransportProtocol.UDT, NioUdtProvider.BYTE_PROVIDER);
// } else {
// log.debug("UDT provider not found on classpath. UDT transport will not be available.");
// }
}
/**
* True when this ServerGroup is stopped.
*/
private final AtomicBoolean stopped = new AtomicBoolean(false);
/**
* Creates a new ServerGroup instance for a proxy. Threads created for this ServerGroup will have the specified
* ServerGroup name in the Thread name. This constructor does not actually initialize any thread pools; instead,
* thread pools for specific transport protocols are lazily initialized as needed.
*
* @param name ServerGroup name to include in thread names
* @param incomingAcceptorThreads number of acceptor threads per protocol
* @param incomingWorkerThreads number of client-to-proxy worker threads per protocol
* @param outgoingWorkerThreads number of proxy-to-server worker threads per protocol
*/
public ServerGroup(String name, int incomingAcceptorThreads, int incomingWorkerThreads, int outgoingWorkerThreads) {
this.name = name;
this.serverGroupId = serverGroupCount.getAndIncrement();
this.incomingAcceptorThreads = incomingAcceptorThreads;
this.incomingWorkerThreads = incomingWorkerThreads;
this.outgoingWorkerThreads = outgoingWorkerThreads;
}
/**
* Lock for initializing any transport protocols.
*/
private final Object THREAD_POOL_INIT_LOCK = new Object();
/**
* Retrieves the {@link ProxyThreadPools} for the specified transport protocol. Lazily initializes the thread pools
* for the transport protocol if they have not yet been initialized. If the protocol has already been initialized,
* this method returns immediately, without synchronization. If initialization is necessary, the initialization
* process creates the acceptor and worker threads necessary to service requests to/from the proxy.
* <p>
* This method is thread-safe; no external locking is necessary.
*
* @param protocol transport protocol to retrieve thread pools for
* @return thread pools for the specified transport protocol
*/
private ProxyThreadPools getThreadPoolsForProtocol(TransportProtocol protocol) {
// if the thread pools have not been initialized for this protocol, initialize them
if (protocolThreadPools.get(protocol) == null) {
synchronized (THREAD_POOL_INIT_LOCK) {
if (protocolThreadPools.get(protocol) == null) {
log.debug("Initializing thread pools for {} with {} acceptor threads, {} incoming worker threads, and {} outgoing worker threads",
protocol, incomingAcceptorThreads, incomingWorkerThreads, outgoingWorkerThreads);
SelectorProvider selectorProvider = TRANSPORT_PROTOCOL_SELECTOR_PROVIDERS.get(protocol);
if (selectorProvider == null) {
throw new UnknownTransportProtocolException(protocol);
}
ProxyThreadPools threadPools = new ProxyThreadPools(selectorProvider,
incomingAcceptorThreads,
incomingWorkerThreads,
outgoingWorkerThreads,
name,
serverGroupId);
protocolThreadPools.put(protocol, threadPools);
}
}
}
return protocolThreadPools.get(protocol);
}
/**
* Lock controlling access to the {@link #registerProxyServer(HttpProxyServer)} and {@link #unregisterProxyServer(HttpProxyServer, boolean)}
* methods.
*/
private final Object SERVER_REGISTRATION_LOCK = new Object();
/**
* Registers the specified proxy server as a consumer of this server group. The server group will not be shut down
* until the proxy unregisters itself.
*
* @param proxyServer proxy server instance to register
*/
public void registerProxyServer(HttpProxyServer proxyServer) {
synchronized (SERVER_REGISTRATION_LOCK) {
registeredServers.add(proxyServer);
}
}
/**
* Unregisters the specified proxy server from this server group. If this was the last registered proxy server, the
* server group will be shut down.
*
* @param proxyServer proxy server instance to unregister
* @param graceful when true, the server group shutdown (if necessary) will be graceful
*/
public void unregisterProxyServer(HttpProxyServer proxyServer, boolean graceful) {
synchronized (SERVER_REGISTRATION_LOCK) {
boolean wasRegistered = registeredServers.remove(proxyServer);
if (!wasRegistered) {
log.warn("Attempted to unregister proxy server from ServerGroup that it was not registered with. Was the proxy unregistered twice?");
}
if (registeredServers.isEmpty()) {
log.debug("Proxy server unregistered from ServerGroup. No proxy servers remain registered, so shutting down ServerGroup.");
shutdown(graceful);
} else {
log.debug("Proxy server unregistered from ServerGroup. Not shutting down ServerGroup ({} proxy servers remain registered).", registeredServers.size());
}
}
}
/**
* Shuts down all event loops owned by this server group.
*
* @param graceful when true, event loops will "gracefully" terminate, waiting for submitted tasks to finish
*/
private void shutdown(boolean graceful) {
if (!stopped.compareAndSet(false, true)) {
log.info("Shutdown requested, but ServerGroup is already stopped. Doing nothing.");
return;
}
log.info("Shutting down server group event loops " + (graceful ? "(graceful)" : "(non-graceful)"));
// loop through all event loops managed by this server group. this includes acceptor and worker event loops
// for both TCP and UDP transport protocols.
List<EventLoopGroup> allEventLoopGroups = new ArrayList<EventLoopGroup>();
for (ProxyThreadPools threadPools : protocolThreadPools.values()) {
allEventLoopGroups.addAll(threadPools.getAllEventLoops());
}
for (EventLoopGroup group : allEventLoopGroups) {
if (graceful) {
group.shutdownGracefully();
} else {
group.shutdownGracefully(0, 0, TimeUnit.SECONDS);
}
}
if (graceful) {
for (EventLoopGroup group : allEventLoopGroups) {
try {
group.awaitTermination(60, TimeUnit.SECONDS);
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
log.warn("Interrupted while shutting down event loop");
}
}
}
log.debug("Done shutting down server group");
}
/**
* Retrieves the client-to-proxy acceptor thread pool for the specified protocol. Initializes the pool if it has not
* yet been initialized.
* <p>
* This method is thread-safe; no external locking is necessary.
*
* @param protocol transport protocol to retrieve the thread pool for
* @return the client-to-proxy acceptor thread pool
*/
public EventLoopGroup getClientToProxyAcceptorPoolForTransport(TransportProtocol protocol) {
return getThreadPoolsForProtocol(protocol).getClientToProxyAcceptorPool();
}
/**
* Retrieves the client-to-proxy acceptor worker pool for the specified protocol. Initializes the pool if it has not
* yet been initialized.
* <p>
* This method is thread-safe; no external locking is necessary.
*
* @param protocol transport protocol to retrieve the thread pool for
* @return the client-to-proxy worker thread pool
*/
public EventLoopGroup getClientToProxyWorkerPoolForTransport(TransportProtocol protocol) {
return getThreadPoolsForProtocol(protocol).getClientToProxyWorkerPool();
}
/**
* Retrieves the proxy-to-server worker thread pool for the specified protocol. Initializes the pool if it has not
* yet been initialized.
* <p>
* This method is thread-safe; no external locking is necessary.
*
* @param protocol transport protocol to retrieve the thread pool for
* @return the proxy-to-server worker thread pool
*/
public EventLoopGroup getProxyToServerWorkerPoolForTransport(TransportProtocol protocol) {
return getThreadPoolsForProtocol(protocol).getProxyToServerWorkerPool();
}
/**
* @return true if this ServerGroup has already been stopped
*/
public boolean isStopped() {
return stopped.get();
}
}
