add ws redis switchover with hysteresis

src/wsConnectionManager.ts

@@ -6,7 +6,161 @@ import { WebSocket, WebSocketServer } from 'ws';
 import { sleep, selectMostRecentBySlot, GROUPING_OPTIONS } from './utils/utils';
 import { DriftEnv, PerpMarkets, SpotMarkets } from '@drift-labs/sdk';
 import { RedisClient, RedisClientPrefix } from '@drift/common/clients';
-import { Metrics } from './core/metricsV2';
+import { Metrics, GaugeValue } from './core/metricsV2';
+
+// Stream selector with hysteresis to prevent flip-flopping between sources
+const STREAM_SWITCH_THRESHOLD_MS = 10_000; // 10 seconds
+const STREAM_STALE_THRESHOLD_MS = 3_000; // Consider stream stale if no message for 3 seconds
+
+class StreamSelector {
+	private activeStream: string | null = null;
+	private streamLastMessageTime: Map<string, number> = new Map();
+	private streamLastSlot: Map<string, number> = new Map();
+	private streamLeadingSince: Map<string, number> = new Map();
+	private streamHealthyGauge: GaugeValue;
+	private activeStreamGauge: GaugeValue;
+	private streams: string[];
+
+	constructor(
+		streams: string[],
+		streamHealthyGauge: GaugeValue,
+		activeStreamGauge: GaugeValue
+	) {
+		this.streams = streams;
+		this.streamHealthyGauge = streamHealthyGauge;
+		this.activeStreamGauge = activeStreamGauge;
+
+		// Initialize all streams as unhealthy
+		for (const stream of streams) {
+			this.streamHealthyGauge.setLatestValue(0, { source: stream });
+			this.activeStreamGauge.setLatestValue(0, { source: stream });
+		}
+	}
+
+	// Record a message from a stream and return whether it should be forwarded
+	recordMessage(stream: string, slot: number): boolean {
+		const now = Date.now();
+		this.streamLastMessageTime.set(stream, now);
+		this.streamLastSlot.set(stream, slot);
+
+		this.streamHealthyGauge.setLatestValue(1, { source: stream });
+
+		if (this.activeStream === null) {
+			this.setActiveStream(stream);
+			return true;
+		}
+
+		// If this is from the active stream, forward it
+		if (stream === this.activeStream) {
+			// Reset leading times for other streams since active stream is still producing
+			for (const s of this.streams) {
+				if (s !== stream) {
+					this.streamLeadingSince.delete(s);
+				}
+			}
+			return true;
+		}
+
+		// This message is from a non-active stream
+		const activeSlot = this.streamLastSlot.get(this.activeStream) || 0;
+		const activeLastMessage =
+			this.streamLastMessageTime.get(this.activeStream) || 0;
+
+		// Check if active stream is stale
+		if (now - activeLastMessage > STREAM_STALE_THRESHOLD_MS) {
+			console.log(
+				`Active stream ${this.activeStream} is stale (no message for ${
+					now - activeLastMessage
+				}ms), switching to ${stream}`
+			);
+			this.setActiveStream(stream);
+			return true;
+		}
+
+		// Check if this stream has a more recent slot
+		if (slot > activeSlot) {
+			// Track how long this stream has been leading
+			if (!this.streamLeadingSince.has(stream)) {
+				this.streamLeadingSince.set(stream, now);
+			}
+
+			const leadingDuration = now - this.streamLeadingSince.get(stream)!;
+			if (leadingDuration >= STREAM_SWITCH_THRESHOLD_MS) {
+				console.log(
+					`Stream ${stream} has been leading for ${leadingDuration}ms, switching from ${this.activeStream}`
+				);
+				this.setActiveStream(stream);
+				return true;
+			}
+		} else {
+			// This stream is not leading, reset its leading time
+			this.streamLeadingSince.delete(stream);
+		}
+
+		// Don't forward messages from non-active streams that haven't proven themselves
+		return false;
+	}
+
+	private setActiveStream(stream: string) {
+		if (this.activeStream) {
+			this.activeStreamGauge.setLatestValue(0, { source: this.activeStream });
+		}
+
+		this.activeStream = stream;
+		this.activeStreamGauge.setLatestValue(1, { source: stream });
+		this.streamLeadingSince.clear();
+
+		console.log(`Active stream set to: ${stream}`);
+	}
+
+	getActiveStream(): string | null {
+		return this.activeStream;
+	}
+
+	checkHealth() {
+		const now = Date.now();
+		let anyHealthy = false;
+
+		for (const stream of this.streams) {
+			const lastMessage = this.streamLastMessageTime.get(stream);
+			const isHealthy =
+				lastMessage && now - lastMessage < STREAM_STALE_THRESHOLD_MS;
+
+			this.streamHealthyGauge.setLatestValue(isHealthy ? 1 : 0, {
+				source: stream,
+			});
+
+			if (isHealthy) {
+				anyHealthy = true;
+			}
+		}
+
+		// If active stream is no longer healthy, try to switch
+		if (this.activeStream) {
+			const activeLastMessage = this.streamLastMessageTime.get(
+				this.activeStream
+			);
+			if (
+				!activeLastMessage ||
+				now - activeLastMessage > STREAM_STALE_THRESHOLD_MS
+			) {
+				// Find a healthy stream to switch to
+				for (const stream of this.streams) {
+					const lastMessage = this.streamLastMessageTime.get(stream);
+					if (lastMessage && now - lastMessage < STREAM_STALE_THRESHOLD_MS) {
+						console.log(
+							`Active stream ${this.activeStream} is unhealthy, switching to ${stream}`
+						);
+						this.setActiveStream(stream);
+						break;
+					}
+				}
+			}
+		}
+
+		return anyHealthy;
+	}
+}
 
 // Set up env constants
 require('dotenv').config();
