Simplify CodecCache

bson/src/main/org/bson/internal/ProvidersCodecRegistry.java

@@ -18,21 +18,24 @@
 
 import org.bson.codecs.Codec;
 import org.bson.codecs.Parameterizable;
+import org.bson.codecs.configuration.CodecConfigurationException;
 import org.bson.codecs.configuration.CodecProvider;
 import org.bson.codecs.configuration.CodecRegistry;
-import org.bson.internal.CodecCache.CodecCacheKey;
 
 import java.lang.reflect.Type;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 
 import static java.lang.String.format;
 import static org.bson.assertions.Assertions.isTrueArgument;
 import static org.bson.assertions.Assertions.notNull;
 
 public final class ProvidersCodecRegistry implements CycleDetectingCodecRegistry {
     private final List<CodecProvider> codecProviders;
-    private final CodecCache codecCache = new CodecCache();
+    private final ConcurrentMap<CodecCacheKey, Codec<?>> codecCache = new ConcurrentHashMap<>();
 
     public ProvidersCodecRegistry(final List<? extends CodecProvider> codecProviders) {
         isTrueArgument("codecProviders must not be null or empty", codecProviders != null && codecProviders.size() > 0);
@@ -67,19 +70,18 @@ public <T> Codec<T> get(final Class<T> clazz, final CodecRegistry registry) {
     @SuppressWarnings({"unchecked"})
     public <T> Codec<T> get(final ChildCodecRegistry<T> context) {
         CodecCacheKey codecCacheKey = new CodecCacheKey(context.getCodecClass(), context.getTypes().orElse(null));
-        if (!codecCache.containsKey(codecCacheKey)) {
+        return (Codec<T>) codecCache.computeIfAbsent(codecCacheKey, k -> {
             for (CodecProvider provider : codecProviders) {
-                Codec<T> codec = provider.get(context.getCodecClass(), context);
+                Codec<?> codec = provider.get(context.getCodecClass(), context);
                 if (codec != null) {
                     if (codec instanceof Parameterizable && context.getTypes().isPresent()) {
-                        codec = (Codec<T>) ((Parameterizable) codec).parameterize(context, context.getTypes().get());
+                        codec = ((Parameterizable) codec).parameterize(context, context.getTypes().get());
                     }
-                    return codecCache.putIfMissing(codecCacheKey, codec);
+                    return codec;
                 }
             }
-            codecCache.put(codecCacheKey, null);
-        }
-        return codecCache.getOrThrow(codecCacheKey);
+            throw new CodecConfigurationException(format("Can't find a codec for %s.", k));
+        });
     }
 
     @Override
@@ -114,4 +116,39 @@ public String toString() {
                 + "codecProviders=" + codecProviders
                 + '}';
     }
+
+    private static final class CodecCacheKey {
+        private final Class<?> clazz;
+        private final List<Type> types;
+
+        CodecCacheKey(final Class<?> clazz, final List<Type> types) {
+            this.clazz = clazz;
+            this.types = types;
+        }
+
+        @Override
+        public boolean equals(final Object o) {
+            if (this == o) {
+                return true;
+            }
+            if (o == null || getClass() != o.getClass()) {
+                return false;
+            }
+            CodecCacheKey that = (CodecCacheKey) o;
+            return clazz.equals(that.clazz) && Objects.equals(types, that.types);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(clazz, types);
+        }
+
+        @Override
+        public String toString() {
+            return "CodecCacheKey{"
+                    + "clazz=" + clazz
+                    + ", types=" + types
+                    + '}';
+        }
+    }
 }

bson/src/test/unit/org/bson/internal/ProvidersCodecRegistrySpecification.groovy

@@ -115,22 +115,32 @@ class ProvidersCodecRegistrySpecification extends Specification {
 
     def 'get should use the codecCache'() {
         given:
-        def provider = Mock(CodecProvider)
+        def codec = Mock(Codec)
+        def provider = new CodecProvider() {
+            private int counter = 0
+
+            @Override
+            Codec get(final Class clazz, final CodecRegistry registry) {
+                if (counter == 0) {
+                    counter++
+                    return codec
+                }
+                throw new AssertionError((Object)'Must not be called more than once.')
+            }
+        }
 
         when:
         def registry = new ProvidersCodecRegistry([provider])
-        registry.get(MinKey)
+        def codecFromRegistry = registry.get(MinKey)
 
         then:
-        thrown(CodecConfigurationException)
-        1 * provider.get(MinKey, _)
+        codecFromRegistry == codec
 
         when:
-        registry.get(MinKey)
+        codecFromRegistry = registry.get(MinKey)
 
         then:
-        thrown(CodecConfigurationException)
-        0 * provider.get(MinKey, _)
+        codecFromRegistry == codec
     }
 
     def 'get with codec registry should return the codec from the first source that has one'() {
@@ -537,4 +547,3 @@ class Nested {
 class Simple {
     int value = 0
 }
-