dnsjava · bhaveshthakker · Jun 11, 2025 · Jun 13, 2025 · Jun 14, 2025
diff --git a/src/main/java/org/xbill/DNS/NioClient.java b/src/main/java/org/xbill/DNS/NioClient.java
@@ -9,6 +9,7 @@
 import java.nio.channels.ClosedSelectorException;
 import java.nio.channels.SelectionKey;
 import java.nio.channels.Selector;
+import java.util.ArrayList;
 import java.util.Iterator;
 import lombok.AccessLevel;
 import lombok.NoArgsConstructor;
@@ -105,12 +106,14 @@ private static void close(boolean fromHook) {
     }
 
     if (localSelector != null) {
-      localSelector.wakeup();
-      try {
-        localSelector.close();
-      } catch (IOException e) {
-        log.warn("Failed to properly close selector, ignoring and continuing close", e);
-      }
+      REGISTRATIONS_TASKS[0] = () -> {
+        try {
+          localSelector.close(); // ✅ runs safely inside the selector thread
+        } catch (IOException e) {
+          log.warn("Failed to properly close selector", e);
+        }
+      };
+      localSelector.wakeup(); // wake up selector thread to execute shutdown
     }
 
     if (localSelectorThread != null) {
@@ -180,10 +183,10 @@ private static synchronized void runTasks(Runnable[] runnables) {
   }
 
   private static void processReadyKeys() {
-    Iterator<SelectionKey> it = selector.selectedKeys().iterator();
-    while (it.hasNext()) {
-      SelectionKey key = it.next();
-      it.remove();
+    // Copy selected keys to avoid ConcurrentModificationException
+    ArrayList<SelectionKey> readyKeys = new ArrayList<>(selector.selectedKeys());
+    for (SelectionKey key : readyKeys) {
+      selector.selectedKeys().remove(key);
       KeyProcessor t = (KeyProcessor) key.attachment();
       t.processReadyKey(key);
     }

diff --git a/src/test/java/org/xbill/DNS/NioTcpClientTest.java b/src/test/java/org/xbill/DNS/NioTcpClientTest.java
@@ -17,6 +17,7 @@
 import java.nio.ByteBuffer;
 import java.time.Duration;
 import java.util.ArrayList;
+import java.util.ConcurrentModificationException;
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
@@ -27,6 +28,14 @@
 import org.opentest4j.AssertionFailedError;
 import org.xbill.DNS.utils.base16;
 
+import java.nio.channels.Selector;
+import java.nio.channels.Pipe;
+import java.nio.channels.SelectionKey;
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.xbill.DNS.NioClient.KeyProcessor;
+import java.lang.reflect.Field;
+
+
 @Slf4j
 class NioTcpClientTest {
   private static final String SELECTOR_TIMEOUT_PROPERTY = "dnsjava.nio.selector_timeout";
@@ -47,6 +56,82 @@ void testSelectorTimeoutLimits(int timeout) {
     }
   }
 
+  /**
+   * Verifies that NioClient.processReadyKeys() does not throw a ConcurrentModificationException
+   * when channels are registered with the selector while processReadyKeys() is iterating over
+   * selector.selectedKeys(). This simulates concurrent modifications that can occur under high load.
+   *
+   * The test starts the selector thread, registers an initial key, and then rapidly registers new
+   * channels from another thread while making them ready. It fails if a ConcurrentModificationException
+   * is observed during the process.
+   * 
+   * Since this involves concurrency timing, the test may be flaky and not fail consistently,
+   * even if the underlying issue exists.
+   */
+  @Test
+  void testProcessReadyKeysShouldNotThrowConcurrentModificationException() throws Exception {
+    // Speed up selector loop
+    System.setProperty("dnsjava.nio.selector_timeout", "10");
+
+    // Start selector thread
+    NioClient.selector();
+
+    // Access the selector
+    Field selectorField = NioClient.class.getDeclaredField("selector");
+    selectorField.setAccessible(true);
+    Selector selector = (Selector) selectorField.get(null);
+
+    // Add initial key to ensure selectedKeys isn't empty
+    Pipe initialPipe = Pipe.open();
+    initialPipe.source().configureBlocking(false);
+    SelectionKey key = initialPipe.source().register(selector, SelectionKey.OP_READ);
+    key.attach((KeyProcessor) (readyKey) -> {
+        try {
+            Thread.sleep(10); // Slow down processing
+        } catch (InterruptedException ignored) {}
+    });
+
+    // Watch for unexpected exceptions
+    AtomicBoolean sawCME = new AtomicBoolean(false);
+    Thread.UncaughtExceptionHandler originalHandler = Thread.getDefaultUncaughtExceptionHandler();
+    Thread.setDefaultUncaughtExceptionHandler((t, e) -> {
+        if (e instanceof ConcurrentModificationException) {
+            sawCME.set(true); // Violation of expected behavior
+        }
+    });
+
+    // Thread that registers new channels and makes them ready
+    Thread writerThread = new Thread(() -> {
+        try {
+            for (int i = 0; i < 100; i++) {
+                Pipe pipe = Pipe.open();
+                pipe.source().configureBlocking(false);
+                SelectionKey sk = pipe.source().register(selector, SelectionKey.OP_READ);
+                sk.attach((KeyProcessor) (readyKey) -> {
+                    try {
+                        Thread.sleep(10);
+                    } catch (InterruptedException ignored) {}
+                });
+
+                // Make the channel ready to trigger selectedKeys modification
+                pipe.sink().write(ByteBuffer.wrap("x".getBytes()));
+                Thread.sleep(2); // Help trigger overlap
+            }
+        } catch (Exception ignored) {}
+    });
+
+    writerThread.start();
+    writerThread.join(5000);
+    Thread.sleep(2000); // Let selector process mutations
+    NioClient.close();
+    Thread.setDefaultUncaughtExceptionHandler(originalHandler); // restore
+
+    // ✅ Assume correctness: fail if any exception was observed
+    if (sawCME.get()) {
+        fail("processReadyKeys() is not thread-safe: Unexpected ConcurrentModificationException was thrown during concurrent channel activation.");
+    }
+  }
+
   @Test
   void testResponseStream() throws InterruptedException, IOException {
     try {