From 669d4a48fdb2085c2c36a1c4a199581a15646dd1 Mon Sep 17 00:00:00 2001
From: Nour Alharithi <nour.alharithi@gmail.com>
Date: Tue, 3 Dec 2024 14:33:21 -0800
Subject: [PATCH] server lite uses new rotation

---
 src/serverLite.ts | 120 +++++++---------------------------------------
 1 file changed, 16 insertions(+), 104 deletions(-)

diff --git a/src/serverLite.ts b/src/serverLite.ts
index f79044c..7fd921f 100644
--- a/src/serverLite.ts
+++ b/src/serverLite.ts
@@ -10,7 +10,6 @@ import {
 	DriftEnv,
 	SlotSubscriber,
 	initialize,
-	isVariant,
 	MarketType,
 	getVariant,
 } from '@drift-labs/sdk';
@@ -26,7 +25,7 @@ import {
 	runtimeSpecsGauge,
 } from './core/metrics';
 import { handleResponseTime } from './core/middleware';
-import { errorHandler, sleep } from './utils/utils';
+import { errorHandler, selectMostRecentBySlot, sleep } from './utils/utils';
 import { setGlobalDispatcher, Agent } from 'undici';
 
 setGlobalDispatcher(
@@ -60,8 +59,7 @@ const serverPort = process.env.PORT || 6969;
 export const ORDERBOOK_UPDATE_INTERVAL = 1000;
 const WS_FALLBACK_FETCH_INTERVAL = ORDERBOOK_UPDATE_INTERVAL * 60;
 const SLOT_STALENESS_TOLERANCE =
-	parseInt(process.env.SLOT_STALENESS_TOLERANCE) || 35;
-const ROTATION_COOLDOWN = parseInt(process.env.ROTATION_COOLDOWN) || 5000;
+	parseInt(process.env.SLOT_STALENESS_TOLERANCE) || 100000;
 
 const logFormat =
 	':remote-addr - :remote-user [:date[clf]] ":method :url HTTP/:http-version" :status :res[content-length] ":referrer" ":user-agent" :req[x-forwarded-for]';
@@ -132,98 +130,20 @@ const main = async (): Promise<void> => {
 
 	// Handle redis client initialization and rotation maps
 	const redisClients: Array<RedisClient> = [];
-	const spotMarketRedisMap: Map<
-		number,
-		{
-			client: RedisClient;
-			clientIndex: number;
-			lastRotationTime: number;
-			lock: boolean;
-		}
-	> = new Map();
-	const perpMarketRedisMap: Map<
-		number,
-		{
-			client: RedisClient;
-			clientIndex: number;
-			lastRotationTime: number;
-			lock: boolean;
-		}
-	> = new Map();
 	logger.info('Connecting to redis');
 	for (let i = 0; i < REDIS_CLIENTS.length; i++) {
 		redisClients.push(new RedisClient({ prefix: REDIS_CLIENTS[i] }));
 	}
 
-	for (let i = 0; i < sdkConfig.SPOT_MARKETS.length; i++) {
-		spotMarketRedisMap.set(sdkConfig.SPOT_MARKETS[i].marketIndex, {
-			client: redisClients[0],
-			clientIndex: 0,
-			lastRotationTime: 0,
-			lock: false,
-		});
-	}
-	for (let i = 0; i < sdkConfig.PERP_MARKETS.length; i++) {
-		perpMarketRedisMap.set(sdkConfig.PERP_MARKETS[i].marketIndex, {
-			client: redisClients[0],
-			clientIndex: 0,
-			lastRotationTime: 0,
-			lock: false,
-		});
-	}
-
-	function canRotate(marketType: MarketType, marketIndex: number) {
-		if (isVariant(marketType, 'spot')) {
-			const state = spotMarketRedisMap.get(marketIndex);
-			if (state) {
-				const now = Date.now();
-				if (now - state.lastRotationTime > ROTATION_COOLDOWN && !state.lock) {
-					state.lastRotationTime = now;
-					return true;
-				}
-			}
-		} else {
-			const state = perpMarketRedisMap.get(marketIndex);
-			if (state) {
-				const now = Date.now();
-				if (now - state.lastRotationTime > ROTATION_COOLDOWN && !state.lock) {
-					state.lastRotationTime = now;
-					return true;
-				}
-			}
-		}
-		return false;
-	}
-
-	function rotateClient(marketType: MarketType, marketIndex: number) {
-		if (isVariant(marketType, 'spot')) {
-			const state = spotMarketRedisMap.get(marketIndex);
-			if (state) {
-				state.lock = true;
-				const nextClientIndex = (state.clientIndex + 1) % redisClients.length;
-				state.client = redisClients[nextClientIndex];
-				state.clientIndex = nextClientIndex;
-				logger.info(
-					`Rotated redis client to index ${nextClientIndex} for spot market ${marketIndex}`
-				);
-				state.lastRotationTime = Date.now();
-				state.lock = false;
-			}
-		} else {
-			const state = perpMarketRedisMap.get(marketIndex);
-			if (state) {
-				state.lock = true;
-				const nextClientIndex = (state.clientIndex + 1) % redisClients.length;
-				state.client = redisClients[nextClientIndex];
-				state.clientIndex = nextClientIndex;
-				logger.info(
-					`Rotated redis client to index ${nextClientIndex} for perp market ${marketIndex}`
-				);
-				state.lastRotationTime = Date.now();
-				state.lock = false;
-			}
-		}
-	}
+	const fetchFromRedis = async (
+		key: string,
+		selectionCriteria: (responses: any) => any
+	): Promise<JSON> => {
+		const redisResponses = await Promise.all(
+			redisClients.map((client) => client.getRaw(key))
+		);
+		return selectionCriteria(redisResponses);
+	};
 
 	const handleStartup = async (_req, res, _next) => {
 		if (slotSubscriber.currentSlot && !redisClients.some((c) => !c.connected)) {
@@ -253,32 +173,24 @@ const main = async (): Promise<void> => {
 			const normedMarketIndex = parseInt(marketIndex as string);
 			const normedMarketType = isSpot ? MarketType.SPOT : MarketType.PERP;
 
-			const redisClient = (
-				isSpot ? spotMarketRedisMap : perpMarketRedisMap
-			).get(normedMarketIndex).client;
-			const redisL3 = await redisClient.getRaw(
+			const redisL3 = await fetchFromRedis(
 				`last_update_orderbook_l3_${getVariant(
 					normedMarketType
-				)}_${normedMarketIndex}`
+				)}_${normedMarketIndex}`,
+				selectMostRecentBySlot
 			);
 			if (
 				redisL3 &&
-				slotSubscriber.getSlot() - parseInt(JSON.parse(redisL3).slot) <
-					SLOT_STALENESS_TOLERANCE
+				slotSubscriber.getSlot() - redisL3['slot'] < SLOT_STALENESS_TOLERANCE
 			) {
 				cacheHitCounter.add(1, {
 					miss: false,
 					path: req.baseUrl + req.path,
 				});
 				res.writeHead(200);
-				res.end(redisL3);
+				res.end(JSON.stringify(redisL3));
 				return;
 			} else {
-				if (redisL3 && redisClients.length > 1) {
-					if (canRotate(normedMarketType, normedMarketIndex)) {
-						rotateClient(normedMarketType, normedMarketIndex);
-					}
-				}
 				cacheHitCounter.add(1, {
 					miss: true,
 					path: req.baseUrl + req.path,