@nevuamarkets/poly-websockets 0.0.1

package/src/WSSubscriptionManager.ts

@@ -0,0 +1,201 @@
+import ms from 'ms';
+import _ from 'lodash';
+import Bottleneck from 'bottleneck';
+import {
+    WebSocketHandlers,
+    PriceChangeEvent,
+    BookEvent,
+    LastTradePriceEvent,
+    TickSizeChangeEvent,
+    PolymarketWSEvent,
+    PolymarketPriceUpdateEvent
+} from './types/PolymarketWebSocket';
+import { SubscriptionManagerOptions } from './types/WebSocketSubscriptions';
+
+import { GroupRegistry } from './modules/GroupRegistry';
+import { OrderBookCache } from './modules/OrderBookCache';
+import { GroupSocket } from './modules/GroupSocket';
+
+import { logger } from './logger';
+
+
+// Keeping a burst limit under 10/s to avoid rate limiting
+// See https://docs.polymarket.com/quickstart/introduction/rate-limits#api-rate-limits
+const BURST_LIMIT_PER_SECOND = 5;
+
+const DEFAULT_RECONNECT_AND_CLEANUP_INTERVAL_MS = ms('10s');
+const DEFAULT_MAX_MARKETS_PER_WS = 100;
+
+class WSSubscriptionManager {
+    private handlers: WebSocketHandlers;
+    private burstLimiter: Bottleneck;
+    private groupRegistry: GroupRegistry;
+    private bookCache: OrderBookCache;
+    private reconnectAndCleanupIntervalMs: number;
+    private maxMarketsPerWS: number;
+
+    constructor(userHandlers: WebSocketHandlers, options?: SubscriptionManagerOptions) {
+        this.groupRegistry = new GroupRegistry();
+        this.bookCache = new OrderBookCache();
+        this.burstLimiter = options?.burstLimiter || new Bottleneck({
+            reservoir: BURST_LIMIT_PER_SECOND,
+            reservoirRefreshAmount: BURST_LIMIT_PER_SECOND,
+            reservoirRefreshInterval: ms('1s'),
+            maxConcurrent: BURST_LIMIT_PER_SECOND
+        });
+
+        this.reconnectAndCleanupIntervalMs = options?.reconnectAndCleanupIntervalMs || DEFAULT_RECONNECT_AND_CLEANUP_INTERVAL_MS;
+        this.maxMarketsPerWS = options?.maxMarketsPerWS || DEFAULT_MAX_MARKETS_PER_WS;
+
+        this.handlers = {
+            onBook: async (events: BookEvent[]) => {
+                await this.actOnSubscribedEvents(events, userHandlers.onBook);
+            },
+            onLastTradePrice: async (events: LastTradePriceEvent[]) => {
+                await this.actOnSubscribedEvents(events, userHandlers.onLastTradePrice);
+            },
+            onTickSizeChange: async (events: TickSizeChangeEvent[]) => {
+                await this.actOnSubscribedEvents(events, userHandlers.onTickSizeChange);
+            },
+            onPriceChange: async (events: PriceChangeEvent[]) => {
+                await this.actOnSubscribedEvents(events, userHandlers.onPriceChange);
+            },
+            onPolymarketPriceUpdate: async (events: PolymarketPriceUpdateEvent[]) => {
+                await this.actOnSubscribedEvents(events, userHandlers.onPolymarketPriceUpdate);
+            },
+            onWSClose: userHandlers.onWSClose,
+            onWSOpen: userHandlers.onWSOpen,
+            onError: userHandlers.onError
+        };
+
+        this.burstLimiter.on('error', (err: Error) => {
+            this.handlers.onError?.(err);
+        });
+
+        // Check for dead groups every 10s and reconnect them if needed
+        setInterval(() => {
+            this.reconnectAndCleanupGroups();
+        }, this.reconnectAndCleanupIntervalMs);
+    }
+
+    /*
+        Clears all WebSocket subscriptions and state.
+
+        This will:
+
+        1. Remove all subscriptions and groups
+        2. Close all WebSocket connections
+        3. Clear the order book cache
+    */
+    public async clearState() {
+        const previousGroups = await this.groupRegistry.clearAllGroups();
+
+        // Close sockets outside the lock
+        for (const group of previousGroups) {
+            this.groupRegistry.disconnectGroup(group);
+        }
+
+        // Also clear the order book cache
+        this.bookCache.clear();
+    }
+
+    /* 
+        This function is called when:
+        - a websocket event is received from the Polymarket WS
+        - a price update event detected, either by after a 'last_trade_price' event or a 'price_change' event
+        depending on the current bid-ask spread (see https://docs.polymarket.com/polymarket-learn/trading/how-are-prices-calculated)
+
+        The user handlers will be called **ONLY** for assets that are actively subscribed to by any groups.
+    */
+    private async actOnSubscribedEvents<T extends PolymarketWSEvent | PolymarketPriceUpdateEvent>(events: T[], action?: (events: T[]) => Promise<void>) {
+
+        // Filter out events that are not subscribed to by any groups
+        events = _.filter(events, (event: T) => {
+            const groupIndices = this.groupRegistry.getGroupIndicesForAsset(event.asset_id);
+
+            if (groupIndices.length > 1) {
+                logger.warn({
+                    message: 'Found multiple groups for asset',
+                    asset_id: event.asset_id,
+                    group_indices: groupIndices
+                });
+            }
+            return groupIndices.length > 0;
+        });
+
+        await action?.(events);
+    }
+
+    /*  
+        Edits wsGroups: Adds new subscriptions.
+
+        - Filters out assets that are already subscribed
+        - Finds a group with capacity or creates a new one
+        - Creates a new WebSocket client and adds it to the group
+    */
+    public async addSubscriptions(assetIdsToAdd: string[]) {
+        try {
+            const groupIdsToConnect = await this.groupRegistry.addAssets(assetIdsToAdd, this.maxMarketsPerWS);
+            for (const groupId of groupIdsToConnect) {
+                await this.createWebSocketClient(groupId, this.handlers);
+            }
+        } catch (error) {
+            const msg = `Error adding subscriptions: ${error instanceof Error ? error.message : String(error)}`;
+            await this.handlers.onError?.(new Error(msg));
+        }
+    }
+
+    /*  
+        Edits wsGroups: Removes subscriptions.
+        The group will use the updated subscriptions when it reconnects.
+        We do that because we don't want to miss events by reconnecting.
+    */
+    public async removeSubscriptions(assetIdsToRemove: string[]) {
+        try {
+            await this.groupRegistry.removeAssets(assetIdsToRemove, this.bookCache);
+        } catch (error) {
+            const errMsg = `Error removing subscriptions: ${error instanceof Error ? error.message : String(error)}`;
+            await this.handlers.onError?.(new Error(errMsg));
+        }
+    }
+
+    /*
+        This function runs periodically and:
+
+        - Tries to reconnect groups that have assets and are disconnected
+        - Cleans up groups that have no assets
+    */
+    private async reconnectAndCleanupGroups() {
+        try {
+            const reconnectIds = await this.groupRegistry.getGroupsToReconnectAndCleanup();
+
+            for (const groupId of reconnectIds) {
+                await this.createWebSocketClient(groupId, this.handlers);
+            }
+        } catch (err) {
+            await this.handlers.onError?.(err as Error);
+        }
+    }
+
+    private async createWebSocketClient(groupId: string, handlers: WebSocketHandlers) {
+        const group = this.groupRegistry.findGroupById(groupId);
+
+        /*
+            Should never happen, but just in case.
+        */
+        if (!group) {
+            await handlers.onError?.(new Error(`Group ${groupId} not found in registry`));
+            return;
+        }
+
+        const groupSocket = new GroupSocket(group, this.burstLimiter, this.bookCache, handlers);
+        try {
+            await groupSocket.connect();
+        } catch (error) {
+            const errorMessage = `Error creating WebSocket client for group ${groupId}: ${error instanceof Error ? error.message : String(error)}`;
+            await handlers.onError?.(new Error(errorMessage));
+        }
+    }
+}
+
+export { WSSubscriptionManager, WebSocketHandlers };

