feat: Add keepAliveTimeoutMs config for AppSync WebSocket link (#724)

manueliglesias · web-flow · commit 74b8351209ec · 2022-06-24T10:58:38.000-06:00
* feat: Add keepAliveTimeoutMs config for AppSync WebSocket link

* fix typo

* validate keepAliveTimeoutMs
diff --git a/packages/aws-appsync-subscription-link/__tests__/link/realtime-subscription-handshake-link-test.ts b/packages/aws-appsync-subscription-link/__tests__/link/realtime-subscription-handshake-link-test.ts
@@ -2,11 +2,14 @@ import { AUTH_TYPE } from "aws-appsync-auth-link";
 import { execute } from "@apollo/client/core";
 import gql from 'graphql-tag';
 import { AppSyncRealTimeSubscriptionHandshakeLink } from '../../src/realtime-subscription-handshake-link';
+import { MESSAGE_TYPES } from "../../src/types";
+import { v4 as uuid } from "uuid";
+jest.mock('uuid', () => ({ v4: jest.fn() }));
 
 const query = gql`subscription { someSubscription { aField } }`
 
 class myWebSocket implements WebSocket {
-    binaryType: BinaryType; 
+    binaryType: BinaryType;
     bufferedAmount: number;
     extensions: string;
     onclose: (this: WebSocket, ev: CloseEvent) => any;
@@ -359,6 +362,164 @@ describe("RealTime subscription link", () => {
             }
 
         });
-    })
+    });
+
+    test("Can use a custom keepAliveTimeoutMs", (done) => {
+        const id = "abcd-efgh-ijkl-mnop";
+        uuid.mockImplementationOnce(() => id);
+
+        expect.assertions(5);
+        jest.spyOn(Date.prototype, 'toISOString').mockImplementation(jest.fn(() => {
+            return "2019-11-13T18:47:04.733Z";
+        }));
+        AppSyncRealTimeSubscriptionHandshakeLink.createWebSocket = jest.fn((url, protocol) => {
+            expect(url).toBe('wss://apikeytest.testcustomdomain.com/graphql/realtime?header=eyJob3N0IjoiYXBpa2V5dGVzdC50ZXN0Y3VzdG9tZG9tYWluLmNvbSIsIngtYW16LWRhdGUiOiIyMDE5MTExM1QxODQ3MDRaIiwieC1hcGkta2V5IjoieHh4eHgifQ==&payload=e30=');
+            expect(protocol).toBe('graphql-ws');
+            const socket = new myWebSocket();
+
+            setTimeout(() => {
+                socket.close = () => {};
+                socket.onopen.call(socket, (undefined as unknown as Event));
+                socket.send = (msg: string) => {
+                    const { type } = JSON.parse(msg);
+
+                    switch (type) {
+                        case MESSAGE_TYPES.GQL_CONNECTION_INIT:
+                            socket.onmessage.call(socket, {
+                                data: JSON.stringify({
+                                    type: MESSAGE_TYPES.GQL_CONNECTION_ACK,
+                                    payload: {
+                                        connectionTimeoutMs: 99999,
+                                    },
+                                })
+                            } as MessageEvent);
+                            setTimeout(() => {
+                                socket.onmessage.call(socket, {
+                                    data: JSON.stringify({
+                                        id,
+                                        type: MESSAGE_TYPES.GQL_DATA,
+                                        payload: {
+                                            data: { something: 123 },
+                                        },
+                                    })
+                                } as MessageEvent);
+
+                            }, 100);
+                            break;
+                    }
+                };
+            }, 100);
+
+            return socket;
+        });
+        const link = new AppSyncRealTimeSubscriptionHandshakeLink({
+            auth: {
+                type: AUTH_TYPE.API_KEY,
+                apiKey: 'xxxxx'
+            },
+            region: 'us-west-2',
+            url: 'https://apikeytest.testcustomdomain.com/graphql',
+            keepAliveTimeoutMs: 123456,
+        });
+
+        expect(link).toBeInstanceOf(AppSyncRealTimeSubscriptionHandshakeLink);
+        expect((link as any).keepAliveTimeout).toBe(123456);
+
+        const sub = execute(link, { query }).subscribe({
+            error: (err) => {
+                console.log(JSON.stringify(err));
+                fail();
+            },
+            next: (data) => {
+                expect((link as any).keepAliveTimeout).toBe(123456);
+                done();
+                sub.unsubscribe();
+            },
+            complete: () => {
+                console.log('done with this');
+                fail();
+            }
+
+        });
+    });
+
+    test("Uses service-provided timeout when no custom keepAliveTimeoutMs is configured", (done) => {
+        const id = "abcd-efgh-ijkl-mnop";
+        uuid.mockImplementationOnce(() => id);
+
+        expect.assertions(5);
+        jest.spyOn(Date.prototype, 'toISOString').mockImplementation(jest.fn(() => {
+            return "2019-11-13T18:47:04.733Z";
+        }));
+        AppSyncRealTimeSubscriptionHandshakeLink.createWebSocket = jest.fn((url, protocol) => {
+            expect(url).toBe('wss://apikeytest.testcustomdomain.com/graphql/realtime?header=eyJob3N0IjoiYXBpa2V5dGVzdC50ZXN0Y3VzdG9tZG9tYWluLmNvbSIsIngtYW16LWRhdGUiOiIyMDE5MTExM1QxODQ3MDRaIiwieC1hcGkta2V5IjoieHh4eHgifQ==&payload=e30=');
+            expect(protocol).toBe('graphql-ws');
+            const socket = new myWebSocket();
+
+            setTimeout(() => {
+                socket.close = () => {};
+                socket.onopen.call(socket, (undefined as unknown as Event));
+                socket.send = (msg: string) => {
+                    const { type } = JSON.parse(msg);
+
+                    switch (type) {
+                        case MESSAGE_TYPES.GQL_CONNECTION_INIT:
+                            socket.onmessage.call(socket, {
+                                data: JSON.stringify({
+                                    type: MESSAGE_TYPES.GQL_CONNECTION_ACK,
+                                    payload: {
+                                        connectionTimeoutMs: 99999,
+                                    },
+                                })
+                            } as MessageEvent);
+                            setTimeout(() => {
+                                socket.onmessage.call(socket, {
+                                    data: JSON.stringify({
+                                        id,
+                                        type: MESSAGE_TYPES.GQL_DATA,
+                                        payload: {
+                                            data: { something: 123 },
+                                        },
+                                    })
+                                } as MessageEvent);
+
+                            }, 100);
+                            break;
+                    }
+                };
+            }, 100);
+
+            return socket;
+        });
+        const link = new AppSyncRealTimeSubscriptionHandshakeLink({
+            auth: {
+                type: AUTH_TYPE.API_KEY,
+                apiKey: 'xxxxx'
+            },
+            region: 'us-west-2',
+            url: 'https://apikeytest.testcustomdomain.com/graphql',
+        });
+
+        expect(link).toBeInstanceOf(AppSyncRealTimeSubscriptionHandshakeLink);
+        expect((link as any).keepAliveTimeout).toBeUndefined();
+
+        const sub = execute(link, { query }).subscribe({
+            error: (err) => {
+                console.log(JSON.stringify(err));
+                fail();
+            },
+            next: (data) => {
+                expect((link as any).keepAliveTimeout).toBe(99999);
+                done();
+                sub.unsubscribe();
+            },
+            complete: () => {
+                console.log('done with this');
+                fail();
+            }
+
+        });
+    });
+
 
 });
