refactor(sdk): add OneShotUserStatsAccountSubscriber

Author: ChesterSim

sdk/src/accounts/basicUserStatsAccountSubscriber.ts

@@ -0,0 +1,65 @@
+import {
+	DataAndSlot,
+	UserStatsAccountEvents,
+	UserStatsAccountSubscriber,
+} from './types';
+import { PublicKey } from '@solana/web3.js';
+import StrictEventEmitter from 'strict-event-emitter-types';
+import { EventEmitter } from 'events';
+import { UserStatsAccount } from '../types';
+
+/**
+ * Basic implementation of UserStatsAccountSubscriber. It will only take in UserStatsAccount
+ * data during initialization and will not fetch or subscribe to updates.
+ */
+export class BasicUserStatsAccountSubscriber
+	implements UserStatsAccountSubscriber
+{
+	isSubscribed: boolean;
+	eventEmitter: StrictEventEmitter<EventEmitter, UserStatsAccountEvents>;
+	userStatsAccountPublicKey: PublicKey;
+
+	callbackId?: string;
+	errorCallbackId?: string;
+
+	userStats: DataAndSlot<UserStatsAccount>;
+
+	public constructor(
+		userStatsAccountPublicKey: PublicKey,
+		data?: UserStatsAccount,
+		slot?: number
+	) {
+		this.isSubscribed = true;
+		this.eventEmitter = new EventEmitter();
+		this.userStatsAccountPublicKey = userStatsAccountPublicKey;
+		this.userStats = { data, slot };
+	}
+
+	async subscribe(_userStatsAccount?: UserStatsAccount): Promise<boolean> {
+		return true;
+	}
+
+	async addToAccountLoader(): Promise<void> {}
+
+	async fetch(): Promise<void> {}
+
+	doesAccountExist(): boolean {
+		return this.userStats !== undefined;
+	}
+
+	async unsubscribe(): Promise<void> {}
+
+	assertIsSubscribed(): void {}
+
+	public getUserStatsAccountAndSlot(): DataAndSlot<UserStatsAccount> {
+		return this.userStats;
+	}
+
+	public updateData(userStatsAccount: UserStatsAccount, slot: number): void {
+		if (!this.userStats || slot >= (this.userStats.slot ?? 0)) {
+			this.userStats = { data: userStatsAccount, slot };
+			this.eventEmitter.emit('userStatsAccountUpdate', userStatsAccount);
+			this.eventEmitter.emit('update');
+		}
+	}
+}

sdk/src/accounts/oneShotUserStatsAccountSubscriber.ts

@@ -0,0 +1,69 @@
+import { Commitment, PublicKey } from '@solana/web3.js';
+import { UserStatsAccount } from '../types';
+import { BasicUserStatsAccountSubscriber } from './basicUserStatsAccountSubscriber';
+import { Program } from '@coral-xyz/anchor';
+import { UserStatsAccountSubscriber } from './types';
+
+/**
+ * Simple implementation of UserStatsAccountSubscriber. It will fetch the UserStatsAccount
+ * data on subscribe (or call to fetch) if no account data is provided on init.
+ * Expect to use only 1 RPC call unless you call fetch repeatedly.
+ */
+export class OneShotUserStatsAccountSubscriber
+	extends BasicUserStatsAccountSubscriber
+	implements UserStatsAccountSubscriber
+{
+	program: Program;
+	commitment: Commitment;
+
+	public constructor(
+		program: Program,
+		userStatsAccountPublicKey: PublicKey,
+		data?: UserStatsAccount,
+		slot?: number,
+		commitment?: Commitment
+	) {
+		super(userStatsAccountPublicKey, data, slot);
+		this.program = program;
+		this.commitment = commitment ?? 'confirmed';
+	}
+
+	async subscribe(userStatsAccount?: UserStatsAccount): Promise<boolean> {
+		if (userStatsAccount) {
+			this.userStats = { data: userStatsAccount, slot: this.userStats.slot };
+			return true;
+		}
+
+		await this.fetchIfUnloaded();
+		if (this.doesAccountExist()) {
+			this.eventEmitter.emit('update');
+		}
+		return true;
+	}
+
+	async fetchIfUnloaded(): Promise<void> {
+		if (this.userStats.data === undefined) {
+			await this.fetch();
+		}
+	}
+
+	async fetch(): Promise<void> {
+		try {
+			const dataAndContext =
+				await this.program.account.userStats.fetchAndContext(
+					this.userStatsAccountPublicKey,
+					this.commitment
+				);
+			if (dataAndContext.context.slot > (this.userStats?.slot ?? 0)) {
+				this.userStats = {
+					data: dataAndContext.data as UserStatsAccount,
+					slot: dataAndContext.context.slot,
+				};
+			}
+		} catch (e) {
+			console.error(
+				`OneShotUserStatsAccountSubscriber.fetch() UserStatsAccount does not exist: ${e.message}`
+			);
+		}
+	}
+}

