JAVA-2934: Handle empty non-final pages in ReactiveResultSetSubscription (#1544)

Author: dssysolyatin

changelog/README.md

@@ -2,6 +2,10 @@
 
 <!-- Note: contrary to 3.x, insert new entries *first* in their section -->
 
+### 4.12.0 (in progress)
+
+- [bug] JAVA-2934: Handle empty non-final pages in ReactiveResultSetSubscription
+
 ### 4.11.0
 
 - [improvement] JAVA-2930: Allow Micrometer to record histograms for timers

core/src/main/java/com/datastax/dse/driver/internal/core/cql/reactive/ReactiveResultSetSubscription.java

@@ -277,13 +277,10 @@ private Object tryNext() {
         if (pages.poll() == null) {
           throw new AssertionError("Queue is empty, this should not happen");
         }
-        current = pages.peek();
         // if the next page is readily available,
         // serve its first row now, no need to wait
         // for the next drain.
-        if (current != null && current.hasMoreRows()) {
-          return current.nextRow();
-        }
+        return tryNext();
       }
     }
     // No items available right now.

core/src/test/java/com/datastax/dse/driver/internal/core/cql/reactive/ReactiveResultSetSubscriptionTest.java

@@ -145,4 +145,34 @@ public void should_report_error_on_intermediary_page() {
     assertThat(wasAppliedSubscriber.getElements()).hasSize(1).containsExactly(true);
     assertThat(wasAppliedSubscriber.getError()).isNull();
   }
+
+  @Test
+  public void should_handle_empty_non_final_pages() {
+    CompletableFuture<AsyncResultSet> future1 = new CompletableFuture<>();
+    CompletableFuture<AsyncResultSet> future2 = new CompletableFuture<>();
+    CompletableFuture<AsyncResultSet> future3 = new CompletableFuture<>();
+    MockAsyncResultSet page1 = new MockAsyncResultSet(10, future2);
+    MockAsyncResultSet page2 = new MockAsyncResultSet(0, future3);
+    MockAsyncResultSet page3 = new MockAsyncResultSet(10, null);
+    TestSubscriber<ReactiveRow> mainSubscriber = new TestSubscriber<>(1);
+    TestSubscriber<ColumnDefinitions> colDefsSubscriber = new TestSubscriber<>();
+    TestSubscriber<ExecutionInfo> execInfosSubscriber = new TestSubscriber<>();
+    TestSubscriber<Boolean> wasAppliedSubscriber = new TestSubscriber<>();
+    ReactiveResultSetSubscription<AsyncResultSet> subscription =
+        new ReactiveResultSetSubscription<>(
+            mainSubscriber, colDefsSubscriber, execInfosSubscriber, wasAppliedSubscriber);
+    mainSubscriber.onSubscribe(subscription);
+    subscription.start(() -> future1);
+    future1.complete(page1);
+    future2.complete(page2);
+    // emulate backpressure
+    subscription.request(1);
+    future3.complete(page3);
+    subscription.request(Long.MAX_VALUE);
+    mainSubscriber.awaitTermination();
+    assertThat(mainSubscriber.getError()).isNull();
+    List<Row> expected = new ArrayList<>(page1.currentPage());
+    expected.addAll(page3.currentPage());
+    assertThat(mainSubscriber.getElements()).hasSize(20).extracting("row").isEqualTo(expected);
+  }
 }

core/src/test/java/com/datastax/dse/driver/internal/core/cql/reactive/TestSubscriber.java

@@ -31,16 +31,25 @@ public class TestSubscriber<T> implements Subscriber<T> {
 
   private final List<T> elements = new ArrayList<>();
   private final CountDownLatch latch = new CountDownLatch(1);
+  private final long demand;
   private Subscription subscription;
   private Throwable error;
 
+  public TestSubscriber() {
+    this.demand = Long.MAX_VALUE;
+  }
+
+  public TestSubscriber(long demand) {
+    this.demand = demand;
+  }
+
   @Override
   public void onSubscribe(Subscription s) {
     if (subscription != null) {
       fail("already subscribed");
     }
     subscription = s;
-    s.request(Long.MAX_VALUE);
+    subscription.request(demand);
   }
 
   @Override
@@ -71,5 +80,6 @@ public List<T> getElements() {
 
   public void awaitTermination() {
     Uninterruptibles.awaitUninterruptibly(latch, 1, TimeUnit.MINUTES);
+    if (latch.getCount() > 0) fail("subscriber not terminated");
   }
 }