Merge pull request #58 from drift-labs/wphan/dlob_publisher_limit_accounts_tracked

take in perp and spot markets from env var

src/dlob-subscriber/DLOBSubscriberIO.ts

@@ -2,16 +2,11 @@ import {
 	BN,
 	DLOBSubscriber,
 	DLOBSubscriptionConfig,
-	DevnetPerpMarkets,
-	DevnetSpotMarkets,
 	L2OrderBookGenerator,
-	MainnetPerpMarkets,
-	MainnetSpotMarkets,
 	MarketType,
 	groupL2,
 	isVariant,
 } from '@drift-labs/sdk';
-import { driftEnv } from '../publishers/dlobPublisher';
 import { RedisClient } from '../utils/redisClient';
 import {
 	SubscriberLookup,
@@ -33,6 +28,11 @@ type wsMarketL2Args = {
 	updateOnChange?: boolean;
 };
 
+export type wsMarketInfo = {
+	marketIndex: number;
+	marketName: string;
+};
+
 export class DLOBSubscriberIO extends DLOBSubscriber {
 	public marketL2Args: wsMarketL2Args[] = [];
 	public lastSeenL2Formatted: Map<MarketType, Map<number, any>>;
@@ -41,6 +41,8 @@ export class DLOBSubscriberIO extends DLOBSubscriber {
 	constructor(
 		config: DLOBSubscriptionConfig & {
 			redisClient: RedisClient;
+			perpMarketInfos: wsMarketInfo[];
+			spotMarketInfos: wsMarketInfo[];
 			spotMarketSubscribers: SubscriberLookup;
 		}
 	) {
@@ -52,17 +54,11 @@ export class DLOBSubscriberIO extends DLOBSubscriber {
 		this.lastSeenL2Formatted.set(MarketType.SPOT, new Map());
 		this.lastSeenL2Formatted.set(MarketType.PERP, new Map());
 
-		// Add all active markets to the market L2Args
-		const perpMarkets =
-			driftEnv === 'devnet' ? DevnetPerpMarkets : MainnetPerpMarkets;
-		const spotMarkets =
-			driftEnv === 'devnet' ? DevnetSpotMarkets : MainnetSpotMarkets;
-
-		for (const market of perpMarkets) {
+		for (const market of config.perpMarketInfos) {
 			this.marketL2Args.push({
 				marketIndex: market.marketIndex,
 				marketType: MarketType.PERP,
-				marketName: market.symbol,
+				marketName: market.marketName,
 				depth: -1,
 				numVammOrders: 100,
 				includeVamm: true,
@@ -70,11 +66,11 @@ export class DLOBSubscriberIO extends DLOBSubscriber {
 				fallbackL2Generators: [],
 			});
 		}
-		for (const market of spotMarkets) {
+		for (const market of config.spotMarketInfos) {
 			this.marketL2Args.push({
 				marketIndex: market.marketIndex,
 				marketType: MarketType.SPOT,
-				marketName: market.symbol,
+				marketName: market.marketName,
 				depth: -1,
 				includeVamm: false,
 				updateOnChange: true,
@@ -92,7 +88,7 @@ export class DLOBSubscriberIO extends DLOBSubscriber {
 			try {
 				this.getL2AndSendMsg(l2Args);
 			} catch (error) {
-				logger.error(error);
+				console.error(error);
 				console.log(`Error getting L2 ${l2Args.marketName}`);
 			}
 		}

src/publishers/dlobPublisher.ts

@@ -14,6 +14,9 @@ import {
 	DriftClientSubscriptionConfig,
 	SlotSubscriber,
 	isVariant,
+	OracleInfo,
+	PerpMarketConfig,
+	SpotMarketConfig,
 } from '@drift-labs/sdk';
 
 import { logger, setLogLevel } from '../utils/logger';
@@ -21,9 +24,13 @@ import {
 	SubscriberLookup,
 	getPhoenixSubscriber,
 	getSerumSubscriber,
+	parsePositiveIntArray,
 	sleep,
 } from '../utils/utils';
-import { DLOBSubscriberIO } from '../dlob-subscriber/DLOBSubscriberIO';
+import {
+	DLOBSubscriberIO,
+	wsMarketInfo,
+} from '../dlob-subscriber/DLOBSubscriberIO';
 import { RedisClient } from '../utils/redisClient';
 import {
 	DLOBProvider,
@@ -64,6 +71,16 @@ const ORDERBOOK_UPDATE_INTERVAL =
 	parseInt(process.env.ORDERBOOK_UPDATE_INTERVAL) || 1000;
 const WS_FALLBACK_FETCH_INTERVAL = 10_000;
 
+// comma separated list of perp market indexes to load: i.e. 0,1,2,3
+const PERP_MARKETS_TO_LOAD = process.env.PERP_MARKETS_TO_LOAD
+	? parsePositiveIntArray(process.env.PERP_MARKETS_TO_LOAD)
+	: undefined;
+
+// comma separated list of spot market indexes to load: i.e. 0,1,2,3
+const SPOT_MARKETS_TO_LOAD = process.env.SPOT_MARKETS_TO_LOAD
+	? parsePositiveIntArray(process.env.SPOT_MARKETS_TO_LOAD)
+	: undefined;
+
 logger.info(`RPC endpoint: ${endpoint}`);
 logger.info(`WS endpoint:  ${wsEndpoint}`);
 logger.info(`DriftEnv:     ${driftEnv}`);
@@ -71,22 +88,117 @@ logger.info(`Commit:       ${commitHash}`);
 
 let MARKET_SUBSCRIBERS: SubscriberLookup = {};
 
+const getMarketsAndOraclesToLoad = (
+	sdkConfig: any
+): {
+	perpMarketInfos: wsMarketInfo[];
+	spotMarketInfos: wsMarketInfo[];
+	oracleInfos?: OracleInfo[];
+} => {
+	const oracleInfos: OracleInfo[] = [];
+	const oraclesTracked = new Set();
+	const perpMarketInfos: wsMarketInfo[] = [];
+	const spotMarketInfos: wsMarketInfo[] = [];
+
+	// only watch all markets if neither env vars are specified
+	const noMarketsSpecified = !PERP_MARKETS_TO_LOAD && !SPOT_MARKETS_TO_LOAD;
+
+	let perpIndexes = PERP_MARKETS_TO_LOAD;
+	if (!perpIndexes) {
+		if (noMarketsSpecified) {
+			perpIndexes = sdkConfig.PERP_MARKETS.map((m) => m.marketIndex);
+		} else {
+			perpIndexes = [];
+		}
+	}
+	let spotIndexes = SPOT_MARKETS_TO_LOAD;
+	if (!spotIndexes) {
+		if (noMarketsSpecified) {
+			spotIndexes = sdkConfig.SPOT_MARKETS.map((m) => m.marketIndex);
+		} else {
+			spotIndexes = [];
+		}
+	}
+
+	if (perpIndexes.length > 0) {
+		for (const idx of perpIndexes) {
+			const perpMarketConfig = sdkConfig.PERP_MARKETS[idx] as PerpMarketConfig;
+			if (!perpMarketConfig) {
+				throw new Error(`Perp market config for ${idx} not found`);
+			}
+			const oracleKey = perpMarketConfig.oracle.toBase58();
+			if (!oraclesTracked.has(oracleKey)) {
+				logger.info(`Tracking oracle ${oracleKey} for perp market ${idx}`);
+				oracleInfos.push({
+					publicKey: perpMarketConfig.oracle,
+					source: perpMarketConfig.oracleSource,
+				});
+				oraclesTracked.add(oracleKey);
+			}
+			perpMarketInfos.push({
+				marketIndex: perpMarketConfig.marketIndex,
+				marketName: perpMarketConfig.symbol,
+			});
+		}
+		logger.info(
+			`DlobPublisher tracking perp markets: ${JSON.stringify(perpMarketInfos)}`
+		);
+	}
+
+	if (spotIndexes.length > 0) {
+		for (const idx of spotIndexes) {
+			const spotMarketConfig = sdkConfig.SPOT_MARKETS[idx] as SpotMarketConfig;
+			if (!spotMarketConfig) {
+				throw new Error(`Spot market config for ${idx} not found`);
+			}
+			const oracleKey = spotMarketConfig.oracle.toBase58();
+			if (!oraclesTracked.has(oracleKey)) {
+				logger.info(`Tracking oracle ${oracleKey} for spot market ${idx}`);
+				oracleInfos.push({
+					publicKey: spotMarketConfig.oracle,
+					source: spotMarketConfig.oracleSource,
+				});
+				oraclesTracked.add(oracleKey);
+			}
+			spotMarketInfos.push({
+				marketIndex: spotMarketConfig.marketIndex,
+				marketName: spotMarketConfig.symbol,
+			});
+		}
+		logger.info(
+			`DlobPublisher tracking spot markets: ${JSON.stringify(spotMarketInfos)}`
+		);
+	}
+
+	return {
+		perpMarketInfos,
+		spotMarketInfos,
+		oracleInfos,
+	};
+};
+
 const initializeAllMarketSubscribers = async (driftClient: DriftClient) => {
 	const markets: SubscriberLookup = {};
 
-	for (const market of sdkConfig.SPOT_MARKETS) {
+	for (const market of driftClient.getSpotMarketAccounts()) {
 		markets[market.marketIndex] = {
 			phoenix: undefined,
 			serum: undefined,
 		};
+		const marketConfig = sdkConfig.SPOT_MARKETS[market.marketIndex];
 
-		if (market.phoenixMarket) {
+		if (marketConfig.phoenixMarket) {
 			const phoenixConfigAccount =
-				await driftClient.getPhoenixV1FulfillmentConfig(market.phoenixMarket);
+				await driftClient.getPhoenixV1FulfillmentConfig(
+					marketConfig.phoenixMarket
+				);
 			if (isVariant(phoenixConfigAccount.status, 'enabled')) {
+				logger.info(
+					`Loading phoenix subscriber for spot market ${market.marketIndex}`
+				);
 				const phoenixSubscriber = getPhoenixSubscriber(
 					driftClient,
-					market,
+					marketConfig,
 					sdkConfig
 				);
 				await phoenixSubscriber.subscribe();
@@ -103,14 +215,17 @@ const initializeAllMarketSubscribers = async (driftClient: DriftClient) => {
 			}
 		}
 
-		if (market.serumMarket) {
+		if (marketConfig.serumMarket) {
 			const serumConfigAccount = await driftClient.getSerumV3FulfillmentConfig(
-				market.serumMarket
+				marketConfig.serumMarket
 			);
 			if (isVariant(serumConfigAccount.status, 'enabled')) {
+				logger.info(
+					`Loading serum subscriber for spot market ${market.marketIndex}`
+				);
 				const serumSubscriber = getSerumSubscriber(
 					driftClient,
-					market,
+					marketConfig,
 					sdkConfig
 				);
 				await serumSubscriber.subscribe();
@@ -177,12 +292,17 @@ const main = async () => {
 		};
 	}
 
+	const { perpMarketInfos, spotMarketInfos, oracleInfos } =
+		getMarketsAndOraclesToLoad(sdkConfig);
 	driftClient = new DriftClient({
 		connection,
 		wallet,
 		programID: clearingHousePublicKey,
 		accountSubscription,
 		env: driftEnv,
+		perpMarketIndexes: perpMarketInfos.map((m) => m.marketIndex),
+		spotMarketIndexes: spotMarketInfos.map((m) => m.marketIndex),
+		oracleInfos,
 	});
 
 	const lamportsBalance = await connection.getBalance(wallet.publicKey);
@@ -271,6 +391,8 @@ const main = async () => {
 		updateFrequency: ORDERBOOK_UPDATE_INTERVAL,
 		redisClient,
 		spotMarketSubscribers: MARKET_SUBSCRIBERS,
+		perpMarketInfos,
+		spotMarketInfos,
 	});
 	await dlobSubscriber.subscribe();
 	if (useWebsocket && !FEATURE_FLAGS.DISABLE_GPA_REFRESH) {

src/utils/utils.ts

@@ -39,6 +39,17 @@ export function sleep(ms: number): Promise<void> {
 	return new Promise((resolve) => setTimeout(resolve, ms));
 }
 
+export function parsePositiveIntArray(
+	intArray: string,
+	separator = ','
+): number[] {
+	return intArray
+		.split(separator)
+		.map((s) => s.trim())
+		.map((s) => parseInt(s))
+		.filter((n) => !isNaN(n) && n > 0);
+}
+
 export const getOracleForMarket = (
 	driftClient: DriftClient,
 	marketType: MarketType,