第83课 - 实战：WebSocket 聊天室

学习目标

完成本课学习后，你将能够：

理解 WebSocket 协议在实时应用（如聊天室）中的作用与优势。
掌握使用 Go 语言标准库 net/http 及第三方库 gorilla/websocket 构建 WebSocket 服务端。
学会利用 Go 的 channel 和 goroutine 高效管理多个并发的 WebSocket 客户端连接。
实现一个包含消息广播、用户上下线通知等核心功能的并发安全的聊天室系统。

核心概念

1. WebSocket：实时双向通信

想象一下传统的 HTTP 请求，就像写信。你（客户端）寄一封信（请求），然后必须等待回信（响应）才能进行下一步。这在聊天这种需要即时性的场景下效率极低。 WebSocket 则像一条电话专线。连接一旦建立，双方就可以随时自由地“说话”（发送数据），而不需要每次都重新拨号（发起新的 HTTP 请求）。这种全双工、持久连接的特性，使其成为实时应用（聊天、直播、协作编辑）的理想选择。

2. Go 的并发原语是绝佳工具

一个聊天室服务需要同时与成百上千个用户保持连接。Go 的 goroutine（轻量级线程）可以轻松地为每一个连接分配一个独立的执行上下文，互不阻塞。而 channel 则像是连接这些 goroutine 的“安全管道”，用于在它们之间安全地传递消息和控制信号，避免复杂繁琐的锁机制。

3. 架构概览：Hub 模式

我们将采用经典的 Hub（中心）模式 来构建聊天室。

Hub：中央协调者，负责管理所有活跃的客户端连接，维护一个在线用户列表，并充当消息广播站。
Client：代表一个用户连接。每个客户端有两个核心任务：1) 持续读取来自浏览器的 WebSocket 消息并转发给 Hub；2) 持续监听来自 Hub 的广播消息并发送回浏览器。
通信流程： 用户A -> 浏览器 -> WebSocket -> Client A -> (通过 channel) -> Hub -> (通过 channel) -> Client B -> 浏览器 -> 用户B。

代码示例

我们将使用 github.com/gorilla/websocket 库来简化 WebSocket 的操作。请先安装：go get github.com/gorilla/websocket。

项目结构：

chatroom/
├── main.go
├── hub.go
└── client.go

1. hub.go - 聊天室的心脏

package main

import "sync"

// Hub 维护所有活跃客户端的集合，并负责广播消息。
type Hub struct {
	// 注册的客户端集合 (map的key是客户端指针，value用于快速查找)
	clients map[*Client]bool

	// 来自客户端的消息，准备广播给所有其他客户端。
	broadcast chan []byte

	// 注册新客户端的请求
	register chan *Client

	// 注销客户端的请求
	unregister chan *Client
	// 使用互斥锁保护 clients map，确保并发安全
	mu sync.RWMutex
}

func newHub() *Hub {
	return &Hub{
		broadcast:  make(chan []byte),
		register:   make(chan *Client),
		unregister: make(chan *Client),
		clients:    make(map[*Client]bool),
	}
}

// run 是Hub的核心事件循环，监听所有通道。
func (h *Hub) run() {
	for {
		select {
		case client := <-h.register:
			h.mu.Lock()
			h.clients[client] = true
			h.mu.Unlock()
			// 通知其他人有新用户加入
			joinMsg := []byte("系统: " + client.nickname + " 加入了聊天室")
			h.broadcast <- joinMsg
		case client := <-h.unregister:
			h.mu.Lock()
			if _, ok := h.clients[client]; ok {
				delete(h.clients, client)
				close(client.send) // 关闭发送通道，通知客户端协程退出
				// 通知其他人有用户离开
				leaveMsg := []byte("系统: " + client.nickname + " 离开了聊天室")
				h.broadcast <- leaveMsg
			}
			h.mu.Unlock()
		case message := <-h.broadcast:
			h.mu.RLock()
			// 将消息发送给每个客户端的 `send` 通道
			for client := range h.clients {
				select {
				case client.send <- message:
				default:
					// 如果客户端的发送缓冲区满了（可能网络慢），则视为掉线
					close(client.send)
					delete(h.clients, client)
				}
			}
			h.mu.RUnlock()
		}
	}
}

2. client.go - 每个用户的代理

package main

