Shutdown NIO loop ASAP

Fixes #11

Author: acogoluegnes

src/main/java/com/rabbitmq/client/impl/nio/NioLoop.java

@@ -66,12 +66,11 @@ public void run() {
         boolean writeRegistered = false;
 
         try {
-            int idlenessCount = 0;
             while (true && !Thread.currentThread().isInterrupted()) {
 
                 for (SelectionKey selectionKey : selector.keys()) {
                     SocketChannelFrameHandlerState state = (SocketChannelFrameHandlerState) selectionKey.attachment();
-                    if (state.getConnection().getHeartbeat() > 0) {
+                    if (state.getConnection() != null && state.getConnection().getHeartbeat() > 0) {
                         long now = System.currentTimeMillis();
                         if ((now - state.getLastActivity()) > state.getConnection().getHeartbeat() * 1000 * 2) {
                             try {
@@ -89,8 +88,7 @@ public void run() {
                 if (!writeRegistered && registrations.isEmpty() && writeRegistrations.isEmpty()) {
                     // we can block, registrations will call Selector.wakeup()
                     select = selector.select(1000);
-                    idlenessCount++;
-                    if (idlenessCount == 10 && selector.keys().size() == 0) {
+                    if (selector.keys().size() == 0) {
                         // we haven't been doing anything for a while, shutdown state
                         boolean clean = context.cleanUp();
                         if (clean) {
@@ -118,7 +116,6 @@ public void run() {
                 }
 
                 if (select > 0) {
-                    idlenessCount = 0;
                     Set<SelectionKey> readyKeys = selector.selectedKeys();
                     Iterator<SelectionKey> iterator = readyKeys.iterator();
                     while (iterator.hasNext()) {
@@ -246,7 +243,7 @@ public void run() {
                 }
             }
         } catch (Exception e) {
-            LOGGER.error("Error in read loop", e);
+            LOGGER.error("Error in NIO loop", e);
         }
     }
 

src/main/java/com/rabbitmq/client/impl/nio/NioLoopContext.java

@@ -9,7 +9,6 @@
 import java.nio.channels.Selector;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ThreadFactory;
-import java.util.concurrent.atomic.AtomicLong;
 
 /**
  *
@@ -29,8 +28,6 @@ public class NioLoopContext {
     SelectorHolder readSelectorState;
     SelectorHolder writeSelectorState;
 
-    private final AtomicLong nioLoopsConnectionCount = new AtomicLong();
-
     public NioLoopContext(SocketChannelFrameHandlerFactory socketChannelFrameHandlerFactory,
         NioParams nioParams) {
         this.socketChannelFrameHandlerFactory = socketChannelFrameHandlerFactory;
@@ -40,10 +37,6 @@ public NioLoopContext(SocketChannelFrameHandlerFactory socketChannelFrameHandler
         this.writeBuffer = ByteBuffer.allocate(nioParams.getWriteByteBufferSize());
     }
 
-    void notifyNewConnection() {
-        nioLoopsConnectionCount.incrementAndGet();
-    }
-
     void initStateIfNecessary() throws IOException {
         if (this.readSelectorState == null) {
             this.readSelectorState = new SelectorHolder(Selector.open());
@@ -67,10 +60,13 @@ private void startIoLoops() {
     }
 
     protected boolean cleanUp() {
-        long connectionCountNow = nioLoopsConnectionCount.get();
+        int readRegistrationsCount = readSelectorState.registrations.size();
+        if(readRegistrationsCount != 0) {
+            return false;
+        }
         socketChannelFrameHandlerFactory.lock();
         try {
-            if (connectionCountNow != nioLoopsConnectionCount.get()) {
+            if (readRegistrationsCount != readSelectorState.registrations.size()) {
                 // a connection request has come in meanwhile, don't do anything
                 return false;
             }

src/main/java/com/rabbitmq/client/impl/nio/SocketChannelFrameHandler.java

@@ -76,7 +76,6 @@ public void sendHeader() throws IOException {
     @Override
     public void initialize(AMQConnection connection) {
         state.setConnection(connection);
-        state.startReading();
     }
 
     @Override

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

@@ -98,20 +98,20 @@ public FrameHandler create(Address addr) throws IOException {
                 long modulo = globalConnectionCount.getAndIncrement() % nioParams.getNbIoThreads();
                 nioLoopContext = nioLoopContexts.get((int) modulo);
                 nioLoopContext.initStateIfNecessary();
-                nioLoopContext.notifyNewConnection();
+                SocketChannelFrameHandlerState state = new SocketChannelFrameHandlerState(
+                    channel,
+                    nioLoopContext,
+                    nioParams,
+                    sslEngine
+                );
+                state.startReading();
+                SocketChannelFrameHandler frameHandler = new SocketChannelFrameHandler(state);
+                return frameHandler;
             } finally {
                 stateLock.unlock();
             }
 
-            SocketChannelFrameHandlerState state = new SocketChannelFrameHandlerState(
-                channel,
-                nioLoopContext,
-                nioParams,
-                sslEngine
-            );
 
-            SocketChannelFrameHandler frameHandler = new SocketChannelFrameHandler(state);
-            return frameHandler;
         } catch(IOException e) {
             try {
                 if(sslEngine != null && channel != null) {

src/test/java/com/rabbitmq/client/test/JavaNioTest.java

@@ -92,6 +92,16 @@ public void shutdownCompleted(ShutdownSignalException cause) {
         }
     }
 
+    @Test
+    public void nioLoopCleaning() throws Exception {
+        ConnectionFactory connectionFactory = new ConnectionFactory();
+        connectionFactory.useNio();
+        for(int i = 0; i < 10; i++) {
+            Connection connection = connectionFactory.newConnection();
+            connection.abort();
+        }
+    }
+
     private Connection basicGetBasicConsume(ConnectionFactory connectionFactory, String queue, final CountDownLatch latch)
         throws IOException, TimeoutException {
         Connection connection = connectionFactory.newConnection();