Adding fetchSize at driver config and session config.

Author: Zhen Li

driver/src/main/java/org/neo4j/driver/Config.java

@@ -18,19 +18,24 @@
  */
 package org.neo4j.driver;
 
+import org.reactivestreams.Subscription;
+
 import java.io.File;
 import java.net.InetAddress;
 import java.util.Objects;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Level;
 
+import org.neo4j.driver.async.AsyncSession;
 import org.neo4j.driver.exceptions.ServiceUnavailableException;
 import org.neo4j.driver.exceptions.SessionExpiredException;
 import org.neo4j.driver.exceptions.TransientException;
 import org.neo4j.driver.internal.async.pool.PoolSettings;
 import org.neo4j.driver.internal.cluster.RoutingSettings;
+import org.neo4j.driver.internal.handlers.pulln.FetchSizeUtil;
 import org.neo4j.driver.internal.retry.RetrySettings;
 import org.neo4j.driver.net.ServerAddressResolver;
+import org.neo4j.driver.reactive.RxSession;
 import org.neo4j.driver.util.Immutable;
 import org.neo4j.driver.util.Resource;
 
@@ -88,6 +93,7 @@ public class Config
 
     private final int routingFailureLimit;
     private final long routingRetryDelayMillis;
+    private final long fetchSize;
     private long routingTablePurgeDelayMillis;
 
     private final int connectionTimeoutMillis;
@@ -114,6 +120,7 @@ private Config( ConfigBuilder builder )
         this.routingTablePurgeDelayMillis = builder.routingTablePurgeDelayMillis;
         this.retrySettings = builder.retrySettings;
         this.resolver = builder.resolver;
+        this.fetchSize = builder.fetchSize;
 
         this.isMetricsEnabled = builder.isMetricsEnabled;
     }
@@ -230,6 +237,11 @@ RetrySettings retrySettings()
         return retrySettings;
     }
 
+    public long fetchSize()
+    {
+        return fetchSize;
+    }
+
     /**
      * @return if the metrics is enabled or not on this driver.
      */
@@ -258,7 +270,7 @@ public static class ConfigBuilder
         private RetrySettings retrySettings = RetrySettings.DEFAULT;
         private ServerAddressResolver resolver;
         private boolean isMetricsEnabled = false;
-
+        private long fetchSize = FetchSizeUtil.DEFAULT_FETCH_SIZE;
 
         private ConfigBuilder() {}
 
@@ -566,6 +578,26 @@ public ConfigBuilder withRoutingTablePurgeDelay( long delay, TimeUnit unit )
             return this;
         }
 
+        /**
+         * Specify how many records to fetch in each batch.
+         * This config is only valid when the driver is used with servers that support Bolt V4 (Server version 4.0 and later).
+         *
+         * Bolt V4 enables pulling records in batches to allow client to take control of data population and apply back pressure to server.
+         * This config specifies the default fetch size for all query runs using {@link Session} and {@link AsyncSession}.
+         * By default, the value is set to {@code 1000}.
+         * Use {@code -1} to disables back pressure and config client to pull all records at once after each run.
+         *
+         * This config only applies to run result obtained via {@link Session} and {@link AsyncSession}.
+         * As with {@link RxSession}, the batch size is provided via {@link Subscription#request(long)} instead.
+         * @param size the default record fetch size when pulling records in batches using Bolt V4.
+         * @return this builder
+         */
+        public ConfigBuilder withFetchSize( long size )
+        {
+            this.fetchSize = FetchSizeUtil.assertValidFetchSize( size );
+            return this;
+        }
+
         /**
          * Specify socket connection timeout.
          * <p>

driver/src/main/java/org/neo4j/driver/SessionConfig.java

@@ -18,6 +18,8 @@
  */
 package org.neo4j.driver;
 
+import org.reactivestreams.Subscription;
+
 import java.util.Arrays;
 import java.util.Objects;
 import java.util.Optional;
@@ -26,6 +28,7 @@
 import org.neo4j.driver.reactive.RxSession;
 
 import static java.util.Objects.requireNonNull;
