fite/lib/ws.js

/**
 * 一个简单到发指的 WebSocket 类
 * @author yutent<yutent.io@gmail.com>
 * @date 2023/01/17 17:32:51
 */
import { WebSocketServer } from 'ws'

class WebSocket {
  #clients = new Map()
  #queue = [] // 消息队列

  constructor(server) {
    if (server.listening) {
      let conn = new WebSocketServer({ server, path: '/ws-fite-hmr' })
      conn.on('connection', (client, req) => {
        let params = new URLSearchParams(req.url.slice(req.url.indexOf('?')))
        let session = params.get('session')
        let lock = +params.get('lock')

        if (lock === 1) {
          client.close()
        } else {
          this.#clients.set(session, client)

          client.once('close', _ => {
            this.#clients.delete(session)
          })

          while (this.#queue.length) {
            let msg = this.#queue.shift()
            this.send(msg)
          }
        }
      })
    }
  }

  send(msg = {}) {
    if (this.#clients.size) {
      for (let [key, client] of this.#clients) {
        client.send(JSON.stringify(msg))
      }
    } else {
      this.#queue.push(msg)
    }
  }
}

export default function (server) {
  return new WebSocket(server)
}