Resolve discriminators to codecs once per codec (plan item 3.3)

core/src/main/java/dev/morphia/mapping/codec/pojo/EntityDecoder.java

@@ -3,7 +3,6 @@
 import com.mongodb.lang.Nullable;
 
 import dev.morphia.annotations.internal.MorphiaInternal;
-import dev.morphia.mapping.DiscriminatorLookup;
 import dev.morphia.mapping.codec.DecodeSession;
 import dev.morphia.mapping.codec.MorphiaInstanceCreator;
 
@@ -15,7 +14,6 @@
 import org.bson.codecs.Decoder;
 import org.bson.codecs.DecoderContext;
 import org.bson.codecs.configuration.CodecConfigurationException;
-import org.bson.codecs.configuration.CodecRegistry;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -31,9 +29,12 @@
 public class EntityDecoder<T> implements Decoder<T> {
     private static final Logger LOG = LoggerFactory.getLogger(EntityDecoder.class);
     private static final DecoderContext CHECKED_DISCRIMINATOR = DecoderContext.builder().checkedDiscriminator(true).build();
+    private static final DecoderContext DEFAULT_CONTEXT = DecoderContext.builder().build();
 
     private final MorphiaCodec<T> morphiaCodec;
     private final EntityModel classModel;
+    // decodes a pre-read _id the way references read theirs, so the session keys match; resolved on first use
+    private Codec<Object> idKeyCodec;
 
     protected EntityDecoder(MorphiaCodec<T> morphiaCodec) {
         this.morphiaCodec = morphiaCodec;
@@ -89,8 +90,7 @@ public T decode(BsonReader reader, DecoderContext decoderContext) {
 
             return instance;
         } else {
-            entity = getCodecFromDocument(reader, classModel.useDiscriminator(), classModel.discriminatorKey(),
-                    morphiaCodec.getRegistry(), morphiaCodec.getDiscriminatorLookup(), morphiaCodec)
+            entity = getCodecFromDocument(reader, classModel.useDiscriminator(), classModel.discriminatorKey(), morphiaCodec)
                     .decode(reader, CHECKED_DISCRIMINATOR);
         }
 
@@ -169,7 +169,7 @@ protected void decodeProperties(BsonReader reader, DecoderContext decoderContext
         while (reader.readBsonType() != BsonType.END_OF_DOCUMENT) {
             String name = reader.readName();
             if (classModel.useDiscriminator() && classModel.discriminatorKey().equals(name)) {
-                reader.readString();
+                reader.skipValue();
             } else {
                 decodeModel(reader, decoderContext, instanceCreator, plan, classModel.getProperty(name));
             }
@@ -178,7 +178,6 @@ protected void decodeProperties(BsonReader reader, DecoderContext decoderContext
     }
 
     protected Codec<T> getCodecFromDocument(BsonReader reader, boolean useDiscriminator, String discriminatorKey,
-            CodecRegistry registry, DiscriminatorLookup discriminatorLookup,
             Codec<T> defaultCodec) {
         Codec<T> codec = null;
         if (useDiscriminator) {
@@ -187,7 +186,7 @@ protected Codec<T> getCodecFromDocument(BsonReader reader, boolean useDiscrimina
                 reader.readStartDocument();
                 while (codec == null && reader.readBsonType() != BsonType.END_OF_DOCUMENT) {
                     if (discriminatorKey.equals(reader.readName())) {
-                        codec = (Codec<T>) registry.get(discriminatorLookup.lookup(reader.readString()));
+                        codec = (Codec<T>) morphiaCodec.codecForDiscriminator(reader.readString());
                     } else {
                         reader.skipValue();
                     }
@@ -211,9 +210,10 @@ private Object peekId(BsonReader reader) {
             while (reader.readBsonType() != BsonType.END_OF_DOCUMENT) {
                 String name = reader.readName();
                 if ("_id".equals(name) || name.equals(idName)) {
-                    return morphiaCodec.getRegistry()
-                            .get(Object.class)
-                            .decode(reader, DecoderContext.builder().build());
+                    if (idKeyCodec == null) {
+                        idKeyCodec = morphiaCodec.getRegistry().get(Object.class);
+                    }
+                    return idKeyCodec.decode(reader, DEFAULT_CONTEXT);
                 } else {
                     reader.skipValue();
                 }

core/src/main/java/dev/morphia/mapping/codec/pojo/LifecycleDecoder.java

@@ -41,9 +41,7 @@ public T decode(BsonReader reader, DecoderContext decoderContext) {
         if (model.useDiscriminator()) {
             String discriminator = document.getString(model.discriminatorKey());
             if (discriminator != null) {
-                Class<?> discriminatorClass = getMorphiaCodec().getDiscriminatorLookup().lookup(discriminator);
-                // need to load the codec to initialize cachedCodecs in field models
-                Codec<?> codec = getMorphiaCodec().getRegistry().get(discriminatorClass);
+                Codec<?> codec = getMorphiaCodec().codecForDiscriminator(discriminator);
                 if (codec instanceof MorphiaCodec) {
                     model = ((MorphiaCodec<?>) codec).getEntityModel();
                     plan = ((MorphiaCodec<?>) codec).getPlan();

core/src/main/java/dev/morphia/mapping/codec/pojo/MorphiaCodec.java

@@ -1,6 +1,8 @@
 package dev.morphia.mapping.codec.pojo;
 
 import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 import dev.morphia.MorphiaDatastore;
 import dev.morphia.annotations.internal.MorphiaInternal;
@@ -48,6 +50,8 @@ public class MorphiaCodec<T> implements CollectibleCodec<T> {
     private final DiscriminatorLookup discriminatorLookup;
     private final Conversions conversions;
     private final CodecPlan plan;
+    // discriminator -> codec for the other types a document read through this codec can name; package-visible for tests
+    final Map<String, Codec<?>> discriminatorCodecs = new ConcurrentHashMap<>();
     private EntityEncoder<T> encoder;
     private EntityDecoder<T> decoder;
     private MorphiaDatastore datastore;
@@ -175,6 +179,31 @@ public EntityModel getEntityModel() {
         return entityModel;
     }
 
+    /**
+     * @param discriminator a discriminator value read from a document
+     * @return the codec for the type it names: this codec when it names this codec's own type, which is the usual
+     *         case, otherwise the registry's codec for that type, looked up once per discriminator
+     * @hidden
+     * @morphia.internal
+     */
+    @MorphiaInternal
+    @SuppressWarnings("unchecked")
+    public Codec<? extends T> codecForDiscriminator(String discriminator) {
+        if (discriminator.equals(entityModel.discriminator())) {
+            return this;
+        }
+        Codec<?> codec = discriminatorCodecs.get(discriminator);
+        if (codec == null) {
+            // not computeIfAbsent: building the codec can look up other codecs, which may land back here
+            codec = registry.get(discriminatorLookup.lookup(discriminator));
+            Codec<?> existing = discriminatorCodecs.putIfAbsent(discriminator, codec);
+            if (existing != null) {
+                codec = existing;
+            }
+        }
+        return (Codec<? extends T>) codec;
+    }
+
     /**
      * @return the plan this codec encodes and decodes with
      * @hidden

core/src/test/java/dev/morphia/mapping/codec/pojo/TestCodecForDiscriminator.java

@@ -0,0 +1,58 @@
+package dev.morphia.mapping.codec.pojo;
+
+import java.util.List;
+
+import com.mongodb.client.MongoClient;
+import com.mongodb.client.MongoClients;
+
+import dev.morphia.MorphiaDatastore;
+import dev.morphia.config.ManualMorphiaConfig;
+import dev.morphia.mapping.DiscriminatorFunction;
+import dev.morphia.mapping.MapperType;
+import dev.morphia.test.mapping.codec.harness.models.HarnessCircle;
+import dev.morphia.test.mapping.codec.harness.models.HarnessShape;
+import dev.morphia.test.mapping.codec.harness.models.HarnessSquare;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+/**
+ * Checks what {@link MorphiaCodec#codecForDiscriminator(String)} caches: nothing for its own discriminator, which it
+ * answers without a lookup, and one entry per other discriminator however often it's asked.
+ */
+public class TestCodecForDiscriminator {
+    @ParameterizedTest
+    @EnumSource(MapperType.class)
+    @SuppressWarnings("unchecked")
+    void cachesOnlyOtherTypes(MapperType mapper) {
+        // nothing here talks to a server
+        try (MongoClient client = MongoClients.create("mongodb://localhost:1")) {
+            MorphiaDatastore datastore = new MorphiaDatastore(client, ManualMorphiaConfig.configure()
+                    .database("codec_for_discriminator")
+                    .mapper(mapper)
+                    .packages(List.of(HarnessShape.class.getPackageName()))
+                    .discriminator(DiscriminatorFunction.className())
+                    .applyIndexes(false)
+                    .applyCaps(false)
+                    .applyDocumentValidations(false));
+            MorphiaCodec<HarnessShape> shapes = (MorphiaCodec<HarnessShape>) datastore.getCodecRegistry().get(HarnessShape.class);
+            MorphiaCodec<?> circles = (MorphiaCodec<?>) datastore.getCodecRegistry().get(HarnessCircle.class);
+            String circle = circles.getEntityModel().discriminator();
+            String square = datastore.getMapper().getEntityModel(HarnessSquare.class).discriminator();
+
+            for (int i = 0; i < 3; i++) {
+                Assertions.assertSame(circles, circles.codecForDiscriminator(circle));
+            }
+            Assertions.assertTrue(circles.discriminatorCodecs.isEmpty(), mapper + ": a self-match needs no lookup");
+
+            for (int i = 0; i < 3; i++) {
+                shapes.codecForDiscriminator(circle);
+                shapes.codecForDiscriminator(square);
+            }
+            Assertions.assertEquals(2, shapes.discriminatorCodecs.size(),
+                    mapper + ": one cached entry per other discriminator: " + shapes.discriminatorCodecs.keySet());
+            Assertions.assertSame(circles, shapes.discriminatorCodecs.get(circle));
+        }
+    }
+}

core/src/test/java/dev/morphia/test/mapping/codec/harness/CodecHarnessTest.java

@@ -39,6 +39,7 @@
 import dev.morphia.test.mapping.codec.harness.models.HarnessCallbackCircle;
 import dev.morphia.test.mapping.codec.harness.models.HarnessCallbackHolder;
 import dev.morphia.test.mapping.codec.harness.models.HarnessCallbackShape;
+import dev.morphia.test.mapping.codec.harness.models.HarnessCircle;
 import dev.morphia.test.mapping.codec.harness.models.HarnessCollections;
 import dev.morphia.test.mapping.codec.harness.models.HarnessConstructed;
 import dev.morphia.test.mapping.codec.harness.models.HarnessDrawing;
@@ -52,6 +53,7 @@
 import dev.morphia.test.mapping.codec.harness.models.HarnessReferences;
 import dev.morphia.test.mapping.codec.harness.models.HarnessRenames;
 import dev.morphia.test.mapping.codec.harness.models.HarnessScalars;
+import dev.morphia.test.mapping.codec.harness.models.HarnessShape;
 import dev.morphia.test.mapping.codec.harness.models.HarnessStringListBox;
 
 import org.bson.BsonBinaryReader;
@@ -292,6 +294,22 @@ void codecPlanResolvesEveryPropertyBySlot() {
         }
     }
 
+    @Test
+    void codecForDiscriminatorResolvesEachTypeOnce() {
+        for (MapperType mapper : MapperType.values()) {
+            MorphiaDatastore datastore = datastore(Variant.DEFAULTS, mapper);
+            MorphiaCodec<HarnessShape> shapes = (MorphiaCodec<HarnessShape>) codec(datastore, HarnessShape.class);
+            MorphiaCodec<?> circles = (MorphiaCodec<?>) codec(datastore, HarnessCircle.class);
+
+            assertSame(circles, circles.codecForDiscriminator(circles.getEntityModel().discriminator()),
+                    mapper + ": a codec's own discriminator resolves to itself");
+            Codec<?> circle = shapes.codecForDiscriminator(circles.getEntityModel().discriminator());
+            assertSame(circles, circle, mapper + ": a subtype's discriminator resolves to the subtype's codec");
+            assertSame(circle, shapes.codecForDiscriminator(circles.getEntityModel().discriminator()),
+                    mapper + ": and the same codec every time");
+        }
+    }
+
     @Test
     void readsLoadOnlyProperties() {
         // @LoadOnly values are never written, so the fixtures can't show that they are read

design/mapping-game-plan.md

@@ -217,7 +217,7 @@ JMH numbers on both mappers.
   `Conversions`.
 - Use a static `Pattern` for `arg[0-9]+`.
 
-### 3.3 Discriminator self-match, step 1 (#3, B2)
+### ~~3.3 Discriminator self-match, step 1 (#3, B2)~~
 - Add `codecForDiscriminator(d)`, which returns `this` on a self-match and otherwise uses a per-codec CHM.
 - Use `skipValue()` instead of `readString()` on `_t`.
 - Step 2, a single pass with no rewind, is moved to 3.17.