Merge pull request #433 from drift-labs/jack/dlob-offload

Jack/dlob offload

package.json

@@ -5,6 +5,8 @@
 	"main": "lib/index.js",
 	"license": "Apache-2.0",
 	"dependencies": {
+		"@aws-sdk/client-kinesis": "^3.830.0",
+		"@aws-sdk/util-retry": "^3.374.0",
 		"@coral-xyz/anchor": "^0.29.0",
 		"@drift-labs/sdk": "file:./drift-common/protocol/sdk",
 		"@drift/common": "file:./drift-common/common-ts",

src/dlob-subscriber/DLOBSubscriberIO.ts

@@ -25,6 +25,7 @@ import {
 	parsePositiveIntArray,
 } from '../utils/utils';
 import { setHealthStatus, HEALTH_STATUS } from '../core/metrics';
+import { OffloadQueue } from '../utils/offload';
 
 type wsMarketArgs = {
 	marketIndex: number;
@@ -69,12 +70,15 @@ export class DLOBSubscriberIO extends DLOBSubscriber {
 		MarketType,
 		Map<number, { slot: number; ts: number }>
 	>;
+	public offloadQueue?: ReturnType<typeof OffloadQueue>;
+	public enableOffload: boolean;
 
 	constructor(
 		config: DLOBSubscriptionConfig & {
 			env: DriftEnv;
 			redisClient: RedisClient;
 			indicativeQuotesRedisClient?: RedisClient;
+			enableOffloadQueue?: boolean;
 			perpMarketInfos: wsMarketInfo[];
 			spotMarketInfos: wsMarketInfo[];
 			spotMarketSubscribers: SubscriberLookup;
@@ -88,6 +92,11 @@ export class DLOBSubscriberIO extends DLOBSubscriber {
 		this.killSwitchSlotDiffThreshold =
 			config.killSwitchSlotDiffThreshold || 200;
 
+		this.enableOffload = config.enableOffloadQueue || false;
+		if (this.enableOffload) {
+			this.offloadQueue = OffloadQueue();
+		}
+
 		// Set up appropriate maps
 		this.lastSeenL2Formatted = new Map();
 		this.lastSeenL2Formatted.set(MarketType.SPOT, new Map());
@@ -392,6 +401,21 @@ export class DLOBSubscriberIO extends DLOBSubscriber {
 			}`,
 			l2Formatted
 		);
+
+		if (this.offloadQueue) {
+			try {
+				this.offloadQueue.addMessage(
+					Object.assign({}, l2Formatted, {
+						bids: l2Formatted.bids.slice(0, 20),
+						asks: l2Formatted.asks.slice(0, 20),
+						ts: Math.floor(new Date().getTime() / 1000),
+					})
+				);
+			} catch (error) {
+				logger.error('Error adding message to offload queue:', error);
+			}
+		}
+
 		this.redisClient.set(
 			`last_update_orderbook_${marketType}_${marketArgs.marketIndex}${
 				this.indicativeQuotesRedisClient ? '_indicative' : ''
@@ -476,6 +500,19 @@ export class DLOBSubscriberIO extends DLOBSubscriber {
 			marketArgs.marketIndex
 		);
 
+		if (this.offloadQueue) {
+			try {
+				this.offloadQueue.addMessage(
+					Object.assign({}, l3, {
+						ts: Math.floor(new Date().getTime() / 1000),
+					}),
+					{ eventType: 'DLOBL3Snapshot', throttle: true }
+				);
+			} catch (error) {
+				logger.error('Error adding message to offload queue:', error);
+			}
+		}
+
 		this.redisClient.set(
 			`last_update_orderbook_l3_${marketType}_${marketArgs.marketIndex}${
 				this.indicativeQuotesRedisClient ? '_indicative' : ''

src/publishers/dlobPublisher.ts

@@ -125,6 +125,8 @@ const SPOT_MARKETS_TO_LOAD =
 		? parsePositiveIntArray(process.env.SPOT_MARKETS_TO_LOAD)
 		: undefined;
 
+const enableOffloadQueue = process.env.ENABLE_OFFLOAD === 'true';
+
 logger.info(`RPC endpoint:  ${endpoint}`);
 logger.info(`WS endpoint:   ${wsEndpoint}`);
 logger.info(`GRPC endpoint: ${grpcEndpoint}`);
@@ -522,6 +524,7 @@ const main = async () => {
 		killSwitchSlotDiffThreshold: KILLSWITCH_SLOT_DIFF_THRESHOLD,
 		protectedMakerView: false,
 		indicativeQuotesRedisClient: indicativeRedisClient,
+		enableOffloadQueue
 	});
 	await dlobSubscriberIndicative.subscribe();
 

src/utils/offload.ts

@@ -0,0 +1,156 @@
+import {
+	KinesisClient,
+	PutRecordsCommand,
+	PutRecordsRequestEntry,
+} from '@aws-sdk/client-kinesis';
+import { ConfiguredRetryStrategy } from '@aws-sdk/util-retry';
+import { logger } from '../utils/logger';
+
+type EventType = 'DLOBSnapshot' | 'DLOBL3Snapshot';
+
+interface QueueMessage {
+	data: any;
+}
+
+interface ThrottleInfo {
+	lastSentTime: number;
+}
+
+const kinesis = new KinesisClient({
+	retryStrategy: new ConfiguredRetryStrategy(
+		5,
+		(attempt: number) => 100 + attempt * 500
+	),
+});
+
+const kinesisStream = process.env.KINESIS_STREAM_NAME;
+const batchInterval = process.env.KINESIS_BATCH_INTERVAL
+	? parseInt(process.env.KINESIS_BATCH_INTERVAL)
+	: 10000;
+
+export const OffloadQueue = () => {
+	const queue: QueueMessage[] = [];
+	const throttleMap: Map<string, ThrottleInfo> = new Map();
+	let isProcessing = false;
+
+	setInterval(async () => {
+		if (!isProcessing && queue.length > 0) {
+			await processQueue();
+		}
+	}, batchInterval);
+
+	const processQueue = async (): Promise<void> => {
+		if (isProcessing || queue.length === 0) {
+			return;
+		}
+
+		isProcessing = true;
+
+		try {
+			const batch = queue.splice(0, 500); // Kinesis limit
+
+			const records: PutRecordsRequestEntry[] = batch.map((msg) => ({
+				Data: Buffer.from(
+					JSON.stringify({
+						...msg.data,
+						source: 'dlob',
+					})
+				),
+				PartitionKey: Date.now().toString(),
+			}));
+
+			const command = new PutRecordsCommand({
+				StreamName: kinesisStream,
+				Records: records,
+			});
+
+			const response = await kinesis.send(command);
+
+			if (response.FailedRecordCount && response.FailedRecordCount > 0) {
+				logger.warn(`${response.FailedRecordCount} records failed to send`);
+
+				const failedMessages: QueueMessage[] = [];
+				response.Records?.forEach((record, index) => {
+					if (record.ErrorCode) {
+						failedMessages.push(batch[index]);
+						logger.error(
+							`Record ${index} failed: ${record.ErrorCode} - ${record.ErrorMessage}`
+						);
+					}
+				});
+
+				queue.unshift(...failedMessages);
+			}
+
+			const successCount = batch.length - (response.FailedRecordCount || 0);
+			logger.info(
+				`Successfully sent ${successCount} records to Kinesis stream: ${kinesisStream}`
+			);
+		} catch (error) {
+			logger.error('Error processing queue:', error);
+		} finally {
+			isProcessing = false;
+		}
+	};
+
+	const shouldSendThrottled = (
+		marketKey: string,
+		currentTime: number
+	): boolean => {
+		const throttleInfo = throttleMap.get(marketKey);
+		const lastSent = throttleInfo?.lastSentTime || 0;
+		const timeSinceLastSent = currentTime - lastSent;
+
+		if (timeSinceLastSent >= 60000) {
+			throttleMap.set(marketKey, { lastSentTime: currentTime });
+			return true;
+		}
+
+		if (timeSinceLastSent >= 45000) {
+			const timeInWindow = timeSinceLastSent - 45000;
+			const probability = Math.min(timeInWindow / 22500, 0.8);
+
+			if (Math.random() < probability) {
+				throttleMap.set(marketKey, { lastSentTime: currentTime });
+				return true;
+			}
+		}
+
+		return false;
+	};
+
+	const addMessage = (
+		data: any,
+		options: {
+			eventType?: EventType;
+			throttle?: boolean;
+		} = {}
+	) => {
+		const { eventType = 'DLOBSnapshot', throttle = false } = options;
+		const currentTime = Date.now();
+
+		if (!throttle) {
+			queue.push({
+				data: {
+					...data,
+					eventType,
+				},
+			});
+			return;
+		}
+
+		const marketKey = `${data.marketType}_${data.marketIndex}`;
+		if (shouldSendThrottled(marketKey, currentTime)) {
+			queue.push({
+				data: {
+					...data,
+					eventType,
+				},
+			});
+		}
+	};
+
+	return {
+		addMessage,
+	};
+};