feat: enable write retry and nack pending writes on reconnect

src/managedwriter/stream_connection.ts

@@ -87,6 +87,7 @@ export class StreamConnection extends EventEmitter {
     this._connection.on('error', this.handleError);
     this._connection.on('close', () => {
       this.trace('connection closed');
+      this.close();
     });
     this._connection.on('pause', () => {
       this.trace('connection paused');
@@ -107,20 +108,19 @@ export class StreamConnection extends EventEmitter {
   private handleError = (err: gax.GoogleError) => {
     this.trace('on error', err, JSON.stringify(err));
     if (this.shouldReconnect(err)) {
+      err.message = 'reconnect triggered due to: ' + err.message;
+      this.ackAllPendingWrites(err);
       this.reconnect();
       return;
     }
-    let nextPendingWrite = this.getNextPendingWrite();
     if (this.isPermanentError(err)) {
       this.trace('found permanent error', err);
-      while (nextPendingWrite) {
-        this.ackNextPendingWrite(err);
-        nextPendingWrite = this.getNextPendingWrite();
-      }
+      this.ackAllPendingWrites(err);
       this.emit('error', err);
       return;
     }
-    if (this.isRequestError(err) && nextPendingWrite) {
+    const nextPendingWrite = this.getNextPendingWrite();
+    if (nextPendingWrite) {
       this.trace(
         'found request error with pending write',
         err,
@@ -144,6 +144,13 @@ export class StreamConnection extends EventEmitter {
     return !!err.code && reconnectionErrorCodes.includes(err.code);
   }
 
+  private isConnectionClosed() {
+    if (this._connection) {
+      return this._connection.destroyed || this._connection.closed;
+    }
+    return true;
+  }
+
   private isPermanentError(err: gax.GoogleError): boolean {
     if (err.code === gax.Status.INVALID_ARGUMENT) {
       const storageErrors = parseStorageErrors(err);
@@ -160,10 +167,6 @@ export class StreamConnection extends EventEmitter {
     return false;
   }
 
-  private isRequestError(err: gax.GoogleError): boolean {
-    return err.code === gax.Status.INVALID_ARGUMENT;
-  }
-
   private resolveCallOptions(
     streamId: string,
     options?: gax.CallOptions
@@ -245,6 +248,19 @@ export class StreamConnection extends EventEmitter {
     return null;
   }
 
+  private ackAllPendingWrites(
+    err: Error | null,
+    result?:
+      | protos.google.cloud.bigquery.storage.v1.IAppendRowsResponse
+      | undefined
+  ) {
+    let nextPendingWrite = this.getNextPendingWrite();
+    while (nextPendingWrite) {
+      this.ackNextPendingWrite(err, result);
+      nextPendingWrite = this.getNextPendingWrite();
+    }
+  }
+
   private ackNextPendingWrite(
     err: Error | null,
     result?:
@@ -280,16 +296,12 @@ export class StreamConnection extends EventEmitter {
 
   private send(pw: PendingWrite) {
     const request = pw.getRequest();
-    if (!this._connection) {
-      pw._markDone(new Error('connection closed'));
-      return;
-    }
-    if (this._connection.destroyed || this._connection.closed) {
+    if (this.isConnectionClosed()) {
       this.reconnect();
     }
     this.trace('sending pending write', pw);
     try {
-      this._connection.write(request, err => {
+      this._connection?.write(request, err => {
         this.trace('wrote pending write', err, this._pendingWrites.length);
         if (err) {
           pw._markDone(err); //TODO: add retries
@@ -306,11 +318,11 @@ export class StreamConnection extends EventEmitter {
    * Check if connection is open and ready to send requests.
    */
   isOpen(): boolean {
-    return !!this._connection;
+    return !this.isConnectionClosed();
   }
 
   /**
-   * Reconnect and re send inflight requests.
+   * Re open appendRows BiDi gRPC connection.
    */
   reconnect() {
     this.trace('reconnect called');

system-test/managed_writer_client_test.ts

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 import * as assert from 'assert';
-import {describe, it, xit} from 'mocha';
+import {describe, it} from 'mocha';
 import * as uuid from 'uuid';
 import * as gax from 'google-gax';
 import * as sinon from 'sinon';
@@ -51,6 +51,11 @@ const generateUuid = () =>
   `${GCLOUD_TESTS_PREFIX}_${uuid.v4()}`.replace(/-/gi, '_');
 const datasetId = generateUuid();
 
+const sleep = (ms: number) =>
+  new Promise(resolve => {
+    setTimeout(resolve, ms);
+  });
+
 const root = protobuf.Root.fromJSON(customerRecordProtoJson);
 if (!root) {
   throw Error('Proto must not be undefined');
@@ -1172,7 +1177,7 @@ describe('managedwriter.WriterClient', () => {
       }
     });
 
-    xit('reconnect on idle connection', async () => {
+    it('reconnect on idle connection', async () => {
       bqWriteClient.initialize();
       const client = new WriterClient();
       client.setClient(bqWriteClient);
@@ -1207,15 +1212,17 @@ describe('managedwriter.WriterClient', () => {
         let pw = writer.appendRows([row1, row2], 0);
         await pw.getResult();
 
-        const sleep = (ms: number) =>
-          new Promise(resolve => {
-            setTimeout(resolve, ms);
-          });
-        const minutes = 10;
-        for (let i = 0; i <= minutes; i++) {
-          console.log('sleeping for a minute: ', minutes - i, 'to go');
-          await sleep(60 * 1000);
-        }
+        // Simulate server sending ABORT error as the conn was idle
+        const conn = connection['_connection'] as gax.CancellableStream; // private method
+        const gerr = new gax.GoogleError(
+          'Closing the stream because it has been inactive for 600 seconds'
+        );
+        gerr.code = gax.Status.ABORTED;
+        conn.emit('error', gerr);
+        // simulate server closing conn.
+        await sleep(100);
+        conn.destroy();
+        await sleep(100);
 
         const row3 = {
           customer_name: 'Test',
@@ -1234,7 +1241,70 @@ describe('managedwriter.WriterClient', () => {
       } finally {
         client.close();
       }
-    }).timeout(20 * 60 * 1000);
+    }).timeout(20 * 1000);
+
+    it('should mark any pending writes with error if connection was closed', async () => {
+      bqWriteClient.initialize();
+      const client = new WriterClient();
+      client.setClient(bqWriteClient);
+
+      const row1 = {
+        customer_name: 'Ada Lovelace',
+        row_num: 1,
+      };
+
+      try {
+        const connection = await client.createStreamConnection({
+          streamType: managedwriter.PendingStream,
+          destinationTable: parent,
+        });
+
+        const writer = new JSONWriter({
+          connection,
+          protoDescriptor,
+        });
+
+        const pw1 = writer.appendRows([row1], 0);
+        await pw1.getResult();
+
+        // Try to append a new row
+        const row2 = {
+          customer_name: 'Test',
+          row_num: 2,
+          customer_email: '[email protected]',
+        };
+
+        let foundError: gax.GoogleError | null = null;
+        const pw2 = writer.appendRows([row2], 1);
+        pw2.getResult().catch(err => {
+          foundError = err as gax.GoogleError;
+        });
+
+        // Simulate server sending RESOURCE_EXHAUSTED error on a write
+        const conn = connection['_connection'] as gax.CancellableStream; // private method
+        // swallow ack for the last appendRow call, so we can simulate it failing
+        conn.removeAllListeners('data');
+        await new Promise(resolve => conn.once('data', resolve));
+        conn.addListener('data', connection['handleData']);
+
+        const gerr = new gax.GoogleError('memory limit exceeded');
+        gerr.code = gax.Status.RESOURCE_EXHAUSTED;
+        conn.emit('error', gerr);
+        // simulate server closing conn.
+        await sleep(100);
+        conn.destroy();
+        await sleep(100);
+
+        // should throw error of reconnection
+        assert.notEqual(foundError, null);
+        assert.equal(foundError!.message.includes('reconnect'), true);
+
+        connection.close();
+        writer.close();
+      } finally {
+        client.close();
+      }
+    });
   });
 
   describe('close', () => {