add ws redis switchover with hysteresis

This commit is contained in:
wphan
2026-01-06 11:41:52 -08:00
parent 1db4c8e54a
commit 1e8f06a3e4

View File

@@ -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,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 (