Optimize BSON codec lookup.

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

@@ -24,6 +24,7 @@
 import org.bson.codecs.configuration.CodecRegistry;
 
 import static org.bson.assertions.Assertions.notNull;
+import static org.bson.codecs.BsonValueCodecProvider.getBsonTypeClassMap;
 import static org.bson.codecs.configuration.CodecRegistries.fromProviders;
 
 /**
@@ -34,16 +35,16 @@
 public class BsonArrayCodec implements Codec<BsonArray> {
 
     private static final CodecRegistry DEFAULT_REGISTRY = fromProviders(new BsonValueCodecProvider());
-
-    private final CodecRegistry codecRegistry;
+    private static final BsonTypeCodecMap DEFAULT_BSON_TYPE_CODEC_MAP = new BsonTypeCodecMap(getBsonTypeClassMap(), DEFAULT_REGISTRY);
+    private final BsonTypeCodecMap bsonTypeCodecMap;
 
     /**
      * Creates a new instance with a default codec registry that uses the {@link BsonValueCodecProvider}.
      *
      * @since 3.4
      */
     public BsonArrayCodec() {
-        this(DEFAULT_REGISTRY);
+        this(DEFAULT_BSON_TYPE_CODEC_MAP);
     }
 
     /**
@@ -52,7 +53,11 @@ public BsonArrayCodec() {
      * @param codecRegistry the codec registry
      */
     public BsonArrayCodec(final CodecRegistry codecRegistry) {
-        this.codecRegistry = notNull("codecRegistry", codecRegistry);
+        this(new BsonTypeCodecMap(getBsonTypeClassMap(), codecRegistry));
+    }
+
+    private BsonArrayCodec(final BsonTypeCodecMap bsonTypeCodecMap) {
+        this.bsonTypeCodecMap = notNull("bsonTypeCodecMap", bsonTypeCodecMap);
     }
 
     @Override
