KAFKA-10014 Always try to close all channels in Selector#close

clients/src/main/java/org/apache/kafka/common/network/PlaintextChannelBuilder.java

@@ -55,14 +55,20 @@ public KafkaChannel buildChannel(String id, SelectionKey key, int maxReceiveSize
         try {
             PlaintextTransportLayer transportLayer = buildTransportLayer(key);
             Supplier<Authenticator> authenticatorCreator = () -> new PlaintextAuthenticator(configs, transportLayer, listenerName);
-            return new KafkaChannel(id, transportLayer, authenticatorCreator, maxReceiveSize,
+            return buildChannel(id, transportLayer, authenticatorCreator, maxReceiveSize,
                     memoryPool != null ? memoryPool : MemoryPool.NONE, metadataRegistry);
         } catch (Exception e) {
             log.warn("Failed to create channel due to ", e);
             throw new KafkaException(e);
         }
     }
 
+    // visible for testing
+    KafkaChannel buildChannel(String id, TransportLayer transportLayer, Supplier<Authenticator> authenticatorCreator,
+                              int maxReceiveSize, MemoryPool memoryPool, ChannelMetadataRegistry metadataRegistry) {
+        return new KafkaChannel(id, transportLayer, authenticatorCreator, maxReceiveSize, memoryPool, metadataRegistry);
+    }
+
     protected PlaintextTransportLayer buildTransportLayer(SelectionKey key) throws IOException {
         return new PlaintextTransportLayer(key);
     }

clients/src/main/java/org/apache/kafka/common/network/Selector.java

@@ -363,23 +363,19 @@ public void wakeup() {
     @Override
     public void close() {
         List<String> connections = new ArrayList<>(channels.keySet());
-        try {
-            for (String id : connections)
-                close(id);
-        } finally {
-            // If there is any exception thrown in close(id), we should still be able
-            // to close the remaining objects, especially the sensors because keeping
-            // the sensors may lead to failure to start up the ReplicaFetcherThread if
-            // the old sensors with the same names has not yet been cleaned up.
-            AtomicReference<Throwable> firstException = new AtomicReference<>();
-            Utils.closeQuietly(nioSelector, "nioSelector", firstException);
-            Utils.closeQuietly(sensors, "sensors", firstException);
-            Utils.closeQuietly(channelBuilder, "channelBuilder", firstException);
-            Throwable exception = firstException.get();
-            if (exception instanceof RuntimeException && !(exception instanceof SecurityException)) {
-                throw (RuntimeException) exception;
-            }
-
+        AtomicReference<Throwable> firstException = new AtomicReference<>();
+        Utils.closeAllQuietly(firstException, "release connections",
+                connections.stream().map(id -> (AutoCloseable) () -> close(id)).toArray(AutoCloseable[]::new));
+        // If there is any exception thrown in close(id), we should still be able
+        // to close the remaining objects, especially the sensors because keeping
+        // the sensors may lead to failure to start up the ReplicaFetcherThread if
+        // the old sensors with the same names has not yet been cleaned up.
+        Utils.closeQuietly(nioSelector, "nioSelector", firstException);
+        Utils.closeQuietly(sensors, "sensors", firstException);
+        Utils.closeQuietly(channelBuilder, "channelBuilder", firstException);
+        Throwable exception = firstException.get();
+        if (exception instanceof RuntimeException && !(exception instanceof SecurityException)) {
+            throw (RuntimeException) exception;
         }
     }
 

clients/src/main/java/org/apache/kafka/common/utils/Utils.java

@@ -948,6 +948,16 @@ public static void closeQuietly(AutoCloseable closeable, String name, AtomicRefe
         }
     }
 
+    /**
+     * close all closable objects even if one of them throws exception.
+     * @param firstException keeps the first exception
+     * @param name message of closing those objects
+     * @param closeables closable objects
+     */
+    public static void closeAllQuietly(AtomicReference<Throwable> firstException, String name, AutoCloseable... closeables) {
+        for (AutoCloseable closeable : closeables) closeQuietly(closeable, name, firstException);
+    }
+
     /**
      * A cheap way to deterministically convert a number to a positive value. When the input is
      * positive, the original value is returned. When the input number is negative, the returned

clients/src/test/java/org/apache/kafka/common/network/SelectorTest.java

@@ -55,6 +55,8 @@
 import java.util.Random;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 import static java.util.Arrays.asList;
@@ -401,6 +403,33 @@ public void testMute() throws Exception {
         assertEquals("The response should be from the previously muted node", "1", selector.completedReceives().iterator().next().source());
     }
 
+    @Test
+    public void testCloseAllChannels() throws Exception {
+        AtomicInteger closedChannelsCount = new AtomicInteger(0);
+        ChannelBuilder channelBuilder = new PlaintextChannelBuilder(null) {
+            private int channelIndex = 0;
+            @Override
+            KafkaChannel buildChannel(String id, TransportLayer transportLayer, Supplier<Authenticator> authenticatorCreator,
+                                      int maxReceiveSize, MemoryPool memoryPool, ChannelMetadataRegistry metadataRegistry) {
+                return new KafkaChannel(id, transportLayer, authenticatorCreator, maxReceiveSize, memoryPool, metadataRegistry) {
+                    private final int index = channelIndex++;
+                    @Override
+                    public void close() throws IOException {
+                        closedChannelsCount.getAndIncrement();
+                        if (index == 0) throw new RuntimeException("you should fail");
+                        else super.close();
+                    }
+                };
+            }
+        };
+        channelBuilder.configure(clientConfigs());
+        Selector selector = new Selector(5000, new Metrics(), new MockTime(), "MetricGroup", channelBuilder, new LogContext());
+        selector.connect("0", new InetSocketAddress("localhost", server.port), BUFFER_SIZE, BUFFER_SIZE);
+        selector.connect("1", new InetSocketAddress("localhost", server.port), BUFFER_SIZE, BUFFER_SIZE);
+        assertThrows(RuntimeException.class, selector::close);
+        assertEquals(2, closedChannelsCount.get());
+    }
+
     @Test
     public void registerFailure() throws Exception {
         ChannelBuilder channelBuilder = new PlaintextChannelBuilder(null) {

clients/src/test/java/org/apache/kafka/common/utils/UtilsTest.java

@@ -39,8 +39,11 @@
 import java.util.Random;
 import java.util.Set;
 import java.util.TreeSet;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
+import java.util.stream.Stream;
 
 import static java.util.Arrays.asList;
 import static java.util.Collections.emptySet;
@@ -705,4 +708,18 @@ private static void assertValue(Object value) {
         props.put("key", value);
         assertEquals(Utils.propsToMap(props).get("key"), value);
     }
+
+    @Test
+    public void testCloseAllQuietly() {
+        AtomicReference<Throwable> exception = new AtomicReference<>();
+        String msg = "you should fail";
+        AtomicInteger count = new AtomicInteger(0);
+        AutoCloseable c0 = () -> {
+            throw new RuntimeException(msg);
+        };
+        AutoCloseable c1 = count::incrementAndGet;
+        Utils.closeAllQuietly(exception, "test", Stream.of(c0, c1).toArray(AutoCloseable[]::new));
+        assertEquals(msg, exception.get().getMessage());
+        assertEquals(1, count.get());
+    }
 }