Skip to content

Commit 35f52e9

Browse files
jyeminjyemin
committed
JAVA-2112: Eliminate double decoding when checking for command failure
1 parent 027d120 commit 35f52e9

File tree

4 files changed

+110
-32
lines changed

4 files changed

+110
-32
lines changed

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

Lines changed: 24 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2008-2015 MongoDB, Inc.
2+
* Copyright 2008-2016 MongoDB, Inc.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -22,12 +22,13 @@
2222
import com.mongodb.diagnostics.logging.Logger;
2323
import com.mongodb.diagnostics.logging.Loggers;
2424
import com.mongodb.event.CommandListener;
25+
import org.bson.BsonBinaryReader;
2526
import org.bson.BsonDocument;
26-
import org.bson.BsonDocumentReader;
2727
import org.bson.FieldNameValidator;
2828
import org.bson.codecs.BsonDocumentCodec;
2929
import org.bson.codecs.Decoder;
30-
import org.bson.codecs.DecoderContext;
30+
import org.bson.codecs.RawBsonDocumentCodec;
31+
import org.bson.io.ByteBufferBsonInput;
3132

3233
import java.util.HashSet;
3334
import java.util.Set;
@@ -105,25 +106,21 @@ public T execute(final InternalConnection connection) {
105106
long startTimeNanos = System.nanoTime();
106107
CommandMessage commandMessage = new CommandMessage(namespace.getFullName(), command, slaveOk, fieldNameValidator,
107108
ProtocolHelper.getMessageSettings(connection.getDescription()));
109+
ResponseBuffers responseBuffers = null;
108110
try {
109111
sendMessage(commandMessage, connection);
110-
ResponseBuffers responseBuffers = connection.receiveMessage(commandMessage.getId());
111-
ReplyMessage<BsonDocument> replyMessage;
112-
try {
113-
replyMessage = new ReplyMessage<BsonDocument>(responseBuffers, new BsonDocumentCodec(), commandMessage.getId());
114-
} finally {
115-
responseBuffers.close();
112+
responseBuffers = connection.receiveMessage(commandMessage.getId());
113+
if (!ProtocolHelper.isCommandOk(new BsonBinaryReader(new ByteBufferBsonInput(responseBuffers.getBodyByteBuffer())))) {
114+
throw getCommandFailureException(getResponseDocument(responseBuffers, commandMessage, new BsonDocumentCodec()),
115+
connection.getDescription().getServerAddress());
116116
}
117117

118-
BsonDocument response = replyMessage.getDocuments().get(0);
119-
if (!ProtocolHelper.isCommandOk(response)) {
120-
throw getCommandFailureException(response, connection.getDescription().getServerAddress());
121-
}
118+
T retval = getResponseDocument(responseBuffers, commandMessage, commandResultDecoder);
122119

123-
T retval = commandResultDecoder.decode(new BsonDocumentReader(response), DecoderContext.builder().build());
124120
if (commandListener != null) {
125121
BsonDocument responseDocumentForEvent = (SECURITY_SENSITIVE_COMMANDS.contains(getCommandName()))
126-
? new BsonDocument() : response;
122+
? new BsonDocument()
123+
: getResponseDocument(responseBuffers, commandMessage, new RawBsonDocumentCodec());
127124
sendCommandSucceededEvent(commandMessage, getCommandName(), responseDocumentForEvent, connection.getDescription(),
128125
startTimeNanos, commandListener);
129126
}
@@ -139,9 +136,21 @@ public T execute(final InternalConnection connection) {
139136
commandListener);
140137
}
141138
throw e;
139+
} finally {
140+
if (responseBuffers != null) {
141+
responseBuffers.close();
142+
}
142143
}
143144
}
144145

