Merge pull request #288 from drift-labs/chore/grpc_publisher

remove grpc impl, use sdk grpc OrderSubscriber

src/dlobProvider.ts

@@ -1,6 +1,5 @@
 import { DLOB, OrderSubscriber, UserAccount, UserMap } from '@drift-labs/sdk';
 import { PublicKey } from '@solana/web3.js';
-import { GeyserOrderSubscriber } from './grpc/OrderSubscriberGRPC';
 
 export type DLOBProvider = {
 	subscribe(): Promise<void>;
@@ -91,45 +90,3 @@ export function getDLOBProviderFromOrderSubscriber(
 		},
 	};
 }
-
-export function getDLOBProviderFromGrpcOrderSubscriber(
-	orderSubscriber: GeyserOrderSubscriber
-): DLOBProvider {
-	return {
-		subscribe: async () => {
-			await orderSubscriber.subscribe();
-		},
-		getDLOB: async (slot: number) => {
-			return await orderSubscriber.getDLOB(slot);
-		},
-		getUniqueAuthorities: () => {
-			const authorities = new Set<string>();
-			for (const { userAccount } of orderSubscriber.usersAccounts.values()) {
-				authorities.add(userAccount.authority.toBase58());
-			}
-			const pubkeys = Array.from(authorities).map((a) => new PublicKey(a));
-			return pubkeys;
-		},
-		getUserAccounts: function* () {
-			for (const [
-				key,
-				{ userAccount },
-			] of orderSubscriber.usersAccounts.entries()) {
-				yield { userAccount: userAccount, publicKey: new PublicKey(key) };
-			}
-		},
-		getUserAccount: (publicKey) => {
-			return orderSubscriber.usersAccounts.get(publicKey.toString())
-				?.userAccount;
-		},
-		size(): number {
-			return orderSubscriber.usersAccounts.size;
-		},
-		fetch() {
-			return orderSubscriber.fetch();
-		},
-		getSlot: () => {
-			return orderSubscriber.getSlot();
-		},
-	};
-}

src/grpc/OrderSubscriberGRPC.ts

