@drift-labs/sdk 2.48.0-beta.2 → 2.48.0-beta.21

package/src/userMap/userMap.ts

@@ -32,33 +32,59 @@ export interface UserMapInterface {
 	values(): IterableIterator<User>;
 }
 
+// filter users that meet these criteria when passing into syncCallback
+export type SyncCallbackCriteria = {
+	// only sync users that have open orders
+	hasOpenOrders: boolean;
+};
+
 export class UserMap implements UserMapInterface {
 	private userMap = new Map<string, User>();
 	private driftClient: DriftClient;
 	private accountSubscription: UserSubscriptionConfig;
 	private includeIdle: boolean;
 	private lastNumberOfSubAccounts;
-	private syncCallback = async (state: StateAccount) => {
+	private stateAccountUpdateCallback = async (state: StateAccount) => {
 		if (state.numberOfSubAccounts !== this.lastNumberOfSubAccounts) {
 			await this.sync();
 			this.lastNumberOfSubAccounts = state.numberOfSubAccounts;
 		}
 	};
+	private syncCallback: (authorities: PublicKey[]) => Promise<void>;
+	private syncCallbackCriteria: SyncCallbackCriteria;
+
+	private syncPromise?: Promise<void>;
+	private syncPromiseResolver: () => void;
 
 	/**
+	 * Constructs a new UserMap instance.
 	 *
-	 * @param driftClient
-	 * @param accountSubscription
-	 * @param includeIdle whether idle users are subscribed to. defaults to false to decrease # of user subscriptions
+	 * @param {DriftClient} driftClient - The DriftClient instance.
+	 * @param {UserSubscriptionConfig} accountSubscription - The UserSubscriptionConfig instance.
+	 * @param {boolean} includeIdle - Whether idle users are subscribed to. Defaults to false to decrease # of user subscriptions.
+	 * @param {(authorities: PublicKey[]) => Promise<void>} syncCallback - Called after `sync` completes, will pas in unique list of authorities. Useful for using it to sync UserStatsMap.
+	 * @param {SyncCallbackCriteria} syncCallbackCriteria - The criteria for the sync callback. Defaults to having no filters
 	 */
 	constructor(
 		driftClient: DriftClient,
 		accountSubscription: UserSubscriptionConfig,
-		includeIdle = false
+		includeIdle = false,
+		syncCallback?: (authorities: PublicKey[]) => Promise<void>,
+		syncCallbackCriteria: SyncCallbackCriteria = { hasOpenOrders: false }
 	) {
 		this.driftClient = driftClient;
 		this.accountSubscription = accountSubscription;
 		this.includeIdle = includeIdle;
+		this.syncCallback = syncCallback;
+		this.syncCallbackCriteria = syncCallbackCriteria;
+	}
+
+	public addSyncCallback(
+		syncCallback?: (authorities: PublicKey[]) => Promise<void>,
+		syncCallbackCriteria: SyncCallbackCriteria = { hasOpenOrders: false }
+	) {
+		this.syncCallback = syncCallback;
+		this.syncCallbackCriteria = syncCallbackCriteria;
 	}
 
 	public async subscribe() {
@@ -69,7 +95,10 @@ export class UserMap implements UserMapInterface {
 		await this.driftClient.subscribe();
 		this.lastNumberOfSubAccounts =
 			this.driftClient.getStateAccount().numberOfSubAccounts;
-		this.driftClient.eventEmitter.on('stateAccountUpdate', this.syncCallback);
+		this.driftClient.eventEmitter.on(
+			'stateAccountUpdate',
+			this.stateAccountUpdateCallback
+		);
 
 		await this.sync();
 	}
@@ -187,74 +216,118 @@ export class UserMap implements UserMapInterface {
 		return this.userMap.size;
 	}
 
+	public getUniqueAuthorities(useSyncCallbackCriteria = true): PublicKey[] {
+		const usersMeetingCriteria = Array.from(this.userMap.values()).filter(
+			(user) => {
+				let pass = true;
+				if (
+					useSyncCallbackCriteria &&
+					this.syncCallbackCriteria.hasOpenOrders
+				) {
+					pass = pass && user.getUserAccount().hasOpenOrder;
+				}
+				return pass;
+			}
+		);
+		const userAuths = new Set(
+			usersMeetingCriteria.map((user) =>
+				user.getUserAccount().authority.toBase58()
+			)
+		);
+		const userAuthKeys = Array.from(userAuths).map(
+			(userAuth) => new PublicKey(userAuth)
+		);
+		return userAuthKeys;
+	}
+
 	public async sync() {
-		const filters = [getUserFilter()];
-		if (!this.includeIdle) {
-			filters.push(getNonIdleUserFilter());
+		if (this.syncPromise) {
+			return this.syncPromise;
 		}
+		this.syncPromise = new Promise((resolver) => {
+			this.syncPromiseResolver = resolver;
+		});
 
-		const rpcRequestArgs = [
-			this.driftClient.program.programId.toBase58(),
-			{
-				commitment: this.driftClient.connection.commitment,
-				filters,
-				encoding: 'base64',
-				withContext: true,
-			},
-		];
-
-		// @ts-ignore
-		const rpcJSONResponse: any = await this.driftClient.connection._rpcRequest(
-			'getProgramAccounts',
-			rpcRequestArgs
-		);
+		try {
+			const filters = [getUserFilter()];
+			if (!this.includeIdle) {
+				filters.push(getNonIdleUserFilter());
+			}
 
-		const rpcResponseAndContext: RpcResponseAndContext<
-			Array<{
-				pubkey: PublicKey;
-				account: {
-					data: [string, string];
-				};
-			}>
-		> = rpcJSONResponse.result;
-
-		const slot = rpcResponseAndContext.context.slot;
-
-		const programAccountBufferMap = new Map<string, Buffer>();
-		for (const programAccount of rpcResponseAndContext.value) {
-			programAccountBufferMap.set(
-				programAccount.pubkey.toString(),
+			const rpcRequestArgs = [
+				this.driftClient.program.programId.toBase58(),
+				{
+					commitment: this.driftClient.connection.commitment,
+					filters,
+					encoding: 'base64',
+					withContext: true,
+				},
+			];
+
+			const rpcJSONResponse: any =
 				// @ts-ignore
-				Buffer.from(
-					programAccount.account.data[0],
-					programAccount.account.data[1]
-				)
-			);
-		}
+				await this.driftClient.connection._rpcRequest(
+					'getProgramAccounts',
+					rpcRequestArgs
+				);
+
+			const rpcResponseAndContext: RpcResponseAndContext<
+				Array<{
+					pubkey: PublicKey;
+					account: {
+						data: [string, string];
+					};
+				}>
+			> = rpcJSONResponse.result;
+
+			const slot = rpcResponseAndContext.context.slot;
 
-		for (const [key, buffer] of programAccountBufferMap.entries()) {
-			if (!this.has(key)) {
-				const userAccount =
-					this.driftClient.program.account.user.coder.accounts.decode(
-						'User',
-						buffer
-					);
-				await this.addPubkey(new PublicKey(key), userAccount);
+			const programAccountBufferMap = new Map<string, Buffer>();
+			for (const programAccount of rpcResponseAndContext.value) {
+				programAccountBufferMap.set(
+					programAccount.pubkey.toString(),
+					// @ts-ignore
+					Buffer.from(
+						programAccount.account.data[0],
+						programAccount.account.data[1]
+					)
+				);
 			}
-		}
 
-		for (const [key, user] of this.userMap.entries()) {
-			if (!programAccountBufferMap.has(key)) {
-				await user.unsubscribe();
-				this.userMap.delete(key);
-			} else {
-				const userAccount =
-					this.driftClient.program.account.user.coder.accounts.decode(
-						'User',
-						programAccountBufferMap.get(key)
-					);
-				user.accountSubscriber.updateData(userAccount, slot);
+			for (const [key, buffer] of programAccountBufferMap.entries()) {
+				if (!this.has(key)) {
+					const userAccount =
+						this.driftClient.program.account.user.coder.accounts.decode(
+							'User',
+							buffer
+						);
+					await this.addPubkey(new PublicKey(key), userAccount);
+				}
+			}
+
+			for (const [key, user] of this.userMap.entries()) {
+				if (!programAccountBufferMap.has(key)) {
+					await user.unsubscribe();
+					this.userMap.delete(key);
+				} else {
+					const userAccount =
+						this.driftClient.program.account.user.coder.accounts.decode(
+							'User',
+							programAccountBufferMap.get(key)
+						);
+					user.accountSubscriber.updateData(userAccount, slot);
+				}
+			}
+
+			if (this.syncCallback) {
+				await this.syncCallback(this.getUniqueAuthorities());
 			}
+		} catch (e) {
+			console.error(`Error in UserMap.sync()`);
+			console.error(e);
+		} finally {
+			this.syncPromiseResolver();
+			this.syncPromise = undefined;
 		}
 	}
 
@@ -267,7 +340,7 @@ export class UserMap implements UserMapInterface {
 		if (this.lastNumberOfSubAccounts) {
 			this.driftClient.eventEmitter.removeListener(
 				'stateAccountUpdate',
-				this.syncCallback
+				this.stateAccountUpdateCallback
 			);
 			this.lastNumberOfSubAccounts = undefined;
 		}

package/src/userMap/userStatsMap.ts

@@ -4,7 +4,6 @@ import {
 	OrderRecord,
 	UserStatsAccount,
 	UserStats,
-	UserStatsSubscriptionConfig,
 	WrappedEvent,
 	DepositRecord,
 	FundingPaymentRecord,
@@ -14,12 +13,12 @@ import {
 	NewUserRecord,
 	LPRecord,
 	InsuranceFundStakeRecord,
-	StateAccount,
+	BulkAccountLoader,
+	PollingUserStatsAccountSubscriber,
 } from '..';
-import { AccountInfo, PublicKey } from '@solana/web3.js';
+import { PublicKey } from '@solana/web3.js';
 
 import { UserMap } from './userMap';
-import { Buffer } from 'buffer';
 
 export class UserStatsMap {
 	/**
@@ -27,39 +26,45 @@ export class UserStatsMap {
 	 */
 	private userStatsMap = new Map<string, UserStats>();
 	private driftClient: DriftClient;
-	private accountSubscription: UserStatsSubscriptionConfig;
-	private lastNumberOfAuthorities;
-	private syncCallback = async (state: StateAccount) => {
-		if (state.numberOfAuthorities !== this.lastNumberOfAuthorities) {
-			await this.sync();
-			this.lastNumberOfAuthorities = state.numberOfAuthorities;
-		}
-	};
+	private bulkAccountLoader: BulkAccountLoader;
 
-	constructor(
-		driftClient: DriftClient,
-		accountSubscription: UserStatsSubscriptionConfig
-	) {
+	/**
+	 * Creates a new UserStatsMap instance.
+	 *
+	 * @param {DriftClient} driftClient - The DriftClient instance.
+	 * @param {BulkAccountLoader} [bulkAccountLoader] - If not provided, a new BulkAccountLoader with polling disabled will be created.
+	 */
+	constructor(driftClient: DriftClient, bulkAccountLoader?: BulkAccountLoader) {
 		this.driftClient = driftClient;
-		this.accountSubscription = accountSubscription;
+		if (!bulkAccountLoader) {
+			bulkAccountLoader = new BulkAccountLoader(
+				driftClient.connection,
+				driftClient.opts.commitment,
+				0
+			);
+		}
+		this.bulkAccountLoader = bulkAccountLoader;
 	}
 
-	public async subscribe() {
+	public async subscribe(authorities: PublicKey[]) {
 		if (this.size() > 0) {
 			return;
 		}
 
 		await this.driftClient.subscribe();
-		this.lastNumberOfAuthorities =
-			this.driftClient.getStateAccount().numberOfAuthorities;
-		this.driftClient.eventEmitter.on('stateAccountUpdate', this.syncCallback);
-
-		await this.sync();
+		await this.sync(authorities);
 	}
 
+	/**
+	 *
+	 * @param authority that owns the UserStatsAccount
+	 * @param userStatsAccount optional UserStatsAccount to subscribe to, if undefined will be fetched later
+	 * @param skipFetch if true, will not immediately fetch the UserStatsAccount
+	 */
 	public async addUserStat(
 		authority: PublicKey,
-		userStatsAccount?: UserStatsAccount
+		userStatsAccount?: UserStatsAccount,
+		skipFetch?: boolean
 	) {
 		const userStat = new UserStats({
 			driftClient: this.driftClient,
@@ -67,9 +72,18 @@ export class UserStatsMap {
 				this.driftClient.program.programId,
 				authority
 			),
-			accountSubscription: this.accountSubscription,
+			accountSubscription: {
+				type: 'polling',
+				accountLoader: this.bulkAccountLoader,
+			},
 		});
-		await userStat.subscribe(userStatsAccount);
+		if (skipFetch) {
+			await (
+				userStat.accountSubscriber as PollingUserStatsAccountSubscriber
+			).addToAccountLoader();
+		} else {
+			await userStat.subscribe(userStatsAccount);
+		}
 
 		this.userStatsMap.set(authority.toString(), userStat);
 	}
@@ -77,7 +91,7 @@ export class UserStatsMap {
 	public async updateWithOrderRecord(record: OrderRecord, userMap: UserMap) {
 		const user = await userMap.mustGet(record.user.toString());
 		if (!this.has(user.getUserAccount().authority.toString())) {
-			await this.addUserStat(user.getUserAccount().authority);
+			await this.addUserStat(user.getUserAccount().authority, undefined, false);
 		}
 	}
 
@@ -156,9 +170,19 @@ export class UserStatsMap {
 		return this.userStatsMap.get(authorityPublicKey);
 	}
 
+	/**
+	 * Enforce that a UserStats will exist for the given authorityPublicKey,
+	 * reading one from the blockchain if necessary.
+	 * @param authorityPublicKey
+	 * @returns
+	 */
 	public async mustGet(authorityPublicKey: string): Promise<UserStats> {
 		if (!this.has(authorityPublicKey)) {
-			await this.addUserStat(new PublicKey(authorityPublicKey));
+			await this.addUserStat(
+				new PublicKey(authorityPublicKey),
+				undefined,
+				false
+			);
 		}
 		return this.get(authorityPublicKey);
 	}
@@ -171,39 +195,18 @@ export class UserStatsMap {
 		return this.userStatsMap.size;
 	}
 
-	public async sync() {
-		const programAccounts =
-			await this.driftClient.connection.getProgramAccounts(
-				this.driftClient.program.programId,
-				{
-					commitment: this.driftClient.connection.commitment,
-					filters: [
-						{
-							memcmp:
-								this.driftClient.program.coder.accounts.memcmp('UserStats'),
-						},
-					],
-				}
-			);
-
-		const programAccountMap = new Map<string, AccountInfo<Buffer>>();
-		for (const programAccount of programAccounts) {
-			programAccountMap.set(
-				new PublicKey(programAccount.account.data.slice(8, 40)).toString(),
-				programAccount.account
-			);
-		}
-
-		for (const key of programAccountMap.keys()) {
-			if (!this.has(key)) {
-				const userStatsAccount =
-					this.driftClient.program.account.userStats.coder.accounts.decode(
-						'UserStats',
-						programAccountMap.get(key).data
-					);
-				await this.addUserStat(new PublicKey(key), userStatsAccount);
-			}
-		}
+	/**
+	 * Sync the UserStatsMap
+	 * @param authorities list of authorities to derive UserStatsAccount public keys from.
+	 * You may want to get this list from UserMap in order to filter out idle users
+	 */
+	public async sync(authorities: PublicKey[]) {
+		await Promise.all(
+			authorities.map((authority) =>
+				this.addUserStat(authority, undefined, true)
+			)
+		);
+		await this.bulkAccountLoader.load();
 	}
 
 	public async unsubscribe() {
@@ -211,13 +214,5 @@ export class UserStatsMap {
 			await userStats.unsubscribe();
 			this.userStatsMap.delete(key);
 		}
-
-		if (this.lastNumberOfAuthorities) {
-			this.driftClient.eventEmitter.removeListener(
-				'stateAccountUpdate',
-				this.syncCallback
-			);
-			this.lastNumberOfAuthorities = undefined;
-		}
 	}
 }

package/tests/amm/test.ts

@@ -967,8 +967,9 @@ describe('AMM Tests', () => {
 		mockMarket1.amm.maxBaseAssetReserve = mockMarket1.amm.baseAssetReserve.add(
 			new BN(9)
 		);
-		mockMarket1.amm.minBaseAssetReserve =
-			mockMarket1.amm.baseAssetReserve.sub(new BN(9));
+		mockMarket1.amm.minBaseAssetReserve = mockMarket1.amm.baseAssetReserve.sub(
+			new BN(9)
+		);
 		mockMarket1.amm.quoteAssetReserve = new BN(cc).mul(BASE_PRECISION);
 		mockMarket1.amm.pegMultiplier = new BN(18.32 * PEG_PRECISION.toNumber());
 		mockMarket1.amm.sqrtK = new BN(cc).mul(BASE_PRECISION);

package/tests/dlob/test.ts

@@ -27,7 +27,7 @@ import {
 
 import { mockPerpMarkets, mockSpotMarkets, mockStateAccount } from './helpers';
 import { DLOBOrdersCoder } from '../../src/dlob/DLOBOrders';
-import {isAuctionComplete, isRestingLimitOrder} from "../../lib";
+import { isAuctionComplete, isRestingLimitOrder } from '../../lib';
 
 function insertOrderToDLOB(
 	dlob: DLOB,
@@ -2537,18 +2537,21 @@ describe('DLOB Perp Tests', () => {
 			OrderTriggerCondition.TRIGGERED_ABOVE, // triggerCondition: OrderTriggerCondition,
 			vBid,
 			vAsk,
-			new BN(1), // slot
+			new BN(1) // slot
 		);
 
-		const restingLimitBids = Array.from(dlob.getRestingLimitBids(marketIndex, slot, MarketType.PERP, oracle));
+		const restingLimitBids = Array.from(
+			dlob.getRestingLimitBids(marketIndex, slot, MarketType.PERP, oracle)
+		);
 		expect(restingLimitBids.length).to.equal(0);
 
-		const takingBids = Array.from(dlob.getTakingBids(marketIndex, MarketType.PERP, slot, oracle));
+		const takingBids = Array.from(
+			dlob.getTakingBids(marketIndex, MarketType.PERP, slot, oracle)
+		);
 		expect(takingBids.length).to.equal(1);
 		const triggerLimitBid = takingBids[0];
 		expect(isAuctionComplete(triggerLimitBid.order, slot)).to.equal(true);
 		expect(isRestingLimitOrder(triggerLimitBid.order, slot)).to.equal(false);
-
 	});
 
 	it('Test will return expired market orders to fill', () => {