From af51b52ef343590dccc5b87d9e4484cba6ca83bd Mon Sep 17 00:00:00 2001
From: wphan <william@drift.trade>
Date: Tue, 26 Dec 2023 10:09:44 -0800
Subject: [PATCH] ts-sdk: allow disabling userMapSync on state account updates

---
 sdk/src/userMap/userMap.ts       | 24 ++++++++++++++++--------
 sdk/src/userMap/userMapConfig.ts |  4 ++++
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/sdk/src/userMap/userMap.ts b/sdk/src/userMap/userMap.ts
index 325e61b5b..b6c30c006 100644
--- a/sdk/src/userMap/userMap.ts
+++ b/sdk/src/userMap/userMap.ts
@@ -51,6 +51,7 @@ export class UserMap implements UserMapInterface {
 	private connection: Connection;
 	private commitment: Commitment;
 	private includeIdle: boolean;
+	private disableSyncOnTotalAccountsChange: boolean;
 	private lastNumberOfSubAccounts: BN;
 	private subscription: PollingSubscription | WebsocketSubscription;
 	private stateAccountUpdateCallback = async (state: StateAccount) => {
@@ -78,6 +79,8 @@ export class UserMap implements UserMapInterface {
 		this.commitment =
 			config.subscriptionConfig.commitment ?? this.driftClient.opts.commitment;
 		this.includeIdle = config.includeIdle ?? false;
+		this.disableSyncOnTotalAccountsChange =
+			config.disableSyncOnTotalAccountsChange ?? false;
 
 		let decodeFn;
 		if (config.fastDecode ?? true) {
@@ -115,10 +118,12 @@ export class UserMap implements UserMapInterface {
 		await this.driftClient.subscribe();
 		this.lastNumberOfSubAccounts =
 			this.driftClient.getStateAccount().numberOfSubAccounts;
-		this.driftClient.eventEmitter.on(
-			'stateAccountUpdate',
-			this.stateAccountUpdateCallback
-		);
+		if (!this.disableSyncOnTotalAccountsChange) {
+			this.driftClient.eventEmitter.on(
+				'stateAccountUpdate',
+				this.stateAccountUpdateCallback
+			);
+		}
 
 		await this.subscription.subscribe();
 	}
@@ -363,10 +368,13 @@ export class UserMap implements UserMapInterface {
 		}
 
 		if (this.lastNumberOfSubAccounts) {
-			this.driftClient.eventEmitter.removeListener(
-				'stateAccountUpdate',
-				this.stateAccountUpdateCallback
-			);
+			if (!this.disableSyncOnTotalAccountsChange) {
+				this.driftClient.eventEmitter.removeListener(
+					'stateAccountUpdate',
+					this.stateAccountUpdateCallback
+				);
+			}
+
 			this.lastNumberOfSubAccounts = undefined;
 		}
 	}
diff --git a/sdk/src/userMap/userMapConfig.ts b/sdk/src/userMap/userMapConfig.ts
index 826231c57..df78714e6 100644
--- a/sdk/src/userMap/userMapConfig.ts
+++ b/sdk/src/userMap/userMapConfig.ts
@@ -31,4 +31,8 @@ export type UserMapConfig = {
 
 	// Whether to skip loading available perp/spot positions and open orders
 	fastDecode?: boolean;
+
+	// If true, will not do a full sync whenever StateAccount.numberOfSubAccounts changes.
+	// default behavior is to do a full sync on changes.
+	disableSyncOnTotalAccountsChange?: boolean;
 };