Repository: michalkvasnicak/aws-lambda-graphql
Pull request #68 (2020-02-06 12:25:27)

Fix timing issues with rapid protocol calls and missing dynamodb database fetch calls
When running locally on a very powerful machine like the one used in continuous integration the calls coming from the graphql client of:
1. $connect
2. $default -> gql_connect
3. $default -> gql_start

Can be quicker in the handler than the createConnection entry dynamodb call is taking to complete. This results into unexpected errors in the connection for local test runs and potentially also on production.

I fixed this by adding a retry mechanism to the connection hydration that is used for the initial call so it retrys once when the connection is not found.


Patch:
diff --git a/packages/aws-lambda-graphql/src/DynamoDBConnectionManager.ts b/packages/aws-lambda-graphql/src/DynamoDBConnectionManager.ts
index ef5472d..019f81c 100644
--- a/packages/aws-lambda-graphql/src/DynamoDBConnectionManager.ts
+++ b/packages/aws-lambda-graphql/src/DynamoDBConnectionManager.ts
@@ -6,6 +6,7 @@ import {
   IConnectionManager,
   ISubscriptionManager,
   IConnectionData,
+  HydrateConnectionOptions,
 } from './types';
 
 export class ConnectionNotFoundError extends ExtendableError {}
@@ -54,22 +55,37 @@ export class DynamoDBConnectionManager implements IConnectionManager {
     this.subscriptions = subscriptions;
   }
 
