googleapis · danieljbruce · Jul 11, 2024 · Jul 9, 2024 · Jul 2, 2024 · Jul 2, 2024
@@ -726,7 +726,7 @@ Please use the format 'prezzy' or '${instance.name}/tables/prezzy'.`);
     let filter: {} | null;
     const rowsLimit = options.limit || 0;
     const hasLimit = rowsLimit !== 0;
-    let rowsRead = 0;
+
     let numConsecutiveErrors = 0;
     let numRequestsMade = 0;
     let retryTimer: NodeJS.Timeout | null;
@@ -749,14 +749,22 @@ Please use the format 'prezzy' or '${instance.name}/tables/prezzy'.`);
     let rowStream: Duplex;
 
     let userCanceled = false;
+    // The key of the last row that was emitted by the per attempt pipeline
+    // Note: this must be updated from the operation level userStream to avoid referencing buffered rows that will be
+    // discarded in the per attempt subpipeline (rowStream)
+    let lastRowKey = '';
+    let rowsRead = 0;
     const userStream = new PassThrough({
       objectMode: true,
-      readableHighWaterMark: 0,
+      readableHighWaterMark: 0, // We need to disable readside buffering to allow for acceptable behavior when the end user cancels the stream early.
+      writableHighWaterMark: 0, // We need to disable writeside buffering because in nodejs 14 the call to _transform happens after write buffering. This creates problems for tracking the last seen row key.
       transform(row, _encoding, callback) {
         if (userCanceled) {
           callback();
           return;
         }
+        lastRowKey = row.id;
+        rowsRead++;
         callback(null, row);
       },
     });
@@ -796,7 +804,6 @@ Please use the format 'prezzy' or '${instance.name}/tables/prezzy'.`);
       // cancelled the stream in the middle of a retry
       retryTimer = null;
 
-      const lastRowKey = chunkTransformer ? chunkTransformer.lastRowKey : '';
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       chunkTransformer = new ChunkTransformer({decode: options.decode} as any);
 
@@ -918,7 +925,6 @@ Please use the format 'prezzy' or '${instance.name}/tables/prezzy'.`);
           ) {
             return next();
           }
-          rowsRead++;
           const row = this.row(rowData.key);
           row.data = rowData.data;
           next(null, row);

@@ -37,7 +37,6 @@ export class MockServer {
       `localhost:${this.port}`,
       grpc.ServerCredentials.createInsecure(),
       () => {
-        server.start();
         callback ? callback(portString) : undefined;
       }
     );

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 import {before, describe, it} from 'mocha';
-import {Bigtable, Row, Table} from '../src';
+import {Bigtable, protos, Row, Table} from '../src';
 import * as assert from 'assert';
 import {Transform, PassThrough, pipeline} from 'stream';
 
@@ -22,7 +22,18 @@ import {MockServer} from '../src/util/mock-servers/mock-server';
 import {BigtableClientMockService} from '../src/util/mock-servers/service-implementations/bigtable-client-mock-service';
 import {MockService} from '../src/util/mock-servers/mock-service';
 import {debugLog, readRowsImpl} from './utils/readRowsImpl';
-import {UntypedHandleCall} from '@grpc/grpc-js';
+import {ServerWritableStream, UntypedHandleCall} from '@grpc/grpc-js';
+import {readRowsImpl2} from './utils/readRowsImpl2';
+
+type PromiseVoid = Promise<void>;
+interface ServerImplementationInterface {
+  (
+    server: ServerWritableStream<
+      protos.google.bigtable.v2.IReadRowsRequest,
+      protos.google.bigtable.v2.IReadRowsResponse
+    >
+  ): PromiseVoid;
+}
 
 describe('Bigtable/ReadRows', () => {
   let server: MockServer;
@@ -317,6 +328,47 @@ describe('Bigtable/ReadRows', () => {
     });
   });
 
