Allows publish after unsubscribe

Author: zxl629

packages/api-graphql/src/Providers/AWSAppSyncEventsProvider/index.ts

@@ -53,6 +53,7 @@ export class AWSAppSyncEventProvider extends AWSWebSocketProvider {
 			wsProtocolName: WS_PROTOCOL_NAME,
 			connectUri: CONNECT_URI,
 		});
+		this.allowNoSubscriptions = true;
 	}
 
 	getProviderName() {
@@ -77,6 +78,10 @@ export class AWSAppSyncEventProvider extends AWSWebSocketProvider {
 		return super.publish(options, customUserAgentDetails);
 	}
 
+	public closeIfNoActiveChannel() {
+		setTimeout(this._closeSocketIfRequired.bind(this), 1000);
+	}
+
 	protected async _prepareSubscriptionPayload({
 		options,
 		subscriptionId,

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

@@ -80,6 +80,7 @@ interface AWSWebSocketProviderArgs {
 export abstract class AWSWebSocketProvider {
 	protected logger: ConsoleLogger;
 	protected subscriptionObserverMap = new Map<string, ObserverQuery>();
+	protected allowNoSubscriptions = false;
 
 	protected awsRealTimeSocket?: WebSocket;
 	private socketStatus: SOCKET_STATUS = SOCKET_STATUS.CLOSED;
@@ -543,10 +544,12 @@ export abstract class AWSWebSocketProvider {
 		this.subscriptionObserverMap.delete(subscriptionId);
 
 		// Verifying 1000ms after removing subscription in case there are new subscription unmount/mount
-		setTimeout(this._closeSocketIfRequired.bind(this), 1000);
+		if (!this.allowNoSubscriptions) {
+			setTimeout(this._closeSocketIfRequired.bind(this), 1000);
+		}
 	}
 
-	private _closeSocketIfRequired() {
+	protected _closeSocketIfRequired() {
 		if (this.subscriptionObserverMap.size > 0) {
 			// Active subscriptions on the WebSocket
 			return;

packages/api-graphql/src/internals/events/index.ts

@@ -3,7 +3,7 @@
 
 import { Subscription } from 'rxjs';
 import { Amplify } from '@aws-amplify/core';
-import { DocumentType } from '@aws-amplify/core/internals/utils';
+import { DocumentType, amplifyUuid } from '@aws-amplify/core/internals/utils';
 
 import { AppSyncEventProvider as eventProvider } from '../../Providers/AWSAppSyncEventsProvider';
 
@@ -17,6 +17,9 @@ import type {
 	SubscriptionObserver,
 } from './types';
 
+// Keeps a list of open channels in the websocket
+const openChannels = new Set<string>();
+
 /**
  * @experimental API may change in future versions
  *
@@ -50,12 +53,18 @@ async function connect(
 
 	await eventProvider.connect(providerOptions);
 
+	const channelId = amplifyUuid();
+	openChannels.add(channelId);
+
 	let _subscription: Subscription;
 
 	const sub = (
 		observer: SubscriptionObserver<any>,
 		subOptions?: EventsOptions,
 	): Subscription => {
+		if (!openChannels.has(channelId)) {
+			throw new Error('Channel is closed');
+		}
 		const subscribeOptions = { ...providerOptions, query: channel };
 		subscribeOptions.authenticationType = normalizeAuth(
 			subOptions?.authMode,
@@ -73,6 +82,9 @@ async function connect(
 		event: DocumentType,
 		pubOptions?: EventsOptions,
 	): Promise<any> => {
+		if (!openChannels.has(channelId)) {
+			throw new Error('Channel is closed');
+		}
 		const publishOptions = {
 			...providerOptions,
 			query: channel,
@@ -86,8 +98,12 @@ async function connect(
 		return eventProvider.publish(publishOptions);
 	};
 
-	const close = () => {
+	const close = async () => {
 		_subscription && _subscription.unsubscribe();
+		openChannels.delete(channelId);
+		if (openChannels.size === 0) {
+			eventProvider.closeIfNoActiveChannel();
+		}
 	};
 
 	return {