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.toString()}/`
  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           // eslint-disable-next-line @typescript-eslint/no-base-to-string
  45           message.toString()
  46         ) as MessagePayload<DataPayload>
  47         switch (type) {
  48           case MessageType.echo:
  49             ws.send(
  50               JSON.stringify({
  51                 type: MessageType.echo,
  52                 data,
  53               })
  54             )
  55             break
  56           case MessageType.factorial:
  57             ws.send(
  58               JSON.stringify(
  59                 {
  60                   type: MessageType.factorial,
  61                   data: {
  62                     // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
  63                     number: WebSocketServerWorker.factorial(data.number!),
  64                   },
  65                 },
  66                 (_, v: unknown) => (typeof v === 'bigint' ? v.toString() : v)
  67               )
  68             )
  69             break
  70         }
  71       })
  72     })
  73     return {
  74       status: true,
  75       port: WebSocketServerWorker.wss.options.port,
  76     }
  77   }
  78
  79   public constructor () {
  80     super(WebSocketServerWorker.startWebSocketServer, {
  81       killHandler: () => {
  82         WebSocketServerWorker.wss.close()
  83       },
  84     })
  85   }
  86 }
  87
  88 export const webSocketServerWorker = new WebSocketServerWorker()