146+
private static <D> D getResponseDocument(final ResponseBuffers responseBuffers, final CommandMessage commandMessage,
147+
final Decoder<D> decoder) {
148+
responseBuffers.reset();
149+
ReplyMessage<D> replyMessage = new ReplyMessage<D>(responseBuffers, decoder, commandMessage.getId());
150+
151+
return replyMessage.getDocuments().get(0);
152+
}
153+
145154
@Override
146155
public void executeAsync(final InternalConnection connection, final SingleResultCallback<T> callback) {
147156
try {

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,14 +36,22 @@
3636
import org.bson.BsonBoolean;
3737
import org.bson.BsonDocument;
3838
import org.bson.BsonInt32;
39+
import org.bson.BsonReader;
3940
import org.bson.BsonString;
41+
import org.bson.BsonType;
4042
import org.bson.BsonValue;
43+
import org.bson.codecs.BsonValueCodecProvider;
44+
import org.bson.codecs.DecoderContext;
45+
import org.bson.codecs.configuration.CodecRegistry;
4146
import org.bson.io.BsonOutput;
4247

4348
import static java.lang.String.format;
49+
import static org.bson.codecs.BsonValueCodecProvider.getClassForBsonType;
50+
import static org.bson.codecs.configuration.CodecRegistries.fromProviders;
4451

4552
final class ProtocolHelper {
4653
private static final Logger PROTOCOL_EVENT_LOGGER = Loggers.getLogger("protocol.event");
54+
private static final CodecRegistry REGISTRY = fromProviders(new BsonValueCodecProvider());
4755

4856
static WriteConcernResult getWriteResult(final BsonDocument result, final ServerAddress serverAddress) {
4957
if (!isCommandOk(result)) {
@@ -67,6 +75,27 @@ private static WriteConcernResult createWriteResult(final BsonDocument result) {
6775

6876
static boolean isCommandOk(final BsonDocument response) {
6977
BsonValue okValue = response.get("ok");
78+
return isCommandOk(okValue);
79+
}
80+
81+
static boolean isCommandOk(final BsonReader bsonReader) {
82+
return isCommandOk(getField(bsonReader, "ok"));
83+
}
84+
85+
private static BsonValue getField(final BsonReader bsonReader, final String fieldName) {
86+
bsonReader.readStartDocument();
87+
while (bsonReader.readBsonType() != BsonType.END_OF_DOCUMENT) {
88+
if (bsonReader.readName().equals(fieldName)) {
89+
return REGISTRY.get(getClassForBsonType(bsonReader.getCurrentBsonType())).decode(bsonReader,
90+
DecoderContext.builder().build());
91+
}
92+
bsonReader.skipValue();
93+
}
94+
bsonReader.readEndDocument();
95+
return null;
96+
}
97+
98+
private static boolean isCommandOk(final BsonValue okValue) {
7099
if (okValue == null) {
71100
return false;
72101
} else if (okValue.isBoolean()) {

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,10 @@ public ByteBuf getBodyByteBuffer() {
5555
return bodyByteBuffer.asReadOnly();
5656
}
5757

58+
public void reset() {
59+
bodyByteBuffer.position(0);
60+
}
61+
5862
@Override
5963
public void close() {
6064
if (!isClosed) {

‎driver/src/test/functional/com/mongodb/client/CommandMonitoringTest.java

Lines changed: 53 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright (c) 2015 MongoDB, Inc.
2+
* Copyright 2015-2016 MongoDB, Inc.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -32,13 +32,21 @@
3232
import org.bson.BsonArray;
3333
import org.bson.BsonBoolean;
3434
import org.bson.BsonDocument;
35+
import org.bson.BsonDocumentWriter;
3536
import org.bson.BsonDouble;
3637
import org.bson.BsonInt32;
3738
import org.bson.BsonInt64;
3839
import org.bson.BsonString;
3940
import org.bson.BsonValue;
4041
import org.bson.Document;
42+
import org.bson.codecs.BsonDocumentCodec;
43+
import org.bson.codecs.BsonValueCodecProvider;
44+
import org.bson.codecs.Codec;
4145
import org.bson.codecs.DocumentCodec;
46+
import org.bson.codecs.EncoderContext;
47+
import org.bson.codecs.configuration.CodecProvider;
48+
import org.bson.codecs.configuration.CodecRegistries;
49+
import org.bson.codecs.configuration.CodecRegistry;
4250
import org.junit.AfterClass;
4351
import org.junit.Before;
4452
import org.junit.BeforeClass;
@@ -69,6 +77,19 @@
6977
// See https://github.com/mongodb/specifications/tree/master/source/command-monitoring/tests
7078
@RunWith(Parameterized.class)
7179
public class CommandMonitoringTest {
80+
private static final CodecRegistry CODEC_REGISTRY_HACK = CodecRegistries.fromProviders(new BsonValueCodecProvider(),
81+
new CodecProvider() {
82+
@Override
83+
@SuppressWarnings("unchecked")
84+
public <T> Codec<T> get(final Class<T> clazz, final CodecRegistry registry) {
85+
// Use BsonDocumentCodec even for a private sub-class of BsonDocument
86+
if (BsonDocument.class.isAssignableFrom(clazz)) {
87+
return (Codec<T>) new BsonDocumentCodec(registry);
88+
}
89+
return null;
90+
}
91+
});
92+
7293
private static MongoClient mongoClient;
7394
private static TestCommandListener commandListener;
7495
private final String filename;
@@ -210,43 +231,48 @@ private CommandSucceededEvent massageExpectedCommandSucceededEvent(final Command
210231
}
211232

212233
private CommandSucceededEvent massageActualCommandSucceededEvent(final CommandSucceededEvent actual) {
234+
BsonDocument response = getWritableCloneOfCommand(actual.getResponse());
235+
213236
// massage numbers that are the wrong BSON type
214-
actual.getResponse().put("ok", new BsonDouble(actual.getResponse().getNumber("ok").doubleValue()));
215-
if (actual.getResponse().containsKey("n")) {
216-
actual.getResponse().put("n", new BsonInt32(actual.getResponse().getNumber("n").intValue()));
237+
response.put("ok", new BsonDouble(response.getNumber("ok").doubleValue()));
238+
if (response.containsKey("n")) {
239+
response.put("n", new BsonInt32(response.getNumber("n").intValue()));
217240
}
218241

219242
if (actual.getCommandName().equals("find") || actual.getCommandName().equals("getMore")) {
220-
if (actual.getResponse().containsKey("cursor")) {
221-
if (actual.getResponse().getDocument("cursor").containsKey("id")
222-
&& !actual.getResponse().getDocument("cursor").getInt64("id").equals(new BsonInt64(0))) {
223-
actual.getResponse().getDocument("cursor").put("id", new BsonInt64(42));
243+
if (response.containsKey("cursor")) {
244+
if (response.getDocument("cursor").containsKey("id")
245+
&& !response.getDocument("cursor").getInt64("id").equals(new BsonInt64(0))) {
246+
response.getDocument("cursor").put("id", new BsonInt64(42));
224247
}
225248
}
226249
} else if (actual.getCommandName().equals("killCursors")) {
227-
actual.getResponse().getArray("cursorsUnknown").set(0, new BsonInt64(42));
250+
response.getArray("cursorsUnknown").set(0, new BsonInt64(42));
228251
} else if (isWriteCommand(actual.getCommandName())) {
229-
if (actual.getResponse().containsKey("writeErrors")) {
230-
for (Iterator<BsonValue> iter = actual.getResponse().getArray("writeErrors").iterator(); iter.hasNext();) {
252+
if (response.containsKey("writeErrors")) {
253+
for (Iterator<BsonValue> iter = response.getArray("writeErrors").iterator(); iter.hasNext();) {
231254
BsonDocument cur = iter.next().asDocument();
232255
cur.put("code", new BsonInt32(42));
233256
cur.put("errmsg", new BsonString(""));
234257
}
235258
}
236259
if (actual.getCommandName().equals("update")) {
237-
actual.getResponse().remove("nModified");
260+
response.remove("nModified");
238261
}
239262
}
240-
return actual;
263+
return new CommandSucceededEvent(actual.getRequestId(), actual.getConnectionDescription(), actual.getCommandName(), response,
264+
actual.getElapsedTime(TimeUnit.NANOSECONDS));
241265
}
242266

243267
private boolean isWriteCommand(final String commandName) {
244268
return asList("insert", "update", "delete").contains(commandName);
245269
}
246270

247271
private CommandStartedEvent massageActualCommandStartedEvent(final CommandStartedEvent actual) {
272+
BsonDocument command = getWritableCloneOfCommand(actual.getCommand());
273+
248274
if (actual.getCommandName().equals("update")) {
249-
for (Iterator<BsonValue> iter = actual.getCommand().getArray("updates").iterator(); iter.hasNext();) {
275+
for (Iterator<BsonValue> iter = command.getArray("updates").iterator(); iter.hasNext();) {
250276
BsonDocument curUpdate = iter.next().asDocument();
251277
if (!curUpdate.containsKey("multi")) {
252278
curUpdate.put("multi", BsonBoolean.FALSE);
@@ -256,12 +282,13 @@ private CommandStartedEvent massageActualCommandStartedEvent(final CommandStarte
256282
}
257283
}
258284
} else if (actual.getCommandName().equals("getMore")) {
259-
actual.getCommand().put("getMore", new BsonInt64(42));
285+
command.put("getMore", new BsonInt64(42));
260286
} else if (actual.getCommandName().equals("killCursors")) {
261-
actual.getCommand().getArray("cursors").set(0, new BsonInt64(42));
287+
command.getArray("cursors").set(0, new BsonInt64(42));
262288
}
263289

264-
return actual;
290+
return new CommandStartedEvent(actual.getRequestId(), actual.getConnectionDescription(), actual.getDatabaseName(),
291+
actual.getCommandName(), command);
265292
}
266293

267294
private void executeOperation() {
@@ -298,6 +325,15 @@ private List<CommandEvent> getExpectedEvents(final BsonArray expectedEventDocume
298325
return expectedEvents;
299326
}
300327

328+
329+
private BsonDocument getWritableCloneOfCommand(final BsonDocument original) {
330+
BsonDocument clone = new BsonDocument();
331+
BsonDocumentWriter writer = new BsonDocumentWriter(clone);
332+
new BsonDocumentCodec(CODEC_REGISTRY_HACK).encode(writer, original, EncoderContext.builder().build());
333+
return clone;
334+
}
335+
336+
301337
@Parameterized.Parameters(name = "{1}")
302338
public static Collection<Object[]> data() throws URISyntaxException, IOException {
303339
List<Object[]> data = new ArrayList<Object[]>();

0 commit comments

Comments
 (0)