feat(pubsub): add support for streaming pull keepalives from the server

handwritten/pubsub/src/message-stream.ts

@@ -26,6 +26,7 @@ import {defaultOptions} from './default-options';
 import {Duration} from './temporal';
 import {ExponentialRetry} from './exponential-retry';
 import {DebugMessage} from './debug';
+import {randomUUID} from 'crypto';
 import {logs as baseLogs, LoggingFunction} from './logs';
 
 /**
@@ -77,6 +78,8 @@ const DEFAULT_OPTIONS: MessageStreamOptions = {
   retryMaxBackoff: Duration.from({seconds: 60}),
 };
 
+const SERVER_KEEP_ALIVE_INTERVAL = 15000;
+
 interface StreamState {
   highWaterMark: number;
 }
@@ -139,6 +142,9 @@ export class ChannelError extends Error implements grpc.ServiceError {
 interface StreamTracked {
   stream?: PullStream;
   receivedStatus?: boolean;
+  lastPingTime?: number;
+  lastResponseTime?: number;
+  aliveTimer?: NodeJS.Timeout;
 }
 
 /**
@@ -200,7 +206,7 @@ export class MessageStream extends PassThrough {
    */
   setStreamAckDeadline(deadline: Duration) {
     const request: StreamingPullRequest = {
-      streamAckDeadlineSeconds: deadline.totalOf('second'),
+      streamAckDeadlineSeconds: deadline.seconds,
     };
 
     for (const tracker of this._streams) {
@@ -227,6 +233,9 @@ export class MessageStream extends PassThrough {
 
     for (let i = 0; i < this._streams.length; i++) {
       const tracker = this._streams[i];
+      if (tracker.aliveTimer) {
+        this._clearAliveTimer(tracker);
+      }
       if (tracker.stream) {
         this._removeStream(i, 'overall message stream destroyed', 'n/a');
       }
@@ -253,6 +262,7 @@ export class MessageStream extends PassThrough {
     const tracker = this._streams[index];
     tracker.stream = stream;
     tracker.receivedStatus = false;
+    tracker.lastResponseTime = Date.now();
 
     stream
       .on('error', err => this._onError(index, err))
@@ -263,11 +273,46 @@ export class MessageStream extends PassThrough {
   private _onData(index: number, data: PullResponse): void {
     // Mark this stream as alive again. (reset backoff)
     const tracker = this._streams[index];
+    tracker.lastResponseTime = Date.now();
     this._retrier.reset(tracker);
 
     this.emit('data', data);
   }
 
+  private _clearAliveTimer(tracker: StreamTracked): void {
+    if (tracker.aliveTimer) {
+      clearTimeout(tracker.aliveTimer);
+      tracker.aliveTimer = undefined;
+    }
+  }
+
+  private _checkAliveTimer(index: number): void {
+    const tracker = this._streams[index];
+    const lastPingTime = tracker.lastPingTime ?? -1;
+    const lastResponseTime = tracker.lastResponseTime ?? 0;
+    if (lastPingTime <= lastResponseTime) {
+      return;
+    }
+
+    this._removeStream(
+      index,
+      'no keepalive response from server within 15 seconds',
+      'will be retried',
+    );
+    this._retrier.retryLater(tracker, () =>
+      this._fillOne(index, undefined, 'retry'),
+    );
+  }
+
+  private _setAliveTimer(index: number): void {
+    const tracker = this._streams[index];
+    this._clearAliveTimer(tracker);
+
+    tracker.aliveTimer = setTimeout(() => {
+      this._checkAliveTimer(index);
+    }, SERVER_KEEP_ALIVE_INTERVAL);
+  }
+
   /**
    * Attempts to create and cache the desired number of StreamingPull requests.
    * gRPC does not supply a way to confirm that a stream is connected, so our
@@ -347,6 +392,8 @@ export class MessageStream extends PassThrough {
       maxOutstandingBytes: this._subscriber.useLegacyFlowControl
         ? 0
         : this._subscriber.maxBytes,
+      clientId: randomUUID().toString(),
+      protocolVersion: 1, // Set protocol version to enable server keepalives
     };
     const otherArgs = {
       headers: {
@@ -386,12 +433,14 @@ export class MessageStream extends PassThrough {
       'sending keepAlive to %i streams',
       this._streams.length,
     );
-    this._streams.forEach(tracker => {
+    this._streams.forEach((tracker, index) => {
       // It's possible that a status event fires off (signaling the rpc being
       // closed) but the stream hasn't drained yet. Writing to such a stream will
       // result in a `write after end` error.
       if (!tracker.receivedStatus && tracker.stream) {
         tracker.stream.write({});
+        tracker.lastPingTime = Date.now();
+        this._setAliveTimer(index);
       }
     });
   }
@@ -511,6 +560,9 @@ export class MessageStream extends PassThrough {
     whatNext?: string,
   ): void {
     const tracker = this._streams[index];
+    if (tracker.aliveTimer) {
+      this._clearAliveTimer(tracker);
+    }
     if (tracker.stream) {
       logs.subscriberStreams.info(
         'closing stream %i; why: %s; next: %s',

handwritten/pubsub/test/message-stream.ts

@@ -523,6 +523,24 @@ describe('MessageStream', () => {
     });
 
     describe('keeping streams alive', () => {
+      it('should set protocolVersion in the initial packet', async () => {
+        // The special handling for messageStream and the spy below are
+        // so that we can test the initial message.
+        messageStream.destroy();
+
+        const spy = sandbox.spy(FakeGrpcStream.prototype, 'write');
+        const ms = new MessageStream(subscriber);
+        await ms.start();
+
+        assert.strictEqual(spy.callCount, 5);
+        const {args} = spy.firstCall;
+        const request = args[0] as any;
+
+        assert.strictEqual(String(request.protocolVersion), '1');
+
+        ms.destroy();
+      });
+
       it('should keep the streams alive', () => {
         const frequency = 30000;
         const stubs = client.streams.map(stream => {
@@ -536,6 +554,48 @@ describe('MessageStream', () => {
           assert.deepStrictEqual(data, {});
         });
       });
+
+      it('should close stream if no data received for 15 seconds after keepalive', async () => {
+        messageStream.destroy();
+        client.streams.length = 0;
+
+        const ms = new MessageStream(subscriber);
+        await ms.start();
+
+        const cancelSpies = client.streams.map(s => sandbox.spy(s, 'cancel'));
+
+        // wait for keepalive ping (30s) + 15s timeout
+        sandbox.clock.tick(45000);
+
+        cancelSpies.forEach(spy => {
+          assert.strictEqual(spy.callCount, 1);
+        });
+
+        ms.destroy();
+      });
+
+      it('should not close stream if data received within 15 seconds of keepalive', async () => {
+        messageStream.destroy();
+
+        const ms = new MessageStream(subscriber);
+        await ms.start();
+
+        const cancelSpies = client.streams.map(s => sandbox.spy(s, 'cancel'));
+
+        sandbox.clock.tick(30000);
+
+        // Simulating data prevents timeout
+        client.streams.forEach(s => s.emit('data', {}));
+
+        // Wait for 15s timeout to pass
+        sandbox.clock.tick(15000);
+
+        cancelSpies.forEach(spy => {
+          assert.strictEqual(spy.callCount, 0);
+        });
+
+        ms.destroy();
+      });
     });
 
     it('should allow updating the ack deadline', async () => {