Merge pull request #159 from graphql-java/reactive-streams-branch-move-reactive-classes-out-of-dataloader-helper

Reactive streams branch move reactive classes out of dataloader helper

Author: bbakerman

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

@@ -3,7 +3,7 @@
 import org.dataloader.annotations.GuardedBy;
 import org.dataloader.annotations.Internal;
 import org.dataloader.impl.CompletableFutureKit;
-import org.dataloader.impl.DataLoaderAssertionException;
+import org.dataloader.reactive.ReactiveSupport;
 import org.dataloader.scheduler.BatchLoaderScheduler;
 import org.dataloader.stats.StatisticsCollector;
 import org.dataloader.stats.context.IncrementBatchLoadCountByStatisticsContext;
@@ -12,13 +12,11 @@
 import org.dataloader.stats.context.IncrementLoadCountStatisticsContext;
 import org.dataloader.stats.context.IncrementLoadErrorCountStatisticsContext;
 import org.reactivestreams.Subscriber;
-import org.reactivestreams.Subscription;
 
 import java.time.Clock;
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.HashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
@@ -510,7 +508,7 @@ private CompletableFuture<List<V>> invokeMapBatchLoader(List<K> keys, BatchLoade
 
     private CompletableFuture<List<V>> invokeBatchPublisher(List<K> keys, List<Object> keyContexts, List<CompletableFuture<V>> queuedFutures, BatchLoaderEnvironment environment) {
         CompletableFuture<List<V>> loadResult = new CompletableFuture<>();
-        Subscriber<V> subscriber = new DataLoaderSubscriber(loadResult, keys, keyContexts, queuedFutures);
+        Subscriber<V> subscriber = ReactiveSupport.batchSubscriber(loadResult, keys, keyContexts, queuedFutures, helperIntegration());
 
         BatchLoaderScheduler batchLoaderScheduler = loaderOptions.getBatchLoaderScheduler();
         if (batchLoadFunction instanceof BatchPublisherWithContext) {
@@ -537,7 +535,7 @@ private CompletableFuture<List<V>> invokeBatchPublisher(List<K> keys, List<Objec
 
     private CompletableFuture<List<V>> invokeMappedBatchPublisher(List<K> keys, List<Object> keyContexts, List<CompletableFuture<V>> queuedFutures, BatchLoaderEnvironment environment) {
         CompletableFuture<List<V>> loadResult = new CompletableFuture<>();
-        Subscriber<Map.Entry<K, V>> subscriber = new DataLoaderMapEntrySubscriber(loadResult, keys, keyContexts, queuedFutures);
+        Subscriber<Map.Entry<K, V>> subscriber = ReactiveSupport.mappedBatchSubscriber(loadResult, keys, keyContexts, queuedFutures, helperIntegration());
 
         BatchLoaderScheduler batchLoaderScheduler = loaderOptions.getBatchLoaderScheduler();
         if (batchLoadFunction instanceof MappedBatchPublisherWithContext) {
@@ -625,246 +623,22 @@ private static <T> DispatchResult<T> emptyDispatchResult() {
         return (DispatchResult<T>) EMPTY_DISPATCH_RESULT;
     }
 
-    /**********************************************************************************************
-     * ********************************************************************************************
-     * <p>
-     * The reactive support classes start here
-     *
-     * @param <T> for two
-     **********************************************************************************************
-     **********************************************************************************************
-     */
-    private abstract class DataLoaderSubscriberBase<T> implements Subscriber<T> {
-
-        final CompletableFuture<List<V>> valuesFuture;
-        final List<K> keys;
-        final List<Object> callContexts;
-        final List<CompletableFuture<V>> queuedFutures;
-
-        List<K> clearCacheKeys = new ArrayList<>();
-        List<V> completedValues = new ArrayList<>();
-        boolean onErrorCalled = false;
-        boolean onCompleteCalled = false;
-
-        DataLoaderSubscriberBase(
-                CompletableFuture<List<V>> valuesFuture,
-                List<K> keys,
-                List<Object> callContexts,
-                List<CompletableFuture<V>> queuedFutures
-        ) {
-            this.valuesFuture = valuesFuture;
-            this.keys = keys;
-            this.callContexts = callContexts;
-            this.queuedFutures = queuedFutures;
-        }
-
-        @Override
-        public void onSubscribe(Subscription subscription) {
-            subscription.request(keys.size());
-        }
-
-        @Override
-        public void onNext(T v) {
-            assertState(!onErrorCalled, () -> "onError has already been called; onNext may not be invoked.");
-            assertState(!onCompleteCalled, () -> "onComplete has already been called; onNext may not be invoked.");
-        }
-
-        @Override
-        public void onComplete() {
-            assertState(!onErrorCalled, () -> "onError has already been called; onComplete may not be invoked.");
-            onCompleteCalled = true;
-        }
-
-        @Override
-        public void onError(Throwable throwable) {
-            assertState(!onCompleteCalled, () -> "onComplete has already been called; onError may not be invoked.");
-            onErrorCalled = true;
-
-            stats.incrementBatchLoadExceptionCount(new IncrementBatchLoadExceptionCountStatisticsContext<>(keys, callContexts));
-        }
-
-        /*
-         * A value has arrived - how do we complete the future that's associated with it in a common way
-         */
-        void onNextValue(K key, V value, Object callContext, List<CompletableFuture<V>> futures) {
-            if (value instanceof Try) {
-                // we allow the batch loader to return a Try so we can better represent a computation
-                // that might have worked or not.
-                //noinspection unchecked
-                Try<V> tryValue = (Try<V>) value;
-                if (tryValue.isSuccess()) {
-                    futures.forEach(f -> f.complete(tryValue.get()));
-                } else {
-                    stats.incrementLoadErrorCount(new IncrementLoadErrorCountStatisticsContext<>(key, callContext));
-                    futures.forEach(f -> f.completeExceptionally(tryValue.getThrowable()));
-                    clearCacheKeys.add(key);
-                }
-            } else {
-                futures.forEach(f -> f.complete(value));
+    private ReactiveSupport.HelperIntegration<K> helperIntegration() {
+        return new ReactiveSupport.HelperIntegration<>() {
+            @Override
+            public StatisticsCollector getStats() {
+                return stats;
             }
-        }
 
-        Throwable unwrapThrowable(Throwable ex) {
-            if (ex instanceof CompletionException) {
-                ex = ex.getCause();
+            @Override
+            public void clearCacheView(K key) {
+                dataLoader.clear(key);
             }
-            return ex;
-        }
-    }
-
-    private class DataLoaderSubscriber extends DataLoaderSubscriberBase<V> {
-
-        private int idx = 0;
-
-        private DataLoaderSubscriber(
-                CompletableFuture<List<V>> valuesFuture,
-                List<K> keys,
-                List<Object> callContexts,
-                List<CompletableFuture<V>> queuedFutures
-        ) {
-            super(valuesFuture, keys, callContexts, queuedFutures);
-        }
 
-        // onNext may be called by multiple threads - for the time being, we pass 'synchronized' to guarantee
-        // correctness (at the cost of speed).
-        @Override
-        public synchronized void onNext(V value) {
-            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++;
-        }
-
-
-        @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"));
-                    }
-                }
+            @Override
+            public void clearCacheEntriesOnExceptions(List<K> keys) {
+                possiblyClearCacheEntriesOnExceptions(keys);
             }
-            possiblyClearCacheEntriesOnExceptions(clearCacheKeys);
-            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
-                    dataLoader.clear(key);
-                }
-            }
-            valuesFuture.completeExceptionally(ex);
-        }
-
-    }
-
-    private class DataLoaderMapEntrySubscriber extends DataLoaderSubscriberBase<Map.Entry<K, V>> {
-
-        private final Map<K, Object> callContextByKey;
-        private final Map<K, List<CompletableFuture<V>>> queuedFuturesByKey;
-        private final Map<K, V> completedValuesByKey = new HashMap<>();
-
-
-        private DataLoaderMapEntrySubscriber(
-                CompletableFuture<List<V>> valuesFuture,
-                List<K> keys,
-                List<Object> callContexts,
-                List<CompletableFuture<V>> queuedFutures
-        ) {
-            super(valuesFuture, keys, callContexts, queuedFutures);
-            this.callContextByKey = new HashMap<>();
-            this.queuedFuturesByKey = new HashMap<>();
-            for (int idx = 0; idx < queuedFutures.size(); idx++) {
-                K key = keys.get(idx);
-                Object callContext = callContexts.get(idx);
-                CompletableFuture<V> queuedFuture = queuedFutures.get(idx);
-                callContextByKey.put(key, callContext);
-                queuedFuturesByKey.computeIfAbsent(key, k -> new ArrayList<>()).add(queuedFuture);
-            }
-        }
-
-
-        @Override
-        public synchronized void onNext(Map.Entry<K, V> entry) {
-            super.onNext(entry);
-            K key = entry.getKey();
-            V value = entry.getValue();
-
-            Object callContext = callContextByKey.get(key);
-            List<CompletableFuture<V>> futures = queuedFuturesByKey.getOrDefault(key, List.of());
-
-            onNextValue(key, value, callContext, futures);
-
-            // did we have an actual key for this value - ignore it if they send us one outside the key set
-            if (!futures.isEmpty()) {
-                completedValuesByKey.put(key, value);
-            }
-        }
-
-        @Override
-        public synchronized void onComplete() {
-            super.onComplete();
-
-            possiblyClearCacheEntriesOnExceptions(clearCacheKeys);
-            List<V> values = new ArrayList<>(keys.size());
-            for (K key : keys) {
-                V value = completedValuesByKey.get(key);
-                values.add(value);
-
-                List<CompletableFuture<V>> futures = queuedFuturesByKey.getOrDefault(key, List.of());
-                for (CompletableFuture<V> future : futures) {
-                    if (! future.isDone()) {
-                        // we have a future that never came back for that key
-                        // but the publisher is done sending in data - it must be null
-                        // e.g. for key X when found no value
-                        future.complete(null);
-                    }
-                }
-            }
-            valuesFuture.complete(values);
-        }
-
-        @Override
-        public synchronized void onError(Throwable ex) {
-            super.onError(ex);
-            ex = unwrapThrowable(ex);
-            // Complete the futures for the remaining keys with the exception.
-            for (int idx = 0; idx < queuedFutures.size(); idx++) {
-                K key = keys.get(idx);
-                List<CompletableFuture<V>> futures = queuedFuturesByKey.get(key);
-                if (!completedValuesByKey.containsKey(key)) {
-                    for (CompletableFuture<V> future : futures) {
-                        future.completeExceptionally(ex);
-                    }
-                    // clear any cached view of this key because they all failed
-                    dataLoader.clear(key);
-                }
-            }
-            valuesFuture.completeExceptionally(ex);
-        }
+        };
     }
 }