Skip to content

Commit 8b489d1

Browse files
stocaarostocaaro
authored
fix(api-graphql): Data messages should maintain the keep alive status (#14164)
1 parent 86d787b commit 8b489d1

File tree

3 files changed

+122
-37
lines changed

3 files changed

+122
-37
lines changed

‎packages/api-graphql/__tests__/AWSAppSyncRealTimeProvider.test.ts‎

Lines changed: 62 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,10 @@ describe('AWSAppSyncRealTimeProvider', () => {
147147
Object.defineProperty(constants, 'RECONNECT_DELAY', {
148148
value: 100,
149149
});
150+
// Reduce the keep alive heartbeat to 10ms
151+
Object.defineProperty(constants, 'DEFAULT_KEEP_ALIVE_HEARTBEAT_TIMEOUT', {
152+
value: 10,
153+
});
150154
});
151155

152156
afterEach(async () => {
@@ -765,7 +769,7 @@ describe('AWSAppSyncRealTimeProvider', () => {
765769
// Resolve the message delivery actions
766770
await replaceConstant(
767771
'DEFAULT_KEEP_ALIVE_ALERT_TIMEOUT',
768-
5,
772+
10,
769773
async () => {
770774
await fakeWebSocketInterface?.readyForUse;
771775
await fakeWebSocketInterface?.triggerOpen();
@@ -776,17 +780,17 @@ describe('AWSAppSyncRealTimeProvider', () => {
776780
await fakeWebSocketInterface?.startAckMessage();
777781

778782
await fakeWebSocketInterface?.keepAlive();
779-
},
780-
);
781783

782-
await fakeWebSocketInterface?.waitUntilConnectionStateIn([
783-
CS.Connected,
784-
]);
784+
await fakeWebSocketInterface?.waitUntilConnectionStateIn([
785+
CS.Connected,
786+
]);
785787

786-
// Wait until the socket is automatically disconnected
787-
await fakeWebSocketInterface?.waitUntilConnectionStateIn([
788-
CS.ConnectionDisrupted,
789-
]);
788+
// Wait until the socket is automatically disconnected
789+
await fakeWebSocketInterface?.waitUntilConnectionStateIn([
790+
CS.ConnectionDisrupted,
791+
]);
792+
},
793+
);
790794

791795
expect(fakeWebSocketInterface?.observedConnectionStates).toContain(
792796
CS.ConnectedPendingKeepAlive,
@@ -798,6 +802,54 @@ describe('AWSAppSyncRealTimeProvider', () => {
798802
);
799803
});
800804

805+
test('subscription observer ka is cleared if data is received', async () => {
806+
expect.assertions(1);
807+
808+
const observer = provider.subscribe({
809+
appSyncGraphqlEndpoint: 'ws://localhost:8080',
810+
});
811+
812+
observer.subscribe({ error: () => {} });
813+
// Resolve the message delivery actions
814+
await replaceConstant(
815+
'DEFAULT_KEEP_ALIVE_ALERT_TIMEOUT',
816+
5,
817+
async () => {
818+
await fakeWebSocketInterface?.readyForUse;
819+
await fakeWebSocketInterface?.triggerOpen();
820+
await fakeWebSocketInterface?.handShakeMessage({
821+
connectionTimeoutMs: 100,
822+
});
823+
824+
await fakeWebSocketInterface?.startAckMessage();
825+
826+
await fakeWebSocketInterface?.keepAlive();
827+
828+
await fakeWebSocketInterface?.waitUntilConnectionStateIn([
829+
CS.ConnectedPendingKeepAlive,
830+
]);
831+
},
832+
);
833+
834+
// Send message
835+
await fakeWebSocketInterface?.sendDataMessage({
836+
type: MESSAGE_TYPES.DATA,
837+
payload: { data: {} },
838+
});
839+
840+
await fakeWebSocketInterface?.waitUntilConnectionStateIn([
841+
CS.Connected,
842+
]);
843+
844+
expect(fakeWebSocketInterface?.observedConnectionStates).toEqual([
845+
CS.Disconnected,
846+
CS.Connecting,
847+
CS.Connected,
848+
CS.ConnectedPendingKeepAlive,
849+
CS.Connected,
850+
]);
851+
});
852+
801853
test('subscription connection disruption triggers automatic reconnection', async () => {
802854
expect.assertions(1);
803855

‎packages/api-graphql/src/Providers/AWSWebSocketProvider/index.ts‎

Lines changed: 55 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ import {
2323
CONNECTION_INIT_TIMEOUT,
2424
CONNECTION_STATE_CHANGE,
2525
DEFAULT_KEEP_ALIVE_ALERT_TIMEOUT,
26-
DEFAULT_KEEP_ALIVE_TIMEOUT,
26+
DEFAULT_KEEP_ALIVE_HEARTBEAT_TIMEOUT,
2727
MAX_DELAY_MS,
2828
MESSAGE_TYPES,
2929
NON_RETRYABLE_CODES,
@@ -83,9 +83,8 @@ export abstract class AWSWebSocketProvider {
8383

8484
protected awsRealTimeSocket?: WebSocket;
8585
private socketStatus: SOCKET_STATUS = SOCKET_STATUS.CLOSED;
86-
private keepAliveTimeoutId?: ReturnType<typeof setTimeout>;
87-
private keepAliveTimeout = DEFAULT_KEEP_ALIVE_TIMEOUT;
88-
private keepAliveAlertTimeoutId?: ReturnType<typeof setTimeout>;
86+
private keepAliveTimestamp: number = Date.now();
87+
private keepAliveHeartbeatIntervalId?: ReturnType<typeof setInterval>;
8988
private promiseArray: { res(): void; rej(reason?: any): void }[] = [];
9089
private connectionState: ConnectionState | undefined;
9190
private readonly connectionStateMonitor = new ConnectionStateMonitor();
@@ -119,6 +118,7 @@ export abstract class AWSWebSocketProvider {
119118
return new Promise<void>((resolve, reject) => {
120119
if (this.awsRealTimeSocket) {
121120
this.awsRealTimeSocket.onclose = (_: CloseEvent) => {
121+
this._closeSocket();
122122
this.subscriptionObserverMap = new Map();
123123
this.awsRealTimeSocket = undefined;
124124
resolve();
@@ -171,7 +171,7 @@ export abstract class AWSWebSocketProvider {
171171
this.logger.debug(
172172
`${CONTROL_MSG.REALTIME_SUBSCRIPTION_INIT_ERROR}: ${err}`,
173173
);
174-
this.connectionStateMonitor.record(CONNECTION_CHANGE.CLOSED);
174+
this._closeSocket();
175175
})
176176
.finally(() => {
177177
subscriptionStartInProgress = false;
@@ -435,7 +435,7 @@ export abstract class AWSWebSocketProvider {
435435
this.logger.debug({ err });
436436
const message = String(err.message ?? '');
437437
// Resolving to give the state observer time to propogate the update
438-
this.connectionStateMonitor.record(CONNECTION_CHANGE.CLOSED);
438+
this._closeSocket();
439439

440440
// Capture the error only when the network didn't cause disruption
441441
if (
@@ -544,20 +544,15 @@ export abstract class AWSWebSocketProvider {
544544
setTimeout(this._closeSocketIfRequired.bind(this), 1000);
545545
} else {
546546
this.logger.debug('closing WebSocket...');
547-
if (this.keepAliveTimeoutId) {
548-
clearTimeout(this.keepAliveTimeoutId);
549-
}
550-
if (this.keepAliveAlertTimeoutId) {
551-
clearTimeout(this.keepAliveAlertTimeoutId);
552-
}
547+
553548
const tempSocket = this.awsRealTimeSocket;
554549
// Cleaning callbacks to avoid race condition, socket still exists
555550
tempSocket.onclose = null;
556551
tempSocket.onerror = null;
557552
tempSocket.close(1000);
558553
this.awsRealTimeSocket = undefined;
559554
this.socketStatus = SOCKET_STATUS.CLOSED;
560-
this.connectionStateMonitor.record(CONNECTION_CHANGE.CLOSED);
555+
this._closeSocket();
561556
}
562557
}
563558

@@ -577,13 +572,40 @@ export abstract class AWSWebSocketProvider {
577572
errorType: string;
578573
};
579574

575+
private maintainKeepAlive() {
576+
this.keepAliveTimestamp = Date.now();
577+
}
578+
579+
private keepAliveHeartbeat(connectionTimeoutMs: number) {
580+
const currentTime = Date.now();
581+
582+
// Check for missed KA message
583+
if (
584+
currentTime - this.keepAliveTimestamp >
585+
DEFAULT_KEEP_ALIVE_ALERT_TIMEOUT
586+
) {
587+
this.connectionStateMonitor.record(CONNECTION_CHANGE.KEEP_ALIVE_MISSED);
588+
} else {
589+
this.connectionStateMonitor.record(CONNECTION_CHANGE.KEEP_ALIVE);
590+
}
591+
592+
// Recognize we are disconnected if we haven't seen messages in the keep alive timeout period
593+
if (currentTime - this.keepAliveTimestamp > connectionTimeoutMs) {
594+
this._errorDisconnect(CONTROL_MSG.TIMEOUT_DISCONNECT);
595+
}
596+
}
597+
580598
private _handleIncomingSubscriptionMessage(message: MessageEvent) {
581599
if (typeof message.data !== 'string') {
582600
return;
583601
}
584602

585603
const [isData, data] = this._handleSubscriptionData(message);
586-
if (isData) return;
604+
if (isData) {
605+
this.maintainKeepAlive();
606+
607+
return;
608+
}
587609

588610
const { type, id, payload } = data;
589611

@@ -632,16 +654,7 @@ export abstract class AWSWebSocketProvider {
632654
}
633655

634656
if (type === MESSAGE_TYPES.GQL_CONNECTION_KEEP_ALIVE) {
635-
if (this.keepAliveTimeoutId) clearTimeout(this.keepAliveTimeoutId);
636-
if (this.keepAliveAlertTimeoutId)
637-
clearTimeout(this.keepAliveAlertTimeoutId);
638-
this.keepAliveTimeoutId = setTimeout(() => {
639-
this._errorDisconnect(CONTROL_MSG.TIMEOUT_DISCONNECT);
640-
}, this.keepAliveTimeout);
641-
this.keepAliveAlertTimeoutId = setTimeout(() => {
642-
this.connectionStateMonitor.record(CONNECTION_CHANGE.KEEP_ALIVE_MISSED);
643-
}, DEFAULT_KEEP_ALIVE_ALERT_TIMEOUT);
644-
this.connectionStateMonitor.record(CONNECTION_CHANGE.KEEP_ALIVE);
657+
this.maintainKeepAlive();
645658

646659
return;
647660
}
@@ -686,13 +699,21 @@ export abstract class AWSWebSocketProvider {
686699
this.logger.debug(`Disconnect error: ${msg}`);
687700

688701
if (this.awsRealTimeSocket) {
689-
this.connectionStateMonitor.record(CONNECTION_CHANGE.CLOSED);
702+
this._closeSocket();
690703
this.awsRealTimeSocket.close();
691704
}
692705

693706
this.socketStatus = SOCKET_STATUS.CLOSED;
694707
}
695708

709+
private _closeSocket() {
710+
if (this.keepAliveHeartbeatIntervalId) {
711+
clearInterval(this.keepAliveHeartbeatIntervalId);
712+
this.keepAliveHeartbeatIntervalId = undefined;
713+
}
714+
this.connectionStateMonitor.record(CONNECTION_CHANGE.CLOSED);
715+
}
716+
696717
private _timeoutStartSubscriptionAck(subscriptionId: string) {
697718
const subscriptionObserver =
698719
this.subscriptionObserverMap.get(subscriptionId);
@@ -708,7 +729,7 @@ export abstract class AWSWebSocketProvider {
708729
subscriptionState: SUBSCRIPTION_STATUS.FAILED,
709730
});
710731

711-
this.connectionStateMonitor.record(CONNECTION_CHANGE.CLOSED);
732+
this._closeSocket();
712733
this.logger.debug(
713734
'timeoutStartSubscription',
714735
JSON.stringify({ query, variables }),
@@ -820,6 +841,7 @@ export abstract class AWSWebSocketProvider {
820841
this.logger.debug(`WebSocket connection error`);
821842
};
822843
newSocket.onclose = () => {
844+
this._closeSocket();
823845
reject(new Error('Connection handshake error'));
824846
};
825847
newSocket.onopen = () => {
@@ -849,6 +871,7 @@ export abstract class AWSWebSocketProvider {
849871

850872
this.awsRealTimeSocket.onclose = event => {
851873
this.logger.debug(`WebSocket closed ${event.reason}`);
874+
this._closeSocket();
852875
reject(new Error(JSON.stringify(event)));
853876
};
854877

@@ -912,7 +935,11 @@ export abstract class AWSWebSocketProvider {
912935
return;
913936
}
914937

915-
this.keepAliveTimeout = connectionTimeoutMs;
938+
// Set up a keep alive heartbeat for this connection
939+
this.keepAliveHeartbeatIntervalId = setInterval(() => {
940+
this.keepAliveHeartbeat(connectionTimeoutMs);
941+
}, DEFAULT_KEEP_ALIVE_HEARTBEAT_TIMEOUT);
942+
916943
this.awsRealTimeSocket.onmessage =
917944
this._handleIncomingSubscriptionMessage.bind(this);
918945

@@ -923,6 +950,7 @@ export abstract class AWSWebSocketProvider {
923950

924951
this.awsRealTimeSocket.onclose = event => {
925952
this.logger.debug(`WebSocket closed ${event.reason}`);
953+
this._closeSocket();
926954
this._errorDisconnect(CONTROL_MSG.CONNECTION_CLOSED);
927955
};
928956
}

‎packages/api-graphql/src/Providers/constants.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,11 @@ export const START_ACK_TIMEOUT = 15000;
128128
*/
129129
export const DEFAULT_KEEP_ALIVE_TIMEOUT = 5 * 60 * 1000;
130130

131+
/**
132+
* Default Time in milleseconds between monitoring checks of keep alive status
133+
*/
134+
export const DEFAULT_KEEP_ALIVE_HEARTBEAT_TIMEOUT = 5 * 1000;
135+
131136
/**
132137
* Default Time in milleseconds to alert for missed GQL_CONNECTION_KEEP_ALIVE message
133138
*/

0 commit comments

Comments
 (0)