@@ -1,253 +0,0 @@
-import { BorshAccountsCoder } from '@coral-xyz/anchor';
-import { Commitment, PublicKey, RpcResponseAndContext } from '@solana/web3.js';
-import { Buffer } from 'buffer';
-import Client, {
-	CommitmentLevel,
-	SubscribeRequest,
-	SubscribeUpdate,
-} from '@triton-one/yellowstone-grpc';
-import {
-	BN,
-	DLOB,
-	DriftClient,
-	UserAccount,
-	getUserFilter,
-	getUserWithOrderFilter,
-} from '@drift-labs/sdk';
-import { ClientDuplexStream } from '@grpc/grpc-js';
-
-type grpcDlobConfig = {
-	endpoint: string;
-	token: string;
-};
-
-export class GeyserOrderSubscriber {
-	usersAccounts = new Map<string, { slot: number; userAccount: UserAccount }>();
-	commitment: Commitment;
-	driftClient: DriftClient;
-	config: grpcDlobConfig;
-	stream: ClientDuplexStream<SubscribeRequest, SubscribeUpdate>;
-
-	fetchPromise?: Promise<void>;
-	fetchPromiseResolver: () => void;
-	mostRecentSlot: number;
-
-	constructor(driftClient: DriftClient, config: grpcDlobConfig) {
-		this.driftClient = driftClient;
-		this.config = config;
-	}
-
-	public async subscribe(): Promise<void> {
-		const client = new Client(this.config.endpoint, this.config.token);
-		this.stream = await client.subscribe();
-		const request: SubscribeRequest = {
-			slots: {},
-			accounts: {
-				drift: {
-					owner: [this.driftClient.program.programId.toBase58()],
-					filters: [
-						{
-							memcmp: {
-								offset: '0',
-								bytes: BorshAccountsCoder.accountDiscriminator('User'),
-							},
-						},
-						{
-							memcmp: {
-								offset: '4350',
-								bytes: Uint8Array.from([1]),
-							},
-						},
-					],
-					account: [],
-				},
-			},
-			transactions: {},
-			blocks: {},
-			blocksMeta: {},
-			accountsDataSlice: [],
-			commitment: CommitmentLevel.PROCESSED,
-			entry: {},
-		};
-
-		this.stream.on('data', (chunk: any) => {
-			if (!chunk.account) {
-				return;
-			}
-			const slot = Number(chunk.account.slot);
-			this.tryUpdateUserAccount(
-				new PublicKey(chunk.account.account.pubkey).toBase58(),
-				'grpc',
-				chunk.account.account.data,
-				slot
-			);
-		});
-
-		return new Promise<void>((resolve, reject) => {
-			this.stream.write(request, (err) => {
-				if (err === null || err === undefined) {
-					resolve();
-				} else {
-					reject(err);
-				}
-			});
-		}).catch((reason) => {
-			console.error(reason);
-			throw reason;
-		});
-	}
-
-	async fetch(): Promise<void> {
-		if (this.fetchPromise) {
-			return this.fetchPromise;
-		}
-
-		this.fetchPromise = new Promise((resolver) => {
-			this.fetchPromiseResolver = resolver;
-		});
-
-		try {
-			const rpcRequestArgs = [
-				this.driftClient.program.programId.toBase58(),
-				{
-					commitment: this.commitment,
-					filters: [getUserFilter(), getUserWithOrderFilter()],
-					encoding: 'base64',
-					withContext: true,
-				},
-			];
-
-			const rpcJSONResponse: any =
-				// @ts-ignore
-				await this.driftClient.connection._rpcRequest(
-					'getProgramAccounts',
-					rpcRequestArgs
-				);
-
-			const rpcResponseAndContext: RpcResponseAndContext<
-				Array<{
-					pubkey: PublicKey;
-					account: {
-						data: [string, string];
-					};
-				}>
-			> = rpcJSONResponse.result;
-
-			const slot: number = rpcResponseAndContext.context.slot;
-
-			const programAccountSet = new Set<string>();
-			for (const programAccount of rpcResponseAndContext.value) {
-				const key = programAccount.pubkey.toString();
-				programAccountSet.add(key);
-				this.tryUpdateUserAccount(
-					key,
-					'raw',
-					programAccount.account.data,
-					slot
-				);
-				// give event loop a chance to breathe
-				await new Promise((resolve) => setTimeout(resolve, 0));
-			}
-
-			for (const key of this.usersAccounts.keys()) {
-				if (!programAccountSet.has(key)) {
-					this.usersAccounts.delete(key);
-				}
-				// give event loop a chance to breathe
-				await new Promise((resolve) => setTimeout(resolve, 0));
-			}
-		} catch (e) {
-			console.error(e);
-		} finally {
-			this.fetchPromiseResolver();
-			this.fetchPromise = undefined;
-		}
-	}
-
-	tryUpdateUserAccount(
-		key: string,
-		dataType: 'raw' | 'grpc',
-		data: string[] | Buffer | UserAccount,
-		slot: number
-	): void {
-		if (!this.mostRecentSlot || slot > this.mostRecentSlot) {
-			this.mostRecentSlot = slot;
-		}
-
-		const slotAndUserAccount = this.usersAccounts.get(key);
-		if (!slotAndUserAccount || slotAndUserAccount.slot <= slot) {
-			// Polling leads to a lot of redundant decoding, so we only decode if data is from a fresh slot
-			let buffer: Buffer;
-			if (dataType === 'raw') {
-				buffer = Buffer.from(data[0], data[1]);
-			} else {
-				buffer = data as Buffer;
-			}
-
-			const newLastActiveSlot = new BN(
-				buffer.subarray(4328, 4328 + 8),
-				undefined,
-				'le'
-			);
-			if (
-				slotAndUserAccount &&
-				slotAndUserAccount.userAccount.lastActiveSlot.gt(newLastActiveSlot)
-			) {
-				return;
-			}
-
-			const userAccount =
-				this.driftClient.program.account.user.coder.accounts.decodeUnchecked(
-					'User',
-					buffer
-				) as UserAccount;
-
-			if (userAccount.hasOpenOrder) {
-				this.usersAccounts.set(key, { slot, userAccount });
-			} else {
-				this.usersAccounts.delete(key);
-			}
-		}
-	}
-
-	public async getDLOB(slot: number): Promise<DLOB> {
-		const dlob = new DLOB();
-		for (const [key, { userAccount }] of this.usersAccounts.entries()) {
-			const userAccountPubkey = new PublicKey(key);
-			for (const order of userAccount.orders) {
-				dlob.insertOrder(order, userAccountPubkey.toBase58(), slot);
-			}
-		}
-		return dlob;
-	}
-
-	public getSlot(): number {
-		return this.mostRecentSlot ?? 0;
-	}
-
-	public async unsubscribe(): Promise<void> {
-		return new Promise<void>((resolve, reject) => {
-			this.stream.write(
-				{
-					slots: {},
-					accounts: {},
-					transactions: {},
-					blocks: {},
-					blocksMeta: {},
-					accountsDataSlice: [],
-					entry: {},
-				},
-				(err) => {
-					if (err === null || err === undefined) {
-						resolve();
-					} else {
-						reject(err);
-					}
-				}
-			);
-		}).catch((reason) => {
-			console.error(reason);
-			throw reason;
-		});
-	}
-}

src/publishers/dlobPublisher.ts

@@ -20,6 +20,7 @@ import {
 	PhoenixSubscriber,
 	MarketType,
 	OraclePriceData,
+	OrderSubscriberConfig,
 } from '@drift-labs/sdk';
 import { RedisClient, RedisClientPrefix } from '@drift/common/clients';
 
