Shutdown cleanly after Ctrl-C

This commit makes PerfTest shut down cleanly when the process is
Ctrl-C-ed. The closing tasks are accumulated in a shutdown service that
executes them whether the process exits normally or is Ctrl-C-ed.

Note this isn't bullet-proof: if some connections are throttled by the
broker because they publish too fast, closing them can easily time out,
so PerfTest may take some time to terminate and some "client
unexpectedly closed TCP connection" messages can still show up in the
broker.

For cases where PerfTest doesn't saturate the broker, it should
terminate fast and cleanly.

Fixes #126

Author: acogoluegnes

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

@@ -15,11 +15,8 @@
 
 package com.rabbitmq.perf;
 
-import com.rabbitmq.client.AlreadyClosedException;
-import com.rabbitmq.client.Connection;
-import com.rabbitmq.client.ConnectionFactory;
-import com.rabbitmq.client.Recoverable;
-import com.rabbitmq.client.RecoveryListener;
+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;
@@ -32,14 +29,7 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Random;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.Future;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
+import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -70,21 +60,27 @@ public class MulticastSet {
     private final List<String> uris;
     private final Random random = new Random();
     private final CompletionHandler completionHandler;
+    private final ShutdownService shutdownService;
     private ThreadingHandler threadingHandler = new DefaultThreadingHandler();
 
     public MulticastSet(Stats stats, ConnectionFactory factory,
         MulticastParams params, List<String> uris, CompletionHandler completionHandler) {
-        this(stats, factory, params, "perftest", uris, completionHandler);
+        this(stats, factory, params, "perftest", uris, completionHandler, new ShutdownService());
     }
 
     public MulticastSet(Stats stats, ConnectionFactory factory,
-        MulticastParams params, String testID, List<String> uris, CompletionHandler completionHandler) {
+                        MulticastParams params, String testID, List<String> uris, CompletionHandler completionHandler) {
+        this(stats, factory, params, testID, uris, completionHandler, new ShutdownService());
+    }
+    public MulticastSet(Stats stats, ConnectionFactory factory,
+        MulticastParams params, String testID, List<String> uris, CompletionHandler completionHandler, ShutdownService shutdownService) {
         this.stats = stats;
         this.factory = factory;
         this.params = params;
         this.testID = testID;
         this.uris = uris;
         this.completionHandler = completionHandler;
+        this.shutdownService = shutdownService;
         this.params.init();
     }
 
@@ -145,9 +141,16 @@ public void run(boolean announceStartup)
         startConsumers(consumerRunnables);
         startProducers(producerStates);
 
-        this.completionHandler.waitForCompletion();
+        AutoCloseable shutdownSequence = this.shutdownService.wrap(
+                () -> shutdown(configurationConnection, consumerConnections, producerStates, producerConnections)
+        );
 
-        shutdown(configurationConnection, consumerConnections, producerStates, producerConnections);
+        this.completionHandler.waitForCompletion();
+        try {
+            shutdownSequence.close();
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
     }
 
     private Function<Integer, ExecutorService> createConsumersExecutorsFactory() {
@@ -248,38 +251,29 @@ private void shutdown(Connection configurationConnection, Connection[] consumerC
         try {
             LOGGER.debug("Starting test shutdown");
             for (AgentState producerState : producerStates) {
-                producerState.task.cancel(true);
+                boolean cancelled = producerState.task.cancel(true);
+                LOGGER.debug("Producer has been correctly cancelled: {}", cancelled);
             }
 
-            try {
-                LOGGER.debug("Closing configuration connection");
-                configurationConnection.close();
-                LOGGER.debug("Configuration connection closed");
-            } catch (AlreadyClosedException e) {
-                LOGGER.debug("Configuration connection was already closed");
-            } catch (Exception e) {
-                LOGGER.debug("Error while closing configuration connection: {}", e.getMessage());
+            // we do our best to stop producers before closing their connections
+            for (AgentState producerState : producerStates) {
+                if (!producerState.task.isDone()) {
+                    try {
+                        producerState.task.get(10, TimeUnit.SECONDS);
+                    } catch (Exception e) {
+                        LOGGER.debug("Error while waiting for producer to stop: {}. Moving on.", e.getMessage());
+                    }
+                }
             }
 
+            dispose(configurationConnection);
+
             for (Connection producerConnection : producerConnections) {
-                try {
-                    producerConnection.close();
-                } catch (Exception e) {
-                    // don't do anything, we need to close the other connections
-                }
+                dispose(producerConnection);
             }
 
             for (Connection consumerConnection : consumerConnections) {
-                try {
-                    LOGGER.debug("Closing consumer connection {}", consumerConnection.getClientProvidedName());
-                    consumerConnection.close();
-                    LOGGER.debug("Consumer connection {} has been closed", consumerConnection.getClientProvidedName());
-                } catch (AlreadyClosedException e) {
-                    LOGGER.debug("Consumer connection {} already closed", consumerConnection.getClientProvidedName());
-                } catch (Exception e) {
-                    // don't do anything, we need to close the other connections
-                    LOGGER.debug("Error while closing consumer connection {}: {}", consumerConnection.getClientProvidedName(), e.getMessage());
-                }
+                dispose(consumerConnection);
             }
 
             LOGGER.debug("Shutting down threading handler");
@@ -314,6 +308,21 @@ public void handleRecovery(Recoverable recoverable) {
         }
     }
 
+    private static void dispose(Connection connection) {
+        try {
+            LOGGER.debug("Closing connection {}", connection.getClientProvidedName());
+            // we need to shutdown, it should not take forever, so we pick a reasonably
+            // small timeout (3 seconds)
+            connection.close(AMQP.REPLY_SUCCESS, "Closed by PerfTest", 3000);
+            LOGGER.debug("Connection {} has been closed", connection.getClientProvidedName());
+        } catch (AlreadyClosedException e) {
+            LOGGER.debug("Connection {} already closed", connection.getClientProvidedName());
+        } catch (Exception e) {
+            // don't do anything, we need to close the other connections
+            LOGGER.debug("Error while closing connection {}: {}", connection.getClientProvidedName(), e.getMessage());
+        }
+    }
+
     private void setUri() throws URISyntaxException, NoSuchAlgorithmException, KeyManagementException {
         if (uris != null) {
             factory.setUri(uri());

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

@@ -17,7 +17,9 @@
 
 import com.rabbitmq.client.ConnectionFactory;
 import com.rabbitmq.client.DefaultSaslConfig;
+import com.rabbitmq.client.ExceptionHandler;
 import com.rabbitmq.client.impl.ClientVersion;
+import com.rabbitmq.client.impl.DefaultExceptionHandler;
 import com.rabbitmq.client.impl.nio.NioParams;
 import io.micrometer.core.instrument.composite.CompositeMeterRegistry;
 import org.apache.commons.cli.*;
@@ -30,6 +32,7 @@
 import java.security.NoSuchAlgorithmException;
 import java.text.SimpleDateFormat;
 import java.util.*;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.SynchronousQueue;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
@@ -46,9 +49,11 @@ public class PerfTest {
 
     public static void main(String [] args, PerfTestOptions perfTestOptions) {
         SystemExiter systemExiter = perfTestOptions.systemExiter;
+        ShutdownService shutdownService = perfTestOptions.shutdownService;
         Options options = getOptions();
         CommandLineParser parser = new GnuParser();
         CompositeMetrics metrics = new CompositeMetrics();
+        shutdownService.wrap(() -> metrics.close());
         Options metricsOptions = metrics.options();
         forEach(metricsOptions, option -> options.addOption(option));
         try {
@@ -143,12 +148,13 @@ public static void main(String [] args, PerfTestOptions perfTestOptions) {
             factory.setTopologyRecoveryEnabled(false);
 
             CompositeMeterRegistry registry = new CompositeMeterRegistry();
+            shutdownService.wrap(() -> registry.close());
 
             metrics.configure(cmd, registry, factory);
 
             PrintWriter output;
             if (outputFile != null) {
-                output = openCsvFileForWriting(outputFile);
+                output = openCsvFileForWriting(outputFile, shutdownService);
             } else {
                 output = null;
             }
@@ -204,6 +210,10 @@ public static void main(String [] args, PerfTestOptions perfTestOptions) {
             }
 
             factory = configureNioIfRequested(cmd, factory);
+            if (factory.getNioParams().getNioExecutor() != null) {
+                ExecutorService nioExecutor = factory.getNioParams().getNioExecutor();
+                shutdownService.wrap(() -> nioExecutor.shutdownNow());
+            }
 
             MulticastParams p = new MulticastParams();
             p.setAutoAck(               autoAck);
@@ -253,16 +263,12 @@ public static void main(String [] args, PerfTestOptions perfTestOptions) {
 
             MulticastSet.CompletionHandler completionHandler = getCompletionHandler(p);
 
-            MulticastSet set = new MulticastSet(stats, factory, p, testID, uris, completionHandler);
+            factory.setExceptionHandler(perfTestOptions.exceptionHandler);
+
+            MulticastSet set = new MulticastSet(stats, factory, p, testID, uris, completionHandler, shutdownService);
             set.run(true);
 
             stats.printFinal();
-
-            if (factory.getNioParams().getNioExecutor() != null) {
-                factory.getNioParams().getNioExecutor().shutdownNow();
-            }
-
-            registry.close();
         }
         catch (ParseException exp) {
             System.err.println("Parsing failed. Reason: " + exp.getMessage());
@@ -272,13 +278,11 @@ public static void main(String [] args, PerfTestOptions perfTestOptions) {
             LOGGER.error("Main thread caught exception", e);
             systemExiter.exit(1);
         } finally {
-            if (metrics != null) {
-                metrics.close();
-            }
+            shutdownService.close();
         }
     }
 
-    private static PrintWriter openCsvFileForWriting(String outputFile) throws IOException {
+    private static PrintWriter openCsvFileForWriting(String outputFile, ShutdownService shutdownService) throws IOException {
         PrintWriter output;
         File file = new File(outputFile);
         if (file.exists()) {
@@ -288,7 +292,7 @@ private static PrintWriter openCsvFileForWriting(String outputFile) throws IOExc
             }
         }
         output = new PrintWriter(new BufferedWriter(new FileWriter(file, true)), true); //NOSONAR
-        Runtime.getRuntime().addShutdownHook(new Thread(() -> output.close()));
+        shutdownService.wrap(() -> output.close());
         return output;
     }
 
@@ -358,7 +362,13 @@ static MulticastSet.CompletionHandler getCompletionHandler(MulticastParams p) {
 
     public static void main(String[] args) throws IOException {
         Log.configureLog();
-        main(args, new PerfTestOptions().setSystemExiter(new JvmSystemExiter()).setSkipSslContextConfiguration(false));
+        PerfTestOptions perfTestOptions = new PerfTestOptions();
+        Runtime.getRuntime().addShutdownHook(new Thread(() -> perfTestOptions.shutdownService.close()));
+        main(args, perfTestOptions
+                .setSystemExiter(new JvmSystemExiter())
+                .setSkipSslContextConfiguration(false)
+                .setExceptionHandler(new RelaxedExceptionHandler())
+                );
     }
 
     private static SSLContext getSslContextIfNecessary(CommandLineProxy cmd, Properties systemProperties) throws NoSuchAlgorithmException {
@@ -585,6 +595,10 @@ public static class PerfTestOptions {
 
         private boolean skipSslContextConfiguration = false;
 
+        private ShutdownService shutdownService = new ShutdownService();
+
+        private ExceptionHandler exceptionHandler = new DefaultExceptionHandler();
+
         private Function<String, String> argumentLookup = LONG_OPTION_TO_ENVIRONMENT_VARIABLE
             .andThen(ENVIRONMENT_VARIABLE_PREFIX)
             .andThen(ENVIRONMENT_VARIABLE_LOOKUP);
@@ -603,6 +617,16 @@ public PerfTestOptions setArgumentLookup(Function<String, String> argumentLookup
             this.argumentLookup = argumentLookup;
             return this;
         }
+
+        public PerfTestOptions setShutdownService(ShutdownService shutdownService) {
+            this.shutdownService = shutdownService;
+            return this;
+        }
+
+        public PerfTestOptions setExceptionHandler(ExceptionHandler exceptionHandler) {
+            this.exceptionHandler = exceptionHandler;
+            return this;
+        }
     }
 
     /**

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

@@ -20,6 +20,7 @@
 
 import java.io.PrintStream;
 import java.io.PrintWriter;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import static java.lang.String.format;
 
@@ -40,6 +41,8 @@ class PrintlnStats extends Stats {
     private final PrintWriter csvOut;
     private final PrintStream out;
 
+    private final AtomicBoolean printFinalOnGoingOrDone = new AtomicBoolean(false);
+
     public PrintlnStats(String testID, long interval,
                         boolean sendStatsEnabled, boolean recvStatsEnabled,
                         boolean returnStatsEnabled, boolean confirmStatsEnabled,
@@ -100,6 +103,12 @@ private static double rate(long count, long elapsed) {
 
     @Override
     protected void report(long now) {
+        if (!printFinalOnGoingOrDone.get()) {
+            doReport(now);
+        }
+    }
+
+    private void doReport(long now) {
         String output = "id: " + testID + ", ";
 
         double ratePublished = 0.0;
@@ -169,7 +178,10 @@ protected void report(long now) {
             }
         }
 
-        this.out.println(output);
+        if (!printFinalOnGoingOrDone.get()) {
+            this.out.println(output);
+        }
+
         writeToCsvIfNecessary(now, ratePublished, rateReturned, rateConfirmed, rateNacked, rateConsumed, consumerLatencyStats, confirmLatencyStats);
     }
 
@@ -181,7 +193,7 @@ private String legacyMetrics() {
     }
 
     private void writeToCsvIfNecessary(long now, double ratePublished, double rateReturned, double rateConfirmed, double rateNacked, double rateConsumed, long[] consumerLatencyStats, long[] confirmLatencyStats) {
-        if (this.csvOut != null) {
+        if (this.csvOut != null && !printFinalOnGoingOrDone.get()) {
             if (consumerLatencyStats == null) {
                 consumerLatencyStats = getStats(latency);
             }
@@ -248,17 +260,19 @@ private String getRate(String descr, double rate, boolean display) {
     }
 
     public void printFinal() {
-        long now = System.currentTimeMillis();
-
-        System.out.println("id: " + testID + ", sending rate avg: " +
-                formatRate(sendCountTotal * 1000.0 / (now - startTime)) +
-                " " + MESSAGE_RATE_LABEL);
+        if (printFinalOnGoingOrDone.compareAndSet(false, true)) {
+            long now = System.currentTimeMillis();
 
-        long elapsed = now - startTime;
-        if (elapsed > 0) {
-            System.out.println("id: " + testID + ", receiving rate avg: " +
-                    formatRate(recvCountTotal * 1000.0 / elapsed) +
+            System.out.println("id: " + testID + ", sending rate avg: " +
+                    formatRate(sendCountTotal * 1000.0 / (now - startTime)) +
                     " " + MESSAGE_RATE_LABEL);
+
+            long elapsed = now - startTime;
+            if (elapsed > 0) {
+                System.out.println("id: " + testID + ", receiving rate avg: " +
+                        formatRate(recvCountTotal * 1000.0 / elapsed) +
+                        " " + MESSAGE_RATE_LABEL);
+            }
         }
     }
 }