Merge pull request #239 from graphql-java/lock-based-synchronisation

Updating the code to use reentrant locks

Author: bbakerman

src/main/java/org/dataloader/DataLoader.java

@@ -35,6 +35,8 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 
@@ -77,6 +79,7 @@ public class DataLoader<K, V extends @Nullable Object> {
     private final ValueCache<K, V> valueCache;
     private final DataLoaderOptions options;
     private final Object batchLoadFunction;
+    final Lock lock;
 
     @VisibleForTesting
     DataLoader(@Nullable String name, Object batchLoadFunction, @Nullable DataLoaderOptions options) {
@@ -93,7 +96,7 @@ public class DataLoader<K, V extends @Nullable Object> {
         this.batchLoadFunction = nonNull(batchLoadFunction);
         this.options = loaderOptions;
         this.name = name;
-
+        this.lock = new ReentrantLock();
         this.helper = new DataLoaderHelper<>(this, batchLoadFunction, loaderOptions, this.futureCache, this.valueCache, this.stats, clock);
     }
 
@@ -261,18 +264,16 @@ public CompletableFuture<List<V>> loadMany(List<K> keys, List<Object> keyContext
         nonNull(keys);
         nonNull(keyContexts);
 
-        synchronized (this) {
-            List<CompletableFuture<V>> collect = new ArrayList<>(keys.size());
-            for (int i = 0; i < keys.size(); i++) {
-                K key = keys.get(i);
-                Object keyContext = null;
-                if (i < keyContexts.size()) {
-                    keyContext = keyContexts.get(i);
-                }
-                collect.add(load(key, keyContext));
+        List<CompletableFuture<V>> collect = new ArrayList<>(keys.size());
+        for (int i = 0; i < keys.size(); i++) {
+            K key = keys.get(i);
+            Object keyContext = null;
+            if (i < keyContexts.size()) {
+                keyContext = keyContexts.get(i);
             }
-            return CompletableFutureKit.allOf(collect);
+            collect.add(load(key, keyContext));
         }
+        return CompletableFutureKit.allOf(collect);
     }
 
     /**
@@ -292,15 +293,13 @@ public CompletableFuture<List<V>> loadMany(List<K> keys, List<Object> keyContext
     public CompletableFuture<Map<K, V>> loadMany(Map<K, ?> keysAndContexts) {
         nonNull(keysAndContexts);
 
-        synchronized (this) {
-            Map<K, CompletableFuture<V>> collect = new HashMap<>(keysAndContexts.size());
-            for (Map.Entry<K, ?> entry : keysAndContexts.entrySet()) {
-                K key = entry.getKey();
-                Object keyContext = entry.getValue();
-                collect.put(key, load(key, keyContext));
-            }
-            return CompletableFutureKit.allOf(collect);
+        Map<K, CompletableFuture<V>> collect = new HashMap<>(keysAndContexts.size());
+        for (Map.Entry<K, ?> entry : keysAndContexts.entrySet()) {
+            K key = entry.getKey();
+            Object keyContext = entry.getValue();
+            collect.put(key, load(key, keyContext));
         }
+        return CompletableFutureKit.allOf(collect);
     }
 
     /**
@@ -376,9 +375,12 @@ public DataLoader<K, V> clear(K key) {
      */
     public DataLoader<K, V> clear(K key, BiConsumer<Void, Throwable> handler) {
         Object cacheKey = getCacheKey(key);
-        synchronized (this) {
+        try {
+            lock.lock();
             futureCache.delete(cacheKey);
             valueCache.delete(key).whenComplete(handler);
+        } finally {
+            lock.unlock();
         }
         return this;
     }
@@ -400,9 +402,12 @@ public DataLoader<K, V> clearAll() {
      * @return the data loader for fluent coding
      */
     public DataLoader<K, V> clearAll(BiConsumer<Void, Throwable> handler) {
-        synchronized (this) {
+        try {
+            lock.lock();
             futureCache.clear();
             valueCache.clear().whenComplete(handler);
+        } finally {
+            lock.unlock();
         }
         return this;
     }
@@ -442,10 +447,13 @@ public DataLoader<K, V> prime(K key, Exception error) {
      */
     public DataLoader<K, V> prime(K key, CompletableFuture<V> value) {
         Object cacheKey = getCacheKey(key);
-        synchronized (this) {
+        try {
+            lock.lock();
             if (!futureCache.containsKey(cacheKey)) {
                 futureCache.set(cacheKey, value);
             }
+        } finally {
+            lock.unlock();
         }
         return this;
     }

src/main/java/org/dataloader/DataLoaderHelper.java

@@ -28,6 +28,7 @@
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.Lock;
 
 import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
@@ -82,6 +83,7 @@ Object getCallContext() {
     private final StatisticsCollector stats;
     private final Clock clock;
     private final AtomicReference<Instant> lastDispatchTime;
+    private final Lock lock;
 
     DataLoaderHelper(DataLoader<K, V> dataLoader,
                      Object batchLoadFunction,
@@ -95,6 +97,7 @@ Object getCallContext() {
         this.loaderOptions = loaderOptions;
         this.futureCache = futureCache;
         this.valueCache = valueCache;
+        this.lock = dataLoader.lock;
         this.loaderQueue = new ArrayList<>();
         this.stats = stats;
         this.clock = clock;
@@ -111,7 +114,8 @@ public Instant getLastDispatchTime() {
     }
 
     Optional<CompletableFuture<V>> getIfPresent(K key) {
-        synchronized (dataLoader) {
+        try {
+            lock.lock();
             boolean cachingEnabled = loaderOptions.cachingEnabled();
             if (cachingEnabled) {
                 Object cacheKey = getCacheKey(nonNull(key));
@@ -124,26 +128,35 @@ Optional<CompletableFuture<V>> getIfPresent(K key) {
                 } catch (Exception ignored) {
                 }
             }
+        } finally {
+            lock.unlock();
         }
         return Optional.empty();
     }
 
     Optional<CompletableFuture<V>> getIfCompleted(K key) {
-        synchronized (dataLoader) {
+        try {
+            lock.lock();
+
             Optional<CompletableFuture<V>> cachedPromise = getIfPresent(key);
             if (cachedPromise.isPresent()) {
                 CompletableFuture<V> promise = cachedPromise.get();
                 if (promise.isDone()) {
                     return cachedPromise;
                 }
             }
+        } finally {
+            lock.unlock();
         }
         return Optional.empty();
     }
 
 
+    @GuardedBy("lock")
     CompletableFuture<V> load(K key, Object loadContext) {
-        synchronized (dataLoader) {
+        try {
+            lock.lock();
+
             boolean batchingEnabled = loaderOptions.batchingEnabled();
             boolean cachingEnabled = loaderOptions.cachingEnabled();
 
@@ -158,6 +171,8 @@ CompletableFuture<V> load(K key, Object loadContext) {
             ctx.onDispatched();
             cf.whenComplete(ctx::onCompleted);
             return cf;
+        } finally {
+            lock.unlock();
         }
     }
 
@@ -173,14 +188,17 @@ Object getCacheKeyWithContext(K key, Object context) {
                 loaderOptions.cacheKeyFunction().get().getKeyWithContext(key, context) : key;
     }
 
+    @GuardedBy("lock")
     DispatchResult<V> dispatch() {
         DataLoaderInstrumentationContext<DispatchResult<?>> instrCtx = ctxOrNoopCtx(instrumentation().beginDispatch(dataLoader));
 
         boolean batchingEnabled = loaderOptions.batchingEnabled();
         final List<K> keys;
         final List<Object> callContexts;
         final List<CompletableFuture<V>> queuedFutures;
-        synchronized (dataLoader) {
+        try {
+            lock.lock();
+
             int queueSize = loaderQueue.size();
             if (queueSize == 0) {
                 lastDispatchTime.set(now());
@@ -200,6 +218,8 @@ DispatchResult<V> dispatch() {
             });
             loaderQueue.clear();
             lastDispatchTime.set(now());
+        } finally {
+            lock.unlock();
         }
         if (!batchingEnabled) {
             instrCtx.onDispatched();
@@ -334,7 +354,7 @@ private void possiblyClearCacheEntriesOnExceptions(List<K> keys) {
         }
     }
 
-    @GuardedBy("dataLoader")
+    @GuardedBy("lock")
     private CompletableFuture<V> loadFromCache(K key, Object loadContext, boolean batchingEnabled) {
         final Object cacheKey = loadContext == null ? getCacheKey(key) : getCacheKeyWithContext(key, loadContext);
 
@@ -353,7 +373,7 @@ private CompletableFuture<V> loadFromCache(K key, Object loadContext, boolean ba
         return loadCallFuture;
     }
 
-    @GuardedBy("dataLoader")
+    @GuardedBy("lock")
     private CompletableFuture<V> queueOrInvokeLoader(K key, Object loadContext, boolean batchingEnabled, boolean cachingEnabled) {
         if (batchingEnabled) {
             CompletableFuture<V> loadCallFuture = new CompletableFuture<>();
@@ -606,8 +626,11 @@ private DataLoaderInstrumentation instrumentation() {
     }
 
     int dispatchDepth() {
-        synchronized (dataLoader) {
+        try {
+            lock.lock();
             return loaderQueue.size();
+        } finally {
+            lock.unlock();
         }
     }
 

src/main/java/org/dataloader/reactive/AbstractBatchSubscriber.java

@@ -10,6 +10,8 @@
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 
 import static org.dataloader.impl.Assertions.assertState;
 
@@ -25,6 +27,7 @@ abstract class AbstractBatchSubscriber<K, V, T> implements Subscriber<T> {
     final List<Object> callContexts;
     final List<CompletableFuture<V>> queuedFutures;
     final ReactiveSupport.HelperIntegration<K> helperIntegration;
+    final Lock lock = new ReentrantLock();
 
     List<K> clearCacheKeys = new ArrayList<>();
     List<V> completedValues = new ArrayList<>();

src/main/java/org/dataloader/reactive/BatchSubscriberImpl.java

@@ -29,58 +29,74 @@ class BatchSubscriberImpl<K, V> extends AbstractBatchSubscriber<K, V, V> {
         super(valuesFuture, keys, callContexts, queuedFutures, helperIntegration);
     }
 
-    // onNext may be called by multiple threads - for the time being, we pass 'synchronized' to guarantee
+    // onNext may be called by multiple threads - for the time being, we use a ReentrantLock to guarantee
     // correctness (at the cost of speed).
     @Override
-    public synchronized void onNext(V value) {
-        super.onNext(value);
+    public void onNext(V value) {
+        try {
+            lock.lock();
 
-        if (idx >= keys.size()) {
-            // hang on they have given us more values than we asked for in keys
-            // we cant handle this
-            return;
-        }
-        K key = keys.get(idx);
-        Object callContext = callContexts.get(idx);
-        CompletableFuture<V> future = queuedFutures.get(idx);
-        onNextValue(key, value, callContext, List.of(future));
+            super.onNext(value);
+
+            if (idx >= keys.size()) {
+                // hang on they have given us more values than we asked for in keys
+                // we cant handle this
+                return;
+            }
+            K key = keys.get(idx);
+            Object callContext = callContexts.get(idx);
+            CompletableFuture<V> future = queuedFutures.get(idx);
+            onNextValue(key, value, callContext, List.of(future));
 
-        completedValues.add(value);
-        idx++;
+            completedValues.add(value);
+            idx++;
+        } finally {
+            lock.unlock();
+        }
     }
 
 
     @Override
-    public synchronized void onComplete() {
-        super.onComplete();
-        if (keys.size() != completedValues.size()) {
-            // we have more or less values than promised
-            // we will go through all the outstanding promises and mark those that
-            // have not finished as failed
-            for (CompletableFuture<V> queuedFuture : queuedFutures) {
-                if (!queuedFuture.isDone()) {
-                    queuedFuture.completeExceptionally(new DataLoaderAssertionException("The size of the promised values MUST be the same size as the key list"));
+    public void onComplete() {
+        try {
+            lock.lock();
+            super.onComplete();
+            if (keys.size() != completedValues.size()) {
+                // we have more or less values than promised
+                // we will go through all the outstanding promises and mark those that
+                // have not finished as failed
+                for (CompletableFuture<V> queuedFuture : queuedFutures) {
+                    if (!queuedFuture.isDone()) {
+                        queuedFuture.completeExceptionally(new DataLoaderAssertionException("The size of the promised values MUST be the same size as the key list"));
+                    }
                 }
             }
+            possiblyClearCacheEntriesOnExceptions();
+            valuesFuture.complete(completedValues);
+        } finally {
+            lock.unlock();
         }
-        possiblyClearCacheEntriesOnExceptions();
-        valuesFuture.complete(completedValues);
     }
 
     @Override
-    public synchronized void onError(Throwable ex) {
-        super.onError(ex);
-        ex = unwrapThrowable(ex);
-        // Set the remaining keys to the exception.
-        for (int i = idx; i < queuedFutures.size(); i++) {
-            K key = keys.get(i);
-            CompletableFuture<V> future = queuedFutures.get(i);
-            if (!future.isDone()) {
-                future.completeExceptionally(ex);
-                // clear any cached view of this key because it failed
-                helperIntegration.clearCacheView(key);
+    public void onError(Throwable ex) {
+        try {
+            lock.lock();
+            super.onError(ex);
+            ex = unwrapThrowable(ex);
+            // Set the remaining keys to the exception.
+            for (int i = idx; i < queuedFutures.size(); i++) {
+                K key = keys.get(i);
+                CompletableFuture<V> future = queuedFutures.get(i);
+                if (!future.isDone()) {
+                    future.completeExceptionally(ex);
+                    // clear any cached view of this key because it failed
+                    helperIntegration.clearCacheView(key);
+                }
             }
+            valuesFuture.completeExceptionally(ex);
+        } finally {
+            lock.unlock();
         }
-        valuesFuture.completeExceptionally(ex);
     }
 }