package/src/index.ts

@@ -0,0 +1,3 @@
+export { WSSubscriptionManager, WebSocketHandlers } from './WSSubscriptionManager';
+export * from './types/PolymarketWebSocket';
+export * from './types/WebSocketSubscriptions'; 

package/src/logger.ts

@@ -0,0 +1,37 @@
+import winston from 'winston';
+
+
+// Override with LOG_LEVEL environment variable (e.g., LOG_LEVEL=info npm start)
+export const logger = winston.createLogger({
+    level: process.env.LOG_LEVEL || 'error',
+    format: winston.format.combine(
+        winston.format.timestamp(),
+        winston.format.errors({ stack: true }),
+        winston.format.colorize(),
+        winston.format.printf(({ level, message, timestamp, ...rest }) => {
+            // Ensure consistent order: timestamp, level, message, then rest of fields
+            const restString = Object.keys(rest)
+                .filter(key => key !== 'service') // Exclude service since we add it in defaultMeta
+                .sort()
+                .map(key => `${key}: ${JSON.stringify(rest[key])}`)
+                .join(', ');
+            return `${timestamp} ${level}: ${message}${restString ? ` (${restString})` : ''}`;
+        })
+    ),
+    defaultMeta: { service: 'poly-websockets' },
+    transports: [
+        new winston.transports.Console({
+            format: winston.format.combine(
+                winston.format.colorize({
+                    all: true,
+                    colors: {
+                        error: 'red',
+                        warn: 'yellow',
+                        info: 'cyan',
+                        debug: 'green'
+                    }
+                })
+            )
+        })
+    ]
+});