+  it('should return row data in the right order', done => {
+    // 150 rows must be enough to reproduce issues with losing the data and to create backpressure
+    const keyFrom = undefined;
+    const keyTo = undefined;
+    // the server will error after sending this chunk (not row)
+    const errorAfterChunkNo = 100;
+    const dataResults = [];
+
+    // TODO: Do not use `any` here, make it a more specific type and address downstream implications on the mock server.
+    service.setService({
+      ReadRows: readRowsImpl2(
+        keyFrom,
+        keyTo,
+        errorAfterChunkNo
+      ) as ServerImplementationInterface,
+    });
+    const sleep = (ms: number) => {
+      return new Promise(resolve => setTimeout(resolve, ms));
+    };
+    (async () => {
+      try {
+        const stream = table.createReadStream({
+          start: '00000000',
+          end: '00000150',
+        });
+
+        for await (const row of stream) {
+          dataResults.push(row.id);
+          await sleep(50);
+        }
+        const expectedResults = Array.from(Array(150).keys())
+          .map(i => '00000000' + i.toString())
+          .map(i => i.slice(-8));
+        assert.deepStrictEqual(dataResults, expectedResults);
+        done();
+      } catch (error) {
+        done(error);
+      }
+    })();
+  });
+
   after(async () => {
     server.shutdown(() => {});
   });

@@ -1118,6 +1118,20 @@ describe('Bigtable/Table', () => {
       let reqOptsCalls: any[];
       let setTimeoutSpy: sinon.SinonSpy;
 
+      /*
+        setImmediate is required here to correctly mock events as they will
+        come in from the request function. It is required for tests to pass,
+        but it is not a problem that it is required because we never expect
+        a single Node event to emit data and then emit an error. That is,
+        a mock without setImmediate around the last error represents a scenario
+        that will never happen.
+       */
+      function emitRetriableError(stream: Duplex) {
+        setImmediate(() => {
+          stream.emit('error', makeRetryableError());
+        });
+      }
+
       beforeEach(() => {
         FakeChunkTransformer.prototype._transform = function (
           rows: Row[],
@@ -1262,7 +1276,7 @@ describe('Bigtable/Table', () => {
           // eslint-disable-next-line @typescript-eslint/no-explicit-any
           ((stream: any) => {
             stream.push([{key: 'a'}]);
-            stream.emit('error', makeRetryableError());
+            emitRetriableError(stream);
           }) as {} as EventEmitter,
           ((stream: Writable) => {
             stream.end();
@@ -1285,7 +1299,7 @@ describe('Bigtable/Table', () => {
         emitters = [
           ((stream: Duplex) => {
             stream.push([{key: 'a'}]);
-            stream.emit('error', makeRetryableError());
+            emitRetriableError(stream);
           }) as {} as EventEmitter,
           ((stream: Writable) => {
             stream.end();
@@ -1310,7 +1324,7 @@ describe('Bigtable/Table', () => {
           ((stream: Duplex) => {
             stream.push([{key: 'a'}]);
             stream.push([{key: 'b'}]);
-            stream.emit('error', makeRetryableError());
+            emitRetriableError(stream);
           }) as {} as EventEmitter,
           ((stream: Duplex) => {
             stream.push([{key: 'c'}]);
@@ -1343,7 +1357,7 @@ describe('Bigtable/Table', () => {
         emitters = [
           ((stream: Duplex) => {
             stream.push([{key: 'a'}]);
-            stream.emit('error', makeRetryableError());
+            emitRetriableError(stream);
           }) as {} as EventEmitter,
           ((stream: Duplex) => {
             stream.end([{key: 'c'}]);
@@ -1362,7 +1376,7 @@ describe('Bigtable/Table', () => {
           ((stream: Duplex) => {
             stream.push([{key: 'a'}]);
             stream.push([{key: 'b'}]);
-            stream.emit('error', makeRetryableError());
+            emitRetriableError(stream);
           }) as {} as EventEmitter,
         ];
 
@@ -1381,7 +1395,7 @@ describe('Bigtable/Table', () => {
         emitters = [
           ((stream: Duplex) => {
             stream.push([{key: 'a'}]);
-            stream.emit('error', makeRetryableError());
+            emitRetriableError(stream);
           }) as {} as EventEmitter,
         ];
 
@@ -1395,7 +1409,7 @@ describe('Bigtable/Table', () => {
         emitters = [
           ((stream: Duplex) => {
             stream.push([{key: 'c'}]);
-            stream.emit('error', makeRetryableError());
+            emitRetriableError(stream);
           }) as {} as EventEmitter,
         ];
 
@@ -1418,7 +1432,7 @@ describe('Bigtable/Table', () => {
           ((stream: Duplex) => {
             stream.push([{key: 'a1'}]);
             stream.push([{key: 'd'}]);
-            stream.emit('error', makeRetryableError());
+            emitRetriableError(stream);
           }) as {} as EventEmitter,
         ];