From 83e28e1cd6f43404989eede919a1c41c161c63f7 Mon Sep 17 00:00:00 2001
From: Nour Alharithi <nour.alharithi@gmail.com>
Date: Mon, 4 Dec 2023 10:51:09 -0800
Subject: [PATCH] adding new spot markets

---
 src/dlob-subscriber/DLOBSubscriberIO.ts | 18 +++++++--
 src/publishers/dlobPublisher.ts         | 52 ++++++++++++++++++++++++-
 2 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/src/dlob-subscriber/DLOBSubscriberIO.ts b/src/dlob-subscriber/DLOBSubscriberIO.ts
index 039c738..a7cea28 100644
--- a/src/dlob-subscriber/DLOBSubscriberIO.ts
+++ b/src/dlob-subscriber/DLOBSubscriberIO.ts
@@ -13,7 +13,11 @@ import {
 } from '@drift-labs/sdk';
 import { driftEnv } from '../publishers/dlobPublisher';
 import { RedisClient } from '../utils/redisClient';
-import { addOracletoResponse, l2WithBNToStrings } from '../utils/utils';
+import {
+	SubscriberLookup,
+	addOracletoResponse,
+	l2WithBNToStrings,
+} from '../utils/utils';
 
 type wsMarketL2Args = {
 	marketIndex: number;
@@ -32,7 +36,12 @@ export class DLOBSubscriberIO extends DLOBSubscriber {
 	public lastSeenL2Formatted: Map<MarketType, Map<number, any>>;
 	redisClient: RedisClient;
 
-	constructor(config: DLOBSubscriptionConfig & { redisClient: RedisClient }) {
+	constructor(
+		config: DLOBSubscriptionConfig & {
+			redisClient: RedisClient;
+			spotMarketSubscribers: SubscriberLookup;
+		}
+	) {
 		super(config);
 		this.redisClient = config.redisClient;
 
@@ -67,7 +76,10 @@ export class DLOBSubscriberIO extends DLOBSubscriber {
 				depth: -1,
 				includeVamm: false,
 				updateOnChange: true,
-				fallbackL2Generators: [],
+				fallbackL2Generators: [
+					config.spotMarketSubscribers[market.marketIndex].phoenix,
+					config.spotMarketSubscribers[market.marketIndex].serum,
+				].filter((a) => !!a),
 			});
 		}
 	}
diff --git a/src/publishers/dlobPublisher.ts b/src/publishers/dlobPublisher.ts
index 0550c72..2d2aca8 100644
--- a/src/publishers/dlobPublisher.ts
+++ b/src/publishers/dlobPublisher.ts
@@ -16,7 +16,12 @@ import {
 } from '@drift-labs/sdk';
 
 import { logger, setLogLevel } from '../utils/logger';
-import { sleep } from '../utils/utils';
+import {
+	SubscriberLookup,
+	getPhoenixSubscriber,
+	getSerumSubscriber,
+	sleep,
+} from '../utils/utils';
 import { DLOBSubscriberIO } from '../dlob-subscriber/DLOBSubscriberIO';
 import { RedisClient } from '../utils/redisClient';
 import {
@@ -53,6 +58,41 @@ logger.info(`WS endpoint:  ${wsEndpoint}`);
 logger.info(`DriftEnv:     ${driftEnv}`);
 logger.info(`Commit:       ${commitHash}`);
 
+let MARKET_SUBSCRIBERS: SubscriberLookup = {};
+
+const initializeAllMarketSubscribers = async (driftClient: DriftClient) => {
+	const markets: SubscriberLookup = {};
+
+	for (const market of sdkConfig.SPOT_MARKETS) {
+		markets[market.marketIndex] = {
+			phoenix: undefined,
+			serum: undefined,
+		};
+
+		if (market.phoenixMarket) {
+			const phoenixSubscriber = getPhoenixSubscriber(
+				driftClient,
+				market,
+				sdkConfig
+			);
+			await phoenixSubscriber.subscribe();
+			markets[market.marketIndex].phoenix = phoenixSubscriber;
+		}
+
+		if (market.serumMarket) {
+			const serumSubscriber = getSerumSubscriber(
+				driftClient,
+				market,
+				sdkConfig
+			);
+			await serumSubscriber.subscribe();
+			markets[market.marketIndex].serum = serumSubscriber;
+		}
+	}
+
+	return markets;
+};
+
 const main = async () => {
 	const wallet = new Wallet(new Keypair());
 	const clearingHousePublicKey = new PublicKey(sdkConfig.DRIFT_PROGRAM_ID);
@@ -120,6 +160,15 @@ const main = async () => {
 		logger.error(e);
 	});
 
+	logger.info(`Initializing all market subscribers...`);
+	const initAllMarketSubscribersStart = Date.now();
+	MARKET_SUBSCRIBERS = await initializeAllMarketSubscribers(driftClient);
+	logger.info(
+		`All market subscribers initialized in ${
+			Date.now() - initAllMarketSubscribersStart
+		} ms`
+	);
+
 	let dlobProvider: DLOBProvider;
 	if (useOrderSubscriber) {
 		let subscriptionConfig;
@@ -167,6 +216,7 @@ const main = async () => {
 		slotSource,
 		updateFrequency: ORDERBOOK_UPDATE_INTERVAL,
 		redisClient,
+		spotMarketSubscribers: MARKET_SUBSCRIBERS,
 	});
 	await dlobSubscriber.subscribe();
 	if (useWebsocket) {