DATACASS-509 - Request subsequent result pages in ReactiveSession asynchronously.

We now no longer use a Scheduler to offload ResultSet's blocking paging but request the next result page asynchronously by adapting ResultSet.fetchMoreResults(). Result pages are requested once all elements of the previous ResultSet are emitted and the row publisher completes successfully. The Scheduler is no longer required.

The request progress is stored in a MonoProcessor to extend the result stream. Increase visibility of utility methods to avoid synthetic accessor creation.

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/config/AbstractReactiveCassandraConfiguration.java

@@ -15,8 +15,6 @@
  */
 package org.springframework.data.cassandra.config;
 
-import reactor.core.scheduler.Schedulers;
-
 import org.springframework.context.annotation.Bean;
 import org.springframework.data.cassandra.ReactiveSession;
 import org.springframework.data.cassandra.ReactiveSessionFactory;
@@ -47,7 +45,7 @@ public abstract class AbstractReactiveCassandraConfiguration extends AbstractCas
 	 */
 	@Bean
 	public ReactiveSession reactiveSession() {
-		return new DefaultBridgedReactiveSession(getRequiredSession(), Schedulers.elastic());
+		return new DefaultBridgedReactiveSession(getRequiredSession());
 	}
 
 	/**

spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/cql/session/DefaultBridgedReactiveSession.java

@@ -17,8 +17,8 @@ package org.springframework.data.cassandra.core.cql.session;
 
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.core.publisher.MonoProcessor;
 import reactor.core.scheduler.Scheduler;
-import reactor.core.scheduler.Schedulers;
 
 import java.util.List;
 import java.util.Map;
@@ -40,13 +40,12 @@ import com.google.common.util.concurrent.ListenableFuture;
  * Calls are deferred until a subscriber subscribes to the resulting {@link org.reactivestreams.Publisher}. The calls
  * are executed by subscribing to {@link ListenableFuture} and returning the result as calls complete.
  * <p>
- * {@link ResultSet} implements transparent paging that invokes in the middle of result streaming blocking calls to
- * Cassandra. {@link DefaultBridgedReactiveSession} uses therefore {@link ReactiveResultSet} to avoid client thread
- * blocking. Elements are emitted on netty EventLoop threads and transported by the provided {@link Scheduler}. However,
- * this is an intermediate solution until Datastax can provide a fully reactive driver.
+ * Elements are emitted on netty EventLoop threads. {@link ResultSet} allows {@link ResultSet#fetchMoreResults()
+ * asynchronous requesting} of subsequent pages. The next page is requested after emitting all elements of the previous
+ * page. However, this is an intermediate solution until Datastax can provide a fully reactive driver.
  * <p>
  * All CQL operations performed by this class are logged at debug level, using
- * "org.springframework.data.cassandra.core.cql.DefaultBridgedReactiveSession" as log category.
+ * {@code org.springframework.data.cassandra.core.cql.DefaultBridgedReactiveSession} as log category.
  * <p>
  *
  * @author Mark Paluch
@@ -61,21 +60,31 @@ public class DefaultBridgedReactiveSession implements ReactiveSession {
 	private final Logger logger = LoggerFactory.getLogger(getClass());
 
 	private final Session session;
-	private final Scheduler scheduler;
+
+	/**
+	 * Create a new {@link DefaultBridgedReactiveSession} for a {@link Session}.
+	 *
+	 * @param session must not be {@literal null}.
+	 * @since 2.1
+	 */
+	public DefaultBridgedReactiveSession(Session session) {
+
+		Assert.notNull(session, "Session must not be null");
+
+		this.session = session;
+	}
 
 	/**
 	 * Create a new {@link DefaultBridgedReactiveSession} for a {@link Session} and {@link Scheduler}.
 	 *
 	 * @param session must not be {@literal null}.
 	 * @param scheduler must not be {@literal null}.
+	 * @deprecated since 2.1. Use {@link #DefaultBridgedReactiveSession(Session)} as a {@link Scheduler} is no longer
+	 *             required to off-load {@link ResultSet}'s blocking behavior.
 	 */
