Make the shutdown timeout configurable.

Author: Simon MacMullen

src/com/rabbitmq/client/ConnectionFactory.java

@@ -75,6 +75,9 @@ public class ConnectionFactory implements Cloneable {
     /** The default connection timeout;
      *  zero means wait indefinitely */
     public static final int    DEFAULT_CONNECTION_TIMEOUT = 0;
+    /** The default shutdown timeout;
+     *  zero means wait indefinitely */
+    public static final int    DEFAULT_SHUTDOWN_TIMEOUT = 10000;
 
     /** The default SSL protocol */
     private static final String DEFAULT_SSL_PROTOCOL = "SSLv3";
@@ -88,6 +91,7 @@ public class ConnectionFactory implements Cloneable {
     private int requestedFrameMax                 = DEFAULT_FRAME_MAX;
     private int requestedHeartbeat                = DEFAULT_HEARTBEAT;
     private int connectionTimeout                 = DEFAULT_CONNECTION_TIMEOUT;
+    private int shutdownTimeout                   = DEFAULT_SHUTDOWN_TIMEOUT;
     private Map<String, Object> _clientProperties = AMQConnection.defaultClientProperties();
     private SocketFactory factory                 = SocketFactory.getDefault();
     private SaslConfig saslConfig                 = DefaultSaslConfig.PLAIN;
@@ -333,6 +337,26 @@ public int getConnectionTimeout() {
         return this.connectionTimeout;
     }
 
+    /**
+     * Set the shutdown timeout. This is the amount of time that Consumer implementations have to
+     * continue working through deliveries (and other Consumer callbacks) <b>after</b> the connection
+     * has closed but before the ConsumerWorkService is torn down. If consumers exceed this timeout
+     * then any remaining queued deliveries (and other Consumer callbacks, <b>including</b>
+     * the Consumer's handleShutdownSignal() invocation) will be lost.
+     * @param shutdownTimeout shutdown timeout in milliseconds; zero for infinite; default 10000
+     */
+    public void setShutdownTimeout(int shutdownTimeout) {
+        this.shutdownTimeout = shutdownTimeout;
+    }
+
+    /**
+     * Retrieve the shutdown timeout.
+     * @return the shutdown timeout, in milliseconds; zero for infinite
+     */
+    public int getShutdownTimeout() {
+        return shutdownTimeout;
+    }
+
     /**
      * Set the requested heartbeat.
      * @param requestedHeartbeat the initially requested heartbeat interval, in seconds; zero for none
@@ -602,7 +626,7 @@ public Connection newConnection(ExecutorService executor, Address[] addrs)
 
     public ConnectionParams params(ExecutorService executor) {
         return new ConnectionParams(username, password, executor, virtualHost, getClientProperties(),
-                                    requestedFrameMax, requestedChannelMax, requestedHeartbeat, saslConfig,
+                                    requestedFrameMax, requestedChannelMax, requestedHeartbeat, shutdownTimeout, saslConfig,
                                     networkRecoveryInterval, topologyRecovery, exceptionHandler, threadFactory);
     }
 

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

@@ -139,6 +139,7 @@ public static final Map<String, Object> defaultClientProperties() {
     private final int requestedHeartbeat;
     private final int requestedChannelMax;
     private final int requestedFrameMax;
+    private final int shutdownTimeout;
     private final String username;
     private final String password;
     private final Collection<BlockedListener> blockedListeners = new CopyOnWriteArrayList<BlockedListener>();
@@ -217,6 +218,7 @@ public AMQConnection(ConnectionParams params, FrameHandler frameHandler)
         this.requestedFrameMax = params.getRequestedFrameMax();
         this.requestedChannelMax = params.getRequestedChannelMax();
         this.requestedHeartbeat = params.getRequestedHeartbeat();
+        this.shutdownTimeout = params.getShutdownTimeout();
         this.saslConfig = params.getSaslConfig();
         this.executor = params.getExecutor();
         this.threadFactory = params.getThreadFactory();
@@ -229,7 +231,7 @@ public AMQConnection(ConnectionParams params, FrameHandler frameHandler)
     }
 
     private void initializeConsumerWorkService() {
-        this._workService  = new ConsumerWorkService(executor, threadFactory);
+        this._workService  = new ConsumerWorkService(executor, threadFactory, shutdownTimeout);
     }
 
     private void initializeHeartbeatSender() {

src/com/rabbitmq/client/impl/ChannelManager.java

@@ -33,8 +33,6 @@
  * Manages a set of channels, indexed by channel number (<code><b>1.._channelMax</b></code>).
  */
 public class ChannelManager {
-    private static final int SHUTDOWN_TIMEOUT_SECONDS = 10;
-
     /** Monitor for <code>_channelMap</code> and <code>channelNumberAllocator</code> */
     private final Object monitor = new Object();
         /** Mapping from <code><b>1.._channelMax</b></code> to {@link ChannelN} instance */
@@ -108,7 +106,13 @@ private void scheduleShutdownProcessing() {
         Runnable target = new Runnable() {
             public void run() {
                 for (CountDownLatch latch : sdSet) {
-                    try { latch.await(SHUTDOWN_TIMEOUT_SECONDS, TimeUnit.SECONDS); } catch (Throwable e) { /*ignored*/ }
+                    try {
+                        int shutdownTimeout = ssWorkService.getShutdownTimeout();
+                        if (shutdownTimeout == 0) latch.await();
+                        else                      latch.await(shutdownTimeout, TimeUnit.MILLISECONDS);
+                    } catch (Throwable e) {
+                         /*ignored*/
+                    }
                 }
                 ssWorkService.shutdown();
             }

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

@@ -16,6 +16,7 @@ public class ConnectionParams {
     private final int requestedFrameMax;
     private final int requestedChannelMax;
     private final int requestedHeartbeat;
+    private final int shutdownTimeout;
     private final SaslConfig saslConfig;
     private final long networkRecoveryInterval;
     private final boolean topologyRecovery;
@@ -41,7 +42,7 @@ public class ConnectionParams {
     public ConnectionParams(String username, String password, ExecutorService executor,
                             String virtualHost, Map<String, Object> clientProperties,
                             int requestedFrameMax, int requestedChannelMax, int requestedHeartbeat,
-                            SaslConfig saslConfig, long networkRecoveryInterval,
+                            int shutdownTimeout, SaslConfig saslConfig, long networkRecoveryInterval,
                             boolean topologyRecovery, ExceptionHandler exceptionHandler, ThreadFactory threadFactory) {
         this.username = username;
         this.password = password;
@@ -51,6 +52,7 @@ public ConnectionParams(String username, String password, ExecutorService execut
         this.requestedFrameMax = requestedFrameMax;
         this.requestedChannelMax = requestedChannelMax;
         this.requestedHeartbeat = requestedHeartbeat;
+        this.shutdownTimeout = shutdownTimeout;
         this.saslConfig = saslConfig;
         this.networkRecoveryInterval = networkRecoveryInterval;
         this.topologyRecovery = topologyRecovery;
@@ -90,6 +92,10 @@ public int getRequestedHeartbeat() {
         return requestedHeartbeat;
     }
 
+    public int getShutdownTimeout() {
+        return shutdownTimeout;
+    }
+
     public SaslConfig getSaslConfig() {
         return saslConfig;
     }

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

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

test/src/com/rabbitmq/client/test/server/ChannelLimitNegotiation.java

@@ -80,7 +80,8 @@ public void testOpeningTooManyChannels() throws Exception {
 
             // Construct a channel directly
             final ChannelN ch = new ChannelN((AMQConnection) conn, n + 1,
-                                             new ConsumerWorkService(Executors.newSingleThreadExecutor(), Executors.defaultThreadFactory()));
+                                             new ConsumerWorkService(Executors.newSingleThreadExecutor(),
+                                                     Executors.defaultThreadFactory(), ConnectionFactory.DEFAULT_SHUTDOWN_TIMEOUT));
             conn.addShutdownListener(new ShutdownListener() {
                 public void shutdownCompleted(ShutdownSignalException cause) {
                     // make sure channel.open continuation is released

test/src/com/rabbitmq/examples/PerfTest.java

@@ -82,6 +82,7 @@ public static void main(String[] args) {
                                     confirm != -1);
 
             ConnectionFactory factory = new ConnectionFactory();
+            factory.setShutdownTimeout(0); // So we still shut down even with slow consumers
             factory.setUri(uri);
             factory.setRequestedFrameMax(frameMax);
             factory.setRequestedHeartbeat(heartbeat);