package/src/modules/GroupRegistry.ts

@@ -0,0 +1,274 @@
+import { Mutex } from 'async-mutex';
+import _ from 'lodash';
+import { v4 as uuidv4 } from 'uuid';
+import { WebSocketGroup, WebSocketStatus } from '../types/WebSocketSubscriptions';
+import { OrderBookCache } from './OrderBookCache';
+import { logger } from '../logger';
+
+/*
+ * Global group store and mutex, intentionally **not** exported anymore to prevent
+ * accidental external mutation.  All access should go through the helper methods
+ * on GroupRegistry instead. 
+ */
+const wsGroups: WebSocketGroup[] = [];
+const wsGroupsMutex = new Mutex();
+
+export class GroupRegistry {
+
+    /** 
+     * Atomic mutate helper.
+     * 
+     * @param fn - The function to run atomically.
+     * @returns The result of the function.
+     */
+    public async mutate<T>(fn: (groups: WebSocketGroup[]) => T | Promise<T>): Promise<T> {
+        const release = await wsGroupsMutex.acquire();
+        try { return await fn(wsGroups); }
+        finally { release(); }
+    }
+
+    /** 
+     * Read-only copy of the registry.
+     * 
+     * Only to be used in test suite.
+     */
+    public snapshot(): WebSocketGroup[] {
+        return wsGroups.map(group => ({
+            ...group,
+            assetIds: new Set(group.assetIds),
+        }));
+    }
+
+    /**
+     * Find the first group with capacity to hold new assets.
+     * 
+     * Returns the groupId if found, otherwise null.
+     */
+    public findGroupWithCapacity(newAssetLen: number, maxPerWS: number): string | null {
+        for (const group of wsGroups) {
+            if (group.assetIds.size === 0) continue;
+            if (group.assetIds.size + newAssetLen <= maxPerWS) return group.groupId;
+        }
+        return null;
+    }
+
+    /**
+     * Get the indices of all groups that contain the asset.
+     * 
+     * Returns an array of indices.
+     */
+    public getGroupIndicesForAsset(assetId: string): number[] {
+        const indices: number[] = [];
+        for (let i = 0; i < wsGroups.length; i++) {
+            if (wsGroups[i]?.assetIds.has(assetId)) indices.push(i);
+        }
+        return indices;
+    }
+
+    /**
+     * Check if any group contains the asset.
+     */
+    public hasAsset(assetId: string): boolean {
+        return wsGroups.some(group => group.assetIds.has(assetId));
+    }
+
+    /**
+     * Find the group by groupId.
+     * 
+     * Returns the group if found, otherwise undefined.
+     */
+    public findGroupById(groupId: string): WebSocketGroup | undefined {
+        return wsGroups.find(g => g.groupId === groupId);
+    }
+
+    /**
+     * Atomically remove **all** groups from the registry and return them so the
+     * caller can perform any asynchronous cleanup (closing sockets, etc.)
+     * outside the lock. 
+     * 
+     * Returns the removed groups.
+     */
+    public async clearAllGroups(): Promise<WebSocketGroup[]> {
+        let removed: WebSocketGroup[] = [];
+        await this.mutate(groups => {
+            removed = [...groups];
+            groups.length = 0;
+        });
+        return removed;
+    }
+
+    /**
+     * Add new asset subscriptions.
+     * 
+     * – Ignores assets that are already subscribed.
+     * – Either reuses an existing group with capacity or creates new groups (size ≤ maxPerWS).
+     * – If appending to a group:
+     *  - A new group is created with the updated assetIds.
+     *  - The old group is marked for cleanup.
+     *  - The group is added to the list of groups to connect.
+     * 
+     * @param assetIds - The assetIds to add.
+     * @param maxPerWS - The maximum number of assets per WebSocket group.
+     * @returns An array of *new* groupIds that need websocket connections.
+     */
+    public async addAssets(assetIds: string[], maxPerWS: number): Promise<string[]> {
+        const groupIdsToConnect: string[] = [];
+        let newAssetIds: string[] = []
+
+        await this.mutate(groups => {
+            newAssetIds = assetIds.filter(id => !groups.some(g => g.assetIds.has(id)));
+            if (newAssetIds.length === 0) return;
+
+            const existingGroupId = this.findGroupWithCapacity(newAssetIds.length, maxPerWS);
+
+            /*
+                If no existing group with capacity is found, create new groups.
+            */
+            if (existingGroupId === null) {
+                const chunks = _.chunk(newAssetIds, maxPerWS);
+                for (const chunk of chunks) {
+                    const groupId = uuidv4();
+                    groups.push(
+                        { 
+                            groupId, 
+                            assetIds: new Set(chunk), 
+                            wsClient: null, 
+                            status: WebSocketStatus.PENDING 
+                        }
+                    );
+                    groupIdsToConnect.push(groupId);
+                }
+
+            /*
+                If an existing group with capacity is found, update the group.
+            */
+            } else {
+                const existingGroup = groups.find(g => g.groupId === existingGroupId);
+                if (!existingGroup) {
+                    // Should never happen
+                    throw new Error(`Group with capacity not found for ${newAssetIds.join(', ')}`);
+                }
+
+                const updatedAssetIds = new Set([...existingGroup.assetIds, ...newAssetIds]);
+
+                // Mark old group ready for cleanup
+                existingGroup.assetIds = new Set();
+                existingGroup.status = WebSocketStatus.CLEANUP;
+
+                const groupId = uuidv4();
+                groups.push(
+                    { 
+                        groupId, 
+                        assetIds: updatedAssetIds, 
+                        wsClient: null, 
+                        status: WebSocketStatus.PENDING 
+                    }
+                );
+                groupIdsToConnect.push(groupId);
+            }
+        });
+
+        if (newAssetIds.length > 0) {
+            logger.info({
+                message: `Added ${newAssetIds.length} new asset(s)`
+            })
+        }
+        return groupIdsToConnect;
+    }
+
+    /**
+     * Remove asset subscriptions from every group that contains the asset.
+     * 
+     * It should be only one group that contains the asset, we search all of them
+     * regardless.
+     * 
+     * Returns the list of assetIds that were removed.
+     */
+    public async removeAssets(assetIds: string[], bookCache: OrderBookCache): Promise<string[]> {
+        const removedAssetIds: string[] = [];
+        await this.mutate(groups => {
+            groups.forEach(group => {
+                if (group.assetIds.size === 0) return;
+
+                assetIds.forEach(id => {
+                    if (group.assetIds.delete(id)) {
+                        bookCache.clear(id);
+                        removedAssetIds.push(id)
+                    }
+                });
+            });
+        });
+        if (removedAssetIds.length > 0) {
+            logger.info({
+                message: `Removed ${removedAssetIds.length} asset(s)`
+            })
+        }
+        return removedAssetIds;
+    }
+
+    /**
+     * Disconnect a group.
+     */
+    public disconnectGroup(group: WebSocketGroup) {
+        group.wsClient?.close();
+        group.wsClient = null;
+
+        logger.info({
+            message: 'Disconnected group',
+            groupId: group.groupId,
+            assetIds: Array.from(group.assetIds),
+        });
+
+    };
+    /**
+     * Check status of groups and reconnect or cleanup as needed.
+     * 
+     * – Empty groups are removed from the global array and returned.
+     * – Dead (but non-empty) groups are reset so that caller can reconnect them.
+     * – Pending groups are returned so that caller can connect them.
+     * 
+     * Returns an array of group IDs that need to be reconnected, after cleaning up empty and cleanup-marked groups.
+     */
+    public async getGroupsToReconnectAndCleanup(): Promise<string[]> {
+        const reconnectIds: string[] = [];
+
+        await this.mutate(groups => {
+            const groupsToRemove = new Set<string>();
+
+            for (const group of groups) {
+                if (group.assetIds.size === 0) {
+                    groupsToRemove.add(group.groupId);
+                    continue;
+                }
+
+                if (group.status === WebSocketStatus.ALIVE) {
+                    continue;
+                }
+
+                if (group.status === WebSocketStatus.DEAD) {
+                    this.disconnectGroup(group);
+                    reconnectIds.push(group.groupId);
+                }
+                if (group.status === WebSocketStatus.CLEANUP) {
+                    groupsToRemove.add(group.groupId);
+                    group.assetIds = new Set();
+                    continue;
+                }
+
+                if (group.status === WebSocketStatus.PENDING) {
+                    reconnectIds.push(group.groupId);
+                }
+            }
+            if (groupsToRemove.size > 0) {
+                groups.forEach(group => {
+                    if (groupsToRemove.has(group.groupId)) {
+                        this.disconnectGroup(group);
+                    }
+                });
+                const remaining = groups.filter(group => !groupsToRemove.has(group.groupId));
+                groups.splice(0, groups.length, ...remaining);
+            }
+        });
+        return reconnectIds;
+    }
+}