+	@Deprecated
 	public DefaultBridgedReactiveSession(Session session, Scheduler scheduler) {
-
-		Assert.notNull(session, "Session must not be null");
-		Assert.notNull(scheduler, "Scheduler must not be null");
-
-		this.session = session;
-		this.scheduler = scheduler;
+		this(session);
 	}
 
 	/* (non-Javadoc)
@@ -133,7 +142,7 @@ public class DefaultBridgedReactiveSession implements ReactiveSession {
 					if (resultSetFuture.isDone()) {
 
 						try {
-							future.complete(new DefaultReactiveResultSet(resultSetFuture.getUninterruptibly(), scheduler));
+							future.complete(new DefaultReactiveResultSet(resultSetFuture.getUninterruptibly()));
 						} catch (Exception e) {
 							future.completeExceptionally(e);
 						}
@@ -144,8 +153,7 @@ public class DefaultBridgedReactiveSession implements ReactiveSession {
 			} catch (Exception e) {
 				return Mono.error(e);
 			}
-
-		}).subscribeOn(scheduler);
+		});
 	}
 
 	/* (non-Javadoc)
@@ -191,8 +199,7 @@ public class DefaultBridgedReactiveSession implements ReactiveSession {
 			} catch (Exception e) {
 				return Mono.error(e);
 			}
-
-		}).subscribeOn(scheduler);
+		});
 	}
 
 	/* (non-Javadoc)
@@ -219,14 +226,12 @@ public class DefaultBridgedReactiveSession implements ReactiveSession {
 		return session.getCluster();
 	}
 
-	private static class DefaultReactiveResultSet implements ReactiveResultSet {
+	static class DefaultReactiveResultSet implements ReactiveResultSet {
 
 		private final ResultSet resultSet;
-		private final Scheduler scheduler;
 
-		DefaultReactiveResultSet(ResultSet resultSet, Scheduler scheduler) {
+		DefaultReactiveResultSet(ResultSet resultSet) {
 			this.resultSet = resultSet;
-			this.scheduler = scheduler;
 		}
 
 		/* (non-Javadoc)
@@ -234,12 +239,50 @@ public class DefaultBridgedReactiveSession implements ReactiveSession {
 		 */
 		@Override
 		public Flux<Row> rows() {
+			return getRows(Mono.just(resultSet));
+		}
+
+		Flux<Row> getRows(Mono<ResultSet> nextResults) {
+
+			return nextResults.flatMapMany(it -> {
+
+				Flux<Row> rows = toRows(it);
+
+				if (it.isFullyFetched()) {
+					return rows;
+				}
+
+				MonoProcessor<ResultSet> processor = MonoProcessor.create();
+				return rows //
+						.doOnComplete(() -> fetchMore(it.fetchMoreResults(), processor)) //
+						.concatWith(getRows(processor));
+			});
+		}
+
+		static Flux<Row> toRows(ResultSet resultSet) {
 
 			int prefetch = Math.max(1, resultSet.getAvailableWithoutFetching());
+			return Flux.fromIterable(resultSet).take(prefetch);
+		}
 
