Fix RawBsonDocument encoding performance regression (#1888)

Add instanceof check in BsonDocumentCodec to route RawBsonDocument to RawBsonDocumentCodec, restoring efficient byte-copy encoding. Previous BsonType-based lookup led to sub-optimal
performance as it could not distinguish RawBsonDocument from BsonDocument.

JAVA-6101

Author: vbabanin

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;
@@ -101,6 +103,10 @@ protected BsonValue readValue(final BsonReader reader, final DecoderContext deco
 
     @Override
     public void encode(final BsonWriter writer, final BsonDocument value, final EncoderContext encoderContext) {
+        if (value instanceof RawBsonDocument) {
+            RAW_BSON_DOCUMENT_CODEC.encode(writer, (RawBsonDocument) value, encoderContext);
+            return;
+        }
         writer.writeStartDocument();
 
         beforeFields(writer, encoderContext, 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;
+    }
+}