Java tutorial
/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You 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 org.apache.nifi.processor.util.listen.dispatcher; import org.apache.commons.io.IOUtils; import org.apache.nifi.logging.ComponentLog; import org.apache.nifi.processor.util.listen.event.Event; import org.apache.nifi.processor.util.listen.event.EventFactory; import org.apache.nifi.processor.util.listen.handler.ChannelHandlerFactory; import org.apache.nifi.remote.io.socket.ssl.SSLSocketChannel; import org.apache.nifi.security.util.SslContextFactory; import javax.net.ssl.SSLContext; import javax.net.ssl.SSLEngine; import java.io.IOException; import java.net.InetAddress; import java.net.InetSocketAddress; import java.net.StandardSocketOptions; import java.nio.ByteBuffer; import java.nio.channels.Channel; import java.nio.channels.SelectionKey; import java.nio.channels.Selector; import java.nio.channels.ServerSocketChannel; import java.nio.channels.SocketChannel; import java.nio.charset.Charset; import java.util.Iterator; import java.util.concurrent.BlockingQueue; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; /** * Accepts Socket connections on the given port and creates a handler for each connection to * be executed by a thread pool. */ public class SocketChannelDispatcher<E extends Event<SocketChannel>> implements AsyncChannelDispatcher { private final EventFactory<E> eventFactory; private final ChannelHandlerFactory<E, AsyncChannelDispatcher> handlerFactory; private final BlockingQueue<ByteBuffer> bufferPool; private final BlockingQueue<E> events; private final ComponentLog logger; private final int maxConnections; private final SSLContext sslContext; private final SslContextFactory.ClientAuth clientAuth; private final Charset charset; private ExecutorService executor; private volatile boolean stopped = false; private Selector selector; private final BlockingQueue<SelectionKey> keyQueue; private final AtomicInteger currentConnections = new AtomicInteger(0); public SocketChannelDispatcher(final EventFactory<E> eventFactory, final ChannelHandlerFactory<E, AsyncChannelDispatcher> handlerFactory, final BlockingQueue<ByteBuffer> bufferPool, final BlockingQueue<E> events, final ComponentLog logger, final int maxConnections, final SSLContext sslContext, final Charset charset) { this(eventFactory, handlerFactory, bufferPool, events, logger, maxConnections, sslContext, SslContextFactory.ClientAuth.REQUIRED, charset); } public SocketChannelDispatcher(final EventFactory<E> eventFactory, final ChannelHandlerFactory<E, AsyncChannelDispatcher> handlerFactory, final BlockingQueue<ByteBuffer> bufferPool, final BlockingQueue<E> events, final ComponentLog logger, final int maxConnections, final SSLContext sslContext, final SslContextFactory.ClientAuth clientAuth, final Charset charset) { this.eventFactory = eventFactory; this.handlerFactory = handlerFactory; this.bufferPool = bufferPool; this.events = events; this.logger = logger; this.maxConnections = maxConnections; this.keyQueue = new LinkedBlockingQueue<>(maxConnections); this.sslContext = sslContext; this.clientAuth = clientAuth; this.charset = charset; if (bufferPool == null || bufferPool.size() == 0 || bufferPool.size() != maxConnections) { throw new IllegalArgumentException( "A pool of available ByteBuffers equal to the maximum number of connections is required"); } } @Override public void open(final InetAddress nicAddress, final int port, final int maxBufferSize) throws IOException { stopped = false; executor = Executors.newFixedThreadPool(maxConnections); final ServerSocketChannel serverSocketChannel = ServerSocketChannel.open(); serverSocketChannel.configureBlocking(false); if (maxBufferSize > 0) { serverSocketChannel.setOption(StandardSocketOptions.SO_RCVBUF, maxBufferSize); final int actualReceiveBufSize = serverSocketChannel.getOption(StandardSocketOptions.SO_RCVBUF); if (actualReceiveBufSize < maxBufferSize) { logger.warn("Attempted to set Socket Buffer Size to " + maxBufferSize + " bytes but could only set to " + actualReceiveBufSize + "bytes. You may want to consider changing the Operating System's " + "maximum receive buffer"); } } serverSocketChannel.socket().bind(new InetSocketAddress(nicAddress, port)); selector = Selector.open(); serverSocketChannel.register(selector, SelectionKey.OP_ACCEPT); } @Override public void run() { while (!stopped) { try { int selected = selector.select(); // if stopped the selector could already be closed which would result in a ClosedSelectorException if (selected > 0 && !stopped) { Iterator<SelectionKey> selectorKeys = selector.selectedKeys().iterator(); // if stopped we don't want to modify the keys because close() may still be in progress while (selectorKeys.hasNext() && !stopped) { SelectionKey key = selectorKeys.next(); selectorKeys.remove(); if (!key.isValid()) { continue; } if (key.isAcceptable()) { // Handle new connections coming in final ServerSocketChannel channel = (ServerSocketChannel) key.channel(); final SocketChannel socketChannel = channel.accept(); // Check for available connections if (currentConnections.incrementAndGet() > maxConnections) { currentConnections.decrementAndGet(); logger.warn("Rejecting connection from {} because max connections has been met", new Object[] { socketChannel.getRemoteAddress().toString() }); IOUtils.closeQuietly(socketChannel); continue; } logger.debug("Accepted incoming connection from {}", new Object[] { socketChannel.getRemoteAddress().toString() }); // Set socket to non-blocking, and register with selector socketChannel.configureBlocking(false); SelectionKey readKey = socketChannel.register(selector, SelectionKey.OP_READ); // Prepare the byte buffer for the reads, clear it out ByteBuffer buffer = bufferPool.poll(); buffer.clear(); buffer.mark(); // If we have an SSLContext then create an SSLEngine for the channel SSLSocketChannel sslSocketChannel = null; if (sslContext != null) { final SSLEngine sslEngine = sslContext.createSSLEngine(); sslEngine.setUseClientMode(false); switch (clientAuth) { case REQUIRED: sslEngine.setNeedClientAuth(true); break; case WANT: sslEngine.setWantClientAuth(true); break; case NONE: sslEngine.setNeedClientAuth(false); sslEngine.setWantClientAuth(false); break; } sslSocketChannel = new SSLSocketChannel(sslEngine, socketChannel); } // Attach the buffer and SSLSocketChannel to the key SocketChannelAttachment attachment = new SocketChannelAttachment(buffer, sslSocketChannel); readKey.attach(attachment); } else if (key.isReadable()) { // Clear out the operations the select is interested in until done reading key.interestOps(0); // Create a handler based on the protocol and whether an SSLEngine was provided or not final Runnable handler; if (sslContext != null) { handler = handlerFactory.createSSLHandler(key, this, charset, eventFactory, events, logger); } else { handler = handlerFactory.createHandler(key, this, charset, eventFactory, events, logger); } // run the handler executor.execute(handler); } } } // Add back all idle sockets to the select SelectionKey key; while ((key = keyQueue.poll()) != null) { key.interestOps(SelectionKey.OP_READ); } } catch (IOException e) { logger.error("Error accepting connection from SocketChannel", e); } } } @Override public int getPort() { // Return the port for the key listening for accepts for (SelectionKey key : selector.keys()) { if (key.isValid()) { final Channel channel = key.channel(); if (channel instanceof ServerSocketChannel) { return ((ServerSocketChannel) channel).socket().getLocalPort(); } } } return 0; } @Override public void close() { stopped = true; if (selector != null) { selector.wakeup(); } if (executor != null) { executor.shutdown(); try { // Wait a while for existing tasks to terminate if (!executor.awaitTermination(1000L, TimeUnit.MILLISECONDS)) { executor.shutdownNow(); } } catch (InterruptedException ie) { // (Re-)Cancel if current thread also interrupted executor.shutdownNow(); // Preserve interrupt status Thread.currentThread().interrupt(); } } if (selector != null) { synchronized (selector.keys()) { for (SelectionKey key : selector.keys()) { IOUtils.closeQuietly(key.channel()); } } } IOUtils.closeQuietly(selector); } @Override public void completeConnection(SelectionKey key) { // connection is done. Return the buffer to the pool SocketChannelAttachment attachment = (SocketChannelAttachment) key.attachment(); try { bufferPool.put(attachment.getByteBuffer()); } catch (InterruptedException e) { Thread.currentThread().interrupt(); } currentConnections.decrementAndGet(); } @Override public void addBackForSelection(SelectionKey key) { keyQueue.offer(key); selector.wakeup(); } }