Skip to content

Commit 0cda51d

Browse files
ibauersachsbhaveshthakker
ibauersachs
and
bhaveshthakker
committed
Run NIO close in the selector thread to prevent CME
Closes #379 Closes #380 Co-authored-by: bhaveshthakker <bhaveshthakker111@gmail.com>
1 parent 7411e13 commit 0cda51d

File tree

4 files changed

+237
-47
lines changed

4 files changed

+237
-47
lines changed

‎src/main/java/org/xbill/DNS/NioClient.java‎

Lines changed: 92 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import java.nio.channels.SelectionKey;
1111
import java.nio.channels.Selector;
1212
import java.util.Iterator;
13+
import java.util.function.Consumer;
1314
import lombok.AccessLevel;
1415
import lombok.NoArgsConstructor;
1516
import lombok.extern.slf4j.Slf4j;
@@ -24,9 +25,9 @@
2425
* <p>The following configuration parameter is available:
2526
*
2627
* <dl>
27-
* <dt>dnsjava.nio.selector_timeout
28+
* <dt>{@value SELECTOR_TIMEOUT_PROPERTY}
2829
* <dd>Set selector timeout in milliseconds. Default/Max 1000, Min 1.
29-
* <dt>dnsjava.nio.register_shutdown_hook
30+
* <dt>{@value REGISTER_SHUTDOWN_HOOK_PROPERTY}
3031
* <dd>Register Shutdown Hook termination of NIO. Default True.
3132
* </dl>
3233
*
@@ -35,24 +36,32 @@
3536
@Slf4j
3637
@NoArgsConstructor(access = AccessLevel.NONE)
3738
public abstract class NioClient {
39+
static final String SELECTOR_TIMEOUT_PROPERTY = "dnsjava.nio.selector_timeout";
40+
static final String REGISTER_SHUTDOWN_HOOK_PROPERTY = "dnsjava.nio.register_shutdown_hook";
41+
private static final Object NIO_CLIENT_LOCK = new Object();
42+
3843
/** Packet logger, if available. */
3944
private static PacketLogger packetLogger = null;
4045

4146
private static final Runnable[] TIMEOUT_TASKS = new Runnable[2];
42-
private static final Runnable[] REGISTRATIONS_TASKS = new Runnable[2];
47+
48+
@SuppressWarnings("unchecked")
49+
private static final Consumer<Selector>[] REGISTRATIONS_TASKS = new Consumer[2];
50+
4351
private static final Runnable[] CLOSE_TASKS = new Runnable[2];
4452
private static Thread selectorThread;
4553
private static Thread closeThread;
4654
private static volatile Selector selector;
4755
private static volatile boolean run;
56+
private static volatile boolean closeDone;
4857

4958
interface KeyProcessor {
5059
void processReadyKey(SelectionKey key);
5160
}
5261

5362
static Selector selector() throws IOException {
5463
if (selector == null) {
55-
synchronized (NioClient.class) {
64+
synchronized (NIO_CLIENT_LOCK) {
5665
if (selector == null) {
5766
selector = Selector.open();
5867
log.debug("Starting dnsjava NIO selector thread");
@@ -63,8 +72,7 @@ static Selector selector() throws IOException {
6372
selectorThread.start();
6473
closeThread = new Thread(() -> close(true));
6574
closeThread.setName("dnsjava NIO shutdown hook");
66-
if (Boolean.parseBoolean(
67-
System.getProperty("dnsjava.nio.register_shutdown_hook", "true"))) {
75+
if (Boolean.parseBoolean(System.getProperty(REGISTER_SHUTDOWN_HOOK_PROPERTY, "true"))) {
6876
Runtime.getRuntime().addShutdownHook(closeThread);
6977
}
7078
}
@@ -81,6 +89,10 @@ public static void close() {
8189

8290
private static void close(boolean fromHook) {
8391
run = false;
92+
Selector localSelector = selector;
93+
if (localSelector != null) {
94+
selector.wakeup();
95+
}
8496

8597
if (!fromHook) {
8698
try {
@@ -90,40 +102,26 @@ private static void close(boolean fromHook) {
90102
}
91103
}
92104

93-
try {
94-
runTasks(CLOSE_TASKS);
95-
} catch (Exception e) {
96-
log.warn("Failed to execute shutdown task, ignoring and continuing close", e);
97-
}
98-
99-
Selector localSelector = selector;
100-
Thread localSelectorThread = selectorThread;
101-
synchronized (NioClient.class) {
102-
selector = null;
103-
selectorThread = null;
104-
closeThread = null;
105+
if (localSelector == null) {
106+
// Prevent hanging when close() was called without starting
107+
return;
105108
}
106109

107-
if (localSelector != null) {
108-
localSelector.wakeup();
110+
synchronized (NIO_CLIENT_LOCK) {
109111
try {
110-
localSelector.close();
111-
} catch (IOException e) {
112-
log.warn("Failed to properly close selector, ignoring and continuing close", e);
113-
}
114-
}
115-
116-
if (localSelectorThread != null) {
117-
try {
118-
localSelectorThread.join();
112+
while (!closeDone) {
113+
NIO_CLIENT_LOCK.wait();
114+
}
119115
} catch (InterruptedException e) {
120116
Thread.currentThread().interrupt();
117+
} finally {
118+
closeDone = false;
121119
}
122120
}
123121
}
124122

125123
static void runSelector() {
126-
int timeout = Integer.getInteger("dnsjava.nio.selector_timeout", 1000);
124+
int timeout = Integer.getInteger(SELECTOR_TIMEOUT_PROPERTY, 1000);
127125

128126
if (timeout <= 0 || timeout > 1000) {
129127
throw new IllegalArgumentException("Invalid selector_timeout, must be between 1 and 1000");
@@ -136,7 +134,7 @@ static void runSelector() {
136134
}
137135

138136
if (run) {
139-
runTasks(REGISTRATIONS_TASKS);
137+
runRegistrationTasks();
140138
processReadyKeys();
141139
}
142140
} catch (IOException e) {
@@ -145,30 +143,74 @@ static void runSelector() {
145143
// ignore
146144
}
147145
}
146+
147+
runClose();
148148
log.debug("dnsjava NIO selector thread stopped");
149149
}
150150

151-
static synchronized void setTimeoutTask(Runnable r, boolean isTcpClient) {
151+
private static void runClose() {
152+
try {
153+
runTasks(CLOSE_TASKS);
154+
} catch (Exception e) {
155+
log.warn("Failed to execute shutdown task, ignoring and continuing close", e);
156+
}
157+
158+
Selector localSelector = selector;
159+
Thread localSelectorThread = selectorThread;
160+
synchronized (NIO_CLIENT_LOCK) {
161+
selector = null;
162+
selectorThread = null;
163+
closeThread = null;
164+
closeDone = true;
165+
NIO_CLIENT_LOCK.notifyAll();
166+
}
167+
168+
if (localSelector != null) {
169+
try {
170+
localSelector.close();
171+
} catch (IOException e) {
172+
log.warn("Failed to properly close selector, ignoring and continuing close", e);
173+
}
174+
}
175+
176+
if (localSelectorThread != null) {
177+
try {
178+
localSelectorThread.join();
179+
} catch (InterruptedException e) {
180+
Thread.currentThread().interrupt();
181+
}
182+
}
183+
}
184+
185+
static void setTimeoutTask(Runnable r, boolean isTcpClient) {
152186
addTask(TIMEOUT_TASKS, r, isTcpClient);
153187
}
154188

155-
static synchronized void setRegistrationsTask(Runnable r, boolean isTcpClient) {
156-
addTask(REGISTRATIONS_TASKS, r, isTcpClient);
189+
static void setRegistrationsTask(Consumer<Selector> r, boolean isTcpClient) {
190+
addRegistrationTask(r, isTcpClient);
157191
}
158192

159-
static synchronized void setCloseTask(Runnable r, boolean isTcpClient) {
193+
static void setCloseTask(Runnable r, boolean isTcpClient) {
160194
addTask(CLOSE_TASKS, r, isTcpClient);
161195
}
162196

163-
private static void addTask(Runnable[] closeTasks, Runnable r, boolean isTcpClient) {
197+
private static void addTask(Runnable[] tasks, Runnable r, boolean isTcpClient) {
164198
if (isTcpClient) {
165-
closeTasks[0] = r;
199+
tasks[0] = r;
166200
} else {
167-
closeTasks[1] = r;
201+
tasks[1] = r;
168202
}
169203
}
170204

171-
private static synchronized void runTasks(Runnable[] runnables) {
205+
private static void addRegistrationTask(Consumer<Selector> r, boolean isTcpClient) {
206+
if (isTcpClient) {
207+
REGISTRATIONS_TASKS[0] = r;
208+
} else {
209+
REGISTRATIONS_TASKS[1] = r;
210+
}
211+
}
212+
213+
private static void runTasks(Runnable[] runnables) {
172214
Runnable r0 = runnables[0];
173215
if (r0 != null) {
174216
r0.run();
@@ -179,6 +221,17 @@ private static synchronized void runTasks(Runnable[] runnables) {
179221
}
180222
}
181223

224+
private static void runRegistrationTasks() {
225+
Consumer<Selector> r0 = REGISTRATIONS_TASKS[0];
226+
if (r0 != null) {
227+
r0.accept(selector);
228+
}
229+
Consumer<Selector> r1 = REGISTRATIONS_TASKS[1];
230+
if (r1 != null) {
231+
r1.accept(selector);
232+
}
233+
}
234+
182235
private static void processReadyKeys() {
183236
Iterator<SelectionKey> it = selector.selectedKeys().iterator();
184237
while (it.hasNext()) {

‎src/main/java/org/xbill/DNS/NioTcpClient.java‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,15 +32,14 @@ final class NioTcpClient extends NioClient implements TcpIoClient {
3232
setCloseTask(this::closeTcp, true);
3333
}
3434

35-
private void processPendingRegistrations() {
35+
private void processPendingRegistrations(Selector selector) {
3636
while (!registrationQueue.isEmpty()) {
3737
ChannelState state = registrationQueue.poll();
3838
if (state == null) {
3939
continue;
4040
}
4141

4242
try {
43-
final Selector selector = selector();
4443
if (!state.channel.isConnected()) {
4544
state.channel.register(selector, SelectionKey.OP_CONNECT, state);
4645
} else {

‎src/main/java/org/xbill/DNS/NioUdpClient.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ final class NioUdpClient extends NioClient implements UdpIoClient {
5454
setCloseTask(this::closeUdp, false);
5555
}
5656

57-
private void processPendingRegistrations() {
57+
private void processPendingRegistrations(Selector selector) {
5858
while (!registrationQueue.isEmpty()) {
5959
Transaction t = registrationQueue.poll();
6060
if (t == null) {
@@ -63,7 +63,7 @@ private void processPendingRegistrations() {
6363

6464
try {
6565
log.trace("Registering OP_READ for transaction with id {}", t.id);
66-
t.channel.register(selector(), SelectionKey.OP_READ, t);
66+
t.channel.register(selector, SelectionKey.OP_READ, t);
6767
t.send();
6868
} catch (IOException e) {
6969
t.completeExceptionally(e);

0 commit comments

Comments
 (0)