NodeJS + WS 访问当前运行的 WS 服务器实例

Question

我已经使用 NodeJS、ExpressJS 和路由控制器实现了一个简单的 REST API。我还实现了一个与 REST API 一起运行并使用 WS 的基本 WebSocket 服务器。

const app = express();

app.use(bodyParser.json({limit: "50mb"}));
app.use(bodyParser.urlencoded({limit: "50mb", extended: true}));

useExpressServer(app, {
    controllers: [
        UserController
    ]
});

const server = app.listen(21443, (err: Error) => {
    console.log("listening on port 21443");
});

const wss = new WebSocket.Server({server});

wss.on("connection", (ws: WebSocket) => {
    ws.on("message", (message: string) => {
        console.log("received: %s", message);
        ws.send(`Hello, you sent -> ${message}`);
    });

    ws.send("Hi there, I am a WebSocket server");
});

我的问题是如何访问当前正在运行的 WS 实例，以便能够从我的控制器方法中访问 send 或 broadcast。我有许多运行长进程的POST 方法，因此向客户端返回HTTP 200，然后我想向所有连接的WS 客户端发送send 或broadcast。

从我的控制器类中访问WebSocket.Server 实例的正确方法是什么？

Answer 1

已连接客户端的列表存储在wss 对象中。你可以像这样接收和循环它们：

wss.clients.forEach((client) => {
    if (client.userId === current_user_id && client.readyState === WebSocket.OPEN) {
        // this is the socket of your current user
    }
})

现在您需要以某种方式识别您的客户。您可以通过在连接时为该客户端分配一些 id 来做到这一点：

wss.on('connection', async (ws, req) => {
    // req.url is the url that user connected with
    // use a query parameter on connection, or an authorization token by which you can identify the user
    // so your connection url will look like
    // http://example.com/socket?token=your_token
    ws.userId = your_user_identifier
    ....
})

广播使用：

wss.clients.forEach((client) => {
    if (client.readyState === WebSocket.OPEN) {
        client.send(data);
    }
});

如果您的控制器和套接字将位于不同的文件中（我相信它们会），您将必须在套接字文件中导出 wss 对象并将其导入控制器中。

Answer 2

您可以更早地创建 websocket 并传递实例：

const notifier = new NotifierService(); 
notifier.connect(http.createServer(app));

app.get("/somethingHappened", () => {
  notifier.broadcast("new notification!!");
});

app.use(routes(notifier))

完整代码：

app.js

将 websocket 传递给其他路由：

const express = require("express");
const http = require("http");
const NotifierService = require("../server/NotifierService.js");
const routes = require("./routes");

const app = express();
const server = http.createServer(app);
const notifier = new NotifierService();
notifier.connect(server);

app.get("/somethingHappened", () => {
  notifier.broadcast("new notification!!");
});

//   to demonstrate how the notifier instance can be
//   passed around to different routes
app.use(routes(notifier));

server
  .listen(4000)
  .on("listening", () =>
    console.log("info", `HTTP server listening on port 4000`)
  );

NotifierService.js 处理 websocket 的类

const url = require("url");
const { Server } = require("ws");

class NotifierService {
  constructor() {
    this.connections = new Map();
  }

  connect(server) {
    this.server = new Server({ noServer: true });
    this.interval = setInterval(this.checkAll.bind(this), 10000);
    this.server.on("close", this.close.bind(this));
    this.server.on("connection", this.add.bind(this));
    server.on("upgrade", (request, socket, head) => {
      console.log("ws upgrade");
      const id = url.parse(request.url, true).query.storeId;

      if (id) {
        this.server.handleUpgrade(request, socket, head, (ws) =>
          this.server.emit("connection", id, ws)
        );
      } else {
        socket.destroy();
      }
    });
  }

  add(id, socket) {
    console.log("ws add");
    socket.isAlive = true;
    socket.on("pong", () => (socket.isAlive = true));
    socket.on("close", this.remove.bind(this, id));
    this.connections.set(id, socket);
  }

  send(id, message) {
    console.log("ws sending message");

    const connection = this.connections.get(id);

    connection.send(JSON.stringify(message));
  }

  broadcast(message) {
    console.log("ws broadcast");
    this.connections.forEach((connection) =>
      connection.send(JSON.stringify(message))
    );
  }

  isAlive(id) {
    return !!this.connections.get(id);
  }

  checkAll() {
    this.connections.forEach((connection) => {
      if (!connection.isAlive) {
        return connection.terminate();
      }

      connection.isAlive = false;
      connection.ping("");
    });
  }

  remove(id) {
    this.connections.delete(id);
  }

  close() {
    clearInterval(this.interval);
  }
}

module.exports = NotifierService;

routes.js

const express = require("express");

const router = express.Router();
module.exports = (webSocketNotifier) => {
  router.post("/newPurchase/:id", (req, res, next) => {
    webSocketNotifier.send(req.params.id, "purchase made");
    res.status(200).send();
  });
  return router;
};