@drift-labs/sdk 2.145.0-beta.2 → 2.145.0

package/src/constituentMap/constituentMap.ts

@@ -0,0 +1,285 @@
+import {
+	Commitment,
+	Connection,
+	MemcmpFilter,
+	PublicKey,
+	RpcResponseAndContext,
+} from '@solana/web3.js';
+import { ConstituentAccountSubscriber, DataAndSlot } from '../accounts/types';
+import { ConstituentAccount } from '../types';
+import { PollingConstituentAccountSubscriber } from './pollingConstituentAccountSubscriber';
+import { WebSocketConstituentAccountSubscriber } from './webSocketConstituentAccountSubscriber';
+import { DriftClient } from '../driftClient';
+import { getConstituentFilter, getConstituentLpPoolFilter } from '../memcmp';
+import { ZSTDDecoder } from 'zstddec';
+import { getLpPoolPublicKey } from '../addresses/pda';
+
+const MAX_CONSTITUENT_SIZE_BYTES = 480; // TODO: update this when account is finalized
+
+export type ConstituentMapConfig = {
+	driftClient: DriftClient;
+	connection?: Connection;
+	subscriptionConfig:
+		| {
+				type: 'polling';
+				frequency: number;
+				commitment?: Commitment;
+		  }
+		| {
+				type: 'websocket';
+				resubTimeoutMs?: number;
+				logResubMessages?: boolean;
+				commitment?: Commitment;
+		  };
+	lpPoolId?: number;
+	// potentially use these to filter Constituent accounts
+	additionalFilters?: MemcmpFilter[];
+};
+
+export interface ConstituentMapInterface {
+	subscribe(): Promise<void>;
+	unsubscribe(): Promise<void>;
+	has(key: string): boolean;
+	get(key: string): ConstituentAccount | undefined;
+	getFromSpotMarketIndex(
+		spotMarketIndex: number
+	): ConstituentAccount | undefined;
+	getFromConstituentIndex(
+		constituentIndex: number
+	): ConstituentAccount | undefined;
+
+	getWithSlot(key: string): DataAndSlot<ConstituentAccount> | undefined;
+	mustGet(key: string): Promise<ConstituentAccount>;
+	mustGetWithSlot(key: string): Promise<DataAndSlot<ConstituentAccount>>;
+}
+
+export class ConstituentMap implements ConstituentMapInterface {
+	private driftClient: DriftClient;
+	private constituentMap = new Map<string, DataAndSlot<ConstituentAccount>>();
+	private constituentAccountSubscriber: ConstituentAccountSubscriber;
+	private additionalFilters?: MemcmpFilter[];
+	private commitment?: Commitment;
+	private connection?: Connection;
+
+	private constituentIndexToKeyMap = new Map<number, string>();
+	private spotMarketIndexToKeyMap = new Map<number, string>();
+
+	private lpPoolId: number;
+
+	constructor(config: ConstituentMapConfig) {
+		this.driftClient = config.driftClient;
+		this.additionalFilters = config.additionalFilters;
+		this.commitment = config.subscriptionConfig.commitment;
+		this.connection = config.connection || this.driftClient.connection;
+		this.lpPoolId = config.lpPoolId ?? 0;
+
+		if (config.subscriptionConfig.type === 'polling') {
+			this.constituentAccountSubscriber =
+				new PollingConstituentAccountSubscriber(
+					this,
+					this.driftClient.program,
+					config.subscriptionConfig.frequency,
+					config.subscriptionConfig.commitment,
+					this.getFilters()
+				);
+		} else if (config.subscriptionConfig.type === 'websocket') {
+			this.constituentAccountSubscriber =
+				new WebSocketConstituentAccountSubscriber(
+					this,
+					this.driftClient.program,
+					config.subscriptionConfig.resubTimeoutMs,
+					config.subscriptionConfig.commitment,
+					this.getFilters()
+				);
+		}
+
+		// Listen for account updates from the subscriber
+		this.constituentAccountSubscriber.eventEmitter.on(
+			'onAccountUpdate',
+			(account: ConstituentAccount, pubkey: PublicKey, slot: number) => {
+				this.updateConstituentAccount(pubkey.toString(), account, slot);
+			}
+		);
+	}
+
+	private getFilters(): MemcmpFilter[] {
+		const filters = [
+			getConstituentFilter(),
+			getConstituentLpPoolFilter(
+				getLpPoolPublicKey(this.driftClient.program.programId, this.lpPoolId)
+			),
+		];
+		if (this.additionalFilters) {
+			filters.push(...this.additionalFilters);
+		}
+		return filters;
+	}
+
+	private decode(name: string, buffer: Buffer): ConstituentAccount {
+		return this.driftClient.program.account.constituent.coder.accounts.decodeUnchecked(
+			name,
+			buffer
+		);
+	}
+
+	public async sync(): Promise<void> {
+		try {
+			const rpcRequestArgs = [
+				this.driftClient.program.programId.toBase58(),
+				{
+					commitment: this.commitment,
+					filters: this.getFilters(),
+					encoding: 'base64+zstd',
+					withContext: true,
+				},
+			];
+
+			// @ts-ignore
+			const rpcJSONResponse: any = await this.connection._rpcRequest(
+				'getProgramAccounts',
+				rpcRequestArgs
+			);
+			const rpcResponseAndContext: RpcResponseAndContext<
+				Array<{ pubkey: PublicKey; account: { data: [string, string] } }>
+			> = rpcJSONResponse.result;
+			const slot = rpcResponseAndContext.context.slot;
+
+			const promises = rpcResponseAndContext.value.map(
+				async (programAccount) => {
+					const compressedUserData = Buffer.from(
+						programAccount.account.data[0],
+						'base64'
+					);
+					const decoder = new ZSTDDecoder();
+					await decoder.init();
+					const buffer = Buffer.from(
+						decoder.decode(compressedUserData, MAX_CONSTITUENT_SIZE_BYTES)
+					);
+					const key = programAccount.pubkey.toString();
+					const currAccountWithSlot = this.getWithSlot(key);
+
+					if (currAccountWithSlot) {
+						if (slot >= currAccountWithSlot.slot) {
+							const constituentAcc = this.decode('Constituent', buffer);
+							this.updateConstituentAccount(key, constituentAcc, slot);
+						}
+					} else {
+						const constituentAcc = this.decode('Constituent', buffer);
+						this.updateConstituentAccount(key, constituentAcc, slot);
+					}
+				}
+			);
+			await Promise.all(promises);
+		} catch (error) {
+			console.log(`ConstituentMap.sync() error: ${error.message}`);
+		}
+	}
+
+	public async subscribe(): Promise<void> {
+		await this.constituentAccountSubscriber.subscribe();
+	}
+
+	public async unsubscribe(): Promise<void> {
+		await this.constituentAccountSubscriber.unsubscribe();
+		this.constituentMap.clear();
+	}
+
+	public has(key: string): boolean {
+		return this.constituentMap.has(key);
+	}
+
+	public get(key: string): ConstituentAccount | undefined {
+		return this.constituentMap.get(key)?.data;
+	}
+
+	public getFromConstituentIndex(
+		constituentIndex: number
+	): ConstituentAccount | undefined {
+		const key = this.constituentIndexToKeyMap.get(constituentIndex);
+		return key ? this.get(key) : undefined;
+	}
+
+	public getFromSpotMarketIndex(
+		spotMarketIndex: number
+	): ConstituentAccount | undefined {
+		const key = this.spotMarketIndexToKeyMap.get(spotMarketIndex);
+		return key ? this.get(key) : undefined;
+	}
+
+	public getWithSlot(key: string): DataAndSlot<ConstituentAccount> | undefined {
+		return this.constituentMap.get(key);
+	}
+
+	public async mustGet(key: string): Promise<ConstituentAccount> {
+		if (!this.has(key)) {
+			await this.sync();
+		}
+		const result = this.constituentMap.get(key);
+		if (!result) {
+			throw new Error(`ConstituentAccount not found for key: ${key}`);
+		}
+		return result.data;
+	}
+
+	public async mustGetWithSlot(
+		key: string
+	): Promise<DataAndSlot<ConstituentAccount>> {
+		if (!this.has(key)) {
+			await this.sync();
+		}
+		const result = this.constituentMap.get(key);
+		if (!result) {
+			throw new Error(`ConstituentAccount not found for key: ${key}`);
+		}
+		return result;
+	}
+
+	public size(): number {
+		return this.constituentMap.size;
+	}
+
+	public *values(): IterableIterator<ConstituentAccount> {
+		for (const dataAndSlot of this.constituentMap.values()) {
+			yield dataAndSlot.data;
+		}
+	}
+
+	public valuesWithSlot(): IterableIterator<DataAndSlot<ConstituentAccount>> {
+		return this.constituentMap.values();
+	}
+
+	public *entries(): IterableIterator<[string, ConstituentAccount]> {
+		for (const [key, dataAndSlot] of this.constituentMap.entries()) {
+			yield [key, dataAndSlot.data];
+		}
+	}
+
+	public entriesWithSlot(): IterableIterator<
+		[string, DataAndSlot<ConstituentAccount>]
+	> {
+		return this.constituentMap.entries();
+	}
+
+	public updateConstituentAccount(
+		key: string,
+		constituentAccount: ConstituentAccount,
+		slot: number
+	): void {
+		const existingData = this.getWithSlot(key);
+		if (existingData) {
+			if (slot >= existingData.slot) {
+				this.constituentMap.set(key, {
+					data: constituentAccount,
+					slot,
+				});
+			}
+		} else {
+			this.constituentMap.set(key, {
+				data: constituentAccount,
+				slot,
+			});
+		}
+		this.constituentIndexToKeyMap.set(constituentAccount.constituentIndex, key);
+		this.spotMarketIndexToKeyMap.set(constituentAccount.spotMarketIndex, key);
+	}
+}

