examples/typescript/websocket-server-pool/ws-cluster/src/worker.ts

   1 import { ClusterWorker } from 'poolifier'
   2 import { type RawData, WebSocketServer } from 'ws'
   3
   4 import {
   5   type DataPayload,
   6   type MessagePayload,
   7   MessageType,
   8   type WorkerData,
   9   type WorkerResponse,
  10 } from './types.js'
  11
  12 class WebSocketServerWorker extends ClusterWorker<WorkerData, WorkerResponse> {
  13   private static wss: WebSocketServer
  14
  15   private static readonly factorial = (n: number | bigint): bigint => {
  16     if (n === 0 || n === 1) {
  17       return 1n
  18     } else {
  19       n = BigInt(n)
  20       let factorial = 1n
  21       for (let i = 1n; i <= n; i++) {
  22         factorial *= i
  23       }
  24       return factorial
  25     }
  26   }
  27
  28   private static readonly startWebSocketServer = (
  29     workerData?: WorkerData
  30   ): WorkerResponse => {
  31     // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
  32     const { port } = workerData!
  33
  34     WebSocketServerWorker.wss = new WebSocketServer({ port }, () => {
  35       console.info(
  36         `⚡️[ws server]: WebSocket server is started in cluster worker at ws://localhost:${port}/`
  37       )
  38     })
  39
  40     WebSocketServerWorker.wss.on('connection', ws => {
  41       ws.on('error', console.error)
  42       ws.on('message', (message: RawData) => {
  43         const { type, data } = JSON.parse(
  44           message.toString()
  45         ) as MessagePayload<DataPayload>
  46         switch (type) {
  47           case MessageType.echo:
  48             ws.send(
  49               JSON.stringify({
  50                 type: MessageType.echo,
  51                 data,
  52               })
  53             )
  54             break
  55           case MessageType.factorial:
  56             ws.send(
  57               JSON.stringify(
  58                 {
  59                   type: MessageType.factorial,
  60                   data: {
  61                     // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
  62                     number: WebSocketServerWorker.factorial(data.number!),
  63                   },
  64                 },
  65                 (_, v) => (typeof v === 'bigint' ? v.toString() : v)
  66               )
  67             )
  68             break
  69         }
  70       })
  71     })
  72     return {
  73       status: true,
  74       port: WebSocketServerWorker.wss.options.port,
  75     }
  76   }
  77
  78   public constructor () {
  79     super(WebSocketServerWorker.startWebSocketServer, {
  80       killHandler: () => {
  81         WebSocketServerWorker.wss.close()
  82       },
  83     })
  84   }
  85 }
  86
  87 export const webSocketServerWorker = new WebSocketServerWorker()