Add shutdown timeout argument, defaults to 5 seconds

The shutdown sequence (mainly connection closing) happens now in a
separate thread to be easily stopped. AMQP connections won't be closed
in case of timeout, but this avoids waiting for a long time when
publisher connections are all blocked.

[#162704749]

References #126

Author: acogoluegnes

src/docs/asciidoc/usage.adoc

@@ -144,3 +144,36 @@ messages from only one queue.
 Note it is possible to link:#customising-queues-and-messages[customise]
 the queue and to work against link:#working-with-many-queues[several queues] as well.
 
+== Stopping PerfTest
+
+There are 2 reasons for a PerfTest run to stop:
+
+ * one of the limits has been reached (time limit, producer or consumer message count)
+ * the process is stopped by the user, e.g. by using Ctrl-C in the terminal
+
+In both cases, PerfTest tries to exit as cleanly as possible, in a reasonable amount of time.
+Nevertheless, when PerfTest AMQP connections are throttled by the broker, because they're
+publishing too fast or because broker http://www.rabbitmq.com/alarms.html[alarms]
+have kicked in, it can take time to close them (several seconds or more for one connection).
+
+If closing connections in the gentle way takes too long (5 seconds by default), PerfTest
+will move on to the most important resources to free and terminates. This can result
+in `client unexpectedly closed TCP connection` messages in the broker logs. Note this
+means the AMQP connection hasn't been closed with the right sequence of AMQP frames,
+but the socket has been closed properly. There's no resource leakage here.
+
+The connection closing timeout can be set up with the `--shutdown-timeout` argument (or `-st`).
+The default timeout can be increased to let more time to close connections, e.g. the
+command below uses a shutdown timeout of 20 seconds:
+
+ bin/runjava com.rabbitmq.perf.PerfTest --shutdown-timeout 20
+
+The connection closing sequence can also be skipped by setting the timeout to 0 or any negative
+value:
+
+ bin/runjava com.rabbitmq.perf.PerfTest --shutdown-timeout -1
+
+With the previous command, PerfTest won't even try to close AMQP connections, it will exit
+as fast as possible, freeing only the most important resources. This is perfectly
+acceptable when performing runs on a test environment.
+

src/main/java/com/rabbitmq/perf/MulticastParams.java

@@ -92,6 +92,7 @@ public class MulticastParams {
     private int producerRandomStartDelayInSeconds;
     private int producerSchedulerThreadCount = -1;
     private int consumersThreadPools = -1;
+    private int shutdownTimeout = 5;
 
     public void setExchangeType(String exchangeType) {
         this.exchangeType = exchangeType;
@@ -234,6 +235,10 @@ public void setConsumersThreadPools(int consumersThreadPools) {
         this.consumersThreadPools = consumersThreadPools;
     }
 
+    public void setShutdownTimeout(int shutdownTimeout) {
+        this.shutdownTimeout = shutdownTimeout;
+    }
+
     public int getConsumerCount() {
         return consumerCount;
     }
@@ -318,6 +323,10 @@ public int getConsumersThreadPools() {
         return consumersThreadPools;
     }
 
+    public int getShutdownTimeout() {
+        return shutdownTimeout;
+    }
+
     public Producer createProducer(Connection connection, Stats stats, MulticastSet.CompletionHandler completionHandler) throws IOException {
         Channel channel = connection.createChannel(); //NOSONAR
         if (producerTxSize > 0) channel.txSelect();

src/main/java/com/rabbitmq/perf/MulticastSet.java

@@ -16,7 +16,6 @@
 package com.rabbitmq.perf;
 
 import com.rabbitmq.client.*;
-import com.rabbitmq.client.impl.DefaultExceptionHandler;
 import com.rabbitmq.client.impl.recovery.AutorecoveringConnection;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -141,11 +140,35 @@ public void run(boolean announceStartup)
         startConsumers(consumerRunnables);
         startProducers(producerStates);
 
-        AutoCloseable shutdownSequence = this.shutdownService.wrap(
-                () -> shutdown(configurationConnection, consumerConnections, producerStates, producerConnections)
-        );
+
+        AutoCloseable shutdownSequence;
+        int shutdownTimeout = this.params.getShutdownTimeout();
+        if (shutdownTimeout > 0) {
+            shutdownSequence = this.shutdownService.wrap(
+                    () -> {
+                        CountDownLatch latch = new CountDownLatch(1);
+                        Thread shutdownThread = new Thread(() -> {
+                            try {
+                                shutdown(configurationConnection, consumerConnections, producerStates, producerConnections);
+                            } finally {
+                                latch.countDown();
+                            }
+                        });
+                        shutdownThread.start();
+                        boolean done = latch.await(shutdownTimeout, TimeUnit.SECONDS);
+                        if (!done) {
+                            LOGGER.debug("Shutdown not completed in {} second(s), aborting.", shutdownTimeout);
+                            shutdownThread.interrupt();
+                        }
+                    }
+            );
+        } else {
+            // no closing timeout, we don't do anything
+            shutdownSequence = () -> { };
+        }
 
         this.completionHandler.waitForCompletion();
+
         try {
             shutdownSequence.close();
         } catch (Exception e) {
@@ -251,6 +274,9 @@ private void shutdown(Connection configurationConnection, Connection[] consumerC
         try {
             LOGGER.debug("Starting test shutdown");
             for (AgentState producerState : producerStates) {
+                if (Thread.interrupted()) {
+                    return;
+                }
                 boolean cancelled = producerState.task.cancel(true);
                 LOGGER.debug("Producer has been correctly cancelled: {}", cancelled);
             }
@@ -259,23 +285,38 @@ private void shutdown(Connection configurationConnection, Connection[] consumerC
             for (AgentState producerState : producerStates) {
                 if (!producerState.task.isDone()) {
                     try {
+                        if (Thread.interrupted()) {
+                            return;
+                        }
                         producerState.task.get(10, TimeUnit.SECONDS);
                     } catch (Exception e) {
                         LOGGER.debug("Error while waiting for producer to stop: {}. Moving on.", e.getMessage());
                     }
                 }
             }
 
+            if (Thread.interrupted()) {
+                return;
+            }
             dispose(configurationConnection);
 
             for (Connection producerConnection : producerConnections) {
+                if (Thread.interrupted()) {
+                    return;
+                }
                 dispose(producerConnection);
             }
 
             for (Connection consumerConnection : consumerConnections) {
+                if (Thread.interrupted()) {
+                    return;
+                }
                 dispose(consumerConnection);
             }
 
+            if (Thread.interrupted()) {
+                return;
+            }
             LOGGER.debug("Shutting down threading handler");
             this.threadingHandler.shutdown();
             LOGGER.debug("Threading handler shut down");
@@ -318,7 +359,7 @@ private static void dispose(Connection connection) {
         } catch (AlreadyClosedException e) {
             LOGGER.debug("Connection {} already closed", connection.getClientProvidedName());
         } catch (Exception e) {
-            // don't do anything, we need to close the other connections
+            // just log, we don't want to stop here
             LOGGER.debug("Error while closing connection {}: {}", connection.getClientProvidedName(), e.getMessage());
         }
     }

src/main/java/com/rabbitmq/perf/PerfTest.java

@@ -135,6 +135,7 @@ public static void main(String [] args, PerfTestOptions perfTestOptions) {
 
             boolean disableConnectionRecovery = hasOption(cmd, "dcr");
             int consumersThreadPools = intArg(cmd, "ctp", -1);
+            int shutdownTimeout = intArg(cmd, "st", 5);
 
             String uri               = strArg(cmd, 'h', "amqp://localhost");
             String urisParameter     = strArg(cmd, 'H', null);
@@ -260,6 +261,7 @@ public static void main(String [] args, PerfTestOptions perfTestOptions) {
             p.setProducerRandomStartDelayInSeconds(producerRandomStartDelayInSeconds);
             p.setProducerSchedulerThreadCount(producerSchedulingThreads);
             p.setConsumersThreadPools(consumersThreadPools);
+            p.setShutdownTimeout(shutdownTimeout);
 
             MulticastSet.CompletionHandler completionHandler = getCompletionHandler(p);
 
@@ -496,6 +498,9 @@ public static Options getOptions() {
 
         options.addOption(new Option("ctp", "consumers-thread-pools",true, "number of thread pools to use for all consumers, "
             + "default is to use a thread pool for each consumer"));
+
+        options.addOption(new Option("st", "shutdown-timeout",true, "shutdown timeout, default is 5 seconds"));
+
         return options;
     }
 

src/test/java/com/rabbitmq/perf/MessageCountTimeLimitAndPublishingIntervalRateTest.java

@@ -189,8 +189,10 @@ public void timeLimit() throws InterruptedException {
             callback("getNextPublishSeqNo", (proxy, method, args) -> 0L)
         );
 
+        AtomicInteger closeCount = new AtomicInteger(0);
         Connection connection = proxy(Connection.class,
-            callback("createChannel", (proxy, method, args) -> channel)
+            callback("createChannel", (proxy, method, args) -> channel),
+            callback("close", (proxy, method, args) -> closeCount.incrementAndGet())
         );
 
         MulticastSet multicastSet = getMulticastSet(connectionFactoryThatReturns(connection));
@@ -202,6 +204,7 @@ public void timeLimit() throws InterruptedException {
         waitAtMost(15, () -> testIsDone.get());
 
         assertThat(testDurationInMs, greaterThanOrEqualTo(3000L));
+        assertThat(closeCount.get(), is(4)); // the configuration connection is actually closed twice
     }
 
     // -y 1 --pmessages 10 -x n -X m
@@ -493,6 +496,99 @@ public void publishingInterval() throws InterruptedException {
         assertThat(testDurationInMs, greaterThan(5000L));
     }
 
+    @Test
+    public void shutdownCalledIfShutdownTimeoutIsGreatherThanZero() throws Exception {
+        countsAndTimeLimit(0, 0, 0);
+
+        Channel channel = proxy(Channel.class,
+                callback("basicPublish", (proxy, method, args) -> null),
+                callback("getNextPublishSeqNo", (proxy, method, args) -> 0L)
+        );
+
+        AtomicInteger connectionCloseCalls = new AtomicInteger(0);
+        Connection connection = proxy(Connection.class,
+                callback("createChannel", (proxy, method, args) -> channel),
+                callback("close", (proxy, method, args) -> connectionCloseCalls.incrementAndGet())
+        );
+
+        MulticastSet multicastSet = getMulticastSet(connectionFactoryThatReturns(connection));
+
+        run(multicastSet);
+
+        waitForRunToStart();
+
+        completionHandler.countDown();
+        waitAtMost(10, () -> testIsDone.get());
+        assertThat(connectionCloseCalls.get(), is(4)); // configuration connection is closed twice
+    }
+
+    @Test
+    public void shutdownNotCalledIfShutdownTimeoutIsZeroOrLess() throws Exception {
+        countsAndTimeLimit(0, 0, 0);
+        params.setShutdownTimeout(-1);
+
+        Channel channel = proxy(Channel.class,
+                callback("basicPublish", (proxy, method, args) -> null),
+                callback("getNextPublishSeqNo", (proxy, method, args) -> 0L)
+        );
+
+        AtomicInteger connectionCloseCalls = new AtomicInteger(0);
+        Connection connection = proxy(Connection.class,
+                callback("createChannel", (proxy, method, args) -> channel),
+                callback("close", (proxy, method, args) -> connectionCloseCalls.incrementAndGet())
+        );
+
+        MulticastSet multicastSet = getMulticastSet(connectionFactoryThatReturns(connection));
+
+        run(multicastSet);
+
+        waitForRunToStart();
+
+        completionHandler.countDown();
+        waitAtMost(10, () -> testIsDone.get());
+        assertThat(connectionCloseCalls.get(), is(1)); // configuration connection is closed after configuration is done
+    }
+
+    @Test
+    public void shutdownNotCompletedIfTimeoutIsReached() throws Exception {
+        countsAndTimeLimit(0, 0, 0);
+        params.setShutdownTimeout(1);
+
+        Channel channel = proxy(Channel.class,
+                callback("basicPublish", (proxy, method, args) -> null),
+                callback("getNextPublishSeqNo", (proxy, method, args) -> 0L)
+        );
+
+        AtomicInteger connectionCloseCalls = new AtomicInteger(0);
+        Connection connection = proxy(Connection.class,
+                callback("createChannel", (proxy, method, args) -> channel),
+                callback("close", (proxy, method, args) -> {
+                    connectionCloseCalls.incrementAndGet();
+                    // the first call is to close the configuration connection at the beginning of the run,
+                    // so we simulate a timeout when closing a connection during the final shutdown
+                    if (connectionCloseCalls.get() == 2) {
+                        try {
+                            Thread.sleep(5000);
+                        } catch (InterruptedException e) {
+                            // the interrupt flag is cleared after an InterruptedException
+                            Thread.currentThread().interrupt();
+                        }
+                    }
+                    return null;
+                })
+        );
+
+        MulticastSet multicastSet = getMulticastSet(connectionFactoryThatReturns(connection));
+
+        run(multicastSet);
+
+        waitForRunToStart();
+
+        completionHandler.countDown();
+        waitAtMost(10, () -> testIsDone.get());
+        assertThat(connectionCloseCalls.get(), is(2));
+    }
+
     private Collection<Future<?>> sendMessagesToConsumer(int messagesCount, Consumer consumer) {
         final Collection<Future<?>> tasks = new ArrayList<>(messagesCount);
         IntStream.range(0, messagesCount).forEach(i -> {
@@ -561,7 +657,11 @@ private void run(MulticastSet multicastSet) {
                 // one of the tests stops the execution, no need to be noisy
                 LOGGER.warn("Run has been interrupted");
             } catch (Exception e) {
-                LOGGER.warn("Error during run", e);
+                if (e.getCause() instanceof InterruptedException) {
+                    LOGGER.warn("Run has been interrupted");
+                } else {
+                    LOGGER.warn("Error during run", e);
+                }
             }
         });
     }