+import static org.neo4j.driver.internal.handlers.pulln.FetchSizeUtil.assertValidFetchSize;
 
 /**
  * The session configurations used to configure a session.
@@ -37,12 +40,14 @@ public class SessionConfig
     private final Iterable<Bookmark> bookmarks;
     private final AccessMode defaultAccessMode;
     private final String database;
+    private final Optional<Long> fetchSize;
 
     private SessionConfig( Builder builder )
     {
         this.bookmarks = builder.bookmarks;
         this.defaultAccessMode = builder.defaultAccessMode;
         this.database = builder.database;
+        this.fetchSize = builder.fetchSize;
     }
 
     /**
@@ -109,6 +114,15 @@ public Optional<String> database()
         return Optional.ofNullable( database );
     }
 
+    /**
+     * This value if set, overrides the default fetch size set on {@link Config#fetchSize()}.
+     * @return an optional value of fetch size.
+     */
+    public Optional<Long> fetchSize()
+    {
+        return fetchSize;
+    }
+
     @Override
     public boolean equals( Object o )
     {
@@ -121,7 +135,8 @@ public boolean equals( Object o )
             return false;
         }
         SessionConfig that = (SessionConfig) o;
-        return Objects.equals( bookmarks, that.bookmarks ) && defaultAccessMode == that.defaultAccessMode && Objects.equals( database, that.database );
+        return Objects.equals( bookmarks, that.bookmarks ) && defaultAccessMode == that.defaultAccessMode && Objects.equals( database, that.database )
+                && Objects.equals( fetchSize, that.fetchSize );
     }
 
     @Override
@@ -133,14 +148,16 @@ public int hashCode()
     @Override
     public String toString()
     {
-        return "SessionParameters{" + "bookmarks=" + bookmarks + ", defaultAccessMode=" + defaultAccessMode + ", database='" + database + '\'' + '}';
+        return "SessionParameters{" + "bookmarks=" + bookmarks + ", defaultAccessMode=" + defaultAccessMode + ", database='" + database + '\'' +
+                ", fetchSize=" + fetchSize + '}';
     }
 
     /**
      * Builder used to configure {@link SessionConfig} which will be used to create a session.
      */
     public static class Builder
     {
+        private Optional<Long> fetchSize = Optional.empty();
         private Iterable<Bookmark> bookmarks = null;
         private AccessMode defaultAccessMode = AccessMode.WRITE;
         private String database = null;
@@ -230,6 +247,27 @@ public Builder withDatabase( String database )
             return this;
         }
 
+        /**
+         * Specify how many records to fetch in each batch for this session.
+         * This config will overrides the default value set on {@link Config#fetchSize()}.
+         * This config is only valid when the driver is used with servers that support Bolt V4 (Server version 4.0 and later).
+         *
+         * Bolt V4 enables pulling records in batches to allow client to take control of data population and apply back pressure to server.
+         * This config specifies the default fetch size for all query runs using {@link Session} and {@link AsyncSession}.
+         * By default, the value is set to {@code 1000}.
+         * Use {@code -1} to disables back pressure and config client to pull all records at once after each run.
+         *
+         * This config only applies to run result obtained via {@link Session} and {@link AsyncSession}.
+         * As with {@link RxSession}, the batch size is provided via {@link Subscription#request(long)} instead.
+         * @param size the default record fetch size when pulling records in batches using Bolt V4.
+         * @return this builder
+         */
+        public Builder withFetchSize( long size )
+        {
+            this.fetchSize = Optional.of( assertValidFetchSize( size ) );
+            return this;
+        }
+
         public SessionConfig build()
         {
             return new SessionConfig( this );

driver/src/main/java/org/neo4j/driver/internal/InternalDriver.java

@@ -164,10 +164,10 @@ private static RuntimeException driverCloseException()
         return new IllegalStateException( "This driver instance has already been closed" );
     }
 
