Merge pull request #349 from rabbitmq/rabbitmq-java-client-341-no-auto-recovery-when-tcp-window-full

Trigger recovery if a socket write fails

Author: michaelklishin

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

@@ -81,6 +81,9 @@ public class ConnectionFactory implements Cloneable {
     /** The default network recovery interval: 5000 millis */
     public static final long   DEFAULT_NETWORK_RECOVERY_INTERVAL = 5000;
 
+    /** The default timeout for work pool enqueueing: no timeout */
+    public static final int    DEFAULT_WORK_POOL_TIMEOUT = -1;
+
     private static final String PREFERRED_TLS_PROTOCOL = "TLSv1.2";
 
     private static final String FALLBACK_TLS_PROTOCOL = "TLSv1";
@@ -138,6 +141,12 @@ public class ConnectionFactory implements Cloneable {
      */
     private boolean channelShouldCheckRpcResponseType = false;
 
+    /**
+     * Timeout in ms for work pool enqueuing.
+     * @since 4.5.0
+     */
+    private int workPoolTimeout = DEFAULT_WORK_POOL_TIMEOUT;
+
     /** @return the default host to use for connections */
     public String getHost() {
         return host;
@@ -974,6 +983,7 @@ public ConnectionParams params(ExecutorService consumerWorkServiceExecutor) {
         result.setHeartbeatExecutor(heartbeatExecutor);
         result.setChannelRpcTimeout(channelRpcTimeout);
         result.setChannelShouldCheckRpcResponseType(channelShouldCheckRpcResponseType);
+        result.setWorkPoolTimeout(workPoolTimeout);
         return result;
     }
 
@@ -1270,4 +1280,25 @@ public void setChannelShouldCheckRpcResponseType(boolean channelShouldCheckRpcRe
     public boolean isChannelShouldCheckRpcResponseType() {
         return channelShouldCheckRpcResponseType;
     }
+
+    /**
+     * Timeout (in ms) for work pool enqueueing.
+     * The {@link WorkPool} dispatches several types of responses
+     * from the broker (e.g. deliveries). A high-traffic
+     * client with slow consumers can exhaust the work pool and
+     * compromise the whole connection (by e.g. letting the broker
+     * saturate the receive TCP buffers). Setting a timeout
+     * would make the connection fail early and avoid hard-to-diagnose
+     * TCP connection failure. Note this shouldn't happen
+     * with clients that set appropriate QoS values.
+     * Default is no timeout.
+     * @param workPoolTimeout timeout in ms
+     */
+    public void setWorkPoolTimeout(int workPoolTimeout) {
+        this.workPoolTimeout = workPoolTimeout;
+    }
+
+    public int getWorkPoolTimeout() {
+        return workPoolTimeout;
+    }
 }

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

@@ -32,6 +32,7 @@
 import java.net.SocketTimeoutException;
 import java.util.*;
 import java.util.concurrent.*;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 final class Copyright {
     final static String COPYRIGHT="Copyright (c) 2007-2017 Pivotal Software, Inc.";
@@ -61,6 +62,12 @@ public class AMQConnection extends ShutdownNotifierComponent implements Connecti
     private final List<RecoveryCanBeginListener> recoveryCanBeginListeners =
             Collections.synchronizedList(new ArrayList<RecoveryCanBeginListener>());
 
+    private final ErrorOnWriteListener errorOnWriteListener;
+
+    private final int workPoolTimeout;
+
+    private final AtomicBoolean finalShutdownStarted = new AtomicBoolean(false);
+
     /**
      * Retrieve a copy of the default table of client properties that
      * will be sent to the server during connection startup. This
@@ -245,10 +252,17 @@ public AMQConnection(ConnectionParams params, FrameHandler frameHandler, Metrics
         this._inConnectionNegotiation = true; // we start out waiting for the first protocol response
 
         this.metricsCollector = metricsCollector;
+
+        this.errorOnWriteListener = params.getErrorOnWriteListener() != null ? params.getErrorOnWriteListener() :
+            new ErrorOnWriteListener() {
+                @Override
+                public void handle(Connection connection, IOException exception) { }
+        };
+        this.workPoolTimeout = params.getWorkPoolTimeout();
     }
 
     private void initializeConsumerWorkService() {
-        this._workService  = new ConsumerWorkService(consumerWorkServiceExecutor, threadFactory, shutdownTimeout);
+        this._workService  = new ConsumerWorkService(consumerWorkServiceExecutor, threadFactory, workPoolTimeout, shutdownTimeout);
     }
 
     private void initializeHeartbeatSender() {
@@ -556,7 +570,11 @@ public void writeFrame(Frame f) throws IOException {
      * Public API - flush the output buffers
      */
     public void flush() throws IOException {
-        _frameHandler.flush();
+        try {
+            _frameHandler.flush();
+        } catch (IOException ioe) {
+            this.errorOnWriteListener.handle(this, ioe);
+        }
     }
 
     private static int negotiatedMaxValue(int clientValue, int serverValue) {
@@ -575,15 +593,24 @@ private class MainLoop implements Runnable {
          */
         @Override
         public void run() {
+            boolean shouldDoFinalShutdown = true;
             try {
                 while (_running) {
                     Frame frame = _frameHandler.readFrame();
                     readFrame(frame);
                 }
             } catch (Throwable ex) {
-                handleFailure(ex);
+                if (ex instanceof InterruptedException) {
+                    // loop has been interrupted during shutdown,
+                    // no need to do it again
+                    shouldDoFinalShutdown = false;
+                } else {
+                    handleFailure(ex);
+                }
             } finally {
-                doFinalShutdown();
+                if (shouldDoFinalShutdown) {
+                    doFinalShutdown();
+                }
             }
         }
     }
@@ -594,6 +621,9 @@ public boolean handleReadFrame(Frame frame) {
             try {
                 readFrame(frame);
                 return true;
+            } catch (WorkPoolFullException e) {
+                // work pool is full, we propagate this one.
+                throw e;
             } catch (Throwable ex) {
                 try {
                     handleFailure(ex);
@@ -686,14 +716,33 @@ private void handleFailure(Throwable ex)  {
 
     /** private API */
     public void doFinalShutdown() {
-        _frameHandler.close();
-        _appContinuation.set(null);
-        notifyListeners();
-        // assuming that shutdown listeners do not do anything
-        // asynchronously, e.g. start new threads, this effectively
-        // guarantees that we only begin recovery when all shutdown
-        // listeners have executed
-        notifyRecoveryCanBeginListeners();
+        if (finalShutdownStarted.compareAndSet(false, true)) {
+            _frameHandler.close();
+            _appContinuation.set(null);
+            closeMainLoopThreadIfNecessary();
+            notifyListeners();
+            // assuming that shutdown listeners do not do anything
+            // asynchronously, e.g. start new threads, this effectively
+            // guarantees that we only begin recovery when all shutdown
+            // listeners have executed
+            notifyRecoveryCanBeginListeners();
+        }
+    }
+
+    private void closeMainLoopThreadIfNecessary() {
+        if (mainLoopReadThreadNotNull() && notInMainLoopThread()) {
+            if (this.mainLoopThread.isAlive()) {
+                this.mainLoopThread.interrupt();
+            }
+        }
+    }
+
+    private boolean notInMainLoopThread() {
+        return Thread.currentThread() != this.mainLoopThread;
+    }
+
+    private boolean mainLoopReadThreadNotNull() {
+        return this.mainLoopThread != null;
     }
 
     private void notifyRecoveryCanBeginListeners() {
@@ -879,7 +928,6 @@ private ShutdownSignalException startShutdown(Method reason,
         _heartbeatSender.shutdown();
 
         _channel0.processShutdownSignal(sse, !initiatedByApplication, notifyRpc);
-
         return sse;
     }
 

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

@@ -387,6 +387,9 @@ private void releaseChannel() {
                 if (callback != null) {
                     try {
                         this.dispatcher.handleCancel(callback, consumerTag);
+                    } catch (WorkPoolFullException e) {
+                        // couldn't enqueue in work pool, propagating
+                        throw e;
                     } catch (Throwable ex) {
                         getConnection().getExceptionHandler().handleConsumerException(this,
                                                                                       ex,
@@ -445,10 +448,13 @@ protected void processDelivery(Command command, Basic.Deliver method) {
             // in case a manual ack in the callback, the stats will be able to record the ack
             metricsCollector.consumedMessage(this, m.getDeliveryTag(), m.getConsumerTag());
             this.dispatcher.handleDelivery(callback,
-                                           m.getConsumerTag(),
-                                           envelope,
-                                           (BasicProperties) command.getContentHeader(),
-                                           command.getContentBody());
+                m.getConsumerTag(),
+                envelope,
+                (BasicProperties) command.getContentHeader(),
+                command.getContentBody());
+        } catch (WorkPoolFullException e) {
+            // couldn't enqueue in work pool, propagating
+            throw e;
         } catch (Throwable ex) {
             getConnection().getExceptionHandler().handleConsumerException(this,
                 ex,

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

@@ -44,6 +44,8 @@ public class ConnectionParams {
     private boolean topologyRecovery;
     private int channelRpcTimeout;
     private boolean channelShouldCheckRpcResponseType;
+    private ErrorOnWriteListener errorOnWriteListener;
+    private int workPoolTimeout = -1;
 
     private ExceptionHandler exceptionHandler;
     private ThreadFactory threadFactory;
@@ -213,4 +215,20 @@ public void setChannelRpcTimeout(int channelRpcTimeout) {
     public void setChannelShouldCheckRpcResponseType(boolean channelShouldCheckRpcResponseType) {
         this.channelShouldCheckRpcResponseType = channelShouldCheckRpcResponseType;
     }
+
+    public void setErrorOnWriteListener(ErrorOnWriteListener errorOnWriteListener) {
+        this.errorOnWriteListener = errorOnWriteListener;
+    }
+
+    public ErrorOnWriteListener getErrorOnWriteListener() {
+        return errorOnWriteListener;
+    }
+
+    public void setWorkPoolTimeout(int workPoolTimeout) {
+        this.workPoolTimeout = workPoolTimeout;
+    }
+
+    public int getWorkPoolTimeout() {
+        return workPoolTimeout;
+    }
 }

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

@@ -31,14 +31,18 @@ final public class ConsumerWorkService {
     private final WorkPool<Channel, Runnable> workPool;
     private final int shutdownTimeout;
 
-    public ConsumerWorkService(ExecutorService executor, ThreadFactory threadFactory, int shutdownTimeout) {
+    public ConsumerWorkService(ExecutorService executor, ThreadFactory threadFactory, int queueingTimeout, int shutdownTimeout) {
         this.privateExecutor = (executor == null);
         this.executor = (executor == null) ? Executors.newFixedThreadPool(DEFAULT_NUM_THREADS, threadFactory)
-                                           : executor;
-        this.workPool = new WorkPool<Channel, Runnable>();
+            : executor;
+        this.workPool = new WorkPool<Channel, Runnable>(queueingTimeout);
         this.shutdownTimeout = shutdownTimeout;
     }
 
+    public ConsumerWorkService(ExecutorService executor, ThreadFactory threadFactory, int shutdownTimeout) {
+        this(executor, threadFactory, -1, shutdownTimeout);
+    }
+
     public int getShutdownTimeout() {
         return shutdownTimeout;
     }

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

@@ -0,0 +1,35 @@
+// Copyright (c) 2018-Present Pivotal Software, Inc.  All rights reserved.
+//
+// This software, the RabbitMQ Java client library, is triple-licensed under the
+// Mozilla Public License 1.1 ("MPL"), the GNU General Public License version 2
+// ("GPL") and the Apache License version 2 ("ASL"). For the MPL, please see
+// LICENSE-MPL-RabbitMQ. For the GPL, please see LICENSE-GPL2.  For the ASL,
+// please see LICENSE-APACHE2.
+//
+// This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND,
+// either express or implied. See the LICENSE file for specific language governing
+// rights and limitations of this software.
+//
+// If you have any questions regarding licensing, please contact us at
+// [email protected].
+
+package com.rabbitmq.client.impl;
+
+import com.rabbitmq.client.Connection;
+
+import java.io.IOException;
+
+/**
+ * Listener called when a connection gets an error trying to write on the socket.
+ * This can be used to trigger connection recovery.
+ */
+public interface ErrorOnWriteListener {
+
+    /**
+     * Called when writing to the socket failed
+     * @param connection the owning connection instance
+     * @param exception the thrown exception
+     */
+    void handle(Connection connection, IOException exception) throws IOException;
+
+}

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

@@ -21,6 +21,9 @@
 import java.util.Iterator;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.Callable;
+import java.util.concurrent.TimeUnit;
 
 /**
  * <p>This is a generic implementation of the channels specification
@@ -61,6 +64,37 @@ public class WorkPool<K, W> {
     private final Map<K, VariableLinkedBlockingQueue<W>> pool = new HashMap<K, VariableLinkedBlockingQueue<W>>();
     /** Those keys which want limits to be removed. We do not limit queue size if this is non-empty. */
     private final Set<K> unlimited = new HashSet<K>();
+    private final EnqueueingCallback<W> enqueueingCallback;
+
+    public WorkPool(final int queueingTimeout) {
+        if (queueingTimeout > 0) {
+            this.enqueueingCallback = new EnqueueingCallback<W>() {
+                @Override
+                public void enqueue(BlockingQueue<W> queue, W item) {
+                    try {
+                        boolean offered = queue.offer(item, queueingTimeout, TimeUnit.MILLISECONDS);
+                        if (!offered) {
+                            throw new WorkPoolFullException("Could not enqueue in work pool after " + queueingTimeout + " ms.");
+                        }
+                    } catch (InterruptedException e) {
+                        Thread.currentThread();
+                    }
+                }
+            };
+        } else {
+            this.enqueueingCallback = new EnqueueingCallback<W>() {
+
+                @Override
+                public void enqueue(BlockingQueue<W> queue, W item) {
+                    try {
+                        queue.put(item);
+                    } catch (InterruptedException e) {
+                        Thread.currentThread().interrupt();
+                    }
+                }
+            };
+        }
+    }
 
     /**
      * Add client <code><b>key</b></code> to pool of item queues, with an empty queue.
@@ -178,11 +212,7 @@ public boolean addWorkItem(K key, W item) {
         }
         // The put operation may block. We need to make sure we are not holding the lock while that happens.
         if (queue != null) {
-            try {
-                queue.put(item);
-            } catch (InterruptedException e) {
-                Thread.currentThread().interrupt();
-            }
+            enqueueingCallback.enqueue(queue, item);
 
             synchronized (this) {
                 if (isDormant(key)) {
@@ -243,4 +273,10 @@ private K readyToInProgress() {
         }
         return key;
     }
+
+    private interface EnqueueingCallback<W> {
+
+        void enqueue(BlockingQueue<W> queue, W item);
+
+    }
 }