perf: release connections before running post commit steps (#274)

Author: ForbesLindesay

packages/pg/src/__tests__/release-connetions-before-post-commit.test.pg.ts

@@ -0,0 +1,36 @@
+import connect, {sql} from '..';
+
+jest.setTimeout(10_000);
+
+const db = connect({
+  bigIntMode: 'bigint',
+  poolSize: 1,
+});
+
+afterAll(async () => {
+  await db.dispose();
+});
+
+test('transaction in task', async () => {
+  await db.task(async (db) => {
+    const taskDb = db;
+    await db.tx(async (db) => {
+      await db.query(sql`SELECT 1 + 1`);
+      await db.addPostCommitStep(async () => {
+        // This code runs after the transaction completes
+        await taskDb.query(sql`SELECT 1 + 1`);
+      });
+    });
+  });
+});
+
+test('tx in pool', async () => {
+  const poolDb = db;
+  await db.tx(async (db) => {
+    await db.query(sql`SELECT 1 + 1`);
+    await db.addPostCommitStep(async () => {
+      // This code runs after the transaction completes
+      await poolDb.query(sql`SELECT 1 + 1`);
+    });
+  });
+});

packages/shared/src/BaseConnection.ts

@@ -59,12 +59,26 @@ export default class BaseConnection<
     options?: TransactionOptions<TDriver>,
   ): Promise<TResult> {
     this._throwIfDisposed();
+    const postCommitSteps: (() => Promise<void>)[] = [];
     await this._lock.acquireLock();
+    let result: TResult;
     try {
-      return await txInternal(this._driver, this._factories, fn, options);
+      result = await txInternal(
+        this._driver,
+        this._factories,
+        fn,
+        options,
+        (fn) => {
+          postCommitSteps.push(fn);
+        },
+      );
     } finally {
       this._lock.releaseLock();
     }
+    for (const step of postCommitSteps) {
+      await step();
+    }
+    return result;
   }
 
   async query(query: SQLQuery): Promise<any[]>;

packages/shared/src/BaseConnectionPool.ts

@@ -4,7 +4,7 @@ import createConnectionPool, {
 } from '@databases/connection-pool';
 import splitSqlQuery from '@databases/split-sql-query';
 import type {SQLQuery} from '@databases/sql';
-import Factory, {Disposable} from './Factory';
+import Factory, {Disposable, TransactionFactory} from './Factory';
 import Driver from './Driver';
 import QueryableType from './QueryableType';
 import {
@@ -94,7 +94,22 @@ export default class BaseConnectionPool<
     options?: TransactionOptions<TDriver>,
   ): Promise<TResult> {
     this._throwIfDisposed();
-    return this._withDriverFromPool(txInternal, this._factories, fn, options);
+    const postCommitSteps: (() => Promise<void>)[] = [];
+    const result = await this._withDriverFromPool<
+      [
+        TransactionFactory<TDriver, TTransaction>,
+        (connection: TTransaction) => Promise<TResult>,
+        TransactionOptions<TDriver> | undefined,
+        (fn: () => Promise<void>) => void,
+      ],
+      TResult
+    >(txInternal, this._factories, fn, options, (fn) => {
+      postCommitSteps.push(fn);
+    });
+    for (const step of postCommitSteps) {
+      await step();
+    }
+    return result;
   }
 
   async query(query: SQLQuery): Promise<any[]>;

packages/shared/src/utils.ts

@@ -85,16 +85,12 @@ export async function txInternal<
   factories: TransactionFactory<TDriver, TTransaction>,
   fn: (connection: TTransaction) => Promise<TResult>,
   options: TTransactionOptions | undefined,
+  addPostCommitStep: (fn: () => Promise<void>) => void,
 ): Promise<TResult> {
   let failureCount = 0;
   while (true) {
     await driver.beginTransaction(options);
-    const postCommitSteps: (() => Promise<void>)[] = [];
-    const tx = factories.createTransaction(driver, {
-      addPostCommitStep: (fn) => {
-        postCommitSteps.push(fn);
-      },
-    });
+    const tx = factories.createTransaction(driver, {addPostCommitStep});
     let result;
     try {
       result = await fn(tx);
@@ -110,9 +106,6 @@ export async function txInternal<
       }
       throw ex;
     }
-    for (const step of postCommitSteps) {
-      await step();
-    }
     return result;
   }
 }