chore: add 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",
@@ -79,5 +81,5 @@
 	},
 	"resolutions": {
 		"rpc-websockets": "^10.0.0"
-    }
+	}
 }

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,19 @@ 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' : ''

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,99 @@
+import {
+	KinesisClient,
+	PutRecordsCommand,
+	PutRecordsRequestEntry,
+} from '@aws-sdk/client-kinesis';
+import { ConfiguredRetryStrategy } from '@aws-sdk/util-retry';
+import { logger } from '../utils/logger';
+
+interface QueueMessage {
+	data: any;
+}
+
+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[] = [];
+	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',
+						eventType: 'DLOBSnapshot',
+					})
+				),
+				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 addMessage = (data: any): void => {
+		queue.push({
+			data,
+		});
+	};
+
+	return {
+		addMessage,
+	};
+};