Merge pull request #234 from drift-labs/master

bump common + nick's script changes

package.json

@@ -20,6 +20,7 @@
 		"@types/ws": "^8.5.8",
 		"async-mutex": "^0.4.0",
 		"axios": "^1.6.2",
+		"bottleneck": "^2.19.5",
 		"commander": "^9.4.0",
 		"compression": "^1.7.4",
 		"cors": "^2.8.5",

src/index.ts

@@ -657,7 +657,10 @@ const main = async (): Promise<void> => {
 				return;
 			}
 
-			const redisClient = perpMarketRedisMap.get(marketIndex).client;
+			const redisClient = redisClients.find(
+				(client) =>
+					client.forceGetClient().options.keyPrefix === RedisClientPrefix.DLOB
+			);
 
 			const redisResponseGainers = await redisClient.getRaw(
 				`perp_market_${marketIndex}_gainers`

src/scripts/publishUnsettledPnlUsers.ts

@@ -1,4 +1,4 @@
-import { RedisClient, RedisClientPrefix } from '@drift/common';
+import { RedisClient, RedisClientPrefix, sleep } from '@drift/common';
 import {
 	BigNum,
 	DriftClient,
@@ -11,11 +11,11 @@ import {
 	Wallet,
 	ZERO,
 	calculateClaimablePnl,
-	calculatePositionPNL,
 	decodeUser,
 } from '@drift-labs/sdk';
 import { Connection, Keypair } from '@solana/web3.js';
 import { logger } from '../utils/logger';
+import Bottleneck from 'bottleneck';
 
 const dotenv = require('dotenv');
 dotenv.config();
@@ -29,10 +29,19 @@ type UserPnlMap = {
 	};
 };
 
+type AllPnlUsers = {
+	[perpMarketIndex: number]: {
+		users: { userPubKey: string; pnl: number }[];
+	};
+};
+
 const startTime = Date.now();
 const endpoint = process.env.ENDPOINT as string;
 const wsEndpoint = process.env.WS_ENDPOINT as string;
 const driftEnv = process.env.ENV as string;
+const chunkSize = Number(process.env.CHUNK_SIZE) || 100;
+const sleepTimeMs = Number(process.env.SLEEP_TIME_MS) || 500;
+const delayMs = Number(process.env.DEFAULT_DELAY_MS) || 100;
 
 const connection = new Connection(endpoint, {
 	commitment: 'confirmed',
@@ -45,6 +54,8 @@ const driftClient = new DriftClient({
 });
 
 const userMapRedisClient = new RedisClient({
+	host: process.env.ELASTICACHE_USERMAP_HOST ?? process.env.ELASTICACHE_HOST,
+	port: process.env.ELASTICACHE_USERMAP_PORT ?? process.env.ELASTICACHE_PORT,
 	prefix: RedisClientPrefix.USER_MAP,
 });
 
@@ -56,14 +67,41 @@ const main = async () => {
 	// get the users from usermap redis client
 	await driftClient.subscribe();
 
+	const limiter = new Bottleneck({
+		maxConcurrent: 10,
+		minTime: delayMs,
+	});
+
 	const userStrings = await userMapRedisClient.lRange('user_pubkeys', 0, -1);
 
-	const redisUsers = await Promise.all(
-		userStrings.map((userStr) => getUserFromRedis(userStr))
+	const totalCount = userStrings.length;
+	let finishedCount = 0;
+	let allNonZeroPnls = {};
+
+	logger.info(`Starting loop for total of ${totalCount} users`);
+
+	while (finishedCount < totalCount) {
+		const userChunk = userStrings.slice(
+			finishedCount,
+			finishedCount + chunkSize
 		);
 
-	// construct an object with the top 20/bottom 20 unsettled in each perp market
-	const userPnlMap = createMarketSpecificPnlLeaderboards(redisUsers);
+		limiter.schedule(async () => {
+			const redisUsers = await Promise.all(
+				userChunk.map((userStr) => getUserFromRedis(userStr))
+			);
+			// add all the nonzero pnl users from each market in chunks
+			allNonZeroPnls = buildUserMarketLists(redisUsers, allNonZeroPnls);
+		});
+
+		finishedCount += chunkSize;
+
+		logger.info(`Wrote ${finishedCount} users, sleeping for ${sleepTimeMs}ms`);
+
+		await sleep(sleepTimeMs);
+	}
+
+	const userPnlMap = createMarketPnlLeaderboards(allNonZeroPnls);
 
 	const success = await writeToDlobRedis(userPnlMap);
 
@@ -73,6 +111,12 @@ const main = async () => {
 		} in ${Date.now() - startTime} ms`
 	);
 
+	await sleep(5000);
+
+	logger.info('==== TEST PRINT LOSERS FROM BTC-PERP ====');
+	const losers = await dlobRedisClient.getRaw(`perp_market_1_losers`);
+	logger.info(losers);
+
 	process.exit();
 };
 
@@ -84,30 +128,66 @@ const writeToDlobRedis = async (userPnlMap: UserPnlMap): Promise<boolean> => {
 				const losersRedisKey = `perp_market_${perpMarketIndex}_losers`;
 
 				// write the new lists
-				await dlobRedisClient.setRaw(
+				await Promise.all([
+					dlobRedisClient.setRaw(
 						gainersRedisKey,
-					JSON.stringify(userPnlMap[perpMarketIndex].gain)
-				);
-				await dlobRedisClient.setRaw(
+						JSON.stringify(userPnlMap[perpMarketIndex].gain ?? [])
+					),
+					dlobRedisClient.setRaw(
 						losersRedisKey,
-					JSON.stringify(userPnlMap[perpMarketIndex].loss)
+						JSON.stringify(userPnlMap[perpMarketIndex].loss ?? [])
+					),
+				]).then(([_winnersResult, _losersResult]) => {
+					logger.info(
+						`Wrote ${userPnlMap[perpMarketIndex].gain.length} users to redis key: ${gainersRedisKey}`
 					);
-
-				return;
+					logger.info(
+						`Wrote ${userPnlMap[perpMarketIndex].loss.length} users to redis key: ${losersRedisKey}`
+					);
+				});
 			})
 		);
 
 		return true;
 	} catch (e) {
-		console.log('Error writing to dlob redis client: ', e);
+		logger.info(`Error writing to dlob redis client: ${e}`);
 		return false;
 	}
 };
 
-const createMarketSpecificPnlLeaderboards = (
-	redisUsers: { user: User; bufferString: string }[]
-): UserPnlMap => {
-	const pnlMap = {};
+const createMarketPnlLeaderboards = (allPnlUsers: AllPnlUsers): UserPnlMap => {
+	let pnlMap = {};
+
+	Object.keys(allPnlUsers).forEach((perpMarketIndex) => {
+		try {
+			const sortedPnls = allPnlUsers[perpMarketIndex].users.sort(
+				(a, b) => b.pnl - a.pnl
+			);
+			// store the top 20 winners and losers in each perp market
+			pnlMap[perpMarketIndex] = {
+				gain: sortedPnls.slice(0, 20),
+				loss: sortedPnls.slice(-20).reverse(),
+			};
+		} catch (e) {
+			logger.info(
+				`Error in ${PerpMarkets[driftEnv][perpMarketIndex].symbol}: ${e}`
+			);
+
+			pnlMap[perpMarketIndex] = {
+				gain: [],
+				loss: [],
+			};
+		}
+	});
+
+	return pnlMap;
+};
+
+const buildUserMarketLists = (
+	redisUsers: { user: User; bufferString: string }[],
+	allPnlUsers: AllPnlUsers
+): AllPnlUsers => {
+	let newPnlUsers = allPnlUsers;
 
 	const usdcSpotMarket = driftClient.getSpotMarketAccount(
 		QUOTE_SPOT_MARKET_INDEX
@@ -123,7 +203,7 @@ const createMarketSpecificPnlLeaderboards = (
 				perpMarket.marketIndex
 			);
 
-			const allNonZeroPnls = redisUsers.flatMap((redisUser) => {
+			const allNonZeroPnlsInMarket = redisUsers.flatMap((redisUser) => {
 				try {
 					const perpPosition = redisUser.user.getPerpPosition(
 						perpMarket.marketIndex
@@ -144,21 +224,12 @@ const createMarketSpecificPnlLeaderboards = (
 							false
 						)[0];
 
-					let marketPnl = calculatePositionPNL(
-						perpMarketAccount,
-						perpPositionWithLpSettle,
-						true,
-						oraclePriceData
-					);
-
-					if (marketPnl.gt(ZERO)) {
-						marketPnl = calculateClaimablePnl(
+					const marketPnl = calculateClaimablePnl(
 						perpMarketAccount,
 						usdcSpotMarket,
 						perpPositionWithLpSettle,
 						oraclePriceData
 					);
-					}
 
 					return marketPnl.eq(ZERO)
 						? []
@@ -175,19 +246,21 @@ const createMarketSpecificPnlLeaderboards = (
 				}
 			});
 
-			const sortedPnls = allNonZeroPnls.sort((a, b) => b.pnl - a.pnl);
-
-			// store the top 20 winners and losers in each perp market
-			pnlMap[perpMarket.marketIndex] = {
-				gain: sortedPnls.slice(0, 20),
-				loss: sortedPnls.slice(-20, -1),
+			if (!allPnlUsers[perpMarket.marketIndex]) {
+				newPnlUsers[perpMarket.marketIndex] = { users: allNonZeroPnlsInMarket };
+			} else {
+				newPnlUsers[perpMarket.marketIndex] = {
+					users: allPnlUsers[perpMarket.marketIndex].users.concat(
+						allNonZeroPnlsInMarket
+					),
 				};
+			}
 		} catch (e) {
 			logger.error(`Could not fetch PnLs for ${perpMarket.symbol}: `, e);
 		}
 	}
 
-	return pnlMap;
+	return newPnlUsers;
 };
 
 const getUserFromRedis = async (userAccountStr: string) => {