Merge pull request #142 from shelfio/implement-backpressure

Author: vladgolubev

package.json

@@ -26,7 +26,7 @@
     "lint": "eslint . --ext .js,.ts,.json --fix --quiet",
     "lint:ci": "eslint . --ext .js,.ts,.json --quiet",
     "prepack": "yarn build",
-    "test": "export ENVIRONMENT=local && jest src --runInBand",
+    "test": "export ENVIRONMENT=local && jest src --runInBand --forceExit",
     "type-check": "tsc --noEmit",
     "type-check:watch": "npm run type-check -- --watch"
   },
@@ -56,7 +56,7 @@
     "@babel/cli": "7.17.6",
     "@babel/core": "7.17.5",
     "@shelf/babel-config": "0.1.8",
-    "@shelf/eslint-config": "2.0.0",
+    "@shelf/eslint-config": "2.16.0",
     "@shelf/jest-dynamodb": "2.1.0",
     "@shelf/prettier-config": "1.0.0",
     "@shelf/tsconfig": "0.0.6",

readme.md

@@ -34,7 +34,7 @@ const {parallelScan} = require('@shelf/dynamodb-parallel-scan');
 
 ### Use as async generator (or streams)
 
-Note: this stream doesn't implement backpressure mechanism just yet, so memory overflow could happen if you don't consume stream fast enough.
+Note: `highWaterMark` determines items count threshold, so Parallel Scan can fetch `concurrency` \* 1MB more data even after highWaterMark was reached.
 
 ```js
 const {parallelScanAsStream} = require('@shelf/dynamodb-parallel-scan');
@@ -49,7 +49,7 @@ const {parallelScanAsStream} = require('@shelf/dynamodb-parallel-scan');
       },
       ProjectionExpression: 'fileSize',
     },
-    {concurrency: 1000, chunkSize: 10000}
+    {concurrency: 1000, chunkSize: 10000, highWaterMark: 10000}
   );
 
   for await (const items of stream) {

src/blocker.test.ts

@@ -0,0 +1,70 @@
+import {inspect} from 'util';
+import {Blocker} from './blocker';
+
+it('should return same promise after sequent block() calls', () => {
+  const blocker = new Blocker();
+
+  blocker.block();
+
+  const oldPromise = blocker.get();
+
+  blocker.block();
+
+  const newPromise = blocker.get();
+
+  expect(newPromise === oldPromise).toBeTruthy();
+});
+
+it('should return pending promise after block()', () => {
+  const blocker = new Blocker();
+
+  blocker.block();
+
+  const inspectedBlocker = inspect(blocker.get());
+  const isInspectedBlockerPending = inspectedBlocker.includes('pending');
+
+  expect(isInspectedBlockerPending).toBeTruthy();
+});
+
+it('should return resolved promise after unblock()', () => {
+  const blocker = new Blocker();
+
+  blocker.unblock();
+
+  const inspectedBlocker = inspect(blocker.get());
+  const isInspectedBlockerResolved = inspectedBlocker.includes('undefined');
+
+  expect(isInspectedBlockerResolved).toBeTruthy();
+});
+
+it('should return resolved promise in default state', () => {
+  const blocker = new Blocker();
+
+  const inspectedBlocker = inspect(blocker.get());
+  const isInspectedBlockerResolved = inspectedBlocker.includes('undefined');
+
+  expect(isInspectedBlockerResolved).toBeTruthy();
+});
+
+it('should be blocked after block() call', () => {
+  const blocker = new Blocker();
+
+  blocker.block();
+
+  expect(blocker.isBlocked()).toBeTruthy();
+});
+
+it('should not be blocked after unblock() call', () => {
+  const blocker = new Blocker();
+
+  blocker.unblock();
+  blocker.unblock();
+
+  expect(blocker.isBlocked()).toBeFalsy();
+});
+
+it('should not be blocked in default state', () => {
+  const blocker = new Blocker();
+
+  expect(blocker.isBlocked()).toBeFalsy();
+});

src/blocker.ts

@@ -0,0 +1,38 @@
+export class Blocker {
+  private _promise: Promise<void>;
+  private _promiseResolver: (...args: any[]) => any;
+  private _isBlocked: boolean;
+
+  constructor() {
+    this._promise = Promise.resolve();
+    this._promiseResolver = () => {}; // eslint-disable-line @typescript-eslint/no-empty-function
+    this._isBlocked = false;
+  }
+
+  block(): void {
+    if (this._isBlocked) {
+      return;
+    }
+
+    this._promise = new Promise(r => {
+      this._promiseResolver = r;
+      setTimeout(r, 2147483647);
+    }); //TODO: Implement endless promise
+
+    this._isBlocked = true;
+  }
+
+  unblock(): void {
+    this._promiseResolver();
+
+    this._isBlocked = false;
+  }
+
+  get(): Promise<void> {
+    return this._promise;
+  }
+
+  isBlocked(): boolean {
+    return this._isBlocked;
+  }
+}

src/ddb.ts

@@ -1,6 +1,6 @@
 import DynamoDB from 'aws-sdk/clients/dynamodb';
-import {DocumentClient} from 'aws-sdk/lib/dynamodb/document_client';
-import {InsertManyParams} from './ddb.types';
+import type {DocumentClient} from 'aws-sdk/lib/dynamodb/document_client';
+import type {InsertManyParams} from './ddb.types';
 
 const isTest = process.env.JEST_WORKER_ID;
 const config = {
@@ -27,11 +27,13 @@ export function insertMany({
 }: InsertManyParams): Promise<DocumentClient.BatchWriteItemOutput> {
   const params: DocumentClient.BatchWriteItemInput = {
     RequestItems: {
-      [tableName]: items.map(item => ({
-        PutRequest: {
-          Item: item,
-        },
-      })),
+      [tableName]: items.map(item => {
+        return {
+          PutRequest: {
+            Item: item,
+          },
+        };
+      }),
     },
   };
 

src/parallel-scan-stream.test.ts

@@ -1,31 +1,44 @@
-import {insertMany} from './ddb';
+jest.setTimeout(25000);
+
+import {uniq} from 'lodash';
+import * as ddbHelpers from './ddb';
 import {parallelScanAsStream} from './parallel-scan-stream';
 
+async function delay(ms: number) {
+  return new Promise(r => {
+    setTimeout(r, ms);
+  });
+}
+
 describe('parallelScanAsStream', () => {
   const files = [
-    {id: 'some-file-id-1'},
-    {id: 'some-file-id-2'},
-    {id: 'some-file-id-3', fileSize: 100},
-    {id: 'some-file-id-4'},
-    {id: 'some-file-id-5'},
-    {id: 'some-file-id-6', fileSize: 200},
-    {id: 'some-file-id-7'},
-    {id: 'some-file-id-8'},
-    {id: 'some-file-id-9', fileSize: 300},
-    {id: 'some-file-id-10'},
+    {id: 'some-file-id-1', isLarge: false},
+    {id: 'some-file-id-2', isLarge: false},
+    {id: 'some-file-id-3', fileSize: 100, isLarge: false},
+    {id: 'some-file-id-4', isLarge: false},
+    {id: 'some-file-id-5', isLarge: false},
+    {id: 'some-file-id-6', fileSize: 200, isLarge: false},
+    {id: 'some-file-id-7', isLarge: false},
+    {id: 'some-file-id-8', isLarge: false},
+    {id: 'some-file-id-9', fileSize: 300, isLarge: false},
+    {id: 'some-file-id-10', isLarge: false},
   ];
 
   beforeAll(async () => {
-    await insertMany({items: files, tableName: 'files'});
+    await ddbHelpers.insertMany({items: files, tableName: 'files'});
   });
 
   it('should stream items with chunks of 2 with concurrency 1', async () => {
     const stream = await parallelScanAsStream(
       {
         TableName: 'files',
-        FilterExpression: 'attribute_exists(#id)',
+        FilterExpression: 'attribute_exists(#id) and #isLarge = :false',
         ExpressionAttributeNames: {
           '#id': 'id',
+          '#isLarge': 'isLarge',
+        },
+        ExpressionAttributeValues: {
+          ':false': false,
         },
       },
       {concurrency: 1, chunkSize: 2}
@@ -40,9 +53,13 @@ describe('parallelScanAsStream', () => {
     const stream = await parallelScanAsStream(
       {
         TableName: 'files',
-        FilterExpression: 'attribute_exists(#id)',
+        FilterExpression: 'attribute_exists(#id) and #isLarge = :false',
         ExpressionAttributeNames: {
           '#id': 'id',
+          '#isLarge': 'isLarge',
+        },
+        ExpressionAttributeValues: {
+          ':false': false,
         },
       },
       {concurrency: 5, chunkSize: 2}
@@ -58,4 +75,50 @@ describe('parallelScanAsStream', () => {
 
     expect(allItems).toHaveLength(10);
   });
+
+  it('should pause calling dynamodb after highWaterMark reached', async () => {
+    const scanSpy = jest.spyOn(ddbHelpers, 'scan');
+
+    const megaByte = Buffer.alloc(1024 * 390); // Maximum allowed item size in ddb is 400KB
+    const megaByteString = megaByte.toString();
+
+    await ddbHelpers.insertMany({
+      items: [
+        {id: 'some-big-file-id-1', isLarge: true, payload: megaByteString},
+        {id: 'some-big-file-id-2', isLarge: true, payload: megaByteString},
+        {id: 'some-big-file-id-3', isLarge: true, payload: megaByteString},
+        {id: 'some-big-file-id-4', isLarge: true, payload: megaByteString},
+        {id: 'some-big-file-id-5', isLarge: true, payload: megaByteString},
+      ],
+      tableName: 'files',
+    });
+
+    const stream = await parallelScanAsStream(
+      {
+        TableName: 'files',
+        FilterExpression: 'attribute_exists(#id) and #isLarge = :true',
+        ExpressionAttributeNames: {
+          '#id': 'id',
+          '#isLarge': 'isLarge',
+        },
+        ExpressionAttributeValues: {
+          ':true': true,
+        },
+      },
+      {concurrency: 1, chunkSize: 1, highWaterMark: 1}
+    );
+
+    const scanCallsByIteration = [];
+    for await (const _ of stream) {
+      expect(_).not.toBeUndefined();
+
+      await delay(1000);
+
+      scanCallsByIteration.push(scanSpy.mock.calls.length);
+    }
+
+    const scanCallsByIterationUniq = uniq(scanCallsByIteration);
+
+    expect(scanCallsByIterationUniq).toEqual([1, 2]);
+  });
 });