package/src/constituentMap/pollingConstituentAccountSubscriber.ts

@@ -0,0 +1,97 @@
+import {
+	NotSubscribedError,
+	ConstituentAccountEvents,
+	ConstituentAccountSubscriber,
+} from '../accounts/types';
+import { Program } from '@coral-xyz/anchor';
+import StrictEventEmitter from 'strict-event-emitter-types';
+import { EventEmitter } from 'events';
+import { Commitment, MemcmpFilter } from '@solana/web3.js';
+import { ConstituentMap } from './constituentMap';
+
+export class PollingConstituentAccountSubscriber
+	implements ConstituentAccountSubscriber
+{
+	isSubscribed: boolean;
+	program: Program;
+	frequency: number;
+	commitment?: Commitment;
+	additionalFilters?: MemcmpFilter[];
+	eventEmitter: StrictEventEmitter<EventEmitter, ConstituentAccountEvents>;
+
+	intervalId?: NodeJS.Timeout;
+	constituentMap: ConstituentMap;
+
+	public constructor(
+		constituentMap: ConstituentMap,
+		program: Program,
+		frequency: number,
+		commitment?: Commitment,
+		additionalFilters?: MemcmpFilter[]
+	) {
+		this.constituentMap = constituentMap;
+		this.isSubscribed = false;
+		this.program = program;
+		this.frequency = frequency;
+		this.commitment = commitment;
+		this.additionalFilters = additionalFilters;
+		this.eventEmitter = new EventEmitter();
+	}
+
+	async subscribe(): Promise<boolean> {
+		if (this.isSubscribed || this.frequency <= 0) {
+			return true;
+		}
+
+		const executeSync = async () => {
+			await this.sync();
+			this.intervalId = setTimeout(executeSync, this.frequency);
+		};
+
+		// Initial sync
+		await this.sync();
+
+		// Start polling
+		this.intervalId = setTimeout(executeSync, this.frequency);
+
+		this.isSubscribed = true;
+		return true;
+	}
+
+	async sync(): Promise<void> {
+		try {
+			await this.constituentMap.sync();
+			this.eventEmitter.emit('update');
+		} catch (error) {
+			console.log(
+				`PollingConstituentAccountSubscriber.sync() error: ${error.message}`
+			);
+			this.eventEmitter.emit('error', error);
+		}
+	}
+
+	async unsubscribe(): Promise<void> {
+		if (!this.isSubscribed) {
+			return;
+		}
+
+		if (this.intervalId) {
+			clearTimeout(this.intervalId);
+			this.intervalId = undefined;
+		}
+
+		this.isSubscribed = false;
+	}
+
+	assertIsSubscribed(): void {
+		if (!this.isSubscribed) {
+			throw new NotSubscribedError(
+				'You must call `subscribe` before using this function'
+			);
+		}
+	}
+
+	didSubscriptionSucceed(): boolean {
+		return this.isSubscribed;
+	}
+}