import (
	"log"
	"time"

	"github.com/gorilla/websocket"
)

const (
	// 允许从客户端读取的最大消息大小
	maxMessageSize = 512
	// 心跳超时时间
	pongWait = 60 * time.Second
	// 心跳发送间隔
	pingInterval = (pongWait * 9) / 10
)

// Client 是一个中间人，在 WebSocket 连接和 Hub 之间。
type Client struct {
	hub *Hub

	// WebSocket 连接
	conn *websocket.Conn

	// 出站消息的缓冲通道
	send chan []byte

	// 用户昵称
	nickname string
}

// readPump 将消息从 WebSocket 连接泵送到 hub。
func (c *Client) readPump() {
	defer func() {
		c.hub.unregister <- c
		c.conn.Close()
	}()
	c.conn.SetReadLimit(maxMessageSize)
	c.conn.SetReadDeadline(time.Now().Add(pongWait))
	c.conn.SetPongHandler(func(string) error {
		c.conn.SetReadDeadline(time.Now().Add(pongWait))
		return nil
	})
	for {
		_, message, err := c.conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
				log.Printf("读取错误: %v", err)
			}
			break
		}
		// 在消息前加上用户昵称
		msgWithNick := []byte(c.nickname + ": " + string(message))
		c.hub.broadcast <- msgWithNick
	}
}

// writePump 将消息从 hub 泵送到 WebSocket 连接。
func (c *Client) writePump() {
	ticker := time.NewTicker(pingInterval)
	defer func() {
		ticker.Stop()
		c.conn.Close()
	}()
	for {
		select {
		case message, ok := <-c.send:
			c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if !ok {
				// Hub 关闭了通道，发送关闭帧。
				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}
			w, err := c.conn.NextWriter(websocket.TextMessage)
			if err != nil {
				return
			}
			w.Write(message)
			// 将队列中剩余的排队消息添加到当前 WebSocket 消息中。
			n := len(c.send)
			for i := 0; i < n; i++ {
				w.Write(<-c.send)
			}
			if err := w.Close(); err != nil {
				return
			}
		case <-ticker.C:
			// 定期发送 Ping 帧，保持连接活跃
			c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

3. main.go - 程序入口与 HTTP 升级处理

package main

import (
	"flag"
	"log"
	"net/http"
)

var addr = flag.String("addr", ":8080", "http 服务地址")

// upgrader 负责将普通 HTTP 连接升级为 WebSocket 连接。
var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	// 允许所有来源的连接（在生产环境中应进行验证）
	CheckOrigin: func(r *http.Request) bool { return true },
}

func serveWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}
	// 从URL查询参数中获取昵称，例如 ws://localhost:8080/ws?nickname=Tom
	nickname := r.URL.Query().Get("nickname")
	if nickname == "" {
		nickname = "匿名"
	}
	client := &Client{hub: hub, conn: conn, send: make(chan []byte, 256), nickname: nickname}
	client.hub.register <- client

	// 启动两个独立的协程，分别用于读写。
	go client.writePump()
	go client.readPump()
}

func serveHome(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.Error(w, "Not found", http.StatusNotFound)
		return
	}
	if r.Method != http.MethodGet {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}
	http.ServeFile(w, r, "home.html")
}

func main() {
	flag.Parse()
	hub := newHub()
	go hub.run() // 启动Hub的主事件循环
	http.HandleFunc("/", serveHome)
	http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
		serveWs(hub, w, r)
	})
	log.Printf("聊天室服务启动在: %s", *addr)
	err := http.ListenAndServe(*addr, nil)
	if err != nil {
		log.Fatal("ListenAndServe: ", err)
	}
}

4. home.html - 简单的前端页面（用于测试）

<!DOCTYPE html>
<html>
<head>
    <title>Go WebSocket 聊天室</title>
