Simplify CodecCache (#1018)

JAVA-4751

Author: stIncMale

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

@@ -17,7 +17,6 @@
 package org.bson.internal;
 
 import org.bson.codecs.Codec;
-import org.bson.codecs.configuration.CodecConfigurationException;
 
 import java.lang.reflect.Type;
 import java.util.List;
@@ -26,7 +25,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 
-import static java.lang.String.format;
+import static org.bson.assertions.Assertions.assertNotNull;
 
 final class CodecCache {
 
@@ -65,29 +64,18 @@ public String toString() {
         }
     }
 
-    private final ConcurrentMap<CodecCacheKey, Optional<Codec<?>>> codecCache = new ConcurrentHashMap<>();
+    private final ConcurrentMap<CodecCacheKey, Codec<?>> codecCache = new ConcurrentHashMap<>();
 
-    public boolean containsKey(final CodecCacheKey codecCacheKey) {
-        return codecCache.containsKey(codecCacheKey);
+    public <T> Codec<T> putIfAbsent(final CodecCacheKey codecCacheKey, final Codec<T> codec) {
+        assertNotNull(codec);
+        @SuppressWarnings("unchecked")
+        Codec<T> prevCodec = (Codec<T>) codecCache.putIfAbsent(codecCacheKey, codec);
+        return prevCodec == null ? codec : prevCodec;
     }
 
-    public void put(final CodecCacheKey codecCacheKey, final Codec<?> codec){
-        codecCache.put(codecCacheKey, Optional.ofNullable(codec));
-    }
-
-    @SuppressWarnings("unchecked")
-    public synchronized <T> Codec<T> putIfMissing(final CodecCacheKey codecCacheKey, final Codec<T> codec) {
-        Optional<Codec<?>> cachedCodec = codecCache.computeIfAbsent(codecCacheKey, clz -> Optional.of(codec));
-        if (cachedCodec.isPresent()) {
-            return (Codec<T>) cachedCodec.get();
-        }
-        codecCache.put(codecCacheKey, Optional.of(codec));
-        return codec;
-    }
-
-    @SuppressWarnings("unchecked")
-    public <T> Codec<T> getOrThrow(final CodecCacheKey codecCacheKey) {
-        return (Codec<T>) codecCache.getOrDefault(codecCacheKey, Optional.empty()).orElseThrow(
-                () -> new CodecConfigurationException(format("Can't find a codec for %s.", codecCacheKey)));
+    public <T> Optional<Codec<T>> get(final CodecCacheKey codecCacheKey) {
+        @SuppressWarnings("unchecked")
+        Codec<T> codec = (Codec<T>) codecCache.get(codecCacheKey);
+        return Optional.ofNullable(codec);
     }
 }

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

@@ -18,6 +18,7 @@
 
 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;
@@ -67,19 +68,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 codecCache.<T>get(codecCacheKey).orElseGet(() -> {
             for (CodecProvider provider : codecProviders) {
                 Codec<T> 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());
                     }
-                    return codecCache.putIfMissing(codecCacheKey, codec);
+                    return codecCache.putIfAbsent(codecCacheKey, codec);
                 }
             }
-            codecCache.put(codecCacheKey, null);
-        }
-        return codecCache.getOrThrow(codecCacheKey);
+            throw new CodecConfigurationException(format("Can't find a codec for %s.", codecCacheKey));
+        });
     }
 
     @Override

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

@@ -17,7 +17,6 @@
 package org.bson.internal
 
 import org.bson.codecs.MinKeyCodec
-import org.bson.codecs.configuration.CodecConfigurationException
 import org.bson.types.MinKey
 import spock.lang.Specification
 
@@ -28,40 +27,38 @@ class CodecCacheSpecification extends Specification {
         def codec = new MinKeyCodec()
         def cache = new CodecCache()
         def cacheKey = new CodecCache.CodecCacheKey(MinKey, null)
-        cache.put(cacheKey, codec)
+        cache.putIfAbsent(cacheKey, codec)
 
         then:
-        cache.getOrThrow(cacheKey).is(codec)
+        cache.get(cacheKey).get().is(codec)
     }
 
-    def 'should throw if codec for class does not exist'() {
+    def 'should return empty if codec for class does not exist'() {
         when:
         def cache = new CodecCache()
         def cacheKey = new CodecCache.CodecCacheKey(MinKey, null)
-        cache.getOrThrow(cacheKey)
 
         then:
-        thrown(CodecConfigurationException)
+        !cache.get(cacheKey).isPresent()
     }
 
     def 'should return the cached codec if a codec for the parameterized class exists'() {
         when:
         def codec = new MinKeyCodec()
         def cache = new CodecCache()
         def cacheKey = new CodecCache.CodecCacheKey(List, [Integer])
-        cache.put(cacheKey, codec)
+        cache.putIfAbsent(cacheKey, codec)
 
         then:
-        cache.getOrThrow(cacheKey).is(codec)
+        cache.get(cacheKey).get().is(codec)
     }
 
-    def 'should throw if codec for the parameterized class does not exist'() {
+    def 'should return empty if codec for the parameterized class does not exist'() {
         when:
         def cache = new CodecCache()
         def cacheKey = new CodecCache.CodecCacheKey(List, [Integer])
-        cache.getOrThrow(cacheKey)
 
         then:
-        thrown(CodecConfigurationException)
+        !cache.get(cacheKey).isPresent()
     }
 }

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
 }
-