Fix RawBsonDocument encoding performance regression

bson/src/main/org/bson/codecs/BsonArrayCodec.java

@@ -21,6 +21,7 @@
 import org.bson.BsonType;
 import org.bson.BsonValue;
 import org.bson.BsonWriter;
+import org.bson.RawBsonDocument;
 import org.bson.codecs.configuration.CodecRegistry;
 
 import static org.bson.assertions.Assertions.notNull;
@@ -36,6 +37,7 @@ public class BsonArrayCodec implements Codec<BsonArray> {
 
     private static final CodecRegistry DEFAULT_REGISTRY = fromProviders(new BsonValueCodecProvider());
     private static final BsonTypeCodecMap DEFAULT_BSON_TYPE_CODEC_MAP = new BsonTypeCodecMap(getBsonTypeClassMap(), DEFAULT_REGISTRY);
+    private static final RawBsonDocumentCodec RAW_BSON_DOCUMENT_CODEC = new RawBsonDocumentCodec();
     private final BsonTypeCodecMap bsonTypeCodecMap;
 
     /**
@@ -77,7 +79,9 @@ public void encode(final BsonWriter writer, final BsonArray array, final Encoder
         writer.writeStartArray();
 
         for (BsonValue value : array) {
-            Codec codec = bsonTypeCodecMap.get(value.getBsonType());
+            Codec codec = value instanceof RawBsonDocument
+                    ? RAW_BSON_DOCUMENT_CODEC
+                    : bsonTypeCodecMap.get(value.getBsonType());
             encoderContext.encodeWithChildContext(codec, writer, value);
         }
 

bson/src/main/org/bson/codecs/BsonDocumentCodec.java

@@ -22,6 +22,7 @@
 import org.bson.BsonType;
 import org.bson.BsonValue;
 import org.bson.BsonWriter;
+import org.bson.RawBsonDocument;
 import org.bson.codecs.configuration.CodecRegistry;
 import org.bson.types.ObjectId;
 
@@ -40,6 +41,7 @@ public class BsonDocumentCodec implements CollectibleCodec<BsonDocument> {
     private static final String ID_FIELD_NAME = "_id";
     private static final CodecRegistry DEFAULT_REGISTRY = fromProviders(new BsonValueCodecProvider());
     private static final BsonTypeCodecMap DEFAULT_BSON_TYPE_CODEC_MAP = new BsonTypeCodecMap(getBsonTypeClassMap(), DEFAULT_REGISTRY);
+    private static final RawBsonDocumentCodec RAW_BSON_DOCUMENT_CODEC = new RawBsonDocumentCodec();
 
     private final CodecRegistry codecRegistry;
     private final BsonTypeCodecMap bsonTypeCodecMap;
@@ -130,7 +132,9 @@ private boolean skipField(final EncoderContext encoderContext, final String key)
 
     @SuppressWarnings({"unchecked", "rawtypes"})
     private void writeValue(final BsonWriter writer, final EncoderContext encoderContext, final BsonValue value) {
-        Codec codec = bsonTypeCodecMap.get(value.getBsonType());
+        Codec codec = value instanceof RawBsonDocument
+                ? RAW_BSON_DOCUMENT_CODEC
+                : bsonTypeCodecMap.get(value.getBsonType());
         encoderContext.encodeWithChildContext(codec, writer, value);
     }
 

driver-benchmarks/src/main/com/mongodb/benchmark/benchmarks/AbstractBsonDocumentBenchmark.java

@@ -61,7 +61,7 @@ public int getBytesPerRun() {
         return fileLength * NUM_INTERNAL_ITERATIONS;
     }
 
-    private byte[] getDocumentAsBuffer(final T document) throws IOException {
+    protected byte[] getDocumentAsBuffer(final T document) throws IOException {
         BasicOutputBuffer buffer = new BasicOutputBuffer();
         codec.encode(new BsonBinaryWriter(buffer), document, EncoderContext.builder().build());
 

driver-benchmarks/src/main/com/mongodb/benchmark/benchmarks/BenchmarkSuite.java

@@ -71,6 +71,9 @@ private static void runBenchmarks()
         runBenchmark(new BsonDecodingBenchmark<>("Deep", "extended_bson/deep_bson.json", DOCUMENT_CODEC));
         runBenchmark(new BsonDecodingBenchmark<>("Full", "extended_bson/full_bson.json", DOCUMENT_CODEC));
 
+        runBenchmark(new RawBsonNestedEncodingBenchmark("Full RawBsonDocument in BsonDocument BSON Encoding", "extended_bson/full_bson.json"));
+        runBenchmark(new RawBsonArrayEncodingBenchmark("Full RawBsonDocument Array in BsonDocument BSON Encoding", "extended_bson/full_bson.json", 10));
+
         runBenchmark(new RunCommandBenchmark<>(DOCUMENT_CODEC));
         runBenchmark(new FindOneBenchmark<Document>("single_and_multi_document/tweet.json", BenchmarkSuite.DOCUMENT_CLASS));
 

driver-benchmarks/src/main/com/mongodb/benchmark/benchmarks/RawBsonArrayEncodingBenchmark.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2016-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.mongodb.benchmark.benchmarks;
+
+import org.bson.BsonArray;import org.bson.BsonDocument;
+import org.bson.RawBsonDocument;
+import org.bson.codecs.BsonDocumentCodec;
+
+import java.io.IOException;
+
+public class RawBsonArrayEncodingBenchmark extends BsonEncodingBenchmark<BsonDocument> {
+
+    private final int arraySize;
+
+    public RawBsonArrayEncodingBenchmark(final String name, final String resourcePath, final int arraySize) {
+        super(name, resourcePath, new BsonDocumentCodec());
+        this.arraySize = arraySize;
+    }
+
+    @Override
+    public void setUp() throws IOException {
+        super.setUp();
+        RawBsonDocument rawDoc = new RawBsonDocument(document, codec);
+
+        BsonArray array = new BsonArray();
+        for (int i = 0; i < arraySize; i++) {
+            array.add(rawDoc);
+        }
+        document = new BsonDocument("results", array);
+
+        // Recalculate documentBytes for accurate throughput reporting
+        documentBytes = getDocumentAsBuffer(document);
+
+    }
+
+    @Override
+    public int getBytesPerRun() {
+        return documentBytes.length * NUM_INTERNAL_ITERATIONS;
+    }
+}

driver-benchmarks/src/main/com/mongodb/benchmark/benchmarks/RawBsonNestedEncodingBenchmark.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2016-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.mongodb.benchmark.benchmarks;
+
+import org.bson.BsonDocument;
+import org.bson.RawBsonDocument;
+import org.bson.codecs.BsonDocumentCodec;
+
+import java.io.IOException;
+
+public class RawBsonNestedEncodingBenchmark extends BsonEncodingBenchmark<BsonDocument> {
+
+    public RawBsonNestedEncodingBenchmark(final String name, final String resourcePath) {
+        super(name, resourcePath, new BsonDocumentCodec());
+    }
+
+    @Override
+    public void setUp() throws IOException {
+        super.setUp();
+
+        RawBsonDocument rawDoc = new RawBsonDocument(document, codec);
+        document = new BsonDocument("nested", rawDoc);
+
+        documentBytes = getDocumentAsBuffer(document);
+    }
+
+    @Override
+    public int getBytesPerRun() {
+        return documentBytes.length * NUM_INTERNAL_ITERATIONS;
+    }
+}