-			return Flux.fromIterable(resultSet) //
-					.subscribeOn(scheduler) //
-					.publishOn(Schedulers.immediate(), prefetch); // limit prefetching to available size
+		static void fetchMore(ListenableFuture<ResultSet> future, MonoProcessor<ResultSet> sink) {
+
+			try {
+
+				future.addListener(() -> {
+
+					try {
+
+						sink.onNext(future.get());
+						sink.onComplete();
+					} catch (Exception e) {
+						sink.onError(e);
+					}
+				}, Runnable::run);
+
+			} catch (Exception e) {
+				sink.onError(e);
+			}
 		}
 
 		/* (non-Javadoc)

spring-data-cassandra/src/test/java/org/springframework/data/cassandra/core/cql/DefaultBridgedReactiveSessionIntegrationTests.java

@@ -21,6 +21,11 @@ import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Schedulers;
 import reactor.test.StepVerifier;
 
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.concurrent.ConcurrentLinkedQueue;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.springframework.data.cassandra.ReactiveResultSet;
@@ -28,6 +33,8 @@ import org.springframework.data.cassandra.core.cql.session.DefaultBridgedReactiv
 import org.springframework.data.cassandra.test.util.AbstractKeyspaceCreatingIntegrationTest;
 
 import com.datastax.driver.core.KeyspaceMetadata;
+import com.datastax.driver.core.QueryLogger;
+import com.datastax.driver.core.SimpleStatement;
 import com.datastax.driver.core.exceptions.SyntaxError;
 
 /**
@@ -40,7 +47,7 @@ public class DefaultBridgedReactiveSessionIntegrationTests extends AbstractKeysp
 	DefaultBridgedReactiveSession reactiveSession;
 
 	@Before
-	public void before() throws Exception {
+	public void before() {
 
 		this.session.execute("DROP TABLE IF EXISTS users;");
 
@@ -48,7 +55,7 @@ public class DefaultBridgedReactiveSessionIntegrationTests extends AbstractKeysp
 	}
 
 	@Test // DATACASS-335
-	public void executeShouldExecuteDeferred() throws Exception {
+	public void executeShouldExecuteDeferred() {
 
 		String query = "CREATE TABLE users (\n" + "  userid text PRIMARY KEY,\n" + "  first_name text\n" + ");";
 
@@ -72,7 +79,7 @@ public class DefaultBridgedReactiveSessionIntegrationTests extends AbstractKeysp
 	}
 
 	@Test // DATACASS-335
-	public void executeShouldReturnRows() throws Exception {
+	public void executeShouldReturnRows() {
 
 		session.execute("CREATE TABLE users (\n" + "  userid text PRIMARY KEY,\n" + "  first_name text\n" + ");");
 		session.execute("INSERT INTO users (userid, first_name) VALUES ('White', 'Walter');");
@@ -87,7 +94,7 @@ public class DefaultBridgedReactiveSessionIntegrationTests extends AbstractKeysp
 	}
 
 	@Test // DATACASS-335
-	public void executeShouldPrepareStatement() throws Exception {
+	public void executeShouldPrepareStatement() {
 
 		session.execute("CREATE TABLE users (\n" + "  userid text PRIMARY KEY,\n" + "  first_name text\n" + ");");
 
@@ -98,6 +105,37 @@ public class DefaultBridgedReactiveSessionIntegrationTests extends AbstractKeysp
 				}).verifyComplete();
 	}
 
+	@Test // DATACASS-509
+	public void shouldFetchBatches() {
+
+		String createTable = "CREATE TABLE users (\n" + "  userid text PRIMARY KEY,\n" + "  first_name text\n" + ");";
+		this.session.execute(createTable);
+
+		List<String> keys = new ArrayList<>();
+		for (int i = 0; i < 100; i++) {
+
+			String key = String.format("u-03%d", i);
+			String value = "v-" + i;
+			keys.add(key);
+
+			this.session.execute(String.format("INSERT INTO users (userid,first_name) VALUES ('%s', '%s');", key, value));
+		}
+
+		session.getCluster().register(QueryLogger.builder().build());
+
+		SimpleStatement statement = new SimpleStatement("SELECT * FROM users;");
+		statement.setFetchSize(10);
+
+		Mono<ReactiveResultSet> execution = reactiveSession.execute(statement);
+
+		Collection<String> received = new ConcurrentLinkedQueue<>();
+		StepVerifier.create(execution.flatMapMany(ReactiveResultSet::rows).map(row -> row.getString(0))) //
+				.recordWith(() -> received) //
+				.expectNextCount(100).verifyComplete();
+
+		assertThat(received).containsAll(keys).hasSize(100);
+	}
+
 	private KeyspaceMetadata getKeyspaceMetadata() {
 		return cluster.getMetadata().getKeyspace(this.session.getLoggedKeyspace());
 	}

spring-data-cassandra/src/test/java/org/springframework/data/cassandra/core/cql/DefaultBridgedReactiveSessionUnitTests.java

@@ -18,9 +18,13 @@ package org.springframework.data.cassandra.core.cql;
 import static org.assertj.core.api.Assertions.*;
 import static org.mockito.Mockito.*;
 
-import reactor.core.scheduler.Schedulers;
+import reactor.core.publisher.Flux;
+import reactor.test.StepVerifier;
 
+import java.util.ArrayDeque;
 import java.util.Collections;
+import java.util.Iterator;
+import java.util.Queue;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -28,15 +32,19 @@ import org.junit.runner.RunWith;
 import org.mockito.ArgumentMatchers;
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnitRunner;
+import org.springframework.data.cassandra.ReactiveResultSet;
 import org.springframework.data.cassandra.core.cql.session.DefaultBridgedReactiveSession;
 
 import com.datastax.driver.core.Cluster;
 import com.datastax.driver.core.PreparedStatement;
 import com.datastax.driver.core.RegularStatement;
+import com.datastax.driver.core.ResultSet;
 import com.datastax.driver.core.ResultSetFuture;
+import com.datastax.driver.core.Row;
 import com.datastax.driver.core.Session;
 import com.datastax.driver.core.SimpleStatement;
 import com.datastax.driver.core.Statement;
+import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 
 /**
@@ -54,16 +62,16 @@ public class DefaultBridgedReactiveSessionUnitTests {
 	private DefaultBridgedReactiveSession reactiveSession;
 
 	@Before
-	public void before() throws Exception {
+	public void before() {
 
-		reactiveSession = new DefaultBridgedReactiveSession(sessionMock, Schedulers.immediate());
+		reactiveSession = new DefaultBridgedReactiveSession(sessionMock);
 
 		when(sessionMock.executeAsync(any(Statement.class))).thenReturn(future);
 		when(sessionMock.prepareAsync(any(RegularStatement.class))).thenReturn(preparedStatementFuture);
 	}
 
 	@Test // DATACASS-335
-	public void executeStatementShouldForwardStatementToSession() throws Exception {
+	public void executeStatementShouldForwardStatementToSession() {
 
 		SimpleStatement statement = new SimpleStatement("SELECT *");
 
@@ -73,7 +81,7 @@ public class DefaultBridgedReactiveSessionUnitTests {
 	}
 
 	@Test // DATACASS-335
-	public void executeShouldForwardStatementToSession() throws Exception {
+	public void executeShouldForwardStatementToSession() {
 
 		reactiveSession.execute("SELECT *").subscribe();
 
@@ -81,7 +89,7 @@ public class DefaultBridgedReactiveSessionUnitTests {
 	}
 
 	@Test // DATACASS-335
-	public void executeWithValuesShouldForwardStatementToSession() throws Exception {
+	public void executeWithValuesShouldForwardStatementToSession() {
 
 		reactiveSession.execute("SELECT * WHERE a = ? and b = ?", "A", "B").subscribe();
 
@@ -89,7 +97,7 @@ public class DefaultBridgedReactiveSessionUnitTests {
 	}
 
 	@Test // DATACASS-335
-	public void executeWithValueMapShouldForwardStatementToSession() throws Exception {
+	public void executeWithValueMapShouldForwardStatementToSession() {
 
 		reactiveSession.execute("SELECT * WHERE a = ?", Collections.singletonMap("a", "value")).subscribe();
 
@@ -98,7 +106,7 @@ public class DefaultBridgedReactiveSessionUnitTests {
 	}
 
 	@Test // DATACASS-335
-	public void testPrepareQuery() throws Exception {
+	public void testPrepareQuery() {
 
 		reactiveSession.prepare("SELECT *").subscribe();
 
@@ -106,7 +114,7 @@ public class DefaultBridgedReactiveSessionUnitTests {
 	}
 
 	@Test // DATACASS-335
-	public void testPrepareStatement() throws Exception {
+	public void testPrepareStatement() {
 
 		SimpleStatement statement = new SimpleStatement("SELECT *");
 		reactiveSession.prepare(statement).subscribe();
@@ -115,7 +123,7 @@ public class DefaultBridgedReactiveSessionUnitTests {
 	}
 
 	@Test // DATACASS-335
-	public void testClose() throws Exception {
+	public void testClose() {
 
 		reactiveSession.close();
 
@@ -123,7 +131,7 @@ public class DefaultBridgedReactiveSessionUnitTests {
 	}
 
 	@Test // DATACASS-335
-	public void testIsClosed() throws Exception {
+	public void testIsClosed() {
 
 		when(reactiveSession.isClosed()).thenReturn(true);
 
@@ -134,7 +142,7 @@ public class DefaultBridgedReactiveSessionUnitTests {
 	}
 
 	@Test // DATACASS-335
-	public void testGetCluster() throws Exception {
+	public void testGetCluster() {
 
 		Cluster clusterMock = mock(Cluster.class);
 		when(sessionMock.getCluster()).thenReturn(clusterMock);
@@ -144,6 +152,128 @@ public class DefaultBridgedReactiveSessionUnitTests {
 		assertThat(result).isSameAs(clusterMock);
 	}
 
+	@Test // DATACASS-509
+	public void shouldReadNotMoreThanAvailable() throws Exception {
+
+		Iterator<Row> rows = mockIterator();
+
+		ResultSet resultSet = mock(ResultSet.class);
+		when(resultSet.getAvailableWithoutFetching()).thenReturn(10);
+		when(resultSet.iterator()).thenReturn(rows);
+		doAnswer(invocation -> {
+
+			Runnable listener = invocation.getArgument(0);
+			listener.run();
+
+			return null;
+		}).when(future).addListener(any(), any());
+
+		when(future.getUninterruptibly()).thenReturn(resultSet);
+		when(future.isDone()).thenReturn(true);
+		when(resultSet.isFullyFetched()).thenReturn(true);
+
+		reactiveSession.execute(new SimpleStatement("")).flatMapMany(ReactiveResultSet::rows).collectList().subscribe();
+
+		verify(rows, times(10)).next();
+		verify(resultSet, never()).fetchMoreResults();
+	}
+
+	@Test // DATACASS-509
+	public void shouldFetchMore() throws Exception {
+
+		Iterator<Row> rows = mockIterator();
+
+		ResultSet resultSet = mock(ResultSet.class);
+		when(resultSet.getAvailableWithoutFetching()).thenReturn(10);
+		when(resultSet.iterator()).thenReturn(rows);
+
+		ResultSet emptyResultSet = mock(ResultSet.class);
+		when(emptyResultSet.iterator()).thenReturn(Collections.emptyIterator());
+		when(emptyResultSet.isFullyFetched()).thenReturn(true);
+
+		doAnswer(invocation -> {
+
+			Runnable listener = invocation.getArgument(0);
+			listener.run();
+
+			return null;
+		}).when(future).addListener(any(), any());
+
+		when(future.getUninterruptibly()).thenReturn(resultSet);
+		when(future.isDone()).thenReturn(true);
+		when(resultSet.isFullyFetched()).thenReturn(false, true);
+		when(resultSet.fetchMoreResults()).thenReturn(Futures.immediateFuture(emptyResultSet));
+
+		Flux<Row> flux = reactiveSession.execute(new SimpleStatement("")).flatMapMany(ReactiveResultSet::rows);
+
+		StepVerifier.create(flux, 0).thenRequest(10).expectNextCount(10).then(() -> {
+
+			verify(rows, times(10)).next();
+			verify(resultSet).fetchMoreResults();
+		}).thenRequest(10).verifyComplete();
+	}
+
+	@Test // DATACASS-509
+	public void shouldFetchDependingOfCompletion() throws Exception {
+
+		Iterator<Row> rows = mockIterator();
+
+		Queue<Runnable> runnables = new ArrayDeque<>();
+		ResultSet resultSet = mock(ResultSet.class);
+		when(resultSet.getAvailableWithoutFetching()).thenReturn(10);
+		when(resultSet.iterator()).thenReturn(rows);
+
+		doAnswer(invocation -> {
+			runnables.offer(invocation.getArgument(0));
+			return null;
+		}).when(future).addListener(any(), any());
+
+		when(future.getUninterruptibly()).thenReturn(resultSet);
+		when(future.get()).thenReturn(resultSet);
+		when(future.isDone()).thenReturn(true);
+		when(resultSet.isFullyFetched()).thenReturn(false, false, true);
+		when(resultSet.fetchMoreResults()).thenReturn(future);
+
+		Flux<Row> flux = reactiveSession.execute(new SimpleStatement("")).flatMapMany(ReactiveResultSet::rows);
+
+		StepVerifier.create(flux, 0) //
+				.then(() -> runnables.poll().run()) // complete the first future from executeAsync()
+				.thenRequest(9).expectNextCount(9) //
+				.then(() -> {
+					// feed the 9 elements from the initial ResultSet
+					verify(resultSet, never()).fetchMoreResults();
+				}).thenRequest(1).expectNextCount(1) //
+				.then(() -> {
+
+					// initial ResultSet exhausted, fetch next chunk
+					verify(resultSet).fetchMoreResults();
+					runnables.poll().run();
+				}).thenRequest(1).expectNextCount(1) //
+				.then(() -> {
+
+					// first element from the second ResultSet received, no subsequent fetch
+					assertThat(runnables).isEmpty();
+				}).thenRequest(19).expectNextCount(9) //
+				.then(() -> {
+
+					// second ResultSet exhausted
+					assertThat(runnables).hasSize(1);
+					runnables.poll().run();
+				}) //
+				.thenRequest(10).expectNextCount(10) //
+				.verifyComplete();
+	}
+
+	private static Iterator<Row> mockIterator() {
+
+		Row row = mock(Row.class);
+		Iterator<Row> rows = mock(Iterator.class);
+		when(rows.hasNext()).thenReturn(true);
+		when(rows.next()).thenReturn(row);
+
+		return rows;
+	}
+
 	private static <T extends Statement> T eq(T value) {
 
 		return ArgumentMatchers.argThat(argument -> argument instanceof Statement //