diff --git a/packages/aws-appsync-subscription-link/src/index.ts b/packages/aws-appsync-subscription-link/src/index.ts
@@ -9,20 +9,20 @@ import { NonTerminatingLink } from "./non-terminating-link";
 import type { OperationDefinitionNode } from "graphql";
 
 import {
-  AppSyncRealTimeSubscriptionHandshakeLink
+  AppSyncRealTimeSubscriptionHandshakeLink,
 } from "./realtime-subscription-handshake-link";
-import { UrlInfo } from "./types";
+import { AppSyncRealTimeSubscriptionConfig } from "./types";
 
 function createSubscriptionHandshakeLink(
-  args: UrlInfo,
+  args: AppSyncRealTimeSubscriptionConfig,
   resultsFetcherLink?: ApolloLink
 ): ApolloLink;
 function createSubscriptionHandshakeLink(
   url: string,
   resultsFetcherLink?: ApolloLink
 ): ApolloLink;
 function createSubscriptionHandshakeLink(
-  infoOrUrl: UrlInfo | string,
+  infoOrUrl: AppSyncRealTimeSubscriptionConfig | string,
   theResultsFetcherLink?: ApolloLink
 ) {
   let resultsFetcherLink: ApolloLink, subscriptionLinks: ApolloLink;
@@ -45,7 +45,7 @@ function createSubscriptionHandshakeLink(
 
               observer.next({ [CONTROL_EVENTS_KEY]: controlEvents });
 
-              return () => {};
+              return () => { };
             })
         )
       }),
