Merge pull request #1569 from multiversx/ws-custom-subscriptions

Ws custom subscriptions

Author: stefangutica

config/config.devnet.yaml

@@ -23,6 +23,8 @@ features:
   websocketSubscription:
     enabled: false
     port: 6002
+    maxSubscriptionsPerInstance: 10000
+    maxSubscriptionsPerClient: 5
   eventsNotifier:
     enabled: false
     port: 5674

config/config.e2e-mocked.mainnet.yaml

@@ -8,6 +8,8 @@ features:
   websocketSubscription:
     enabled: false
     port: 6002
+    maxSubscriptionsPerInstance: 10000
+    maxSubscriptionsPerClient: 5
   dataApi:
     enabled: false
     serviceUrl: 'https://data-api.multiversx.com'

config/config.e2e.mainnet.yaml

@@ -21,8 +21,10 @@ flags:
   collectionPropertiesFromGateway: false
 features:
   websocketSubscription:
-    enabled: false
+    enabled: true
     port: 6002
+    maxSubscriptionsPerInstance: 10000
+    maxSubscriptionsPerClient: 5
   eventsNotifier:
     enabled: false
     port: 5674

config/config.mainnet.yaml

@@ -23,6 +23,8 @@ features:
   websocketSubscription:
     enabled: false
     port: 6002
+    maxSubscriptionsPerInstance: 10000
+    maxSubscriptionsPerClient: 5
   eventsNotifier:
     enabled: false
     port: 5674

config/config.testnet.yaml

@@ -23,6 +23,8 @@ features:
   websocketSubscription:
     enabled: false
     port: 6002
+    maxSubscriptionsPerInstance: 10000
+    maxSubscriptionsPerClient: 5
   eventsNotifier:
     enabled: false
     port: 5674

src/common/api-config/api.config.service.ts

@@ -1010,6 +1010,14 @@ export class ApiConfigService {
     return port;
   }
 