@@ -72,7 +77,7 @@ public void encode(final BsonWriter writer, final BsonArray array, final Encoder
         writer.writeStartArray();
 
         for (BsonValue value : array) {
-            Codec codec = codecRegistry.get(value.getClass());
+            Codec codec = bsonTypeCodecMap.get(value.getBsonType());
             encoderContext.encodeWithChildContext(codec, writer, value);
         }
 
@@ -93,7 +98,7 @@ public Class<BsonArray> getEncoderClass() {
      * @return the non-null value read from the reader
      */
     protected BsonValue readValue(final BsonReader reader, final DecoderContext decoderContext) {
-        return codecRegistry.get(BsonValueCodecProvider.getClassForBsonType(reader.getCurrentBsonType())).decode(reader, decoderContext);
+        BsonType currentBsonType = reader.getCurrentBsonType();
+        return (BsonValue) bsonTypeCodecMap.get(currentBsonType).decode(reader, decoderContext);
     }
-
 }

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

@@ -17,7 +17,6 @@
 package org.bson.codecs;
 
 import org.bson.BsonDocument;
-import org.bson.BsonElement;
 import org.bson.BsonObjectId;
 import org.bson.BsonReader;
 import org.bson.BsonType;
@@ -26,8 +25,6 @@
 import org.bson.codecs.configuration.CodecRegistry;
 import org.bson.types.ObjectId;
 
-import java.util.ArrayList;
-import java.util.List;
 import java.util.Map;
 
 import static org.bson.assertions.Assertions.notNull;
@@ -79,17 +76,15 @@ public CodecRegistry getCodecRegistry() {
 
     @Override
     public BsonDocument decode(final BsonReader reader, final DecoderContext decoderContext) {
-        List<BsonElement> keyValuePairs = new ArrayList<>();
-
+        BsonDocument bsonDocument = new BsonDocument();
         reader.readStartDocument();
         while (reader.readBsonType() != BsonType.END_OF_DOCUMENT) {
             String fieldName = reader.readName();
-            keyValuePairs.add(new BsonElement(fieldName, readValue(reader, decoderContext)));
+            bsonDocument.append(fieldName, readValue(reader, decoderContext));
         }
 
         reader.readEndDocument();
-
-        return new BsonDocument(keyValuePairs);
+        return bsonDocument;
     }
 
     /**
@@ -135,7 +130,7 @@ 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 = codecRegistry.get(value.getClass());
+        Codec codec = bsonTypeCodecMap.get(value.getBsonType());
         encoderContext.encodeWithChildContext(codec, writer, value);
     }
 

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

@@ -31,12 +31,11 @@
 import org.bson.types.ObjectId;
 import org.bson.types.Symbol;
 
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.Date;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 
 /**
  * <p>A map from a BSON types to the Class to which it should be decoded.  This class is useful if, for example,
@@ -71,7 +70,7 @@
  */
 public class BsonTypeClassMap {
     static final BsonTypeClassMap DEFAULT_BSON_TYPE_CLASS_MAP = new BsonTypeClassMap();
-    private final Map<BsonType, Class<?>> map = new HashMap<>();
+    private final Class<?>[] bsonTypeOrdinalToClassMap = new Class<?>[BsonType.values().length];
 
     /**
      * Construct an instance with the default mapping, but replacing the default mapping with any values contained in the given map.
@@ -81,7 +80,7 @@ public class BsonTypeClassMap {
      */
     public BsonTypeClassMap(final Map<BsonType, Class<?>> replacementsForDefaults) {
         addDefaults();
-        map.putAll(replacementsForDefaults);
+        replacementsForDefaults.forEach((key, value) -> bsonTypeOrdinalToClassMap[key.ordinal()] = value);
     }
 
     /**
@@ -91,41 +90,37 @@ public BsonTypeClassMap() {
         this(Collections.emptyMap());
     }
 
-    Set<BsonType> keys() {
-        return map.keySet();
-    }
-
     /**
      * Gets the Class that is mapped to the given BSON type.
      *
      * @param bsonType the BSON type
      * @return the Class that is mapped to the BSON type
      */
     public Class<?> get(final BsonType bsonType) {
-        return map.get(bsonType);
+        return bsonTypeOrdinalToClassMap[bsonType.ordinal()];
     }
 
     private void addDefaults() {
-        map.put(BsonType.ARRAY, List.class);
-        map.put(BsonType.BINARY, Binary.class);
-        map.put(BsonType.BOOLEAN, Boolean.class);
-        map.put(BsonType.DATE_TIME, Date.class);
-        map.put(BsonType.DB_POINTER, BsonDbPointer.class);
-        map.put(BsonType.DOCUMENT, Document.class);
-        map.put(BsonType.DOUBLE, Double.class);
-        map.put(BsonType.INT32, Integer.class);
-        map.put(BsonType.INT64, Long.class);
-        map.put(BsonType.DECIMAL128, Decimal128.class);
-        map.put(BsonType.MAX_KEY, MaxKey.class);
-        map.put(BsonType.MIN_KEY, MinKey.class);
-        map.put(BsonType.JAVASCRIPT, Code.class);
-        map.put(BsonType.JAVASCRIPT_WITH_SCOPE, CodeWithScope.class);
-        map.put(BsonType.OBJECT_ID, ObjectId.class);
-        map.put(BsonType.REGULAR_EXPRESSION, BsonRegularExpression.class);
-        map.put(BsonType.STRING, String.class);
-        map.put(BsonType.SYMBOL, Symbol.class);
-        map.put(BsonType.TIMESTAMP, BsonTimestamp.class);
-        map.put(BsonType.UNDEFINED, BsonUndefined.class);
+        bsonTypeOrdinalToClassMap[BsonType.ARRAY.ordinal()] = List.class;
+        bsonTypeOrdinalToClassMap[BsonType.BINARY.ordinal()] = Binary.class;
+        bsonTypeOrdinalToClassMap[BsonType.BOOLEAN.ordinal()] = Boolean.class;
+        bsonTypeOrdinalToClassMap[BsonType.DATE_TIME.ordinal()] = Date.class;
+        bsonTypeOrdinalToClassMap[BsonType.DB_POINTER.ordinal()] = BsonDbPointer.class;
+        bsonTypeOrdinalToClassMap[BsonType.DOCUMENT.ordinal()] = Document.class;
+        bsonTypeOrdinalToClassMap[BsonType.DOUBLE.ordinal()] = Double.class;
+        bsonTypeOrdinalToClassMap[BsonType.INT32.ordinal()] = Integer.class;
+        bsonTypeOrdinalToClassMap[BsonType.INT64.ordinal()] = Long.class;
+        bsonTypeOrdinalToClassMap[BsonType.DECIMAL128.ordinal()] = Decimal128.class;
+        bsonTypeOrdinalToClassMap[BsonType.MAX_KEY.ordinal()] = MaxKey.class;
+        bsonTypeOrdinalToClassMap[BsonType.MIN_KEY.ordinal()] = MinKey.class;
+        bsonTypeOrdinalToClassMap[BsonType.JAVASCRIPT.ordinal()] = Code.class;
+        bsonTypeOrdinalToClassMap[BsonType.JAVASCRIPT_WITH_SCOPE.ordinal()] = CodeWithScope.class;
+        bsonTypeOrdinalToClassMap[BsonType.OBJECT_ID.ordinal()] = ObjectId.class;
+        bsonTypeOrdinalToClassMap[BsonType.REGULAR_EXPRESSION.ordinal()] = BsonRegularExpression.class;
+        bsonTypeOrdinalToClassMap[BsonType.STRING.ordinal()] = String.class;
+        bsonTypeOrdinalToClassMap[BsonType.SYMBOL.ordinal()] = Symbol.class;
+        bsonTypeOrdinalToClassMap[BsonType.TIMESTAMP.ordinal()] = BsonTimestamp.class;
+        bsonTypeOrdinalToClassMap[BsonType.UNDEFINED.ordinal()] = BsonUndefined.class;
     }
 
     @Override
@@ -139,15 +134,11 @@ public boolean equals(final Object o) {
 
         BsonTypeClassMap that = (BsonTypeClassMap) o;
 
-        if (!map.equals(that.map)) {
-            return false;
-        }
-
-        return true;
+        return Arrays.equals(bsonTypeOrdinalToClassMap, that.bsonTypeOrdinalToClassMap);
     }
 
     @Override
     public int hashCode() {
-        return map.hashCode();
+        return Arrays.hashCode(bsonTypeOrdinalToClassMap);
     }
 }

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

@@ -40,7 +40,7 @@ public class BsonTypeCodecMap {
     public BsonTypeCodecMap(final BsonTypeClassMap bsonTypeClassMap, final CodecRegistry codecRegistry) {
         this.bsonTypeClassMap = notNull("bsonTypeClassMap", bsonTypeClassMap);
         notNull("codecRegistry", codecRegistry);
-        for (BsonType cur : bsonTypeClassMap.keys()) {
+        for (BsonType cur : BsonType.values()) {
             Class<?> clazz = bsonTypeClassMap.get(cur);
             if (clazz != null) {
                 try {

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

@@ -156,7 +156,7 @@ public void encode(final BsonWriter writer, final Document document, final Encod
 
         beforeFields(writer, encoderContext, document);
 
-        for (final Map.Entry<String, Object> entry : ((Map<String, Object>) document).entrySet()) {
+        for (final Map.Entry<String, Object> entry : document.entrySet()) {
             if (skipField(encoderContext, entry.getKey())) {
                 continue;
             }

bson/src/test/unit/org/bson/codecs/IterableCodecProviderSpecification.groovy

@@ -16,6 +16,7 @@
 
 package org.bson.codecs
 
+import org.bson.BsonType
 import spock.lang.Specification
 
 import static org.bson.codecs.configuration.CodecRegistries.fromProviders
@@ -57,7 +58,7 @@ class IterableCodecProviderSpecification extends Specification {
     def 'unidentical instances should not be equal'() {
         given:
         def first = new IterableCodecProvider()
-        def second = new IterableCodecProvider(new BsonTypeClassMap([BOOLEAN: String]))
+        def second = new IterableCodecProvider(new BsonTypeClassMap([(BsonType.BOOLEAN): String]))
         def third = new IterableCodecProvider(new BsonTypeClassMap(), { Object from ->
             from
         })

driver-benchmarks/build.gradle

@@ -33,6 +33,16 @@ dependencies {
     api project(':driver-sync')
     api project(':mongodb-crypt')
     implementation "ch.qos.logback:logback-classic:$logbackVersion"
+
+    implementation 'org.openjdk.jmh:jmh-core:1.37'
+    annotationProcessor 'org.openjdk.jmh:jmh-generator-annprocess:1.37'
+}
+
+tasks.register("jmh", JavaExec) {
+    group = 'benchmark'
+    description = 'Run JMH benchmarks.'
+    mainClass = 'org.openjdk.jmh.Main'
+    classpath = sourceSets.main.runtimeClasspath
 }
 
 javadoc {

driver-benchmarks/src/main/com/mongodb/benchmark/jmh/codec/BsonArrayCodecBenchmark.java

@@ -0,0 +1,101 @@
+/*
+ * 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.jmh.codec;
+
+import com.mongodb.internal.connection.ByteBufferBsonOutput;
+import com.mongodb.internal.connection.PowerOfTwoBufferPool;
+import org.bson.BsonArray;
+import org.bson.BsonBinaryReader;
+import org.bson.BsonBinaryWriter;
+import org.bson.BsonDocument;
+import org.bson.BsonDouble;
+import org.bson.codecs.Codec;
+import org.bson.codecs.DecoderContext;
+import org.bson.codecs.EncoderContext;
+import org.jetbrains.annotations.NotNull;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.infra.Blackhole;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.concurrent.TimeUnit;
+
+import static com.mongodb.MongoClientSettings.getDefaultCodecRegistry;
+import static com.mongodb.benchmark.jmh.codec.BsonUtils.getDocumentAsBuffer;
+
+@BenchmarkMode(Mode.Throughput)
+@Warmup(iterations = 20, time = 2, timeUnit = TimeUnit.SECONDS)
+@Measurement(iterations = 20, time = 2, timeUnit = TimeUnit.SECONDS)
+@OutputTimeUnit(TimeUnit.SECONDS)
+@Fork(3)
+public class BsonArrayCodecBenchmark {
+
+    @State(Scope.Benchmark)
+    public static class Input {
+        protected final PowerOfTwoBufferPool bufferPool = PowerOfTwoBufferPool.DEFAULT;
+        protected final Codec<BsonArray> bsonArrayCodec = getDefaultCodecRegistry().get(BsonArray.class);
+        protected BsonDocument document;
+        protected byte[] documentBytes;
+        private BsonBinaryReader reader;
+        private BsonBinaryWriter writer;
+        private BsonArray bsonValues;
+
+        @Setup
+        public void setup() throws IOException {
+            bsonValues = new BsonArray();
+            document = new BsonDocument("array", bsonValues);
+
+            for (int i = 0; i < 1000; i++) {
+                bsonValues.add(new BsonDouble(i));
+            }
+
+            documentBytes = getDocumentAsBuffer(document);
+        }
+
+        @Setup(Level.Invocation)
+        public void beforeIteration() {
+            reader = new BsonBinaryReader(ByteBuffer.wrap(documentBytes));
+            writer = new BsonBinaryWriter(new ByteBufferBsonOutput(bufferPool));
+
+            reader.readStartDocument();
+            writer.writeStartDocument();
+            writer.writeName("array");
+        }
+    }
+
+    @Benchmark
+    public void decode(@NotNull Input input, @NotNull Blackhole blackhole) {
+        blackhole.consume(input.bsonArrayCodec.decode(input.reader, DecoderContext.builder().build()));
+    }
+
+    @Benchmark
+    public void encode(@NotNull Input input, @NotNull Blackhole blackhole) {
+        input.bsonArrayCodec.encode(input.writer, input.bsonValues, EncoderContext.builder().build());
+        blackhole.consume(input);
+    }
+}