diff --git a/packages/aws-appsync-subscription-link/src/realtime-subscription-handshake-link.ts b/packages/aws-appsync-subscription-link/src/realtime-subscription-handshake-link.ts
@@ -17,7 +17,7 @@ import { GraphQLError, print } from "graphql";
 import * as url from "url";
 import { v4 as uuid } from "uuid";
 import {
-  UrlInfo,
+  AppSyncRealTimeSubscriptionConfig,
   SOCKET_STATUS,
   ObserverQuery,
   SUBSCRIPTION_STATUS,
@@ -50,6 +50,11 @@ const CONNECTION_INIT_TIMEOUT = 15000;
  */
 const START_ACK_TIMEOUT = 15000;
 
+/**
+ * Frequency in milliseconds in which the server sends GQL_CONNECTION_KEEP_ALIVE messages
+ */
+const SERVER_KEEP_ALIVE_TIMEOUT = 1 * 60 * 1000;
+
 /**
  * Default Time in milliseconds to wait for GQL_CONNECTION_KEEP_ALIVE message
  */
@@ -66,21 +71,28 @@ export class AppSyncRealTimeSubscriptionHandshakeLink extends ApolloLink {
   private awsRealTimeSocket: WebSocket;
   private socketStatus: SOCKET_STATUS = SOCKET_STATUS.CLOSED;
   private keepAliveTimeoutId;
-  private keepAliveTimeout = DEFAULT_KEEP_ALIVE_TIMEOUT;
+  private keepAliveTimeout?: number = undefined;
   private subscriptionObserverMap: Map<string, ObserverQuery> = new Map();
   private promiseArray: Array<{ res: Function; rej: Function }> = [];
 
-  constructor({ url: theUrl, region: theRegion, auth: theAuth }: UrlInfo) {
+  constructor({ url: theUrl, region: theRegion, auth: theAuth, keepAliveTimeoutMs }: AppSyncRealTimeSubscriptionConfig) {
     super();
     this.url = theUrl;
     this.region = theRegion;
     this.auth = theAuth;
+    this.keepAliveTimeout = keepAliveTimeoutMs;
+
+    if (this.keepAliveTimeout < SERVER_KEEP_ALIVE_TIMEOUT) {
+      let configName: keyof AppSyncRealTimeSubscriptionConfig = 'keepAliveTimeoutMs';
+
+      throw new Error(`${configName} must be greater than or equal to ${SERVER_KEEP_ALIVE_TIMEOUT} (${this.keepAliveTimeout} used).`);
+    }
   }
 
   // Check if url matches standard domain pattern
-	private isCustomDomain(url: string): boolean {
-		return url.match(standardDomainPattern) === null;
-	}
+  private isCustomDomain(url: string): boolean {
+    return url.match(standardDomainPattern) === null;
+  }
 
   request(operation: Operation) {
     const { query, variables } = operation;
@@ -384,7 +396,7 @@ export class AppSyncRealTimeSubscriptionHandshakeLink extends ApolloLink {
               region,
               credentials,
               token,
-              graphql_headers: () => {}
+              graphql_headers: () => { }
             })
           );
           const headerQs = Buffer.from(headerString).toString("base64");
@@ -394,16 +406,16 @@ export class AppSyncRealTimeSubscriptionHandshakeLink extends ApolloLink {
           let discoverableEndpoint = appSyncGraphqlEndpoint;
 
           if (this.isCustomDomain(discoverableEndpoint)) {
-					    discoverableEndpoint = discoverableEndpoint.concat(
-							    customDomainPath
-              );
-					} else {
-						  discoverableEndpoint = discoverableEndpoint.replace('appsync-api', 'appsync-realtime-api').replace('gogi-beta', 'grt-beta');
-					}
+            discoverableEndpoint = discoverableEndpoint.concat(
+              customDomainPath
+            );
+          } else {
+            discoverableEndpoint = discoverableEndpoint.replace('appsync-api', 'appsync-realtime-api').replace('gogi-beta', 'grt-beta');
+          }
 
           discoverableEndpoint = discoverableEndpoint
-              .replace("https://", "wss://")
-              .replace('http://', 'ws://')
+            .replace("https://", "wss://")
+            .replace('http://', 'ws://')
 
           const awsRealTimeUrl = `${discoverableEndpoint}?header=${headerQs}&payload=${payloadQs}`;
 
@@ -607,7 +619,7 @@ export class AppSyncRealTimeSubscriptionHandshakeLink extends ApolloLink {
             } = data;
             if (type === MESSAGE_TYPES.GQL_CONNECTION_ACK) {
               ackOk = true;
-              this.keepAliveTimeout = connectionTimeoutMs;
+              this.keepAliveTimeout = this.keepAliveTimeout ?? connectionTimeoutMs;
               this.awsRealTimeSocket.onmessage = this._handleIncomingSubscriptionMessage.bind(
                 this
               );
diff --git a/packages/aws-appsync-subscription-link/src/types/index.ts b/packages/aws-appsync-subscription-link/src/types/index.ts
@@ -83,6 +83,10 @@ export type UrlInfo = {
   region: string;
 };
 
+export type AppSyncRealTimeSubscriptionConfig = UrlInfo & {
+  keepAliveTimeoutMs?: number;
+};
+
 export type ObserverQuery = {
   observer: ZenObservable.SubscriptionObserver<any>;
   query: string;