+  getWebsocketMaxSubscriptionsPerInstance(): number {
+    return this.configService.get<number>('features.websocketSubscription.maxSubscriptionsPerInstance') ?? 10_000;
+  }
+
+  getWebsocketMaxSubscriptionsPerClient(): number {
+    return this.configService.get<number>('features.websocketSubscription.maxSubscriptionsPerClient') ?? 5;
+  }
+
   getHeadersForCustomUrl(url: string): Record<string, string> | undefined {
     let customUrlConfigs = this.configService.get<any>('customUrlHeaders');
 

src/common/metrics/api.metrics.service.ts

@@ -23,6 +23,7 @@ export class ApiMetricsService {
   private static transactionsPendingResultsCounter: Counter<string>;
   private static batchUpdatesCounter: Counter<string>;
   private static subscriptionsConnectionsGauge: Gauge<string>;
+  private static subscriptionsTopicConnectionsGauge: Gauge<string>;
 
   constructor(
     private readonly apiConfigService: ApiConfigService,
@@ -40,6 +41,14 @@ export class ApiMetricsService {
       });
     }
 
+    if (!ApiMetricsService.subscriptionsTopicConnectionsGauge) {
+      ApiMetricsService.subscriptionsTopicConnectionsGauge = new Gauge({
+        name: 'websocket_subscriptions_topic_connections',
+        help: 'Unique websocket clients per topic',
+        labelNames: ['topic'],
+      });
+    }
+
     if (!ApiMetricsService.vmQueriesHistogram) {
       ApiMetricsService.vmQueriesHistogram = new Histogram({
         name: 'vm_query',
@@ -191,10 +200,16 @@ export class ApiMetricsService {
   }
 
   @OnEvent(MetricsEvents.SetWebsocketMetrics)
-  setWebsocketSubscriptionsMetrics(payload: { connectedClients: number }) {
-    const { connectedClients } = payload;
+  setWebsocketSubscriptionsMetrics(payload: { connectedClients: number; topics?: Record<string, number> }) {
+    const { connectedClients, topics } = payload;
 
     ApiMetricsService.subscriptionsConnectionsGauge.set(connectedClients);
+
+    if (topics) {
+      for (const [topic, count] of Object.entries(topics)) {
+        ApiMetricsService.subscriptionsTopicConnectionsGauge.set({ topic }, count);
+      }
+    }
   }
 
 

src/crons/websocket/blocks.gateway.ts

@@ -4,37 +4,59 @@ import { BlockService } from '../../endpoints/blocks/block.service';
 import { BlockFilter } from '../../endpoints/blocks/entities/block.filter';
 import { QueryPagination } from 'src/common/entities/query.pagination';
 import { BlockSubscribePayload } from '../../endpoints/blocks/entities/block.subscribe';
-import { UseFilters } from '@nestjs/common';
+import { UseFilters, UseGuards } from '@nestjs/common';
 import { WebsocketExceptionsFilter } from 'src/utils/ws-exceptions.filter';
 import { WsValidationPipe } from 'src/utils/ws-validation.pipe';
 import { OriginLogger } from '@multiversx/sdk-nestjs-common';
+import { WsSubscriptionLimiterGuard } from 'src/utils/ws.subscription.limiter';
 
 @UseFilters(WebsocketExceptionsFilter)
 @WebSocketGateway({ cors: { origin: '*' }, path: '/ws/subscription' })
 export class BlocksGateway {
   private readonly logger = new OriginLogger(BlocksGateway.name);
+  static readonly keyPrefix = 'blocks-';
 
   @WebSocketServer()
   server!: Server;
 
   constructor(private readonly blockService: BlockService) { }
 
+  @UseGuards(WsSubscriptionLimiterGuard)
   @SubscribeMessage('subscribeBlocks')
   async handleSubscription(
     @ConnectedSocket() client: Socket,
     @MessageBody(new WsValidationPipe()) payload: BlockSubscribePayload
   ) {
     const filterIdentifier = JSON.stringify(payload);
-    await client.join(`blocks-${filterIdentifier}`);
+    const roomName = `${BlocksGateway.keyPrefix}${filterIdentifier}`;
+
+    if (!client.rooms.has(roomName)) {
+      await client.join(roomName);
+    }
 
     return { status: 'success' };
   }
 
+  @SubscribeMessage('unsubscribeBlocks')
+  async handleUnsubscribe(
+    @ConnectedSocket() client: Socket,
+    @MessageBody(new WsValidationPipe()) payload: BlockSubscribePayload
+  ) {
+    const filterIdentifier = JSON.stringify(payload);
+    const roomName = `${BlocksGateway.keyPrefix}${filterIdentifier}`;
+
+    if (client.rooms.has(roomName)) {
+      await client.leave(roomName);
+    }
+
+    return { status: 'unsubscribed' };
+  }
+
   async pushBlocksForRoom(roomName: string): Promise<void> {
-    if (!roomName.startsWith("blocks-")) return;
+    if (!roomName.startsWith(BlocksGateway.keyPrefix)) return;
 
     try {
-      const filterIdentifier = roomName.replace("blocks-", "");
+      const filterIdentifier = roomName.replace(BlocksGateway.keyPrefix, "");
       const filter: BlockSubscribePayload = JSON.parse(filterIdentifier);
 
       const blockFilter = new BlockFilter({

src/crons/websocket/connection.handler.ts

@@ -15,6 +15,18 @@ export class ConnectionHandler implements OnGatewayDisconnect, OnGatewayConnecti
     handleDisconnect(_client: Socket) { }
 
     handleConnection(client: Socket, ..._args: any[]) {
-        client.setMaxListeners(12);
+        client.setMaxListeners(16);
+    }
+
+    hasSubscriptionsWithPrefixes(prefixes: string[]): boolean {
+        const rooms = this.server.sockets.adapter.rooms;
+
+        for (const roomName of rooms.keys()) {
+            if (prefixes.some(prefix => roomName.startsWith(prefix))) {
+                return true;
+            }
+        }
+
+        return false;
     }
 }

src/crons/websocket/events.custom.gateway.ts

@@ -0,0 +1,98 @@
+import { WebSocketGateway, WebSocketServer, SubscribeMessage, ConnectedSocket, MessageBody } from '@nestjs/websockets';
+import { Server, Socket } from 'socket.io';
+import { UseFilters, UseGuards } from '@nestjs/common';
+import { OriginLogger } from '@multiversx/sdk-nestjs-common';
+
+import { QueryPagination } from 'src/common/entities/query.pagination';
+import { WsValidationPipe } from 'src/utils/ws-validation.pipe';
+import { WebsocketExceptionsFilter } from 'src/utils/ws-exceptions.filter';
+import { WsSubscriptionLimiterGuard } from 'src/utils/ws.subscription.limiter';
+import { RoomKeyGenerator } from './room.key.generator';
+import { EventsService } from 'src/endpoints/events/events.service';
+import { EventsCustomSubscribePayload } from 'src/endpoints/events/entities/events.custom.subscribe';
+import { EventsFilter } from 'src/endpoints/events/entities/events.filter';
+import { Events } from 'src/endpoints/events/entities/events';
+
+
+@UseFilters(WebsocketExceptionsFilter)
+@WebSocketGateway({ cors: { origin: '*' }, path: '/ws/subscription' })
+export class EventsCustomGateway {
+  private readonly logger = new OriginLogger(EventsCustomGateway.name);
+
+  static keyPrefix = 'custom-events-';
+
+  @WebSocketServer()
+  server!: Server;
+
+  constructor(
+    private readonly eventsService: EventsService,
+  ) { }
+
+  @UseGuards(WsSubscriptionLimiterGuard)
+  @SubscribeMessage('subscribeCustomEvents')
+  async handleCustomSubscription(
+    @ConnectedSocket() client: Socket,
+    @MessageBody(new WsValidationPipe()) payload: EventsCustomSubscribePayload
+  ) {
+    const filterIdentifier = RoomKeyGenerator.deterministicStringify(payload);
+    const roomName = `${EventsCustomGateway.keyPrefix}${filterIdentifier}`;
+
+    if (!client.rooms.has(roomName)) {
+      await client.join(roomName);
+    }
+
+    return { status: 'success' };
+  }
+
+  @SubscribeMessage('unsubscribeCustomEvents')
+  async handleCustomUnsubscribe(
+    @ConnectedSocket() client: Socket,
+    @MessageBody(new WsValidationPipe()) payload: EventsCustomSubscribePayload
+  ) {
+    const filterIdentifier = RoomKeyGenerator.deterministicStringify(payload);
+    const roomName = `${EventsCustomGateway.keyPrefix}${filterIdentifier}`;
+
+    if (client.rooms.has(roomName)) {
+      await client.leave(roomName);
+    }
+
+    return { status: 'unsubscribed' };
+  }
+
+  async pushEventsForTimestampMs(timestampMs: number): Promise<void> {
+    try {
+      const allEvents = await this.eventsService.getEvents(
+        new QueryPagination({ size: 10000 }),
+        new EventsFilter({ before: timestampMs, after: timestampMs }),
+      );
+
+      const eventsFilteredForBroadcast: Map<string, Events[]> = new Map();
+
+      for (const event of allEvents) {
+        const roomKeys = RoomKeyGenerator.generate(
+          EventsCustomGateway.keyPrefix,
+          event,
+          EventsCustomSubscribePayload,
+        );
+
+        for (const roomKey of roomKeys) {
+          if (this.server.sockets.adapter.rooms.has(roomKey)) {
+            if (!eventsFilteredForBroadcast.has(roomKey)) {
+              eventsFilteredForBroadcast.set(roomKey, []);
+            }
+            eventsFilteredForBroadcast.get(roomKey)!.push(event);
+          }
+        }
+      }
+
+      for (const [roomName] of eventsFilteredForBroadcast) {
+        this.server.to(roomName).emit("customEventUpdate", {
+          events: eventsFilteredForBroadcast.get(roomName),
+          timestampMs,
+        });
+      }
+    } catch (error) {
+      this.logger.error(error);
+    }
+  }
+}