Handle heartbeat and shutdown

Limit also written frames to the (snapshot) size when iteration starts (instead of
unqueuing during writing).

Fixes #11

Author: acogoluegnes

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

@@ -283,6 +283,7 @@ public void start()
 
         // FIXME properly prepare the connection depending on IO implementation
         if(this._frameHandler instanceof SocketChannelFrameHandler) {
+            // FIXME set the connection earlier
             ((SocketChannelFrameHandler) _frameHandler).getState().setConnection(this);
         } else {
             // start the main loop going
@@ -617,6 +618,17 @@ private void readFrame(Frame frame) throws IOException {
         }
     }
 
+    void handleHeartbeatFailure() {
+        Exception ex = new MissedHeartbeatException("Heartbeat missing with heartbeat = " +
+            _heartbeat + " seconds");
+        try {
+            _exceptionHandler.handleUnexpectedConnectionDriverException(this, ex);
+            shutdown(null, false, ex, true);
+        } finally {
+            doFinalShutdownInRead();
+        }
+    }
+
     private void handleFailureInRead(Throwable ex)  {
         if(ex instanceof EOFException) {
             if (!_brokerInitiatedShutdown)

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

@@ -296,7 +296,7 @@ public void writeTo(WritableByteChannel writableChannel, ByteBuffer buffer) thro
                 buffer.put(b);
             }
 
-            buffer.put((byte) AMQP.FRAME_END);
+            safePut(writableChannel, buffer, (byte) AMQP.FRAME_END);
         } else {
             safePut(writableChannel, buffer, (byte) type);
             safePut(writableChannel, buffer, (byte) ((channel >>> 8) & 0xFF));
@@ -337,7 +337,7 @@ private void safePut(WritableByteChannel channel, ByteBuffer buffer, byte conten
 
     public static void drain(WritableByteChannel channel, ByteBuffer buffer) throws IOException {
         buffer.flip();
-        while(buffer.hasRemaining() && channel.write(buffer) != 0);
+        while(buffer.hasRemaining() && channel.write(buffer) != -1);
         buffer.clear();
     }
 

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

@@ -32,8 +32,6 @@ public class SocketChannelFrameHandler implements FrameHandler {
 
     private final SocketChannelFrameHandlerState state;
 
-    private volatile int readTimeout = ConnectionFactory.DEFAULT_HEARTBEAT * 1000;
-
     public SocketChannelFrameHandler(SocketChannelFrameHandlerState state) {
         this.state = state;
     }
@@ -61,7 +59,6 @@ public int getPort() {
     @Override
     public void setTimeout(int timeoutMs) throws SocketException {
         state.getChannel().socket().setSoTimeout(timeoutMs);
-        this.readTimeout = timeoutMs;
     }
 
     @Override

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

@@ -32,9 +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.*;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 
 /**
  *
@@ -46,24 +47,35 @@ public class SocketChannelFrameHandlerFactory extends FrameHandlerFactory {
     private final SelectorState readSelectorState;
     private final SelectorState writeSelectorState;
 
-    private final ExecutorService executorService = Executors.newFixedThreadPool(2);
+    private ExecutorService executorService;
+
+    private Future<?> readLoop;
+    private Future<?> writeLoop;
+
+    private Lock loopsLock = new ReentrantLock();
+
+    private final ThreadFactory threadFactory = new ThreadFactory() {
+
+        AtomicLong counter = new AtomicLong();
+
+        @Override
+        public Thread newThread(Runnable r) {
+            return new Thread(r, "rabbitmq-nio-"+counter.getAndIncrement());
+        }
+    };
 
     public SocketChannelFrameHandlerFactory(int connectionTimeout, SocketFactory factory, SocketConfigurator configurator,
         boolean ssl) throws IOException {
         super(connectionTimeout, factory, configurator, ssl);
         this.readSelectorState = new SelectorState(Selector.open());
         this.writeSelectorState = new SelectorState(Selector.open());
-        this.executorService.submit(new ReadLoop(readSelectorState));
-        this.executorService.submit(new WriteLoop(writeSelectorState));
     }
 
     public SocketChannelFrameHandlerFactory(int connectionTimeout, SocketFactory factory, SocketConfigurator configurator, boolean ssl,
         ExecutorService shutdownExecutor) throws IOException {
         super(connectionTimeout, factory, configurator, ssl, shutdownExecutor);
         this.readSelectorState = new SelectorState(Selector.open());
         this.writeSelectorState = new SelectorState(Selector.open());
-        this.executorService.submit(new ReadLoop(readSelectorState));
-        this.executorService.submit(new WriteLoop(writeSelectorState));
     }
 
     @Override
@@ -79,7 +91,25 @@ public FrameHandler create(Address addr) throws IOException {
 
         SocketChannelFrameHandlerState state = new SocketChannelFrameHandlerState(channel, writeSelectorState);
 
-        readSelectorState.registerFrameHandlerState(state, SelectionKey.OP_READ);
+        loopsLock.lock();
+        try {
+            readSelectorState.registerFrameHandlerState(state, SelectionKey.OP_READ);
+
+            if(this.executorService == null) {
+
+                this.executorService = Executors.newFixedThreadPool(2, threadFactory);
+            }
+            if(readLoop == null) {
+                readLoop = this.executorService.submit(new ReadLoop(readSelectorState));
+            }
+            if(writeLoop == null) {
+                writeLoop = this.executorService.submit(new WriteLoop(writeSelectorState));
+            }
+        } finally {
+            loopsLock.unlock();
+        }
+
+
 
         SocketChannelFrameHandler frameHandler = new SocketChannelFrameHandler(state);
         return frameHandler;
@@ -102,7 +132,28 @@ public void registerFrameHandlerState(SocketChannelFrameHandlerState state, int
 
     }
 
-    private static class ReadLoop implements Runnable {
+    private boolean cleanLoopsOrKeepRunning() {
+        loopsLock.lock();
+        try {
+            if(readSelectorState.statesToBeRegistered.isEmpty()) {
+                boolean cancelled = writeLoop.cancel(true);
+                if(!cancelled) {
+                    LOGGER.warn("Could not stop write loop");
+                }
+                this.writeLoop = null;
+                this.readLoop = null;
+                this.executorService.shutdownNow();
+                return false;
+            } else {
+                // looks like someone is trying to connect, keep running
+                return true;
+            }
+        } finally {
+            loopsLock.unlock();
+        }
+    }
+
+    private class ReadLoop implements Runnable {
 
         private final SelectorState state;
 
@@ -117,16 +168,45 @@ public void run() {
             ByteBuffer buffer = ByteBuffer.allocate(8192);
             try {
                 while(true) {
+
+                    // if there's no read key anymore
+                    boolean someoneIsStillReading = false;
+                    for (SelectionKey selectionKey : selector.keys()) {
+                        SocketChannelFrameHandlerState state = (SocketChannelFrameHandlerState) selectionKey.attachment();
+                        // FIXME connection should always be here
+                        if(state.getConnection() != null && state.getConnection().getHeartbeat() > 0) {
+                            long now = System.currentTimeMillis();
+                            if((now - state.getLastActivity()) > state.getConnection().getHeartbeat() * 1000 * 2) {
+                                try {
+                                    state.getConnection().handleHeartbeatFailure();
+                                } catch(Exception e) {
+                                    LOGGER.warn("Error after heartbeat failure of connection {}", state.getConnection());
+                                } finally {
+                                    selectionKey.cancel();
+                                }
+                            }
+                        }
+
+                        if(!selectionKey.channel().isOpen()) {
+                            selectionKey.cancel();
+                        } else {
+                            someoneIsStillReading = true;
+                        }
+                    }
+
+                    if(!someoneIsStillReading && state.statesToBeRegistered.isEmpty()) {
+                        boolean keepRunning = cleanLoopsOrKeepRunning();
+                        if(!keepRunning) {
+                            return;
+                        }
+                    }
+
                     int select;
                     if(state.statesToBeRegistered.isEmpty()) {
                         // we can block, registration will call Selector.wakeup()
-
-                        // FIXME check the number of keys and stop the read and write loops
-                        // if there's no read key anymore
-
-                        select = selector.select();
+                        select = selector.select(1000);
                     } else {
-                        // we cannot block, we need to select and clean cancelled keys before registration
+                        // we don't have to block, we need to select and clean cancelled keys before registration
                         select = selector.selectNow();
                     }
 
@@ -153,11 +233,12 @@ public void run() {
                                     buffer.flip();
                                     while(buffer.hasRemaining()) {
                                         Frame frame = Frame.readFrom(channel, buffer);
+
                                         // FIXME the connection may not be there yet (to be checked)
                                         boolean handled = state.getConnection().handleReadFrame(frame);
 
                                         // problem during frame processing, the connection triggered shutdown
-                                        // we can stop
+                                        // we can stop for this channel
                                         if(!handled) {
                                             break;
                                         }
@@ -169,9 +250,10 @@ public void run() {
                                         }
 
                                     }
-
+                                    state.setLastActivity(System.currentTimeMillis());
                                 } catch (Exception e) {
-                                    LOGGER.warn("Error during reading frames: "+e.getMessage());
+                                    LOGGER.warn("Error during reading frames", e);
+                                    key.cancel();
                                 } finally {
                                     buffer.clear();
                                 }
@@ -203,7 +285,7 @@ public void run() {
             ByteBuffer buffer = ByteBuffer.allocate(8192);
 
             try {
-                while(true) {
+                while(true && !Thread.currentThread().isInterrupted()) {
                     int select;
                     if(state.statesToBeRegistered.isEmpty()) {
                         // we can block, registration will call Selector.wakeup()
@@ -228,12 +310,9 @@ public void run() {
                             SelectionKey key = iterator.next();
                             iterator.remove();
                             SocketChannel channel = (SocketChannel) key.channel();
-                            if (key.isConnectable()) {
-                                if (!channel.isConnected()) {
-                                    channel.finishConnect();
-                                }
-                            } else if (key.isWritable()) {
+                            if (key.isWritable()) {
                                 SocketChannelFrameHandlerState state = (SocketChannelFrameHandlerState) key.attachment();
+                                int toBeWritten = state.getWriteQueue().size();
                                 // FIXME property handle header sending request
                                 if(state.isSendHeader()) {
                                     buffer.put("AMQP".getBytes("US-ASCII"));
@@ -247,9 +326,11 @@ public void run() {
                                     state.setSendHeader(false);
                                 }
 
+                                int written = 0;
                                 Frame frame;
-                                while((frame = state.getWriteQueue().poll()) != null) {
+                                while(written <= toBeWritten && (frame = state.getWriteQueue().poll()) != null) {
                                     frame.writeTo(channel, buffer);
+                                    written++;
                                 }
                                 Frame.drain(channel, buffer);
                             }

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

@@ -3,11 +3,14 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.io.IOException;
 import java.nio.channels.SelectionKey;
 import java.nio.channels.SocketChannel;
 import java.util.Queue;
 import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
 
 /**
  *
@@ -19,12 +22,15 @@ public class SocketChannelFrameHandlerState {
     private final SocketChannel channel;
 
     // FIXME find appropriate default for limit in write queue
-    private final BlockingQueue<Frame> writeQueue = new ArrayBlockingQueue<Frame>(1000);
+    private final BlockingQueue<Frame> writeQueue = new LinkedBlockingQueue<Frame>(10000);
 
     private volatile AMQConnection connection;
 
     private volatile boolean sendHeader = false;
 
+    /** should be used only in the NIO read thread */
+    private long lastActivity;
+
     private final SocketChannelFrameHandlerFactory.SelectorState selectorState;
 
     public SocketChannelFrameHandlerState(SocketChannel channel, SocketChannelFrameHandlerFactory.SelectorState selectorState) {
@@ -51,10 +57,15 @@ public void setSendHeader(boolean sendHeader) {
         }
     }
 
-    public void write(Frame frame) {
+    public void write(Frame frame) throws IOException {
         try {
-            this.writeQueue.put(frame);
-            this.selectorState.registerFrameHandlerState(this, SelectionKey.OP_WRITE);
+            boolean offered = this.writeQueue.offer(frame, 10, TimeUnit.SECONDS);
+            if(offered) {
+                this.selectorState.registerFrameHandlerState(this, SelectionKey.OP_WRITE);
+            } else {
+                throw new IOException("Frame enqueuing failed");
+            }
+
         } catch (InterruptedException e) {
             LOGGER.warn("Thread interrupted during enqueuing frame in write queue");
         }
@@ -67,4 +78,12 @@ public AMQConnection getConnection() {
     public void setConnection(AMQConnection connection) {
         this.connection = connection;
     }
+
+    public void setLastActivity(long lastActivity) {
+        this.lastActivity = lastActivity;
+    }
+
+    public long getLastActivity() {
+        return lastActivity;
+    }
 }