Merge pull request #686 from drift-labs/sdk-unsubscribe-fix

fix unsubscribe race condition
drift-labs · Nov 8, 2023 · 6755a9e · 6755a9e
2 parents 54009dc + f50419f
commit 6755a9e
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 12 deletions.
diff --git a/sdk/src/accounts/webSocketAccountSubscriber.ts b/sdk/src/accounts/webSocketAccountSubscriber.ts
@@ -14,6 +14,8 @@ export class WebSocketAccountSubscriber<T> implements AccountSubscriber<T> {
 	onChange: (data: T) => void;
 	listenerId?: number;
 	resubTimeoutMs?: number;
+	isUnsubscribing = false;
+
 	timeoutId?: NodeJS.Timeout;
 
 	receivingData: boolean;
@@ -34,7 +36,7 @@ export class WebSocketAccountSubscriber<T> implements AccountSubscriber<T> {
 	}
 
 	async subscribe(onChange: (data: T) => void): Promise<void> {
-		if (this.listenerId) {
+		if (this.listenerId || this.isUnsubscribing) {
 			return;
 		}
 
@@ -80,6 +82,11 @@ export class WebSocketAccountSubscriber<T> implements AccountSubscriber<T> {
 			throw new Error('onChange callback function must be set');
 		}
 		this.timeoutId = setTimeout(async () => {
+			if (this.isUnsubscribing) {
+				// If we are in the process of unsubscribing, do not attempt to resubscribe
+				return;
+			}
+
 			if (this.receivingData) {
 				console.log(
 					`No ws data from ${this.accountName} in ${this.resubTimeoutMs}ms, resubscribing`
@@ -154,13 +161,20 @@ export class WebSocketAccountSubscriber<T> implements AccountSubscriber<T> {
 	}
 
 	unsubscribe(): Promise<void> {
+		this.isUnsubscribing = true;
+		clearTimeout(this.timeoutId);
+		this.timeoutId = undefined;
+
 		if (this.listenerId) {
-			const promise =
-				this.program.provider.connection.removeAccountChangeListener(
-					this.listenerId
-				);
-			this.listenerId = undefined;
+			const promise = this.program.provider.connection
+				.removeAccountChangeListener(this.listenerId)
+				.then(() => {
+					this.listenerId = undefined;
+					this.isUnsubscribing = false;
+				});
 			return promise;
+		} else {
+			this.isUnsubscribing = false;
 		}
 	}
 }
diff --git a/sdk/src/accounts/webSocketProgramAccountSubscriber.ts b/sdk/src/accounts/webSocketProgramAccountSubscriber.ts
@@ -21,6 +21,7 @@ export class WebSocketProgramAccountSubscriber<T>
 	onChange: (accountId: PublicKey, data: T, context: Context) => void;
 	listenerId?: number;
 	resubTimeoutMs?: number;
+	isUnsubscribing = false;
 	timeoutId?: NodeJS.Timeout;
 	options: { filters: MemcmpFilter[]; commitment?: Commitment };
 
@@ -48,7 +49,7 @@ export class WebSocketProgramAccountSubscriber<T>
 	async subscribe(
 		onChange: (accountId: PublicKey, data: T, context: Context) => void
 	): Promise<void> {
-		if (this.listenerId) {
+		if (this.listenerId || this.isUnsubscribing) {
 			return;
 		}
 
@@ -81,6 +82,11 @@ export class WebSocketProgramAccountSubscriber<T>
 			throw new Error('onChange callback function must be set');
 		}
 		this.timeoutId = setTimeout(async () => {
+			if (this.isUnsubscribing) {
+				// If we are in the process of unsubscribing, do not attempt to resubscribe
+				return;
+			}
+
 			if (this.receivingData) {
 				console.log(
 					`No ws data from ${this.subscriptionName} in ${this.resubTimeoutMs}ms, resubscribing`
@@ -140,13 +146,20 @@ export class WebSocketProgramAccountSubscriber<T>
 	}
 
 	unsubscribe(): Promise<void> {
+		this.isUnsubscribing = true;
+		clearTimeout(this.timeoutId);
+		this.timeoutId = undefined;
+
 		if (this.listenerId) {
-			const promise =
-				this.program.provider.connection.removeAccountChangeListener(
-					this.listenerId
-				);
-			this.listenerId = undefined;
+			const promise = this.program.provider.connection
+				.removeAccountChangeListener(this.listenerId)
+				.then(() => {
+					this.listenerId = undefined;
+					this.isUnsubscribing = false;
+				});
 			return promise;
+		} else {
+			this.isUnsubscribing = false;
 		}
 	}
 }