@@ -37,12 +38,10 @@ import {
 } from '../dlob-subscriber/DLOBSubscriberIO';
 import {
 	DLOBProvider,
-	getDLOBProviderFromGrpcOrderSubscriber,
 	getDLOBProviderFromOrderSubscriber,
 	getDLOBProviderFromUserMap,
 } from '../dlobProvider';
 import FEATURE_FLAGS from '../utils/featureFlags';
-import { GeyserOrderSubscriber } from '../grpc/OrderSubscriberGRPC';
 import express, { Response, Request } from 'express';
 import { handleHealthCheck } from '../core/metrics';
 import { setGlobalDispatcher, Agent } from 'undici';
@@ -76,6 +75,17 @@ const dlobSlotGauge = new Gauge({
 		'redisClient',
 	],
 });
+const oracleSlotGauge = new Gauge({
+	name: 'oracle_slot',
+	help: 'Last updated slot of oracle',
+	labelNames: [
+		'marketIndex',
+		'marketType',
+		'marketName',
+		'redisPrefix',
+		'redisClient',
+	],
+});
 
 //@ts-ignore
 const sdkConfig = initialize({ env: process.env.ENV });
@@ -116,6 +126,10 @@ const SPOT_MARKETS_TO_LOAD =
 
 logger.info(`RPC endpoint: ${endpoint}`);
 logger.info(`WS endpoint:  ${wsEndpoint}`);
+logger.info(`Token:        ${token}`);
+logger.info(
+	`useOrderSubscriber: ${useOrderSubscriber}, useWebsocket: ${useWebsocket}, useGrpc: ${useGrpc}`
+);
 logger.info(`DriftEnv:     ${driftEnv}`);
 logger.info(`Commit:       ${commitHash}`);
 
@@ -377,17 +391,38 @@ const main = async () => {
 
 	let dlobProvider: DLOBProvider;
 	if (useOrderSubscriber) {
-		let subscriptionConfig;
+		let subscriptionConfig: any = {
+			type: 'polling',
+			commitment: stateCommitment,
+			frequency: ORDERBOOK_UPDATE_INTERVAL,
+		};
+
 		if (useWebsocket) {
 			subscriptionConfig = {
 				type: 'websocket',
 				commitment: stateCommitment,
 			};
-		} else {
+		}
+
+		// USE_GRPC=true will override websocket
+		if (useGrpc) {
+			if (!token) {
+				throw new Error('TOKEN is required for grpc');
+			}
+			if (!endpoint) {
+				throw new Error('ENDPOINT is required for grpc');
+			}
+			if (useWebsocket) {
+				logger.warn('USE_GRPC overriding USE_WEBSOCKET');
+			}
 			subscriptionConfig = {
-				type: 'polling',
+				type: 'grpc',
+				grpcConfigs: {
+					endpoint: endpoint,
+					token: token,
+					commitmentLevel: stateCommitment,
+				},
 				commitment: stateCommitment,
-				frequency: ORDERBOOK_UPDATE_INTERVAL,
 			};
 		}
 
@@ -401,17 +436,6 @@ const main = async () => {
 		slotSource = {
 			getSlot: () => orderSubscriber.getSlot(),
 		};
-	} else if (useGrpc) {
-		const grpcOrderSubscriber = new GeyserOrderSubscriber(driftClient, {
-			endpoint: endpoint,
-			token: token,
-		});
-
-		dlobProvider = getDLOBProviderFromGrpcOrderSubscriber(grpcOrderSubscriber);
-
-		slotSource = {
-			getSlot: () => grpcOrderSubscriber.getSlot(),
-		};
 	} else {
 		const userMap = new UserMap({
 			driftClient,
@@ -463,6 +487,9 @@ const main = async () => {
 	setInterval(() => {
 		const slot = slotSource.getSlot();
 		perpMarketInfos.forEach((market) => {
+			const oracleDataAndSlot = driftClient.getOracleDataForPerpMarket(
+				market.marketIndex
+			);
 			dlobSlotGauge.set(
 				{
 					marketIndex: market.marketIndex,
@@ -473,8 +500,21 @@ const main = async () => {
 				},
 				slot
 			);
+			oracleSlotGauge.set(
+				{
+					marketIndex: market.marketIndex,
+					marketType: 'perp',
+					marketName: market.marketName,
+					redisClient: REDIS_CLIENT,
+					redisPrefix: RedisClientPrefix[REDIS_CLIENT],
+				},
+				oracleDataAndSlot.slot.toNumber()
+			);
 		});
 		spotMarketInfos.forEach((market) => {
+			const oracleDataAndSlot = driftClient.getOracleDataForSpotMarket(
+				market.marketIndex
+			);
 			dlobSlotGauge.set(
 				{
 					marketIndex: market.marketIndex,
@@ -483,7 +523,7 @@ const main = async () => {
 					redisClient: REDIS_CLIENT,
 					redisPrefix: RedisClientPrefix[REDIS_CLIENT],
 				},
-				slot
+				oracleDataAndSlot.slot.toNumber()
 			);
 		});
 	}, 10_000);