Merge pull request #279 from zenios/fix/issue-273-reactive-backpressure

Fix reactive data loader hanging on backpressured publishers (#273)

Author: bbakerman

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

@@ -34,6 +34,11 @@ abstract class AbstractBatchSubscriber<K, V, T> implements Subscriber<T> {
     boolean onErrorCalled = false;
     boolean onCompleteCalled = false;
 
+    // the upstream subscription and how much demand we currently have outstanding with it.
+    // guarded by lock (see requestMore / requestMoreIfNeeded).
+    Subscription subscription;
+    long pendingDemand = 0;
+
     AbstractBatchSubscriber(
             CompletableFuture<List<V>> valuesFuture,
             List<K> keys,
@@ -50,7 +55,13 @@ abstract class AbstractBatchSubscriber<K, V, T> implements Subscriber<T> {
 
     @Override
     public void onSubscribe(Subscription subscription) {
-        subscription.request(keys.size());
+        lock.lock();
+        try {
+            this.subscription = subscription;
+            requestMore();
+        } finally {
+            lock.unlock();
+        }
     }
 
     @Override
@@ -59,6 +70,57 @@ public void onNext(T v) {
         assertState(!onCompleteCalled, () -> "onComplete has already been called; onNext may not be invoked.");
     }
 
+    /**
+     * Requests the next window of demand (sized to the key count) from the upstream subscription.
+     * <p>
+     * Must be called while holding {@link #lock}.
+     */
+    private void requestMore() {
+        long n = keys.size();
+        if (n <= 0) {
+            return;
+        }
+        pendingDemand += n;
+        subscription.request(n);
+    }
+
+    /**
+     * Called by the concrete subscribers once they have processed a value in {@code onNext}.
+     * <p>
+     * A reactive publisher only emits while it has outstanding demand. We originally only ever
+     * requested {@code keys.size()} once, so a publisher that emitted values not matching our keys
+     * (or simply emitted them lazily as more demand arrived) could leave us blocked forever waiting
+     * for a value that needed another request to be delivered. So we:
+     * <ol>
+     *     <li>re-request another window whenever the outstanding demand drains to zero, until the
+     *     publisher completes or errors, and</li>
+     *     <li>once every key has a result there is nothing left to wait for, so we cancel the upstream
+     *     subscription and complete ourselves rather than blocking on a publisher that may never call
+     *     {@code onComplete}.</li>
+     * </ol>
+     * <p>
+     * Must be called while holding {@link #lock}.
+     */
+    void requestMoreIfNeeded() {
+        if (onCompleteCalled || onErrorCalled) {
+            return;
+        }
+        if (allResultsReceived()) {
+            subscription.cancel();
+            onComplete();
+            return;
+        }
+        if (--pendingDemand <= 0) {
+            requestMore();
+        }
+    }
+
+    /**
+     * @return true once a result has been received for every key, so that we can complete early
+     * without waiting for the upstream publisher to finish
+     */
+    abstract boolean allResultsReceived();
+
     @Override
     public void onComplete() {
         assertState(!onErrorCalled, () -> "onError has already been called; onComplete may not be invoked.");

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

@@ -50,11 +50,19 @@ public void onNext(V value) {
 
             completedValues.add(value);
             idx++;
+
+            requestMoreIfNeeded();
         } finally {
             lock.unlock();
         }
     }
 
+    @Override
+    boolean allResultsReceived() {
+        // the values come back in key index order so once we have as many as keys we have them all
+        return completedValues.size() >= keys.size();
+    }
+
 
     @Override
     public void onComplete() {

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

@@ -59,12 +59,20 @@ public void onNext(Map.Entry<K, V> entry) {
             if (!futures.isEmpty()) {
                 completedValuesByKey.put(key, value);
             }
+
+            requestMoreIfNeeded();
         } finally {
             lock.unlock();
         }
 
     }
 
+    @Override
+    boolean allResultsReceived() {
+        // once every distinct requested key has a value we have everything we asked for
+        return completedValuesByKey.size() >= queuedFuturesByKey.size();
+    }
+
     @Override
     public void onComplete() {
         try {

src/test/java/org/dataloader/ReactiveBackpressureTest.java

@@ -0,0 +1,157 @@
+package org.dataloader;
+
+import org.awaitility.Duration;
+import org.junit.jupiter.api.Test;
+import org.reactivestreams.Publisher;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+
+import static java.util.Arrays.asList;
+import static org.awaitility.Awaitility.await;
+import static org.dataloader.DataLoaderFactory.newMappedPublisherDataLoader;
+import static org.dataloader.DataLoaderFactory.newPublisherDataLoader;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+
+/**
+ * Reproduces <a href="https://github.com/graphql-java/java-dataloader/issues/273">issue #273</a>.
+ * <p>
+ * A reactive publisher only emits while it has outstanding demand.  The reactive subscribers used
+ * to request {@code keys.size()} exactly once and then wait for {@code onComplete}.  A publisher
+ * that emits values lazily as more demand arrives - for example one that emits an unrelated value
+ * before the one that actually matches a key - would therefore leave the data loader blocked
+ * forever, because the matching value was never requested and the publisher never completed.
+ */
+public class ReactiveBackpressureTest {
+
+    @Test
+    public void mapped_loader_recovers_from_a_lazy_backpressured_publisher_that_never_completes() {
+        List<Collection<String>> loadCalls = new ArrayList<>();
+
+        DataLoader<String, String> loader = newMappedPublisherDataLoader((keys, subscriber) -> {
+            loadCalls.add(new ArrayList<>(keys));
+
+            List<Map.Entry<String, String>> items = new ArrayList<>();
+            // an entry that does not match any requested key - the original code consumed the one
+            // and only window of demand on this value and then blocked
+            items.add(Map.entry("an-unrelated-key", "an-unrelated-value"));
+            for (String key : keys) {
+                items.add(Map.entry(key, "value-" + key));
+            }
+            // note: this publisher NEVER calls onComplete - it only emits as demand arrives
+            new BackpressuredPublisher<>(items, false).subscribe(subscriber);
+        });
+
+        CompletableFuture<String> cfA = loader.load("a");
+        CompletableFuture<String> cfB = loader.load("b");
+        CompletableFuture<List<String>> dispatch = loader.dispatch();
+
+        // before the fix the matching values were never re-requested and dispatch never completed
+        await().atMost(Duration.FIVE_SECONDS).until(() -> cfA.isDone() && cfB.isDone() && dispatch.isDone());
+
+        assertThat(cfA.join(), equalTo("value-a"));
+        assertThat(cfB.join(), equalTo("value-b"));
+        // we got everything we asked for so we completed early rather than waiting on the publisher
+        assertThat(dispatch.join(), equalTo(asList("value-a", "value-b")));
+        assertThat(loadCalls, equalTo(List.of(asList("a", "b"))));
+    }
+
+    @Test
+    public void list_loader_recovers_from_a_lazy_backpressured_publisher_that_never_completes() {
+        List<Collection<String>> loadCalls = new ArrayList<>();
+
+        DataLoader<String, String> loader = newPublisherDataLoader((keys, subscriber) -> {
+            loadCalls.add(new ArrayList<>(keys));
+
+            List<String> items = new ArrayList<>();
+            for (String key : keys) {
+                items.add("value-" + key);
+            }
+            new BackpressuredPublisher<>(items, false).subscribe(subscriber);
+        });
+
+        CompletableFuture<String> cfA = loader.load("a");
+        CompletableFuture<String> cfB = loader.load("b");
+        CompletableFuture<List<String>> dispatch = loader.dispatch();
+
+        await().atMost(Duration.FIVE_SECONDS).until(() -> cfA.isDone() && cfB.isDone() && dispatch.isDone());
+
+        assertThat(cfA.join(), equalTo("value-a"));
+        assertThat(cfB.join(), equalTo("value-b"));
+        assertThat(dispatch.join(), equalTo(asList("value-a", "value-b")));
+        assertThat(loadCalls, equalTo(List.of(asList("a", "b"))));
+    }
+
+    /**
+     * A minimal reactive-streams {@link Publisher} that strictly honours backpressure: it only ever
+     * emits an item when there is outstanding demand for it, and (optionally) never signals
+     * {@code onComplete}.  It uses the standard work-in-progress drain loop so it behaves correctly
+     * even when {@code request} is called re-entrantly from within {@code onNext}.
+     */
+    static final class BackpressuredPublisher<T> implements Publisher<T> {
+        private final List<T> items;
+        private final boolean completeWhenDrained;
+
+        BackpressuredPublisher(List<T> items, boolean completeWhenDrained) {
+            this.items = items;
+            this.completeWhenDrained = completeWhenDrained;
+        }
+
+        @Override
+        public void subscribe(Subscriber<? super T> subscriber) {
+            subscriber.onSubscribe(new Subscription() {
+                final AtomicLong demand = new AtomicLong();
+                final AtomicInteger wip = new AtomicInteger();
+                volatile boolean cancelled;
+                int idx;
+                boolean completed;
+
+                @Override
+                public void request(long n) {
+                    if (n <= 0) {
+                        return;
+                    }
+                    demand.addAndGet(n);
+                    drain();
+                }
+
+                @Override
+                public void cancel() {
+                    cancelled = true;
+                }
+
+                private void drain() {
+                    if (wip.getAndIncrement() != 0) {
+                        // a drain is already in progress (possibly our own re-entrant caller) - it will
+                        // pick up the extra demand we just registered
+                        return;
+                    }
+                    int missed = 1;
+                    for (; ; ) {
+                        while (!cancelled && demand.get() > 0 && idx < items.size()) {
+                            T item = items.get(idx++);
+                            demand.decrementAndGet();
+                            subscriber.onNext(item);
+                        }
+                        if (!cancelled && completeWhenDrained && !completed && idx >= items.size()) {
+                            completed = true;
+                            subscriber.onComplete();
+                        }
+                        missed = wip.addAndGet(-missed);
+                        if (missed == 0) {
+                            break;
+                        }
+                    }
+                }
+            });
+        }
+    }
+}