-    public NetworkSession newSession( SessionConfig parameters )
+    public NetworkSession newSession( SessionConfig config )
     {
         assertOpen();
-        NetworkSession session = sessionFactory.newInstance( parameters );
+        NetworkSession session = sessionFactory.newInstance( config );
         if ( closed.get() )
         {
             // session does not immediately acquire connection, it is fine to just throw

driver/src/main/java/org/neo4j/driver/internal/SessionFactory.java

@@ -25,7 +25,7 @@
 
 public interface SessionFactory
 {
-    NetworkSession newInstance( SessionConfig parameters );
+    NetworkSession newInstance( SessionConfig sessionConfig );
 
     CompletionStage<Void> verifyConnectivity();
 

driver/src/main/java/org/neo4j/driver/internal/SessionFactoryImpl.java

@@ -36,21 +36,28 @@ public class SessionFactoryImpl implements SessionFactory
     private final RetryLogic retryLogic;
     private final Logging logging;
     private final boolean leakedSessionsLoggingEnabled;
+    private final long defaultFetchSize;
 
     SessionFactoryImpl( ConnectionProvider connectionProvider, RetryLogic retryLogic, Config config )
     {
         this.connectionProvider = connectionProvider;
         this.leakedSessionsLoggingEnabled = config.logLeakedSessions();
         this.retryLogic = retryLogic;
         this.logging = config.logging();
+        this.defaultFetchSize = config.fetchSize();
     }
 
     @Override
     public NetworkSession newInstance( SessionConfig sessionConfig )
     {
         BookmarkHolder bookmarkHolder = new DefaultBookmarkHolder( InternalBookmark.from( sessionConfig.bookmarks() ) );
         return createSession( connectionProvider, retryLogic, parseDatabaseName( sessionConfig ),
-                sessionConfig.defaultAccessMode(), bookmarkHolder, logging );
+                sessionConfig.defaultAccessMode(), bookmarkHolder, parseFetchSize( sessionConfig ), logging );
+    }
+
+    private long parseFetchSize( SessionConfig sessionConfig )
+    {
+        return sessionConfig.fetchSize().orElse( defaultFetchSize );
     }
 
     private DatabaseName parseDatabaseName( SessionConfig sessionConfig )
@@ -85,10 +92,10 @@ public ConnectionProvider getConnectionProvider()
     }
 
     private NetworkSession createSession( ConnectionProvider connectionProvider, RetryLogic retryLogic, DatabaseName databaseName, AccessMode mode,
-            BookmarkHolder bookmarkHolder, Logging logging )
+            BookmarkHolder bookmarkHolder, long fetchSize, Logging logging )
     {
         return leakedSessionsLoggingEnabled
-               ? new LeakLoggingNetworkSession( connectionProvider, retryLogic, databaseName, mode, bookmarkHolder, logging )
-               : new NetworkSession( connectionProvider, retryLogic, databaseName, mode, bookmarkHolder, logging );
+               ? new LeakLoggingNetworkSession( connectionProvider, retryLogic, databaseName, mode, bookmarkHolder, fetchSize, logging )
+               : new NetworkSession( connectionProvider, retryLogic, databaseName, mode, bookmarkHolder, fetchSize, logging );
     }
 }

driver/src/main/java/org/neo4j/driver/internal/async/ExplicitTransaction.java

@@ -62,15 +62,17 @@ private enum State
     private final BoltProtocol protocol;
     private final BookmarkHolder bookmarkHolder;
     private final ResultCursorsHolder resultCursors;
+    private final long fetchSize;
 
     private volatile State state = State.ACTIVE;
 
-    public ExplicitTransaction( Connection connection, BookmarkHolder bookmarkHolder )
+    public ExplicitTransaction( Connection connection, BookmarkHolder bookmarkHolder, long fetchSize )
     {
         this.connection = connection;
         this.protocol = connection.protocol();
         this.bookmarkHolder = bookmarkHolder;
         this.resultCursors = new ResultCursorsHolder();
+        this.fetchSize = fetchSize;
     }
 
     public CompletionStage<ExplicitTransaction> beginAsync( InternalBookmark initialBookmark, TransactionConfig config )
@@ -140,7 +142,7 @@ public CompletionStage<StatementResultCursor> runAsync( Statement statement, boo
     {
         ensureCanRunQueries();
         CompletionStage<AsyncStatementResultCursor> cursorStage =
-                protocol.runInExplicitTransaction( connection, statement, this, waitForRunResponse ).asyncResult();
+                protocol.runInExplicitTransaction( connection, statement, this, waitForRunResponse, fetchSize ).asyncResult();
         resultCursors.add( cursorStage );
         return cursorStage.thenApply( cursor -> cursor );
     }
@@ -149,7 +151,7 @@ public CompletionStage<RxStatementResultCursor> runRx( Statement statement )
     {
         ensureCanRunQueries();
         CompletionStage<RxStatementResultCursor> cursorStage =
-                protocol.runInExplicitTransaction( connection, statement, this, false ).rxResult();
+                protocol.runInExplicitTransaction( connection, statement, this, false, fetchSize ).rxResult();
         resultCursors.add( cursorStage );
         return cursorStage;
     }

driver/src/main/java/org/neo4j/driver/internal/async/LeakLoggingNetworkSession.java

@@ -33,9 +33,9 @@ public class LeakLoggingNetworkSession extends NetworkSession
     private final String stackTrace;
 
     public LeakLoggingNetworkSession( ConnectionProvider connectionProvider, RetryLogic retryLogic, DatabaseName databaseName, AccessMode mode,
-            BookmarkHolder bookmarkHolder, Logging logging )
+            BookmarkHolder bookmarkHolder, long fetchSize, Logging logging )
     {
-        super( connectionProvider, retryLogic, databaseName, mode, bookmarkHolder, logging );
+        super( connectionProvider, retryLogic, databaseName, mode, bookmarkHolder, fetchSize, logging );
         this.stackTrace = captureStackTrace();
     }
 

