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

   1 import {
   2   availableParallelism,
   3   ClusterWorker,
   4   DynamicThreadPool
   5 } from 'poolifier'
   6 import { type RawData, WebSocketServer } from 'ws'
   7
   8 import {
   9   type ClusterWorkerData,
  10   type ClusterWorkerResponse,
  11   type DataPayload,
  12   type MessagePayload,
  13   MessageType,
  14   type ThreadWorkerData,
  15   type ThreadWorkerResponse
  16 } from './types.js'
  17
  18 const emptyFunction = (): void => {
  19   /* Intentional */
  20 }
  21
  22 class WebSocketServerWorker extends ClusterWorker<
  23 ClusterWorkerData,
  24 ClusterWorkerResponse
  25 > {
  26   private static wss: WebSocketServer
  27   private static requestHandlerPool: DynamicThreadPool<
  28   ThreadWorkerData<DataPayload>,
  29   ThreadWorkerResponse<DataPayload>
  30   >
  31
  32   private static readonly startWebSocketServer = (
  33     workerData?: ClusterWorkerData
  34   ): ClusterWorkerResponse => {
  35     const { port, workerFile, minWorkers, maxWorkers, ...poolOptions } =
  36       workerData!
  37
  38     WebSocketServerWorker.requestHandlerPool = new DynamicThreadPool<
  39     ThreadWorkerData<DataPayload>,
  40     ThreadWorkerResponse<DataPayload>
  41     >(
  42       minWorkers ?? 1,
  43       maxWorkers ?? availableParallelism(),
  44       workerFile,
  45       poolOptions
  46     )
  47
  48     WebSocketServerWorker.wss = new WebSocketServer({ port }, () => {
  49       console.info(
  50         `⚡️[ws server]: WebSocket server is started in cluster worker at ws://localhost:${port}/`
  51       )
  52     })
  53
  54     WebSocketServerWorker.wss.on('connection', ws => {
  55       ws.on('error', console.error)
  56       ws.on('message', (message: RawData) => {
  57         const { type, data } = JSON.parse(
  58           // eslint-disable-next-line @typescript-eslint/no-base-to-string
  59           message.toString()
  60         ) as MessagePayload<DataPayload>
  61         switch (type) {
  62           case MessageType.echo:
  63             WebSocketServerWorker.requestHandlerPool
  64               .execute({ data }, 'echo')
  65               .then(response => {
  66                 ws.send(
  67                   JSON.stringify({
  68                     type: MessageType.echo,
  69                     data: response.data
  70                   })
  71                 )
  72                 return undefined
  73               })
  74               .catch(emptyFunction)
  75             break
  76           case MessageType.factorial:
  77             WebSocketServerWorker.requestHandlerPool
  78               .execute({ data }, 'factorial')
  79               .then(response => {
  80                 ws.send(
  81                   JSON.stringify(
  82                     {
  83                       type: MessageType.factorial,
  84                       data: response.data
  85                     },
  86                     (_, v) => (typeof v === 'bigint' ? v.toString() : v)
  87                   )
  88                 )
  89                 return undefined
  90               })
  91               .catch(emptyFunction)
  92             break
  93         }
  94       })
  95     })
  96     return {
  97       status: true,
  98       port: WebSocketServerWorker.wss.options.port
  99     }
 100   }
 101
 102   public constructor () {
 103     super(WebSocketServerWorker.startWebSocketServer, {
 104       killHandler: async () => {
 105         await WebSocketServerWorker.requestHandlerPool.destroy()
 106         WebSocketServerWorker.wss.close()
 107       }
 108     })
 109   }
 110 }
 111
 112 export const webSocketServerWorker = new WebSocketServerWorker()