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')

        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.entries()) {
        client.send(JSON.stringify(msg))
      }
    } else {
      this.#queue.push(msg)
    }
  }
}

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