前情提要

各位好！最近做了一个有意思的小案例，分享出来给大家介绍一下 WebSocket 的基本使用。案例虽小，但同时具有客户端和服务端，算是五脏俱全。

本文采取快餐形式，即不介绍 WebSocket 的概念和进阶用法，如果想要了解 WebSocket 概念的推荐阮一峰老师的博客《阮一峰的网络日志-websocket 教程》

客户端

这里我们使用 React 开发一个简单的聊天室页面：

为了方便我们功能复用，我们将使用 WebSocket 的功能封装成一个自定义 Hook，这样在其它地方可以直接导入使用。

在自定义 Hook 中，我们需要注意的是，当我们更新聊天消息时会重新渲染组件，组件中的 Hook 也会再次调用（React 渲染机制这里不再详细描述，感兴趣的可以查看我之前的文章）。为了防止多次创建重复的 WebSocket 实例，我们需要通过 useRef 来缓存 WebSocket 实例，当再次渲染时直接使用缓存的 WebSocket 实例，避免性能浪费。

在 WebSocket 实例中，我们通过监听 “message” 事件来接受服务端传递的聊天消息，通过调用实例的 send 方法来发送消息，当我们卸载组件时需要调用实例的 close 方法来关闭 WebSocket 连接避免性能浪费。

import { useState, useLayoutEffect, useRef } from 'react'

function useWebsocket(url) {
    // 用于缓存 WebSocket 实例
    const ws = useRef(null)
    // 用于接收服务端传递的内容
    const [message, setMsg] = useState('');
    // 发送
    const sendMessage = (msg) => {
        ws.current?.send(msg)
    }
    useLayoutEffect(() => {
        // 判断 WebSocket 实例是否存在，如果存在避免重复创建
        if (!ws.current || ws.current.readyState === 3) {
            // 缓存 WebSocket 实例
            ws.current = new WebSocket(url)
            // 建立 WebSocket 连接后立刻发送一条消息:
            ws.current.addEventListener('open', () => {
                ws.current.send('欢迎 React User 加入聊天室！');
            });
            // 响应收到的消息:
            ws.current.addEventListener('message', (msgObj) => {
                let msg = msgObj.data + '\n'
                setMsg((m) => m + msg)
            })
            // 关闭 WebSocket 连接后立刻发送一条消息:
            ws.current.addEventListener('close', () => {
                ws.current.send('React User 已退出聊天室');
            });
        }
        return () => {
            // 卸载组件时，关闭 WebSocket 连接
            ws.current?.close()
        }
    }, [ws, url])
    return { message, sendMessage }
}

export default useWebsocket

在页面中，我们首先导入封装好的 WebSocket Hook，然后直接使用就可以了！

服务端

人们常说 WebSocket 的工作量全在服务端，这个确实，但如果我们直接简单的使用它，那么服务端的编写也很简单，因为 NodeJS 已经帮我们封装好了模块。

首先我们引入 “ws” 模块，获取其导出对象中的 Server 方法，然后通过 Express 启动监听 8000 端口的服务，并将其注入 Server 方法生成的 WebSocketServer 实例对象中。在 WebSocketServer 中，如果有WebSocket请求接入，可以响应 connection 事件来处理这个WebSocket。在 connection 事件中，回调函数会传入一个 WebSocket 实例，表示这个WebSocket连接，我们可以调用该 WebSocket 实例的 send 方法来向客户端发送消息。

// 引入 NodeJS 自带的 WebSocket 模块
const WebSocket = require('ws')
const express = require('express');
const app = express();

// 获取生成 WebSocket 实例的方法
const WebSocketServer = WebSocket.Server

// 监听 8000 端口
const server = app.listen(8000, () => console.log('Example app listening on port 8000!'));

// 生成 WebSocket 实例，传入服务信息
const wss = new WebSocketServer({
    server
})

// 在 WebSocketServer 中，需要响应 connection 事件
wss.on('connection', function (ws, req) {
    // ws 为 WebSocketServer 的实例，req 为请求信息
    // 当客户端发送数据时，需要响应 message 事件
    ws.on('message', function (message) {
        message = message.toString('utf8')
        // wss.clients 获取所有已请求的浏览器信息
        wss.clients.forEach(client => {
            client.send(message)
        })
    })
});

需要注意的是，我们可以看到 Express 已经监听了 8000 端口，那么我们新加的 WebSocketServer 为什么还是可以使用 8000 端口呢？实际上，8000 端口并非由 Express 监听，而是 Express 调用 NodeJS 的 http.Server 监听的。 Express 只是把响应函数注册到该 http.Server 中了。类似的，WebSocketServer 也可以把自己的响应函数注册到 http.Server 中，这样，同一个端口，根据协议，可以分别由 Express 和 WebSocket 处理：