From 1e8f06a3e442ea86e1c209af4ab8e6e20cd9c357 Mon Sep 17 00:00:00 2001 From: wphan <6348407+wphan@users.noreply.github.com> Date: Tue, 6 Jan 2026 11:41:52 -0800 Subject: [PATCH] add ws redis switchover with hysteresis --- src/wsConnectionManager.ts | 214 +++++++++++++++++++++++++++++++++++-- 1 file changed, 208 insertions(+), 6 deletions(-) diff --git a/src/wsConnectionManager.ts b/src/wsConnectionManager.ts index 68a5d91..7bf2172 100644 --- a/src/wsConnectionManager.ts +++ b/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 = new Map(); + private streamLastSlot: Map = new Map(); + private streamLeadingSince: Map = 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 = 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 = []; const lastMessageClients: Array = []; for (let i = 0; i < REDIS_CLIENTS.length; i++) { @@ -157,12 +345,26 @@ async function main() { source: channelPrefix, }); - const lastMessageSlot = subscribedChannelToSlot.get(sanitizedChannel); - if (!lastMessageSlot || lastMessageSlot <= messageSlot) { - subscribedChannelToSlot.set(sanitizedChannel, messageSlot); - } else if (lastMessageSlot > messageSlot) { - return; + 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 (