fix(serve, ext-driver-pg): destroy the source stream to release connections

Author: oscar60310

packages/extension-driver-pg/src/lib/pgDataSource.ts

@@ -58,10 +58,12 @@ export class PGDataSource extends DataSource<any, PGOptions> {
       const cursor = client.query(
         new Cursor(sql, Array.from(bindParams.values()))
       );
-      cursor.once('done', () => {
+      cursor.once('done', async () => {
         this.logger.debug(
           `Data fetched, release connection from ${profileName}`
         );
+        // It is important to close the cursor before releasing connection, or the connection might not able to handle next request.
+        await cursor.close();
         client.release();
       });
       // All promises MUST fulfilled in this function or we are not able to release the connection when error occurred
@@ -92,34 +94,35 @@ export class PGDataSource extends DataSource<any, PGOptions> {
     const { chunkSize = 100 } = options;
     const cursorRead = this.cursorRead.bind(this);
     const firstChunk = await cursorRead(cursor, chunkSize);
-
+    // save first chunk in buffer for incoming requests
+    let bufferedRows = [...firstChunk.rows];
+    let bufferReadIndex = 0;
+    const fetchNext = async () => {
+      if (bufferReadIndex >= bufferedRows.length) {
+        bufferedRows = (await cursorRead(cursor, chunkSize)).rows;
+        bufferReadIndex = 0;
+      }
+      return bufferedRows[bufferReadIndex++] || null;
+    };
     const stream = new Readable({
       objectMode: true,
       read() {
-        cursorRead(cursor, chunkSize)
-          .then(({ rows }) => {
-            if (rows.length === 0) {
-              this.push(null);
-              // Send done event to notify upstream to release the connection.
-              cursor.emit('done');
-            }
-            for (const row of rows) {
-              this.push(row);
-            }
+        fetchNext()
+          .then((row) => {
+            this.push(row);
           })
           .catch((error) => {
-            this.emit('error', error);
-            // Send done event to notify upstream to release the connection.
-            cursor.emit('done');
+            this.destroy(error);
           });
       },
+      destroy(error: Error | null, cb: (error: Error | null) => void) {
+        // Send done event to notify upstream to release the connection.
+        cursor.emit('done');
+        cb(error);
+      },
+      // automatically destroy() the stream when it emits 'finish' or errors. Node > 10.16
+      autoDestroy: true,
     });
-
-    // Push the first chunk data.
-    for (const row of firstChunk.rows) {
-      stream.push(row);
-    }
-
     return {
       getColumns: () =>
         firstChunk.result.fields.map((field) => ({

packages/extension-driver-pg/test/pgDataSource.spec.ts

@@ -1,6 +1,7 @@
 import { PGServer } from './pgServer';
 import { PGDataSource, PGOptions } from '../src';
 import { streamToArray } from '@vulcan-sql/core';
+import { Writable } from 'stream';
 
 const pg = new PGServer();
 let dataSource: PGDataSource;
@@ -295,10 +296,58 @@ it('Data source should return correct column types', async () => {
     operations: {} as any,
   });
   const column = getColumns();
-  // We need to consume the data stream or the driver waits for us
-  await streamToArray(getData());
+  // We need to destroy the data stream or the driver waits for us
+  const data = getData();
+  data.destroy();
+
   // Assert
   expect(column[0]).toEqual({ name: 'id', type: 'number' });
   expect(column[1]).toEqual({ name: 'name', type: 'string' });
   expect(column[2]).toEqual({ name: 'enabled', type: 'boolean' });
 }, 30000);
+
+it('Data source should release connection when readable stream is destroyed', async () => {
+  // Arrange
+  dataSource = new PGDataSource({}, '', [
+    {
+      name: 'profile1',
+      type: 'pg',
+      connection: {
+        host: pg.host,
+        user: pg.user,
+        password: pg.password,
+        database: pg.database,
+        port: pg.port,
+      } as PGOptions,
+      allow: '*',
+    },
+  ]);
+  await dataSource.activate();
+  // Act
+  const { getData } = await dataSource.execute({
+    statement: 'select * from users limit 100',
+    bindParams: new Map(),
+    profileName: 'profile1',
+    operations: {} as any,
+  });
+  const readStream = getData();
+  const rows: any[] = [];
+  let resolve: any;
+  const waitForStream = () => new Promise((res) => (resolve = res));
+  const writeStream = new Writable({
+    write(chunk, _, cb) {
+      rows.push(chunk);
+      // After read 5 records, destroy the upstream
+      if (rows.length === 5) {
+        readStream.destroy();
+        resolve();
+      } else cb();
+    },
+    objectMode: true,
+  });
+  readStream.pipe(writeStream);
+  await waitForStream();
+  // Assert
+  expect(rows.length).toBe(5);
+  // afterEach hook will timeout if any leak occurred.
+}, 30000);

packages/serve/src/models/extensions/responseFormatter.ts

@@ -37,6 +37,10 @@ export abstract class BaseResponseFormatter
     // if response has data and columns.
     const { data, columns } = ctx.response.body as BodyResponse;
     const formatted = this.format(data, columns);
+    // koa destroy the stream when connection close, we need to destroy our upstream too to notice them to release the resource.
+    formatted.on('close', () => {
+      data.destroy();
+    });
     // set formatted stream to response in context
     this.toResponse(formatted, ctx);
     return;

packages/serve/test/response-formatter/json.spec.ts

@@ -25,6 +25,37 @@ describe('Test to respond to json', () => {
     expect(ctx.response.body).toEqual(expected);
   });
 
+  it('Test to destroy the source when downstream is destroyed', (done) => {
+    // Arrange
+    const stubResponse = sinon.stubInterface<Response>();
+    const source = new Stream.Readable({
+      read() {
+        this.push('some-data');
+        // we don't push pull to simulate an endless stream
+      },
+      destroy(err, cb) {
+        done();
+        cb(err);
+      },
+    });
+    stubResponse.body = {
+      data: source,
+      columns: {},
+    };
+    const ctx = {
+      ...sinon.stubInterface<KoaContext>(),
+      url: faker.internet.url(),
+      response: stubResponse,
+    };
+    // Act
+    const formatter = new JsonFormatter({}, '');
+    formatter.formatToResponse(ctx);
+    (ctx.response.body as Stream.Readable).destroy();
+    // Assert
+    expect(true).toEqual(true);
+    // the done function should be call after destroyed.
+  }, 1000);
+
   it.each([
     {
       input: {