fix: pause request stream on backpressure (#936)

* fix: pause request stream on backpressure

The request stream should be paused if the downstream is indicating
that it cannot handle any more data at the moment. The request stream
should be resumed once the downstream does accept data again. This
reduces memory consumption and potentially out of memory errors when
a result stream is piped into a slow writer.

Fixes #934

* fix: do not retry stream indefinitely

PartialResultSetStream should stop retrying to push data into the
stream after a configurable number of retries have failed.

* fix: process review comments

* fix: remove unused code

* tests: add test for pause/resume

* fix: return after giving up retrying + add test

Author: olavloite

handwritten/spanner/.gitignore

@@ -12,3 +12,4 @@ system-test/*key.json
 .DS_Store
 package-lock.json
 __pycache__
+.idea

handwritten/spanner/src/partial-result-stream.ts

@@ -43,10 +43,16 @@ interface RequestFunction {
  * @property {boolean} [json=false] Indicates if the Row objects should be
  *     formatted into JSON.
  * @property {JSONOptions} [jsonOptions] JSON options.
+ * @property {number} [maxResumeRetries=20] The maximum number of times that the
+ *     stream will retry to push data downstream, when the downstream indicates
+ *     that it is not ready for any more data. Increase this value if you
+ *     experience 'Stream is still not ready to receive data' errors as a
+ *     result of a slow writer in your receiving stream.
  */
 export interface RowOptions {
   json?: boolean;
   jsonOptions?: JSONOptions;
+  maxResumeRetries?: number;
 }
 
 /**
@@ -136,11 +142,12 @@ export class PartialResultStream extends Transform implements ResultEvents {
   private _options: RowOptions;
   private _pendingValue?: p.IValue;
   private _values: p.IValue[];
+  private _numPushFailed = 0;
   constructor(options = {}) {
     super({objectMode: true});
 
     this._destroyed = false;
-    this._options = options;
+    this._options = Object.assign({maxResumeRetries: 20}, options);
     this._values = [];
   }
   /**
@@ -187,20 +194,60 @@ export class PartialResultStream extends Transform implements ResultEvents {
         .fields as google.spanner.v1.StructType.Field[];
     }
 
+    let res = true;
     if (!is.empty(chunk.values)) {
-      this._addChunk(chunk);
+      res = this._addChunk(chunk);
     }
 
-    next();
+    if (res) {
+      next();
+    } else {
+      // Wait a little before we push any more data into the pipeline as a
+      // component downstream has indicated that a break is needed. Pause the
+      // request stream to prevent it from filling up the buffer while we are
+      // waiting.
+      // The stream will initially pause for 2ms, and then double the pause time
+      // for each new pause.
+      const initialPauseMs = 2;
+      setTimeout(() => {
+        this._tryResume(next, 2 * initialPauseMs);
+      }, initialPauseMs);
+    }
+  }
+
+  private _tryResume(next: Function, timeout: number) {
+    // Try to push an empty chunk to check whether more data can be accepted.
+    if (this.push(undefined)) {
+      this._numPushFailed = 0;
+      this.emit('resumed');
+      next();
+    } else {
+      // Downstream returned false indicating that it is still not ready for
+      // more data.
+      this._numPushFailed++;
+      if (this._numPushFailed === this._options.maxResumeRetries) {
+        this.destroy(
+          new Error(
+            `Stream is still not ready to receive data after ${this._numPushFailed} attempts to resume.`
+          )
+        );
+        return;
+      }
+      setTimeout(() => {
+        const nextTimeout = Math.min(timeout * 2, 1024);
+        this._tryResume(next, nextTimeout);
+      }, timeout);
+    }
   }
+
   /**
    * Manages any chunked values.
    *
    * @private
    *
    * @param {object} chunk The partial result set.
    */
-  private _addChunk(chunk: google.spanner.v1.PartialResultSet): void {
+  private _addChunk(chunk: google.spanner.v1.PartialResultSet): boolean {
     const values: Value[] = chunk.values.map(GrpcService.decodeValue_);
 
     // If we have a chunk to merge, merge the values now.
@@ -223,7 +270,14 @@ export class PartialResultStream extends Transform implements ResultEvents {
       this._pendingValue = values.pop();
     }
 
-    values.forEach(value => this._addValue(value));
+    let res = true;
+    values.forEach(value => {
+      res = this._addValue(value) && res;
+      if (!res) {
+        this.emit('paused');
+      }
+    });
+    return res;
   }
   /**
    * Manages complete values, pushing a completed row into the stream once all
@@ -233,25 +287,24 @@ export class PartialResultStream extends Transform implements ResultEvents {
    *
    * @param {*} value The complete value.
    */
-  private _addValue(value: Value): void {
+  private _addValue(value: Value): boolean {
     const values = this._values;
 
     values.push(value);
 
     if (values.length !== this._fields.length) {
-      return;
+      return true;
     }
 
     this._values = [];
 
     const row: Row = this._createRow(values);
 
     if (this._options.json) {
-      this.push(row.toJSON(this._options.jsonOptions));
-      return;
+      return this.push(row.toJSON(this._options.jsonOptions));
     }
 
-    this.push(row);
+    return this.push(row);
   }
   /**
    * Converts an array of values into a row.
@@ -373,7 +426,8 @@ export function partialResultStream(
   // mergeStream allows multiple streams to be connected into one. This is good;
   // if we need to retry a request and pipe more data to the user's stream.
   const requestsStream = mergeStream();
-  const userStream = streamEvents(new PartialResultStream(options));
+  const partialRSStream = new PartialResultStream(options);
+  const userStream = streamEvents(partialRSStream);
   const batchAndSplitOnTokenStream = checkpointStream.obj({
     maxQueued: 10,
     isCheckpointFn: (row: google.spanner.v1.PartialResultSet): boolean => {
@@ -433,5 +487,7 @@ export function partialResultStream(
         lastResumeToken = row.resumeToken;
       })
       .pipe(userStream)
+      .on('paused', () => requestsStream.pause())
+      .on('resumed', () => requestsStream.resume())
   );
 }

handwritten/spanner/src/transaction.ts

@@ -55,6 +55,7 @@ export interface RequestOptions {
   json?: boolean;
   jsonOptions?: JSONOptions;
   gaxOptions?: CallOptions;
+  maxResumeRetries?: number;
 }
 
 export interface Statement {
@@ -484,7 +485,7 @@ export class Snapshot extends EventEmitter {
     table: string,
     request = {} as ReadRequest
   ): PartialResultStream {
-    const {gaxOptions, json, jsonOptions} = request;
+    const {gaxOptions, json, jsonOptions, maxResumeRetries} = request;
     const keySet = Snapshot.encodeKeySet(request);
     const transaction: spannerClient.spanner.v1.ITransactionSelector = {};
 
@@ -499,6 +500,7 @@ export class Snapshot extends EventEmitter {
     delete request.gaxOptions;
     delete request.json;
     delete request.jsonOptions;
+    delete request.maxResumeRetries;
     delete request.keys;
     delete request.ranges;
 
@@ -518,7 +520,11 @@ export class Snapshot extends EventEmitter {
       });
     };
 
-    return partialResultStream(makeRequest, {json, jsonOptions});
+    return partialResultStream(makeRequest, {
+      json,
+      jsonOptions,
+      maxResumeRetries,
+    });
   }
 
   /**
@@ -879,7 +885,7 @@ export class Snapshot extends EventEmitter {
       query.queryOptions
     );
 
-    const {gaxOptions, json, jsonOptions} = query;
+    const {gaxOptions, json, jsonOptions, maxResumeRetries} = query;
     const {params, paramTypes} = Snapshot.encodeParams(query);
     const transaction: spannerClient.spanner.v1.ITransactionSelector = {};
 
@@ -892,6 +898,7 @@ export class Snapshot extends EventEmitter {
     delete query.gaxOptions;
     delete query.json;
     delete query.jsonOptions;
+    delete query.maxResumeRetries;
     delete query.types;
 
     const reqOpts: ExecuteSqlRequest = Object.assign(query, {
@@ -911,7 +918,11 @@ export class Snapshot extends EventEmitter {
       });
     };
 
-    return partialResultStream(makeRequest, {json, jsonOptions});
+    return partialResultStream(makeRequest, {
+      json,
+      jsonOptions,
+      maxResumeRetries,
+    });
   }
 
   /**

handwritten/spanner/test/mockserver/mockspanner.ts

@@ -583,16 +583,24 @@ export class MockSpanner {
     resultSet: protobuf.ResultSet,
     queryMode:
       | google.spanner.v1.ExecuteSqlRequest.QueryMode
-      | keyof typeof google.spanner.v1.ExecuteSqlRequest.QueryMode
+      | keyof typeof google.spanner.v1.ExecuteSqlRequest.QueryMode,
+    rowsPerPartialResultSet = 1
   ): protobuf.PartialResultSet[] {
     const res: protobuf.PartialResultSet[] = [];
     let first = true;
-    for (let i = 0; i < resultSet.rows.length; i++) {
+    for (let i = 0; i < resultSet.rows.length; i += rowsPerPartialResultSet) {
       const token = i.toString().padStart(8, '0');
       const partial = protobuf.PartialResultSet.create({
         resumeToken: Buffer.from(token),
-        values: resultSet.rows[i].values,
+        values: [],
       });
+      for (
+        let row = i;
+        row < Math.min(i + rowsPerPartialResultSet, resultSet.rows.length);
+        row++
+      ) {
+        partial.values.push(...resultSet.rows[row].values!);
+      }
       if (first) {
         partial.metadata = resultSet.metadata;
         first = false;
@@ -937,6 +945,50 @@ export function createSimpleResultSet(): protobuf.ResultSet {
   });
 }
 
+export const NUM_ROWS_LARGE_RESULT_SET = 100;
+
+export function createLargeResultSet(): protobuf.ResultSet {
+  const fields = [
+    protobuf.StructType.Field.create({
+      name: 'NUM',
+      type: protobuf.Type.create({code: protobuf.TypeCode.INT64}),
+    }),
+    protobuf.StructType.Field.create({
+      name: 'NAME',
+      type: protobuf.Type.create({code: protobuf.TypeCode.STRING}),
+    }),
+  ];
+  const metadata = new protobuf.ResultSetMetadata({
+    rowType: new protobuf.StructType({
+      fields,
+    }),
+  });
+  const rows: google.protobuf.IListValue[] = [];
+  for (let num = 1; num <= NUM_ROWS_LARGE_RESULT_SET; num++) {
+    rows.push({
+      values: [
+        {stringValue: `${num}`},
+        {stringValue: generateRandomString(100)},
+      ],
+    });
+  }
+  return protobuf.ResultSet.create({
+    metadata,
+    rows,
+  });
+}
+
+function generateRandomString(length: number) {
+  let result = '';
+  const characters =
+    'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789';
+  const charactersLength = characters.length;
+  for (let i = 0; i < length; i++) {
+    result += characters.charAt(Math.floor(Math.random() * charactersLength));
+  }
+  return result;
+}
+
 /**
  * Returns a protobuf Timestamp containing the current local system time.
  */

handwritten/spanner/test/partial-result-stream.ts

@@ -83,7 +83,7 @@ describe('PartialResultStream', () => {
       it(`should pass acceptance test: ${test.name}`, done => {
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
         const values: any[] = [];
-        const stream = new PartialResultStream();
+        const stream = new PartialResultStream({});
 
         stream
           .on('error', done)
@@ -113,13 +113,13 @@ describe('PartialResultStream', () => {
     let stream: prs.PartialResultStream;
 
     beforeEach(() => {
-      stream = new PartialResultStream();
+      stream = new PartialResultStream({});
     });
 
     afterEach(() => stream.destroy());
 
     it('should emit the response', done => {
-      const stream = new PartialResultStream();
+      const stream = new PartialResultStream({});
 
       stream.on('error', done).on('response', response => {
         assert.strictEqual(response, RESULT);