sdk/src/driftClient.ts

@@ -924,7 +924,7 @@ export class DriftClient {
 			this.userStats = new UserStats({
 				driftClient: this,
 				userStatsAccountPublicKey: this.userStatsAccountPublicKey,
-				accountSubscription: this.userAccountSubscriptionConfig,
+				accountSubscription: this.userStatsAccountSubscriptionConfig,
 			});
 
 			this.userStats.subscribe();

sdk/src/index.ts

@@ -29,6 +29,7 @@ export * from './accounts/pollingInsuranceFundStakeAccountSubscriber';
 export * from './accounts/pollingHighLeverageModeConfigAccountSubscriber';
 export * from './accounts/basicUserAccountSubscriber';
 export * from './accounts/oneShotUserAccountSubscriber';
+export * from './accounts/oneShotUserStatsAccountSubscriber';
 export * from './accounts/types';
 export * from './addresses/pda';
 export * from './adminClient';

sdk/src/userStats.ts

@@ -54,9 +54,14 @@ export class UserStats {
 				},
 				config.accountSubscription.commitment
 			);
+		} else if (config.accountSubscription?.type === 'custom') {
+			this.accountSubscriber =
+				config.accountSubscription.userStatsAccountSubscriber;
 		} else {
+			const exhaustiveCheck: never = config.accountSubscription;
+
 			throw new Error(
-				`Unknown user stats account subscription type: ${config.accountSubscription?.type}`
+				`Unknown user stats account subscription type: ${exhaustiveCheck}`
 			);
 		}
 	}

sdk/src/userStatsConfig.ts

@@ -1,7 +1,7 @@
 import { DriftClient } from './driftClient';
 import { Commitment, PublicKey } from '@solana/web3.js';
 import { BulkAccountLoader } from './accounts/bulkAccountLoader';
-import { GrpcConfigs } from './accounts/types';
+import { GrpcConfigs, UserStatsAccountSubscriber } from './accounts/types';
 
 export type UserStatsConfig = {
 	accountSubscription?: UserStatsSubscriptionConfig;
@@ -22,6 +22,7 @@ export type UserStatsSubscriptionConfig =
 	  }
 	| {
 			type: 'custom';
+			userStatsAccountSubscriber: UserStatsAccountSubscriber;
 	  }
 	| {
 			type: 'grpc';

sdk/src/wallet.ts

@@ -5,6 +5,7 @@ import {
 	VersionedTransaction,
 } from '@solana/web3.js';
 import { IWallet, IVersionedWallet } from './types';
+import nacl from 'tweetnacl';
 
 export class Wallet implements IWallet, IVersionedWallet {
 	constructor(readonly payer: Keypair) {}
@@ -41,3 +42,13 @@ export class Wallet implements IWallet, IVersionedWallet {
 		return this.payer.publicKey;
 	}
 }
+
+export class WalletV2 extends Wallet {
+	constructor(readonly payer: Keypair) {
+		super(payer);
+	}
+
+	async signMessage(message: Uint8Array): Promise<Uint8Array> {
+		return Buffer.from(nacl.sign.detached(message, this.payer.secretKey));
+	}
+}