refactor(subscriber): message stream retry logic (#607)

callmehiphop · web-flow · commit bd5bcad7c21f · 2019-05-15T11:19:50.000-07:00
diff --git a/handwritten/pubsub/src/message-stream.ts b/handwritten/pubsub/src/message-stream.ts
@@ -26,38 +26,15 @@ import {
 import * as isStreamEnded from 'is-stream-ended';
 import {PassThrough} from 'stream';
 
-import {PullResponse, Subscriber} from './subscriber';
+import {PullRetry} from './pull-retry';
+import {Subscriber} from './subscriber';
+import {google} from '../proto/pubsub';
 
 /*!
  * Frequency to ping streams.
  */
 const KEEP_ALIVE_INTERVAL = 30000;
 
-/*!
- * Deadline Exceeded status code
- */
-const DEADLINE: status = 4;
-
-/*!
- * Unknown status code
- */
-const UNKNOWN: status = 2;
-
-/*!
- * codes to retry streams
- */
-const RETRY_CODES: status[] = [
-  0, // ok
-  1, // canceled
-  2, // unknown
-  4, // deadline exceeded
-  8, // resource exhausted
-  10, // aborted
-  13, // internal error
-  14, // unavailable
-  15, // dataloss
-];
-
 /*!
  * Deadline for the stream.
  */
@@ -79,14 +56,8 @@ interface StreamState {
   highWaterMark: number;
 }
 
-interface StreamingPullRequest {
-  subscription?: string;
-  ackIds?: string[];
-  modifyDeadlineSeconds?: number[];
-  modifyDeadlineAckIds?: string[];
-  streamAckDeadlineSeconds?: number;
-}
-
+type StreamingPullRequest = google.pubsub.v1.IStreamingPullRequest;
+type PullResponse = google.pubsub.v1.IPullResponse;
 type PullStream = ClientDuplexStream<StreamingPullRequest, PullResponse> & {
   _readableState: StreamState;
 };
@@ -119,7 +90,9 @@ export class ChannelError extends Error implements ServiceError {
   code: status;
   constructor(err: Error) {
     super(`Failed to connect to channel. Reason: ${err.message}`);
-    this.code = err.message.includes('deadline') ? DEADLINE : UNKNOWN;
+    this.code = err.message.includes('deadline')
+      ? status.DEADLINE_EXCEEDED
+      : status.UNKNOWN;
   }
 }
 
@@ -154,7 +127,9 @@ export interface MessageStreamOptions {
 export class MessageStream extends PassThrough {
   destroyed: boolean;
   private _keepAliveHandle: NodeJS.Timer;
+  private _fillHandle?: NodeJS.Timer;
   private _options: MessageStreamOptions;
+  private _retrier: PullRetry;
   private _streams: Map<PullStream, boolean>;
   private _subscriber: Subscriber;
   constructor(sub: Subscriber, options = {} as MessageStreamOptions) {
@@ -164,6 +139,7 @@ export class MessageStream extends PassThrough {
 
     this.destroyed = false;
     this._options = options;
+    this._retrier = new PullRetry();
     this._streams = new Map();
     this._subscriber = sub;
 
@@ -253,6 +229,8 @@ export class MessageStream extends PassThrough {
       streamAckDeadlineSeconds: this._subscriber.ackDeadline,
     };
 
+    delete this._fillHandle;
+
     for (let i = this._streams.size; i < this._options.maxStreams!; i++) {
       const stream: PullStream = client.streamingPull({deadline});
       this._addStream(stream);
@@ -302,8 +280,13 @@ export class MessageStream extends PassThrough {
   private _onEnd(stream: PullStream, status: StatusObject): void {
     this._removeStream(stream);
 
-    if (RETRY_CODES.includes(status.code)) {
-      this._fillStreamPool();
+    if (this._fillHandle) {
+      return;
+    }
+
+    if (this._retrier.retry(status)) {
+      const delay = this._retrier.createTimeout();
+      this._fillHandle = setTimeout(() => this._fillStreamPool(), delay);
     } else if (!this._streams.size) {
       this.destroy(new StatusError(status));
     }
diff --git a/handwritten/pubsub/src/pull-retry.ts b/handwritten/pubsub/src/pull-retry.ts
@@ -0,0 +1,78 @@
+/*!
+ * Copyright 2019 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+import {StatusObject, status} from 'grpc';
+
+/*!
+ * retryable status codes
+ */
+export const RETRY_CODES: status[] = [
+  status.OK,
+  status.CANCELLED,
+  status.UNKNOWN,
+  status.DEADLINE_EXCEEDED,
+  status.RESOURCE_EXHAUSTED,
+  status.ABORTED,
+  status.INTERNAL,
+  status.UNAVAILABLE,
+  status.DATA_LOSS,
+];
+
+/**
+ * Used to track pull requests and determine if additional requests should be
+ * made, etc.
+ *
+ * @class
+ * @private
+ */
+export class PullRetry {
+  private failures = 0;
+  /**
+   * Generates a timeout that can be used for applying a backoff based on the
+   * current number of failed requests.
+   *
+   * @see {@link https://cloud.google.com/iot/docs/how-tos/exponential-backoff}
+   * @private
+   * @returns {number}
+   */
+  createTimeout(): number {
+    if (this.failures === 0) {
+      return 0;
+    }
+    return Math.pow(2, this.failures) * 1000 + Math.floor(Math.random() * 1000);
+  }
+  /**
+   * Determines if a request status should be retried.
+   *
+   * Deadlines behave kind of unexpectedly on streams, rather than using it as
+   * an indicator of when to give up trying to connect, it actually dictates
+   * how long the stream should stay open. Because of this, it is virtually
+   * impossible to determine whether or not a deadline error is the result of
+   * the server closing the stream or if we timed out waiting for a connection.
+   *
+   * @private
+   * @param {object} status The request status.
+   * @returns {boolean}
+   */
+  retry(err: StatusObject): boolean {
+    if (err.code === status.OK || err.code === status.DEADLINE_EXCEEDED) {
+      this.failures = 0;
+    } else {
+      this.failures += 1;
+    }
+
+    return RETRY_CODES.includes(err.code);
+  }
+}
diff --git a/handwritten/pubsub/test/message-stream.ts b/handwritten/pubsub/test/message-stream.ts
@@ -506,7 +506,7 @@ describe('MessageStream', () => {
 
         stream.push(null);
         setImmediate(() => {
-          assert.strictEqual(client.streams.length, 6);
+          assert.strictEqual(client.streams.length, 5);
           done();
         });
       });
@@ -524,7 +524,7 @@ describe('MessageStream', () => {
           assert.strictEqual(stream.listenerCount('end'), count);
 
           setImmediate(() => {
-            assert.strictEqual(client.streams.length, 6);
+            assert.strictEqual(client.streams.length, 5);
             done();
           });
         });
diff --git a/handwritten/pubsub/test/pull-retry.ts b/handwritten/pubsub/test/pull-retry.ts
@@ -0,0 +1,97 @@
+/**
+ * Copyright 2019 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import assert = require('assert');
+import sinon = require('sinon');
+import {StatusObject, status} from 'grpc';
+import {PullRetry} from '../src/pull-retry';
+
+describe('PullRetry', () => {
+  const sandbox = sinon.createSandbox();
+
+  let retrier: PullRetry;
+
+  beforeEach(() => {
+    retrier = new PullRetry();
+  });
+
+  afterEach(() => {
+    sandbox.restore();
+  });
+
+  describe('createTimeout', () => {
+    it('should return 0 when no failures have occurred', () => {
+      assert.strictEqual(retrier.createTimeout(), 0);
+    });
+
+    it('should use a backoff factoring in the failure count', () => {
+      const random = Math.random();
+      const expected = Math.pow(2, 1) * 1000 + Math.floor(random * 1000);
+
+      sandbox.stub(global.Math, 'random').returns(random);
+
+      retrier.retry({code: status.CANCELLED} as StatusObject);
+      assert.strictEqual(retrier.createTimeout(), expected);
+    });
+  });
+
+  describe('retry', () => {
+    it('should return true for retryable errors', () => {
+      [
+        status.OK,
+        status.CANCELLED,
+        status.UNKNOWN,
+        status.DEADLINE_EXCEEDED,
+        status.RESOURCE_EXHAUSTED,
+        status.ABORTED,
+        status.INTERNAL,
+        status.UNAVAILABLE,
+        status.DATA_LOSS,
+      ].forEach((code: status) => {
+        const shouldRetry = retrier.retry({code} as StatusObject);
+        assert.strictEqual(shouldRetry, true);
+      });
+    });
+
+    it('should return false for non-retryable errors', () => {
+      [
+        status.INVALID_ARGUMENT,
+        status.NOT_FOUND,
+        status.PERMISSION_DENIED,
+        status.FAILED_PRECONDITION,
+        status.OUT_OF_RANGE,
+        status.UNIMPLEMENTED,
+      ].forEach((code: status) => {
+        const shouldRetry = retrier.retry({code} as StatusObject);
+        assert.strictEqual(shouldRetry, false);
+      });
+    });
+
+    it('should reset the failure count on OK', () => {
+      retrier.retry({code: status.CANCELLED} as StatusObject);
+      retrier.retry({code: status.OK} as StatusObject);
+
+      assert.strictEqual(retrier.createTimeout(), 0);
+    });
+
+    it('should reset the failure count on DEADLINE_EXCEEDED', () => {
+      retrier.retry({code: status.CANCELLED} as StatusObject);
+      retrier.retry({code: status.DEADLINE_EXCEEDED} as StatusObject);
+
+      assert.strictEqual(retrier.createTimeout(), 0);
+    });
+  });
+});