-  hydrateConnection = async (connectionId: string): Promise<IConnection> => {
+  hydrateConnection = async (
+    connectionId: string,
+    options: HydrateConnectionOptions,
+  ): Promise<IConnection> => {
+    const { retryCount = 0, timeout = 50 } = options || {};
     // if connection is not found, throw so we can terminate connection
-    const result = await this.db
-      .get({
-        TableName: this.connectionsTable,
-        Key: {
-          id: connectionId,
-        },
-      })
-      .promise();
+    let connection;
+
+    for (let i = 0; i <= retryCount; i++) {
+      const result = await this.db
+        .get({
+          TableName: this.connectionsTable,
+          Key: {
+            id: connectionId,
+          },
+        })
+        .promise();
+      if (result.Item) {
+        // Jump out of loop
+        connection = result.Item as IConnection;
+        break;
+      }
+      // wait for another round
+      await new Promise(r => setTimeout(r, timeout));
+    }
 
-    if (result.Item == null) {
+    if (!connection) {
       throw new ConnectionNotFoundError(`Connection ${connectionId} not found`);
     }
 
-    return result.Item as IConnection;
+    return connection as IConnection;
   };
 
   setConnectionData = async (
diff --git a/packages/aws-lambda-graphql/src/Server.ts b/packages/aws-lambda-graphql/src/Server.ts
index be2e89a..4e030cc 100644
--- a/packages/aws-lambda-graphql/src/Server.ts
+++ b/packages/aws-lambda-graphql/src/Server.ts
@@ -275,14 +275,19 @@ export class Server<
                 timeout: waitTimeout = 50,
               } = {},
             } = this.subscriptionOptions || {};
-            // hydrate connectiont
-            let connection = await this.connectionManager.hydrateConnection(
-              connectionId,
-            );
 
             // parse operation from body
             const operation = parseOperationFromEvent(event);
 
+            // hydrate connection
+            let connection = await this.connectionManager.hydrateConnection(
+              connectionId,
+              {
+                retryCount: 1,
+                timeout: waitTimeout,
+              },
+            );
+
             if (isGQLConnectionInit(operation)) {
               let newConnectionContext = operation.payload;
 
@@ -350,7 +355,7 @@ export class Server<
                 return freshConnection;
               }
 
-              for (let i = 0; i < waitRetryCount; i++) {
+              for (let i = 0; i <= waitRetryCount; i++) {
                 freshConnection = await this.connectionManager.hydrateConnection(
                   connectionId,
                 );
diff --git a/packages/aws-lambda-graphql/src/types/connections.ts b/packages/aws-lambda-graphql/src/types/connections.ts
index 6ab2e73..8c83d35 100644
--- a/packages/aws-lambda-graphql/src/types/connections.ts
+++ b/packages/aws-lambda-graphql/src/types/connections.ts
@@ -27,13 +27,31 @@ export interface IConnectionData {
   readonly isInitialized: boolean;
 }
 
+export interface HydrateConnectionOptions {
+  /**
+   * How many times should we retry the connection query in case it fails for timing issues
+   *
+   * Default is 0
+   */
+  retryCount?: number;
+  /**
+   * How long should we wait until we try determine connection state again?
+   *
+   * Default is 50ms
+   */
+  timeout?: number;
+}
+
 export interface IConnectEvent {
   connectionId: string;
   endpoint: string;
 }
 
 export interface IConnectionManager {
-  hydrateConnection(connectionId: string): Promise<IConnection>;
+  hydrateConnection(
+    connectionId: string,
+    options?: HydrateConnectionOptions,
+  ): Promise<IConnection>;
   setConnectionData(data: Object, connection: IConnection): Promise<void>;
   registerConnection(event: IConnectEvent): Promise<IConnection>;
   sendToConnection(


Test patch:
diff --git a/packages/aws-lambda-graphql/src/__tests__/DynamoDBConnectionManager.test.ts b/packages/aws-lambda-graphql/src/__tests__/DynamoDBConnectionManager.test.ts
index 1589435..93660e1 100644
--- a/packages/aws-lambda-graphql/src/__tests__/DynamoDBConnectionManager.test.ts
+++ b/packages/aws-lambda-graphql/src/__tests__/DynamoDBConnectionManager.test.ts
@@ -99,6 +99,25 @@ describe('DynamoDBConnectionManager', () => {
 
       expect(getMock as jest.Mock).toHaveBeenCalledTimes(1);
     });
+
+    it('hydrates connection with retry', async () => {
+      (getPromiseMock as jest.Mock)
+        .mockResolvedValueOnce({})
+        .mockResolvedValueOnce({
+          Item: { id: 'id', data: { endpoint: '' } },
+        });
+
+      await expect(
+        manager.hydrateConnection('id', { retryCount: 1 }),
+      ).resolves.toEqual({
+        id: 'id',
+        data: {
+          endpoint: '',
+        },
+      });
+
+      expect(getMock as jest.Mock).toHaveBeenCalledTimes(2);
+    });
   });
 
   describe('setConnectionData', () => {
diff --git a/packages/aws-lambda-graphql/src/__tests__/Server.test.ts b/packages/aws-lambda-graphql/src/__tests__/Server.test.ts
index 0095da2..7b87d54 100644
--- a/packages/aws-lambda-graphql/src/__tests__/Server.test.ts
+++ b/packages/aws-lambda-graphql/src/__tests__/Server.test.ts
@@ -2,9 +2,14 @@ import { APIGatewayProxyEvent } from 'aws-lambda';
 import { ulid } from 'ulid';
 import { createSchema } from '../fixtures/schema';
 import { Server } from '../Server';
-import { IConnectionManager, ISubscriptionManager } from '../types';
+import {
+  IConnection,
+  IConnectionManager,
+  ISubscriptionManager,
+} from '../types';
 import { formatMessage } from '../formatMessage';
 import { SERVER_EVENT_TYPES, CLIENT_EVENT_TYPES } from '../protocol';
+import { ConnectionNotFoundError } from '../DynamoDBConnectionManager';
 
 describe('Server', () => {
   describe('createHttpHandler()', () => {
@@ -160,6 +165,67 @@ describe('Server', () => {
           connectionId: '1',
         });
       });
+
+      it('recover from missing connection on quick succession of $connect and GQL_CONNECTION_INIT message on offline server', async () => {
+        // Mock the connection database flow
+        let connection;
+        (connectionManager.registerConnection as jest.Mock).mockImplementationOnce(
+          c => {
+            connection = c;
+          },
+        );
+        (connectionManager.hydrateConnection as jest.Mock).mockImplementationOnce(
+          async connectionId => {
+            for (let i = 0; i <= 1; i++) {
+              if (connection) {
+                return connection as IConnection;
+              }
+              // wait for another round
+              await new Promise(r => setTimeout(r, 50));
+            }
+            throw new ConnectionNotFoundError(
+              `Connection ${connectionId} not found`,
+            );
+          },
+        );
+        // Call connection request before connect to over emphasize the case
+        const gqlConnectionRequest = handler(
+          {
+            body: formatMessage({
+              payload: {
+                contextAttribute: 'contextAttributeValue',
+              },
+              type: CLIENT_EVENT_TYPES.GQL_CONNECTION_INIT,
+            }),
+            requestContext: {
+              connectionId: '1',
+              domainName: 'domain',
+              routeKey: '$default',
+              stage: 'stage',
+            } as any,
+          } as any,
+          {} as any,
+        );
+        await handler(
+          {
+            requestContext: {
+              connectionId: '1',
+              domainName: 'domain',
+              routeKey: '$connect',
+              stage: 'stage',
+            } as any,
+          } as any,
+          {} as any,
+        );
+        await expect(gqlConnectionRequest).resolves.toEqual(
+          expect.objectContaining({
+            body: formatMessage({
+              type: SERVER_EVENT_TYPES.GQL_CONNECTION_ACK,
+            }),
+            statusCode: 200,
+          }),
+        );
+      });
     });
 
     describe('disconnect phase', () => {
@@ -273,7 +339,10 @@ describe('Server', () => {
         });
 
         expect(connectionManager.hydrateConnection).toHaveBeenCalledTimes(1);
-        expect(connectionManager.hydrateConnection).toHaveBeenCalledWith('1');
+        expect(connectionManager.hydrateConnection).toHaveBeenCalledWith('1', {
+          retryCount: 1,
+          timeout: 50,
+        });
       });
 
       it('returns http 200 with GQL_CONNECTION_ACK on connection_init operation and sets context', async () => {
@@ -309,7 +378,10 @@ describe('Server', () => {
         );
 
         expect(connectionManager.hydrateConnection).toHaveBeenCalledTimes(1);
-        expect(connectionManager.hydrateConnection).toHaveBeenCalledWith('1');
+        expect(connectionManager.hydrateConnection).toHaveBeenCalledWith('1', {
+          retryCount: 1,
+          timeout: 50,
+        });
         expect(connectionManager.sendToConnection).toHaveBeenCalledTimes(1);
         expect(connectionManager.setConnectionData).toHaveBeenCalledTimes(1);
         expect(connectionManager.setConnectionData).toHaveBeenCalledWith(
@@ -527,7 +599,7 @@ describe('Server', () => {
           }),
         );
 
-        expect(connectionManager.hydrateConnection).toHaveBeenCalledTimes(3);
+        expect(connectionManager.hydrateConnection).toHaveBeenCalledTimes(4);
         expect(connectionManager.hydrateConnection).toHaveBeenCalledWith('1');
         expect(connectionManager.sendToConnection).toHaveBeenCalledTimes(1);
       });
@@ -585,7 +657,10 @@ describe('Server', () => {
         );
 
         expect(connectionManager.hydrateConnection).toHaveBeenCalledTimes(1);
-        expect(connectionManager.hydrateConnection).toHaveBeenCalledWith('1');
+        expect(connectionManager.hydrateConnection).toHaveBeenCalledWith('1', {
+          retryCount: 1,
+          timeout: 50,
+        });
         expect(connectionManager.sendToConnection).toHaveBeenCalledTimes(1);
         expect(onOperationComplete).toHaveBeenCalledTimes(1);
         expect(onOperationComplete).toHaveBeenCalledWith(
@@ -641,7 +716,10 @@ describe('Server', () => {
         );
 
         expect(connectionManager.hydrateConnection).toHaveBeenCalledTimes(1);
-        expect(connectionManager.hydrateConnection).toHaveBeenCalledWith('1');
+        expect(connectionManager.hydrateConnection).toHaveBeenCalledWith('1', {
+          retryCount: 1,
+          timeout: 50,
+        });
         expect(connectionManager.sendToConnection).toHaveBeenCalledTimes(1);
       });
 
@@ -753,7 +831,10 @@ describe('Server', () => {
         );
 
         expect(connectionManager.hydrateConnection).toHaveBeenCalledTimes(1);
-        expect(connectionManager.hydrateConnection).toHaveBeenCalledWith('1');
+        expect(connectionManager.hydrateConnection).toHaveBeenCalledWith('1', {
+          retryCount: 1,
+          timeout: 50,
+        });
         expect(subscriptionManager.unsubscribeOperation).toHaveBeenCalledTimes(
           1,
         );
@@ -810,7 +891,10 @@ describe('Server', () => {
         );
 
         expect(connectionManager.hydrateConnection).toHaveBeenCalledTimes(1);
-        expect(connectionManager.hydrateConnection).toHaveBeenCalledWith('1');
+        expect(connectionManager.hydrateConnection).toHaveBeenCalledWith('1', {
+          retryCount: 1,
+          timeout: 50,
+        });
         expect(connectionManager.sendToConnection).toHaveBeenCalledTimes(1);
       });
     });