@@ -34,6 +188,14 @@ const wsOrderbookSourceLastSlotGauge = metricsV2.addGauge(
 	'websocket_orderbook_source_last_slot',
 	'Last slot of orderbook messages from a source'
 );
+const wsStreamHealthyGauge = metricsV2.addGauge(
+	'websocket_stream_healthy',
+	'Whether a Redis stream source is healthy (1) or not (0). Alert if all sources are 0.'
+);
+const wsActiveStreamGauge = metricsV2.addGauge(
+	'websocket_active_stream',
+	'Whether a Redis stream source is currently active (1) or not (0). Only one source should be active at a time.'
+);
 metricsV2.finalizeObservables();
 
 const server = http.createServer(app);
@@ -125,6 +287,32 @@ const getRedisChannelFromMessage = (message: any): string => {
 async function main() {
 	const subscribedChannelToSlot: Map<string, number> = new Map();
 
+	// Create stream selector for managing failover between Redis sources
+	const streamSelector =
+		REDIS_CLIENTS.length > 1
+			? new StreamSelector(
+					REDIS_CLIENTS,
+					wsStreamHealthyGauge,
+					wsActiveStreamGauge
+			  )
+			: null;
+
+	// If only one client, mark it as healthy and active
+	if (!streamSelector && REDIS_CLIENTS.length === 1) {
+		wsStreamHealthyGauge.setLatestValue(1, { source: REDIS_CLIENTS[0] });
+		wsActiveStreamGauge.setLatestValue(1, { source: REDIS_CLIENTS[0] });
+	}
+
+	// Periodic health check for stream selector
+	if (streamSelector) {
+		setInterval(() => {
+			const anyHealthy = streamSelector.checkHealth();
+			if (!anyHealthy) {
+				console.error('WARNING: No healthy Redis streams available!');
+			}
+		}, 1000);
+	}
+
 	const redisClients: Array<RedisClient> = [];
 	const lastMessageClients: Array<RedisClient> = [];
 	for (let i = 0; i < REDIS_CLIENTS.length; i++) {
@@ -157,13 +345,27 @@ async function main() {
 						source: channelPrefix,
 					});
 
-					const lastMessageSlot = subscribedChannelToSlot.get(sanitizedChannel);
+					if (streamSelector) {
+						const shouldForward = streamSelector.recordMessage(
+							channelPrefix,
+							messageSlot
+						);
+						if (!shouldForward) {
+							return;
+						}
+					} else {
+						const lastMessageSlot =
+							subscribedChannelToSlot.get(sanitizedChannel);
 						if (!lastMessageSlot || lastMessageSlot <= messageSlot) {
 							subscribedChannelToSlot.set(sanitizedChannel, messageSlot);
 						} else if (lastMessageSlot > messageSlot) {
 							return;
 						}
 					}
+
+					// Update slot tracking for the forwarded message
+					subscribedChannelToSlot.set(sanitizedChannel, messageSlot);
+				}
 				subscribers.forEach((ws) => {
 					if (
 						ws.readyState === WebSocket.OPEN &&