Java tutorial
/** * Copyright 2016 Yahoo 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.yahoo.pulsar.client.impl; import java.io.Closeable; import java.io.File; import java.io.IOException; import java.net.InetSocketAddress; import java.security.cert.X509Certificate; import java.util.Random; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import org.apache.commons.lang3.SystemUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import com.yahoo.pulsar.client.api.AuthenticationDataProvider; import com.yahoo.pulsar.client.api.ClientConfiguration; import com.yahoo.pulsar.client.api.PulsarClientException; import com.yahoo.pulsar.common.api.PulsarLengthFieldFrameDecoder; import io.netty.bootstrap.Bootstrap; import io.netty.buffer.PooledByteBufAllocator; import io.netty.channel.ChannelFuture; import io.netty.channel.ChannelInitializer; import io.netty.channel.ChannelOption; import io.netty.channel.EventLoopGroup; import io.netty.channel.epoll.EpollEventLoopGroup; import io.netty.channel.epoll.EpollSocketChannel; import io.netty.channel.socket.SocketChannel; import io.netty.channel.socket.nio.NioSocketChannel; import io.netty.handler.ssl.SslContext; import io.netty.handler.ssl.SslContextBuilder; import io.netty.handler.ssl.util.InsecureTrustManagerFactory; public class ConnectionPool implements Closeable { private final ConcurrentHashMap<InetSocketAddress, ConcurrentMap<Integer, CompletableFuture<ClientCnx>>> pool; private final Bootstrap bootstrap; private final EventLoopGroup eventLoopGroup; private final int maxConnectionsPerHosts; private static final int MaxMessageSize = 5 * 1024 * 1024; public static final String TLS_HANDLER = "tls"; public ConnectionPool(final PulsarClientImpl client, EventLoopGroup eventLoopGroup) { this.eventLoopGroup = eventLoopGroup; this.maxConnectionsPerHosts = client.getConfiguration().getConnectionsPerBroker(); pool = new ConcurrentHashMap<>(); bootstrap = new Bootstrap(); bootstrap.group(eventLoopGroup); if (SystemUtils.IS_OS_LINUX && eventLoopGroup instanceof EpollEventLoopGroup) { bootstrap.channel(EpollSocketChannel.class); } else { bootstrap.channel(NioSocketChannel.class); } bootstrap.option(ChannelOption.CONNECT_TIMEOUT_MILLIS, 10000); bootstrap.option(ChannelOption.TCP_NODELAY, client.getConfiguration().isUseTcpNoDelay()); bootstrap.option(ChannelOption.ALLOCATOR, PooledByteBufAllocator.DEFAULT); bootstrap.handler(new ChannelInitializer<SocketChannel>() { public void initChannel(SocketChannel ch) throws Exception { ClientConfiguration clientConfig = client.getConfiguration(); if (clientConfig.isUseTls()) { SslContextBuilder builder = SslContextBuilder.forClient(); if (clientConfig.isTlsAllowInsecureConnection()) { builder.trustManager(InsecureTrustManagerFactory.INSTANCE); } else { if (clientConfig.getTlsTrustCertsFilePath().isEmpty()) { // Use system default builder.trustManager((File) null); } else { File trustCertCollection = new File(clientConfig.getTlsTrustCertsFilePath()); builder.trustManager(trustCertCollection); } } // Set client certificate if available AuthenticationDataProvider authData = clientConfig.getAuthentication().getAuthData(); if (authData.hasDataForTls()) { builder.keyManager(authData.getTlsPrivateKey(), (X509Certificate[]) authData.getTlsCertificates()); } SslContext sslCtx = builder.build(); ch.pipeline().addLast(TLS_HANDLER, sslCtx.newHandler(ch.alloc())); } ch.pipeline().addLast("frameDecoder", new PulsarLengthFieldFrameDecoder(MaxMessageSize, 0, 4, 0, 4)); ch.pipeline().addLast("handler", new ClientCnx(client)); } }); } private static final Random random = new Random(); public CompletableFuture<ClientCnx> getConnection(final InetSocketAddress address) { if (maxConnectionsPerHosts == 0) { // Disable pooling return createConnection(address, -1); } final int randomKey = signSafeMod(random.nextInt(), maxConnectionsPerHosts); return pool.computeIfAbsent(address, a -> new ConcurrentHashMap<>()) // .computeIfAbsent(randomKey, k -> createConnection(address, randomKey)); } private CompletableFuture<ClientCnx> createConnection(InetSocketAddress address, int connectionKey) { if (log.isDebugEnabled()) { log.debug("Connection for {} not found in cache", address); } final CompletableFuture<ClientCnx> cnxFuture = new CompletableFuture<ClientCnx>(); // Trigger async connect to broker bootstrap.connect(address).addListener((ChannelFuture future) -> { if (!future.isSuccess()) { cnxFuture.completeExceptionally(new PulsarClientException(future.cause())); cleanupConnection(address, connectionKey, cnxFuture); return; } log.info("[{}] Connected to server", future.channel()); future.channel().closeFuture().addListener(v -> { // Remove connection from pool when it gets closed cleanupConnection(address, connectionKey, cnxFuture); }); // We are connected to broker, but need to wait until the connect/connected handshake is // complete final ClientCnx cnx = (ClientCnx) future.channel().pipeline().get("handler"); cnx.connectionFuture().thenRun(() -> { if (log.isDebugEnabled()) { log.debug("[{}] Connection handshake completed", cnx.channel()); } cnxFuture.complete(cnx); }).exceptionally(exception -> { log.warn("[{}] Connection handshake failed: {}", cnx.channel(), exception.getMessage()); cnxFuture.completeExceptionally(exception); cleanupConnection(address, connectionKey, cnxFuture); cnx.ctx().close(); return null; }); }); return cnxFuture; } @Override public void close() throws IOException { eventLoopGroup.shutdownGracefully(); } private void cleanupConnection(InetSocketAddress address, int connectionKey, CompletableFuture<ClientCnx> connectionFuture) { ConcurrentMap<Integer, CompletableFuture<ClientCnx>> map = pool.get(address); if (map != null) { map.remove(connectionKey, connectionFuture); } } public static int signSafeMod(long dividend, int divisor) { int mod = (int) (dividend % (long) divisor); if (mod < 0) { mod += divisor; } return mod; } private static final Logger log = LoggerFactory.getLogger(ConnectionPool.class); }