driver/src/main/java/org/neo4j/driver/internal/async/NetworkSession.java

@@ -57,21 +57,23 @@ public class NetworkSession
     protected final Logger logger;
 
     private final BookmarkHolder bookmarkHolder;
+    private final long fetchSize;
     private volatile CompletionStage<ExplicitTransaction> transactionStage = completedWithNull();
     private volatile CompletionStage<Connection> connectionStage = completedWithNull();
     private volatile CompletionStage<? extends FailableCursor> resultCursorStage = completedWithNull();
 
     private final AtomicBoolean open = new AtomicBoolean( true );
 
     public NetworkSession( ConnectionProvider connectionProvider, RetryLogic retryLogic, DatabaseName databaseName, AccessMode mode,
-            BookmarkHolder bookmarkHolder, Logging logging )
+            BookmarkHolder bookmarkHolder, long fetchSize, Logging logging )
     {
         this.connectionProvider = connectionProvider;
         this.mode = mode;
         this.retryLogic = retryLogic;
         this.logger = new PrefixedLogger( "[" + hashCode() + "]", logging.getLog( LOG_NAME ) );
         this.bookmarkHolder = bookmarkHolder;
         this.connectionContext = new NetworkSessionConnectionContext( databaseName, bookmarkHolder.getBookmark() );
+        this.fetchSize = fetchSize;
     }
 
     public CompletionStage<StatementResultCursor> runAsync( Statement statement, TransactionConfig config, boolean waitForRunResponse )
@@ -106,7 +108,7 @@ public CompletionStage<ExplicitTransaction> beginTransactionAsync( AccessMode mo
                 .thenCompose( ignore -> acquireConnection( mode ) )
                 .thenCompose( connection ->
                 {
-                    ExplicitTransaction tx = new ExplicitTransaction( connection, bookmarkHolder );
+                    ExplicitTransaction tx = new ExplicitTransaction( connection, bookmarkHolder, fetchSize );
                     return tx.beginAsync( bookmarkHolder.getBookmark(), config );
                 } );
 
@@ -231,7 +233,7 @@ private CompletionStage<StatementResultCursorFactory> buildResultCursorFactory(
                     try
                     {
                         StatementResultCursorFactory factory = connection.protocol()
-                                .runInAutoCommitTransaction( connection, statement, bookmarkHolder, config, waitForRunResponse );
+                                .runInAutoCommitTransaction( connection, statement, bookmarkHolder, config, waitForRunResponse, fetchSize );
                         return completedFuture( factory );
                     }
                     catch ( Throwable e )

driver/src/main/java/org/neo4j/driver/internal/cluster/RoutingProcedureRunner.java

@@ -39,6 +39,7 @@
 
 import static org.neo4j.driver.Values.parameters;
 import static org.neo4j.driver.internal.DatabaseNameUtil.defaultDatabase;
+import static org.neo4j.driver.internal.handlers.pulln.FetchSizeUtil.UNLIMITED_FETCH_SIZE;
 
 public class RoutingProcedureRunner
 {
@@ -86,7 +87,7 @@ BookmarkHolder bookmarkHolder( InternalBookmark ignored )
     CompletionStage<List<Record>> runProcedure( Connection connection, Statement procedure, BookmarkHolder bookmarkHolder )
     {
         return connection.protocol()
-                .runInAutoCommitTransaction( connection, procedure, bookmarkHolder, TransactionConfig.empty(), true )
+                .runInAutoCommitTransaction( connection, procedure, bookmarkHolder, TransactionConfig.empty(), true, UNLIMITED_FETCH_SIZE )
                 .asyncResult().thenCompose( StatementResultCursor::listAsync );
     }
 

driver/src/main/java/org/neo4j/driver/internal/handlers/PullHandlers.java

@@ -47,11 +47,11 @@ public static PullAllResponseHandler newBoltV3PullAllHandler( Statement statemen
     }
 
     public static PullAllResponseHandler newBoltV4AutoPullHandler( Statement statement, RunResponseHandler runHandler, Connection connection,
-            BookmarkHolder bookmarkHolder, ExplicitTransaction tx )
+            BookmarkHolder bookmarkHolder, ExplicitTransaction tx, long fetchSize )
     {
         PullResponseCompletionListener completionListener = createPullResponseCompletionListener( connection, bookmarkHolder, tx );
 
-        return new AutoPullResponseHandler( statement, runHandler, connection, BoltProtocolV3.METADATA_EXTRACTOR, completionListener );
+        return new AutoPullResponseHandler( statement, runHandler, connection, BoltProtocolV3.METADATA_EXTRACTOR, completionListener, fetchSize );
     }