Skip to content

Commit c6bb3aa

Browse files
rozzarozza
committed
PR updates - remove NPE races and release buffers
1 parent 2df4d98 commit c6bb3aa

2 files changed

Lines changed: 71 additions & 23 deletions

File tree

‎driver-core/src/main/com/mongodb/internal/connection/DefaultServerMonitor.java‎

Lines changed: 53 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -225,8 +225,9 @@ public void run() {
225225
logStateChange(previousServerDescription, currentServerDescription);
226226
sdamProvider.get().monitorUpdate(currentServerDescription);
227227

228+
InternalConnection localConnection = connection;
228229
if ((shouldStreamResponses && currentServerDescription.getType() != UNKNOWN)
229-
|| (connection != null && connection.hasMoreToCome())
230+
|| (localConnection != null && localConnection.hasMoreToCome())
230231
|| (currentServerDescription.getException() instanceof MongoSocketException
231232
&& previousServerDescription.getType() != UNKNOWN)) {
232233
continue;
@@ -239,8 +240,9 @@ public void run() {
239240
LOGGER.error(format("%s for %s stopped working. You may want to recreate the MongoClient", this, serverId), t);
240241
throw t;
241242
} finally {
242-
if (connection != null) {
243-
connection.close();
243+
InternalConnection localConnection = connection;
244+
if (localConnection != null) {
245+
localConnection.close();
244246
}
245247
}
246248
}
@@ -255,7 +257,8 @@ private ServerDescription lookupServerDescription(final ServerDescription curren
255257
lookupStartTimeNanos = System.nanoTime();
256258

257259
// Handle connection setup
258-
if (connection == null || connection.isClosed()) {
260+
InternalConnection localConnection = connection;
261+
if (localConnection == null || localConnection.isClosed()) {
259262
return setupNewConnectionAndGetInitialDescription(shouldStreamResponses);
260263
}
261264

@@ -281,17 +284,47 @@ private ServerDescription lookupServerDescription(final ServerDescription curren
281284
}
282285

283286
private ServerDescription setupNewConnectionAndGetInitialDescription(final boolean shouldStreamResponses) {
284-
connection = internalConnectionFactory.create(serverId);
287+
InternalConnection newConnection = internalConnectionFactory.create(serverId);
288+
289+
// Publish the connection to the field under the lock so that heartbeat
290+
// started logging (which reads the field) can see it, but only if the
291+
// monitor has not been closed in the meantime.
292+
boolean published = withLock(lock, () -> {
293+
if (!isClosed) {
294+
connection = newConnection;
295+
return true;
296+
}
297+
return false;
298+
});
299+
300+
if (!published) {
301+
newConnection.close();
302+
throw new MongoSocketException("Monitor closed", serverId.getAddress());
303+
}
304+
285305
logAndNotifyHeartbeatStarted(shouldStreamResponses);
286306

287307
try {
288-
connection.open(operationContextFactory.create());
289-
roundTripTimeSampler.addSample(connection.getInitialServerDescription().getRoundTripTimeNanos());
290-
return connection.getInitialServerDescription();
308+
newConnection.open(operationContextFactory.create());
291309
} catch (Exception e) {
292310
logAndNotifyHeartbeatFailed(shouldStreamResponses, e);
293311
throw e;
294312
}
313+
314+
// After the potentially long open(), verify the monitor is still open
315+
// before using the connection. If close() ran during open(), it already
316+
// nulled the field and closed the connection, so we must not use it.
317+
boolean stillValid = withLock(lock, () -> !isClosed && connection == newConnection);
318+
319+
if (!stillValid) {
320+
// close() may or may not have closed newConnection already;
321+
// closing an already-closed connection is a safe no-op.
322+
newConnection.close();
323+
throw new MongoSocketException("Monitor closed during connection open", serverId.getAddress());
324+
}
325+
326+
roundTripTimeSampler.addSample(newConnection.getInitialServerDescription().getRoundTripTimeNanos());
327+
return newConnection.getInitialServerDescription();
295328
}
296329

297330
/**
@@ -586,15 +619,20 @@ public void run() {
586619
try {
587620
while (!isClosed) {
588621
try {
589-
if (connection == null) {
622+
InternalConnection localConnection = connection;
623+
if (localConnection == null) {
590624
initialize();
591625
} else {
592-
pingServer(connection);
626+
pingServer(localConnection);
593627
}
594628
} catch (Exception t) {
595-
if (connection != null) {
596-
connection.close();
629+
InternalConnection localConnection = withLock(lock, () -> {
630+
InternalConnection result = connection;
597631
connection = null;
632+
return result;
633+
});
634+
if (localConnection != null) {
635+
localConnection.close();
598636
}
599637
}
600638
waitForNext();
@@ -605,8 +643,9 @@ public void run() {
605643
LOGGER.error(format("%s for %s stopped working. You may want to recreate the MongoClient", this, serverId), t);
606644
throw t;
607645
} finally {
608-
if (connection != null) {
609-
connection.close();
646+
InternalConnection localConnection = connection;
647+
if (localConnection != null) {
648+
localConnection.close();
610649
}
611650
}
612651
}

‎driver-core/src/test/functional/com/mongodb/internal/connection/ReplyHeaderTest.java‎

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -76,14 +76,16 @@ void testParseReplyHeaderWithCompressedHeader(final int responseFlags) {
7676
outputBuffer.writeInt(4);
7777
outputBuffer.writeInt(1);
7878

79-
ByteBuf byteBuf = outputBuffer.getByteBuffers().get(0);
79+
List<ByteBuf> byteBuffers = outputBuffer.getByteBuffers();
80+
ByteBuf byteBuf = byteBuffers.get(0);
8081
CompressedHeader compressedHeader = new CompressedHeader(byteBuf,
8182
new MessageHeader(byteBuf, getDefaultMaxMessageSize()));
8283
ReplyHeader replyHeader = new ReplyHeader(byteBuf, compressedHeader);
8384

8485
assertEquals(274, replyHeader.getMessageLength());
8586
assertEquals(45, replyHeader.getRequestId());
8687
assertEquals(23, replyHeader.getResponseTo());
88+
byteBuffers.forEach(ByteBuf::release);
8789
}
8890
}
8991

@@ -100,12 +102,14 @@ void testThrowExceptionOnIncorrectOpCode() {
100102
outputBuffer.writeInt(0);
101103
outputBuffer.writeInt(0);
102104

103-
ByteBuf byteBuf = outputBuffer.getByteBuffers().get(0);
105+
List<ByteBuf> byteBuffers = outputBuffer.getByteBuffers();
106+
ByteBuf byteBuf = byteBuffers.get(0);
104107

105108
MongoInternalException ex = assertThrows(MongoInternalException.class,
106109
() -> new ReplyHeader(byteBuf, new MessageHeader(byteBuf, getDefaultMaxMessageSize())));
107110

108111
assertEquals("Unexpected reply message opCode 2", ex.getMessage());
112+
byteBuffers.forEach(ByteBuf::release);
109113
}
110114
}
111115

@@ -122,12 +126,13 @@ void testThrowExceptionOnMessageSizeLessThan36() {
122126
outputBuffer.writeInt(0);
123127
outputBuffer.writeInt(0);
124128

125-
ByteBuf byteBuf = outputBuffer.getByteBuffers().get(0);
126-
129+
List<ByteBuf> byteBuffers = outputBuffer.getByteBuffers();
130+
ByteBuf byteBuf = byteBuffers.get(0);
127131
MongoInternalException ex = assertThrows(MongoInternalException.class,
128132
() -> new ReplyHeader(byteBuf, new MessageHeader(byteBuf, getDefaultMaxMessageSize())));
129133

130134
assertEquals("The reply message length 35 is less than the minimum message length 36", ex.getMessage());
135+
byteBuffers.forEach(ByteBuf::release);
131136
}
132137
}
133138

@@ -144,12 +149,13 @@ void testThrowExceptionOnMessageSizeExceedingMax() {
144149
outputBuffer.writeInt(0);
145150
outputBuffer.writeInt(0);
146151

147-
ByteBuf byteBuf = outputBuffer.getByteBuffers().get(0);
148-
152+
List<ByteBuf> byteBuffers = outputBuffer.getByteBuffers();
153+
ByteBuf byteBuf = byteBuffers.get(0);
149154
MongoInternalException ex = assertThrows(MongoInternalException.class,
150155
() -> new ReplyHeader(byteBuf, new MessageHeader(byteBuf, 399)));
151156

152157
assertEquals("The reply message length 400 is greater than the maximum message length 399", ex.getMessage());
158+
byteBuffers.forEach(ByteBuf::release);
153159
}
154160
}
155161

@@ -166,12 +172,13 @@ void testThrowExceptionOnNegativeNumberOfDocuments() {
166172
outputBuffer.writeInt(4);
167173
outputBuffer.writeInt(-1);
168174

169-
ByteBuf byteBuf = outputBuffer.getByteBuffers().get(0);
170-
175+
List<ByteBuf> byteBuffers = outputBuffer.getByteBuffers();
176+
ByteBuf byteBuf = byteBuffers.get(0);
171177
MongoInternalException ex = assertThrows(MongoInternalException.class,
172178
() -> new ReplyHeader(byteBuf, new MessageHeader(byteBuf, getDefaultMaxMessageSize())));
173179

174180
assertEquals("The reply message number of returned documents, -1, is expected to be 1", ex.getMessage());
181+
byteBuffers.forEach(ByteBuf::release);
175182
}
176183
}
177184

@@ -191,14 +198,16 @@ void testThrowExceptionOnNegativeNumberOfDocumentsWithCompressedHeader() {
191198
outputBuffer.writeInt(4);
192199
outputBuffer.writeInt(-1);
193200

194-
ByteBuf byteBuf = outputBuffer.getByteBuffers().get(0);
201+
List<ByteBuf> byteBuffers = outputBuffer.getByteBuffers();
202+
ByteBuf byteBuf = byteBuffers.get(0);
195203
CompressedHeader compressedHeader = new CompressedHeader(byteBuf,
196204
new MessageHeader(byteBuf, getDefaultMaxMessageSize()));
197205

198206
MongoInternalException ex = assertThrows(MongoInternalException.class,
199207
() -> new ReplyHeader(byteBuf, compressedHeader));
200208

201209
assertEquals("The reply message number of returned documents, -1, is expected to be 1", ex.getMessage());
210+
byteBuffers.forEach(ByteBuf::release);
202211
}
203212
}
204213
}

0 commit comments

Comments
 (0)