package/src/constituentMap/webSocketConstituentAccountSubscriber.ts

@@ -0,0 +1,112 @@
+import {
+	NotSubscribedError,
+	ConstituentAccountEvents,
+	ConstituentAccountSubscriber,
+} from '../accounts/types';
+import { Program } from '@coral-xyz/anchor';
+import StrictEventEmitter from 'strict-event-emitter-types';
+import { EventEmitter } from 'events';
+import { Commitment, Context, MemcmpFilter, PublicKey } from '@solana/web3.js';
+import { ConstituentAccount } from '../types';
+import { WebSocketProgramAccountSubscriber } from '../accounts/webSocketProgramAccountSubscriber';
+import { getConstituentFilter } from '../memcmp';
+import { ConstituentMap } from './constituentMap';
+
+export class WebSocketConstituentAccountSubscriber
+	implements ConstituentAccountSubscriber
+{
+	isSubscribed: boolean;
+	resubTimeoutMs?: number;
+	commitment?: Commitment;
+	program: Program;
+	eventEmitter: StrictEventEmitter<EventEmitter, ConstituentAccountEvents>;
+
+	constituentDataAccountSubscriber: WebSocketProgramAccountSubscriber<ConstituentAccount>;
+	constituentMap: ConstituentMap;
+	private additionalFilters?: MemcmpFilter[];
+
+	public constructor(
+		constituentMap: ConstituentMap,
+		program: Program,
+		resubTimeoutMs?: number,
+		commitment?: Commitment,
+		additionalFilters?: MemcmpFilter[]
+	) {
+		this.constituentMap = constituentMap;
+		this.isSubscribed = false;
+		this.program = program;
+		this.eventEmitter = new EventEmitter();
+		this.resubTimeoutMs = resubTimeoutMs;
+		this.commitment = commitment;
+		this.additionalFilters = additionalFilters;
+	}
+
+	async subscribe(): Promise<boolean> {
+		if (this.isSubscribed) {
+			return true;
+		}
+		this.constituentDataAccountSubscriber =
+			new WebSocketProgramAccountSubscriber<ConstituentAccount>(
+				'LpPoolConstituent',
+				'Constituent',
+				this.program,
+				this.program.account.constituent.coder.accounts.decode.bind(
+					this.program.account.constituent.coder.accounts
+				),
+				{
+					filters: [getConstituentFilter(), ...(this.additionalFilters || [])],
+					commitment: this.commitment,
+				}
+			);
+
+		await this.constituentDataAccountSubscriber.subscribe(
+			(accountId: PublicKey, account: ConstituentAccount, context: Context) => {
+				this.constituentMap.updateConstituentAccount(
+					accountId.toBase58(),
+					account,
+					context.slot
+				);
+				this.eventEmitter.emit(
+					'onAccountUpdate',
+					account,
+					accountId,
+					context.slot
+				);
+			}
+		);
+
+		this.eventEmitter.emit('update');
+		this.isSubscribed = true;
+		return true;
+	}
+
+	async sync(): Promise<void> {
+		try {
+			await this.constituentMap.sync();
+			this.eventEmitter.emit('update');
+		} catch (error) {
+			console.log(
+				`WebSocketConstituentAccountSubscriber.sync() error: ${error.message}`
+			);
+			this.eventEmitter.emit('error', error);
+		}
+	}
+
+	async unsubscribe(): Promise<void> {
+		if (!this.isSubscribed) {
+			return;
+		}
+
+		await Promise.all([this.constituentDataAccountSubscriber.unsubscribe()]);
+
+		this.isSubscribed = false;
+	}
+
+	assertIsSubscribed(): void {
+		if (!this.isSubscribed) {
+			throw new NotSubscribedError(
+				'You must call `subscribe` before using this function'
+			);
+		}
+	}
+}