</head>
<body>
    <h1>WebSocket 聊天室</h1>
    <input id="nickname" type="text" placeholder="输入昵称" value="用户A">
    <button onclick="connect()">连接</button>
    <br><br>
    <div id="output" style="border:1px solid #ccc; height:300px; overflow-y: scroll; padding:10px; width:500px;"></div>
    <br>
    <input id="message" type="text" placeholder="输入消息" style="width:400px;" onkeydown="if(event.key==='Enter') sendMessage()">
    <button onclick="sendMessage()">发送</button>

    <script>
        let ws;
        function connect() {
            const nickname = document.getElementById('nickname').value;
            ws = new WebSocket(`ws://${location.host}/ws?nickname=${encodeURIComponent(nickname)}`);
            ws.onopen = function() {
                appendOutput("系统: 已连接到聊天室");
            };
            ws.onmessage = function(event) {
                appendOutput(event.data);
            };
            ws.onclose = function() {
                appendOutput("系统: 连接已关闭");
            };
        }
        function sendMessage() {
            const input = document.getElementById('message');
            if (ws && ws.readyState === WebSocket.OPEN && input.value) {
                ws.send(input.value);
                input.value = '';
            }
        }
        function appendOutput(message) {
            const output = document.getElementById('output');
            output.innerHTML += message + '<br>';
            output.scrollTop = output.scrollHeight;
        }
    </script>
</body>
</html>

运行步骤：

创建上述文件。
在项目目录下运行 go run .。
在浏览器中打开 http://localhost:8080。
在多个浏览器标签页中打开此页面，输入不同的昵称并点击“连接”，即可开始聊天。

实践练习

练习 1：在线用户列表

要求：扩展 Hub 功能，使其能够维护一个在线用户的昵称列表。当用户加入或离开时，将更新后的列表广播给所有客户端。前端页面应能显示当前在线的所有用户。提示：

在 Hub 中增加一个字段，如 onlineUsers map[string]bool 或使用 sync.Map。
在 register 和 unregister 操作时更新该列表。
广播一个结构化的 JSON 消息，包含 type（如 "system" 或 "userlist"）和 data（如用户列表）。 预期输出：用户加入/离开时，所有客户端的界面上能实时看到在线用户数量的变化。

练习 2：私聊功能

要求：允许用户通过特定命令（如 /w 用户名消息内容）给另一个在线用户发送私聊消息。只有目标用户和发送者能看到该消息。提示：

修改客户端发送的消息格式，识别以 /w 开头的指令。
在 client.go 的 readPump 中解析指令，提取目标用户名和私聊内容。
Hub 需要维护一个 nickname 到 client 的映射，以便能根据用户名找到对应的客户端连接。
私聊消息不通过 broadcast 通道，而是直接通过 Hub 找到目标客户端的 send 通道发送。 预期输出：执行 /w Bob 你好 后，只有用户“Bob”和发送者自己能看到“你悄悄对 Bob 说: 你好”这样的消息，其他用户不可见。

常见错误

并发访问 map：直接在多个 goroutine 中读写 Hub.clients map 会导致程序崩溃（fatal error: concurrent map read and map write）。必须使用 sync.RWMutex 进行保护，或者使用并发安全的 sync.Map。
Channel 死锁：如果 Client.send 通道已满（例如，客户端网络极慢），Hub 在向其发送消息时会发生阻塞，导致整个 Hub 的 run() 循环卡住。解决方案是使用 select 语句和 default 分支，当无法立即发送时，断开该慢速客户端的连接（如代码示例中所示）。
WebSocket 连接未正确关闭：在客户端协程退出时，如果没有通过 Hub.unregister 正确注销，会导致资源泄漏。务必确保 readPump 的 defer 函数中执行注销操作。
心跳机制缺失：如果客户端异常断开（如断网），服务端可能无法立即察觉，从而保留无效的连接。实现 Ping/Pong 心跳机制（如代码中的 pingInterval 和 PongHandler）是保持连接活跃和检测僵尸连接的关键。
前端重连逻辑：示例前端代码没有实现断线重连。在生产环境中，应在 ws.onclose 事件中添加重连逻辑（如指数退避重试）。

小结

本课我们通过构建一个实战性的 WebSocket 聊天室，深入理解了 Go 并发编程在实时 Web 应用中的强大威力。关键要点如下：

WebSocket 提供了低延迟、双向的通信通道，是构建实时应用的基石。
Hub 模式 是一个清晰、可扩展的并发服务架构，通过中心化的协调者管理所有连接和消息流。
Goroutine 为每个连接提供了轻量级的并发执行单元，使代码逻辑清晰，性能高效。
Channel 是连接各个并发单元的“安全桥梁”，用于传递消息

实战：WebSocket 聊天室