Driver.executeQuery optimization (#1124)

bigmontz · web-flow · commit 5efc5c7e0d08 · 2023-08-21T17:18:33.000+02:00
Pipeline begin and run when using `executeQuery` to reduce the number of round trips.
diff --git a/packages/core/src/driver.ts b/packages/core/src/driver.ts
@@ -444,7 +444,7 @@ class Driver {
     config: DriverConfig = {},
     createConnectionProvider: CreateConnectionProvider,
     createSession: CreateSession = args => new Session(args),
-    createQueryExecutor: CreateQueryExecutor = createQuery => new QueryExecutor(createQuery)
+    createQueryExecutor: CreateQueryExecutor = createSession => new QueryExecutor(createSession)
   ) {
     sanitizeConfig(config)
 
diff --git a/packages/core/src/internal/query-executor.ts b/packages/core/src/internal/query-executor.ts
@@ -46,6 +46,10 @@ export default class QueryExecutor {
       bookmarkManager: config.bookmarkManager,
       impersonatedUser: config.impersonatedUser
     })
+
+    // @ts-expect-error The method is private for external users
+    session._setTxExecutorToPipelineBegin(true)
+
     try {
       const executeInTransaction: TransactionFunction<T> = config.routing === 'READ'
         ? session.executeRead.bind(session)
diff --git a/packages/core/src/internal/transaction-executor.ts b/packages/core/src/internal/transaction-executor.ts
@@ -39,6 +39,7 @@ export class TransactionExecutor {
   private readonly _multiplier: number
   private readonly _jitterFactor: number
   private _inFlightTimeoutIds: Timeout[]
+  public pipelineBegin: boolean
 
   constructor (
     maxRetryTimeMs?: number | null,
@@ -64,6 +65,7 @@ export class TransactionExecutor {
     )
 
     this._inFlightTimeoutIds = []
+    this.pipelineBegin = false
 
     this._verifyAfterConstruction()
   }
@@ -154,7 +156,8 @@ export class TransactionExecutor {
   ): Promise<void> {
     let tx: Transaction
     try {
-      tx = await transactionCreator()
+      const txPromise = transactionCreator()
+      tx = this.pipelineBegin ? txPromise : await txPromise
     } catch (error) {
       // failed to create a transaction
       reject(error)
@@ -192,7 +195,7 @@ export class TransactionExecutor {
 
   _handleTransactionWorkSuccess<T>(
     result: T,
-    tx: Transaction,
+    tx: Transaction | TransactionPromise,
     resolve: Resolve<T>,
     reject: Reject
   ): void {
@@ -215,7 +218,7 @@ export class TransactionExecutor {
     }
   }
 
-  _handleTransactionWorkFailure (error: any, tx: Transaction, reject: Reject): void {
+  _handleTransactionWorkFailure (error: any, tx: Transaction | TransactionPromise, reject: Reject): void {
     if (tx.isOpen()) {
       // transaction work failed and the transaction is still open, roll it back and propagate the failure
       tx.rollback()
diff --git a/packages/core/src/session.ts b/packages/core/src/session.ts
@@ -587,6 +587,15 @@ class Session {
     }
   }
 
+  /**
+   * Configure the transaction executor to pipeline transaction begin.
+   *
+   * @private
+   */
+  private _setTxExecutorToPipelineBegin (pipelined: boolean): void {
+    this._transactionExecutor.pipelineBegin = pipelined
+  }
+
   /**
    * @protected
    */
diff --git a/packages/core/test/internal/query-executor.test.ts b/packages/core/test/internal/query-executor.test.ts
@@ -86,6 +86,21 @@ describe('QueryExecutor', () => {
       expect(spyOnExecuteRead).toHaveBeenCalled()
     })
 
+    it('should configure the session with pipeline begin', async () => {
+      const { queryExecutor, sessionsCreated } = createExecutor()
+
+      await queryExecutor.execute(baseConfig, 'query')
+
+      expect(sessionsCreated.length).toBe(1)
+      const [{ spyOnSetTxExecutorToPipelineBegin, spyOnExecuteRead }] = sessionsCreated
+
+      expect(spyOnSetTxExecutorToPipelineBegin).toHaveBeenCalledTimes(1)
+      expect(spyOnSetTxExecutorToPipelineBegin).toHaveBeenCalledWith(true)
+      expect(spyOnExecuteRead.mock.invocationCallOrder[0]).toBeGreaterThan(
+        spyOnSetTxExecutorToPipelineBegin.mock.invocationCallOrder[0]
+      )
+    })
+
     it('should call not call executeWrite', async () => {
       const { queryExecutor, sessionsCreated } = createExecutor()
 
@@ -213,6 +228,21 @@ describe('QueryExecutor', () => {
       expect(spyOnExecuteWrite).toHaveBeenCalled()
     })
 
+    it('should configure the session with pipeline begin', async () => {
+      const { queryExecutor, sessionsCreated } = createExecutor()
+
+      await queryExecutor.execute(baseConfig, 'query')
+
+      expect(sessionsCreated.length).toBe(1)
+      const [{ spyOnSetTxExecutorToPipelineBegin, spyOnExecuteWrite }] = sessionsCreated
+
+      expect(spyOnSetTxExecutorToPipelineBegin).toHaveBeenCalledTimes(1)
+      expect(spyOnSetTxExecutorToPipelineBegin).toHaveBeenCalledWith(true)
+      expect(spyOnExecuteWrite.mock.invocationCallOrder[0]).toBeGreaterThan(
+        spyOnSetTxExecutorToPipelineBegin.mock.invocationCallOrder[0]
+      )
+    })
+
     it('should call not call executeRead', async () => {
       const { queryExecutor, sessionsCreated } = createExecutor()
 
@@ -316,7 +346,7 @@ describe('QueryExecutor', () => {
         spyOnExecuteRead: jest.SpyInstance<any>
         spyOnExecuteWrite: jest.SpyInstance<any>
         spyOnClose: jest.SpyInstance<Promise<void>>
-
+        spyOnSetTxExecutorToPipelineBegin: jest.SpyInstance<void>
       }>
       createSession: jest.Mock<Session, [args: any]>
     } {
@@ -329,15 +359,17 @@ describe('QueryExecutor', () => {
       spyOnExecuteRead: jest.SpyInstance<any>
       spyOnExecuteWrite: jest.SpyInstance<any>
       spyOnClose: jest.SpyInstance<Promise<void>>
-
+      spyOnSetTxExecutorToPipelineBegin: jest.SpyInstance<void>
     }> = []
     const createSession = jest.fn((args) => {
       const session = new Session(args)
       const sessionCreated = {
         session,
         spyOnExecuteRead: jest.spyOn(session, 'executeRead'),
         spyOnExecuteWrite: jest.spyOn(session, 'executeWrite'),
-        spyOnClose: jest.spyOn(session, 'close')
+        spyOnClose: jest.spyOn(session, 'close'),
+        // @ts-expect-error
+        spyOnSetTxExecutorToPipelineBegin: jest.spyOn(session, '_setTxExecutorToPipelineBegin')
       }
       sessionsCreated.push(sessionCreated)
       _mockSessionExecuteRead(sessionCreated.spyOnExecuteRead)
diff --git a/packages/core/test/session.test.ts b/packages/core/test/session.test.ts
@@ -20,6 +20,7 @@ import { ConnectionProvider, Session, Connection, TransactionPromise, Transactio
 import { bookmarks } from '../src/internal'
 import { ACCESS_MODE_READ, FETCH_ALL } from '../src/internal/constants'
 import { Logger } from '../src/internal/logger'
+import { TransactionExecutor } from '../src/internal/transaction-executor'
 import ManagedTransaction from '../src/transaction-managed'
 import { AuthToken, LoggerFunction } from '../src/types'
 import FakeConnection from './utils/connection.fake'
@@ -1150,6 +1151,25 @@ describe('session', () => {
       )
     })
   })
+
+  describe('Pipeline Begin on TxFunc', () => {
+    it('session should not change the default on session creation', () => {
+      const session = newSessionWithConnection(new FakeConnection())
+
+      // @ts-expect-error
+      expect(session._transactionExecutor.pipelineBegin).toEqual(new TransactionExecutor().pipelineBegin)
+    })
+
+    it.each([true, false])('_setTxExecutorToPipelineBegin(%s) => configure executor', (pipelined) => {
+      const session = newSessionWithConnection(new FakeConnection())
+
+      // @ts-expect-error
+      session._setTxExecutorToPipelineBegin(pipelined)
+
+      // @ts-expect-error
+      expect(session._transactionExecutor.pipelineBegin).toEqual(pipelined)
+    })
+  })
 })
 
 function mockBeginWithSuccess (connection: FakeConnection): FakeConnection {
diff --git a/packages/neo4j-driver-deno/lib/core/driver.ts b/packages/neo4j-driver-deno/lib/core/driver.ts
@@ -444,7 +444,7 @@ class Driver {
     config: DriverConfig = {},
     createConnectionProvider: CreateConnectionProvider,
     createSession: CreateSession = args => new Session(args),
-    createQueryExecutor: CreateQueryExecutor = createQuery => new QueryExecutor(createQuery)
+    createQueryExecutor: CreateQueryExecutor = createSession => new QueryExecutor(createSession)
   ) {
     sanitizeConfig(config)
 
diff --git a/packages/neo4j-driver-deno/lib/core/internal/query-executor.ts b/packages/neo4j-driver-deno/lib/core/internal/query-executor.ts
@@ -46,6 +46,10 @@ export default class QueryExecutor {
       bookmarkManager: config.bookmarkManager,
       impersonatedUser: config.impersonatedUser
     })
+
+    // @ts-expect-error The method is private for external users
+    session._setTxExecutorToPipelineBegin(true)
+
     try {
       const executeInTransaction: TransactionFunction<T> = config.routing === 'READ'
         ? session.executeRead.bind(session)
diff --git a/packages/neo4j-driver-deno/lib/core/internal/transaction-executor.ts b/packages/neo4j-driver-deno/lib/core/internal/transaction-executor.ts
@@ -39,6 +39,7 @@ export class TransactionExecutor {
   private readonly _multiplier: number
   private readonly _jitterFactor: number
   private _inFlightTimeoutIds: Timeout[]
+  public pipelineBegin: boolean
 
   constructor (
     maxRetryTimeMs?: number | null,
@@ -64,6 +65,7 @@ export class TransactionExecutor {
     )
 
     this._inFlightTimeoutIds = []
+    this.pipelineBegin = false
 
     this._verifyAfterConstruction()
   }
@@ -154,7 +156,8 @@ export class TransactionExecutor {
   ): Promise<void> {
     let tx: Transaction
     try {
-      tx = await transactionCreator()
+      const txPromise = transactionCreator()
+      tx = this.pipelineBegin ? txPromise : await txPromise
     } catch (error) {
       // failed to create a transaction
       reject(error)
@@ -192,7 +195,7 @@ export class TransactionExecutor {
 
   _handleTransactionWorkSuccess<T>(
     result: T,
-    tx: Transaction,
+    tx: Transaction | TransactionPromise,
     resolve: Resolve<T>,
     reject: Reject
   ): void {
@@ -215,7 +218,7 @@ export class TransactionExecutor {
     }
   }
 
-  _handleTransactionWorkFailure (error: any, tx: Transaction, reject: Reject): void {
+  _handleTransactionWorkFailure (error: any, tx: Transaction | TransactionPromise, reject: Reject): void {
     if (tx.isOpen()) {
       // transaction work failed and the transaction is still open, roll it back and propagate the failure
       tx.rollback()
diff --git a/packages/neo4j-driver-deno/lib/core/session.ts b/packages/neo4j-driver-deno/lib/core/session.ts
@@ -587,6 +587,15 @@ class Session {
     }
   }
 
+  /**
+   * Configure the transaction executor to pipeline transaction begin.
+   *
+   * @private
+   */
+  private _setTxExecutorToPipelineBegin (pipelined: boolean): void {
+    this._transactionExecutor.pipelineBegin = pipelined
+  }
+
   /**
    * @protected
    */
diff --git a/packages/neo4j-driver/test/internal/transaction-executor.test.js b/packages/neo4j-driver/test/internal/transaction-executor.test.js
diff --git a/packages/testkit-backend/src/feature/common.js b/packages/testkit-backend/src/feature/common.js
