Clean NIO state automatically

Use NIO by default (except in SSLTests).

Fixes #11

Author: acogoluegnes

src/main/java/com/rabbitmq/client/ConnectionFactory.java

@@ -108,8 +108,9 @@ public class ConnectionFactory implements Cloneable {
 
     private MetricsCollector metricsCollector;
 
-    private boolean nio = false;
+    private boolean nio = true;
     private FrameHandlerFactory frameHandlerFactory;
+    private ExecutorService nioExecutor;
 
     /** @return the default host to use for connections */
     public String getHost() {
@@ -647,7 +648,11 @@ public MetricsCollector getMetricsCollector() {
     protected synchronized FrameHandlerFactory createFrameHandlerFactory() throws IOException {
         if(nio) {
             if(this.frameHandlerFactory == null) {
-                this.frameHandlerFactory = new SocketChannelFrameHandlerFactory(connectionTimeout, socketConf, isSSL());
+                if(this.nioExecutor == null) {
+                    this.frameHandlerFactory = new SocketChannelFrameHandlerFactory(connectionTimeout, socketConf, isSSL(), this.threadFactory);
+                } else {
+                    this.frameHandlerFactory = new SocketChannelFrameHandlerFactory(connectionTimeout, socketConf, isSSL(), this.nioExecutor);
+                }
             }
             return this.frameHandlerFactory;
         } else {
@@ -1034,4 +1039,8 @@ public void setNetworkRecoveryInterval(long networkRecoveryInterval) {
     public void setNio(boolean nio) {
         this.nio = nio;
     }
+
+    public void setNioExecutor(ExecutorService nioExecutor) {
+        this.nioExecutor = nioExecutor;
+    }
 }

src/main/java/com/rabbitmq/client/impl/AMQConnection.java

@@ -577,6 +577,14 @@ boolean handleReadFrame(Frame frame) {
         return false;
     }
 
+    public boolean isRunning() {
+        return _running;
+    }
+
+    public boolean hasBrokerInitiatedShutdown() {
+        return _brokerInitiatedShutdown;
+    }
+
     private void readFrame(Frame frame) throws IOException {
         if (frame != null) {
             _missedHeartbeats = 0;
@@ -646,7 +654,7 @@ private void handleFailure(Throwable ex)  {
         }
     }
 
-    private void doFinalShutdown() {
+    void doFinalShutdown() {
         _frameHandler.close();
         _appContinuation.set(null);
         notifyListeners();
@@ -798,7 +806,7 @@ public void run() {
                 // this releases the thread
             } finally {
                 _running = false;
-                _channel0.notifyOutstandingRpc(cause);
+                    _channel0.notifyOutstandingRpc(cause);
             }
         }
     }

src/main/java/com/rabbitmq/client/impl/ForgivingExceptionHandler.java

@@ -20,7 +20,6 @@
 
 import java.io.IOException;
 import java.net.ConnectException;
-import java.net.SocketException;
 
 /**
  * An implementation of {@link com.rabbitmq.client.ExceptionHandler} that does not
@@ -120,7 +119,9 @@ protected void log(String message, Throwable e) {
     }
 
     private boolean isSocketClosedOrConnectionReset(Throwable e) {
-        return e instanceof SocketException &&
-            ("Connection reset".equals(e.getMessage()) || "Socket closed".equals(e.getMessage()));
+        return e instanceof IOException &&
+            ("Connection reset".equals(e.getMessage()) || "Socket closed".equals(e.getMessage()) ||
+             "Connection reset by peer".equals(e.getMessage())
+            );
     }
 }

src/main/java/com/rabbitmq/client/impl/SocketChannelFrameHandlerFactory.java

@@ -32,10 +32,10 @@
 import java.util.Iterator;
 import java.util.Queue;
 import java.util.Set;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.LinkedBlockingQueue;
-import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.*;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 
 /**
  *
@@ -44,19 +44,31 @@ public class SocketChannelFrameHandlerFactory extends AbstractFrameHandlerFactor
 
     private static final Logger LOGGER = LoggerFactory.getLogger(SocketChannelFrameHandlerFactory.class);
 
-    private final SelectorState readSelectorState;
-    private final SelectorState writeSelectorState;
+    private SelectorState readSelectorState;
+    private SelectorState writeSelectorState;
 
-    // FIXME provide constructor with executorservice
-    private final ExecutorService executorService = null;
+    private final ExecutorService executorService;
 
-    private final ThreadFactory threadFactory = Executors.defaultThreadFactory();
+    private final ThreadFactory threadFactory;
 
-    public SocketChannelFrameHandlerFactory(int connectionTimeout, SocketConfigurator configurator, boolean ssl) throws IOException {
+    private final Lock stateLock = new ReentrantLock();
+
+    private final AtomicLong connectionCount = new AtomicLong();
+
+    private Thread readThread, writeThread;
+
+    private Future<?> readTask, writeTask;
+
+    public SocketChannelFrameHandlerFactory(int connectionTimeout, SocketConfigurator configurator, boolean ssl, ExecutorService executorService) throws IOException {
+        super(connectionTimeout, configurator, ssl);
+        this.executorService = executorService;
+        this.threadFactory = null;
+    }
+
+    public SocketChannelFrameHandlerFactory(int connectionTimeout, SocketConfigurator configurator, boolean ssl, ThreadFactory threadFactory) throws IOException {
         super(connectionTimeout, configurator, ssl);
-        this.readSelectorState = new SelectorState(Selector.open());
-        this.writeSelectorState = new SelectorState(Selector.open());
-        startIoLoops();
+        this.executorService = null;
+        this.threadFactory = threadFactory;
     }
 
     @Override
@@ -67,6 +79,15 @@ public FrameHandler create(Address addr) throws IOException {
         SocketChannel channel = SocketChannel.open();
         configurator.configure(channel.socket());
 
+        // lock
+        stateLock.lock();
+        try {
+            connectionCount.incrementAndGet();
+            initStateIfNecessary();
+        } finally {
+            stateLock.unlock();
+        }
+
         // FIXME handle connection failure
         channel.connect(address);
 
@@ -78,15 +99,68 @@ public FrameHandler create(Address addr) throws IOException {
         return frameHandler;
     }
 
+    protected boolean cleanUp() {
+        // get connection count
+        // lock
+        // if connection count has changed, do nothing
+        // if connection count hasn't changed, clean
+        long connectionCountNow = connectionCount.get();
+        stateLock.lock();
+        try {
+            if(connectionCountNow != connectionCount.get()) {
+                // a connection request has come in meanwhile, don't do anything
+                return false;
+            }
+
+            if(this.executorService == null) {
+                this.writeThread.interrupt();
+            } else {
+                boolean canceled = this.writeTask.cancel(true);
+                if(!canceled) {
+                    LOGGER.info("Could not stop write NIO task");
+                }
+            }
+
+            try {
+                readSelectorState.selector.close();
+            } catch (IOException e) {
+                LOGGER.warn("Could not close read selector: {}", e.getMessage());
+            }
+            try {
+                writeSelectorState.selector.close();
+            } catch (IOException e) {
+                LOGGER.warn("Could not close write selector: {}", e.getMessage());
+            }
+
+            this.readSelectorState = null;
+            this.writeSelectorState = null;
+
+        } finally {
+            stateLock.unlock();
+        }
+        return true;
+    }
+
+    protected void initStateIfNecessary() throws IOException {
+        if(this.readSelectorState == null) {
+            // create selectors
+            this.readSelectorState = new SelectorState(Selector.open());
+            this.writeSelectorState = new SelectorState(Selector.open());
+
+            // create threads/tasks
+            startIoLoops();
+        }
+    }
+
     protected void startIoLoops() {
         if(executorService == null) {
-            Thread readThread = Environment.newThread(threadFactory, new ReadLoop(readSelectorState), "rabbitmq-nio-read");
-            Thread writeThread = Environment.newThread(threadFactory, new WriteLoop(writeSelectorState), "rabbitmq-nio-write");
+            this.readThread = Environment.newThread(threadFactory, new ReadLoop(this.readSelectorState), "rabbitmq-nio-read");
+            this.writeThread = Environment.newThread(threadFactory, new WriteLoop(this.writeSelectorState), "rabbitmq-nio-write");
             readThread.start();
             writeThread.start();
         } else {
-            this.executorService.submit(new ReadLoop(readSelectorState));
-            this.executorService.submit(new WriteLoop(writeSelectorState));
+            this.readTask = this.executorService.submit(new ReadLoop(this.readSelectorState));
+            this.writeTask = this.executorService.submit(new WriteLoop(this.writeSelectorState));
         }
     }
 
@@ -104,9 +178,9 @@ public void run() {
             // FIXME find a better default?
             ByteBuffer buffer = ByteBuffer.allocate(8192);
             try {
+                int idlenessCount = 0;
                 while(true && !Thread.currentThread().isInterrupted()) {
 
-                    // if there's no read key anymore
                     for (SelectionKey selectionKey : selector.keys()) {
                         SocketChannelFrameHandlerState state = (SocketChannelFrameHandlerState) selectionKey.attachment();
                         // FIXME connection should always be here
@@ -123,18 +197,30 @@ public void run() {
                             }
                         }
 
+                        // FIXME really necessary? key are supposed to be removed when channel is closed
+                        /*
                         if(!selectionKey.channel().isOpen()) {
-                            // FIXME maybe better to wait for an exception and to trigger AMQConnection shutdown?
-                            // or check if AMQConnection is closed and init shutdown if appropriate
                             LOGGER.warn("Channel for connection {} closed, removing it from IO thread", state.getConnection());
                             selectionKey.cancel();
                         }
+                        */
                     }
 
                     int select;
                     if(state.statesToBeRegistered.isEmpty()) {
                         // we can block, registration will call Selector.wakeup()
                         select = selector.select(1000);
+                        idlenessCount++;
+                        if(idlenessCount == 10 && selector.keys().size() == 0) {
+                        //if(false) {
+                            // we haven't been doing anything for a while, shutdown state
+                            boolean clean = cleanUp();
+                            if(clean) {
+                                // we stop this thread
+                                return;
+                            }
+                            // there may be incoming connections, keep going
+                        }
                     } else {
                         // we don't have to block, we need to select and clean cancelled keys before registration
                         select = selector.selectNow();
@@ -149,6 +235,7 @@ public void run() {
                     }
 
                     if (select > 0) {
+                        idlenessCount = 0;
                         Set<SelectionKey> readyKeys = selector.selectedKeys();
                         Iterator<SelectionKey> iterator = readyKeys.iterator();
                         while (iterator.hasNext()) {
@@ -165,11 +252,19 @@ public void run() {
                                         Frame frame = Frame.readFrom(channel, buffer);
 
                                         try {
-                                            state.getConnection().handleReadFrame(frame);
+                                            boolean noProblem = state.getConnection().handleReadFrame(frame);
+                                            if(noProblem && (!state.getConnection().isRunning() || state.getConnection().hasBrokerInitiatedShutdown())) {
+                                                // looks like the frame was Close-Ok or Close
+                                                dispatchShutdownToConnection(state);
+                                                key.cancel();
+                                                break;
+                                            }
+
                                         } catch(Throwable ex) {
                                             // problem during frame processing, tell connection, and
                                             // we can stop for this channel
                                             handleIoError(state, ex);
+                                            key.cancel();
                                             break;
                                         }
 
@@ -231,7 +326,13 @@ public void run() {
                     RegistrationState registration;
                     while((registration = state.statesToBeRegistered.poll()) != null) {
                         int operations = registration.operations;
-                        registration.state.getChannel().register(selector, operations, registration.state);
+                        try {
+                            registration.state.getChannel().register(selector, operations, registration.state);
+                        } catch(Exception e) {
+                            // can happen if the channel has been closed since the operation has been enqueued
+                            LOGGER.info("Error while registering socket channel for write: {}", e.getMessage());
+                        }
+
                     }
 
                     if(select > 0) {
@@ -264,10 +365,11 @@ public void run() {
                                         frame.writeTo(channel, buffer);
                                         written++;
                                     }
+                                    Frame.drain(channel, buffer);
                                 } catch(Exception e) {
                                     handleIoError(state, e);
                                 } finally {
-                                    Frame.drain(channel, buffer);
+                                    buffer.clear();
                                     key.cancel();
                                 }
                             }
@@ -310,6 +412,22 @@ public void run() {
         }
     }
 
+    protected void dispatchShutdownToConnection(final SocketChannelFrameHandlerState state) {
+        Runnable shutdown = new Runnable() {
+            @Override
+            public void run() {
+                state.getConnection().doFinalShutdown();
+            }
+        };
+        if(this.executorService == null) {
+            String name = "rabbitmq-connection-shutdown-" + state.getConnection();
+            Thread shutdownThread = Environment.newThread(threadFactory, shutdown, name);
+            shutdownThread.start();
+        } else {
+            this.executorService.submit(shutdown);
+        }
+    }
+
     public static class RegistrationState {
 
         private final SocketChannelFrameHandlerState state;