第 45 课 - WebSocket 实时通信

学习目标

通过本课学习，你将能够：

理解 WebSocket 协议与 HTTP 轮询的区别
掌握使用 gorilla/websocket 库创建 WebSocket 服务器
实现基本的双向实时通信功能
构建简单的多人聊天室应用
处理 WebSocket 连接中的常见并发问题

核心概念

什么是 WebSocket？

WebSocket 是一种在单个 TCP 连接上进行全双工通信的协议。它解决了传统 HTTP 协议只能从客户端发送请求的局限，允许服务器主动向客户端推送数据。

与 HTTP 轮询的对比

// 传统 HTTP 轮询（客户端定期请求）
// 客户端每 5 秒向服务器请求一次新消息
// 问题：资源浪费、延迟高、实时性差

// WebSocket 方式
// 连接建立后，双方可以随时发送消息
// 优点：实时性高、资源占用少、支持二进制数据

gorilla/websocket 库

Go 生态中最流行的 WebSocket 实现库，提供了完整的 WebSocket 协议支持：

支持文本和二进制消息
内置消息压缩
优雅的连接关闭机制
完善的错误处理

代码示例

1. 基础 WebSocket 服务器

package main

import (
	"fmt"
	"log"
	"net/http"

	"github.com/gorilla/websocket"
)

// 升级 HTTP 连接为 WebSocket 连接
var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	// 允许所有 CORS 请求（生产环境应限制）
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func wsHandler(w http.ResponseWriter, r *http.Request) {
	// 升级连接
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("升级失败: %v", err)
		return
	}
	defer conn.Close()

	fmt.Println("新客户端连接建立")

	// 循环读取消息
	for {
		// 读取客户端消息
		messageType, message, err := conn.ReadMessage()
		if err != nil {
			if websocket.IsCloseError(err, websocket.CloseNormalClosure, 
				websocket.CloseGoingAway) {
				fmt.Println("客户端正常断开连接")
			} else {
				log.Printf("读取错误: %v", err)
			}
			break
		}

		// 打印收到的消息
		fmt.Printf("收到消息: %s\n", string(message))

		// 回显消息（回声服务器）
		err = conn.WriteMessage(messageType, message)
		if err != nil {
			log.Printf("写入错误: %v", err)
			break
		}
	}
}

func main() {
	http.HandleFunc("/ws", wsHandler)
	
	fmt.Println("WebSocket 服务器启动在 :8080")
	log.Fatal(http.ListenAndServe(":8080", nil))
}

2. 简单聊天室实现

package main

import (
	"fmt"
	"log"
	"net/http"
	"sync"

	"github.com/gorilla/websocket"
)

// 客户端结构体
type Client struct {
	conn *websocket.Conn
	send chan []byte
}

// 聊天室
type ChatRoom struct {
	clients    map[*Client]bool
	broadcast  chan []byte
	register   chan *Client
	unregister chan *Client
	mutex      sync.RWMutex
}

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

func (cr *ChatRoom) Run() {
	for {
		select {
		case client := <-cr.register:
			cr.mutex.Lock()
			cr.clients[client] = true
			cr.mutex.Unlock()
			fmt.Printf("新客户端加入，当前在线: %d\n", len(cr.clients))

		case client := <-cr.unregister:
			cr.mutex.Lock()
			if _, ok := cr.clients[client]; ok {
				delete(cr.clients, client)
				close(client.send)
			}
			cr.mutex.Unlock()
			fmt.Printf("客户端离开，当前在线: %d\n", len(cr.clients))

		case message := <-cr.broadcast:
			cr.mutex.RLock()
			for client := range cr.clients {
				select {
				case client.send <- message:
				default:
					// 如果发送失败，关闭连接
					close(client.send)
					delete(cr.clients, client)
				}
			}
			cr.mutex.RUnlock()
		}
	}
}

// 处理单个客户端连接
func (cr *ChatRoom) HandleClient(conn *websocket.Conn) {
	client := &Client{
		conn: conn,
		send: make(chan []byte, 256),
	}

	// 注册客户端
	cr.register <- client

	// 启动读写协程
	go client.writePump()
	go client.readPump(cr)
}

// 客户端读取消息
func (c *Client) readPump(cr *ChatRoom) {
	defer func() {
		cr.unregister <- c
		c.conn.Close()
	}()

	for {
		_, message, err := c.conn.ReadMessage()
		if err != nil {
			break
		}
		// 广播消息
		cr.broadcast <- message
	}
}

// 客户端写入消息
func (c *Client) writePump() {
	defer c.conn.Close()

	for message := range c.send {
		err := c.conn.WriteMessage(websocket.TextMessage, message)
		if err != nil {
			break
		}
	}
}

var chatRoom = NewChatRoom()

func wsHandler(w http.ResponseWriter, r *http.Request) {
	upgrader := websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("升级失败: %v", err)
		return
	}

	chatRoom.HandleClient(conn)
}

func main() {
	go chatRoom.Run()

	http.HandleFunc("/ws", wsHandler)
	http.Handle("/", http.FileServer(http.Dir("./public")))

	fmt.Println("聊天室服务器启动在 :8080")
	log.Fatal(http.ListenAndServe(":8080", nil))
}

3. 前端客户端 (public/index.html)

<!DOCTYPE html>
<html>
<head>
    <title>WebSocket 聊天室</title>
    <style>
        body { font-family: Arial, sans-serif; max-width: 600px; margin: 50px auto; }
        #messages { height: 400px; border: 1px solid #ccc; overflow-y: scroll; padding: 10px; margin-bottom: 20px; }
        #input { width: 80%; padding: 10px; }
        button { padding: 10px 20px; background: #007bff; color: white; border: none; cursor: pointer; }
    </style>
</head>
<body>
    <h1>WebSocket 聊天室</h1>
    <div id="messages"></div>
    <input type="text" id="input" placeholder="输入消息..." autofocus>
    <button onclick="sendMessage()">发送</button>

    <script>
        const messages = document.getElementById('messages');
        const input = document.getElementById('input');
        let ws;

        function connect() {
            // 建立 WebSocket 连接
            ws = new WebSocket('ws://localhost:8080/ws');

            ws.onopen = function() {
                addMessage('系统', '已连接到服务器');
            };

            ws.onmessage = function(event) {
                addMessage('其他用户', event.data);
            };

            ws.onclose = function() {
                addMessage('系统', '连接已断开，3秒后重连...');
                setTimeout(connect, 3000);
            };

            ws.onerror = function(error) {
                addMessage('错误', '连接错误');
            };
        }

        function addMessage(sender, content) {
            const div = document.createElement('div');
            div.innerHTML = `<strong>${sender}:</strong> ${content}`;
            messages.appendChild(div);
            messages.scrollTop = messages.scrollHeight;
        }

        function sendMessage() {
            const message = input.value.trim();
            if (message && ws.readyState === WebSocket.OPEN) {
                ws.send(message);
                addMessage('我', message);
                input.value = '';
            }
        }

        // 按回车发送
        input.addEventListener('keypress', function(e) {
            if (e.key === 'Enter') {
                sendMessage();
            }
        });

        // 建立连接
        connect();
    </script>
</body>
</html>

实践练习

练习 1: 回声服务器增强

创建一个 WebSocket 服务器，实现以下功能：

接收客户端消息
将消息转为大写后返回
在消息前添加时间戳
如果收到 "quit" 消息，主动关闭连接

预期输出：

客户端发送: hello
服务器返回: [2024-01-15 10:30:45] HELLO

练习 2: 实时通知系统

实现一个简单的实时通知系统：

服务器维护一个在线用户列表
当有新用户加入时，通知所有在线用户
用户可以发送 "broadcast:xxx" 格式的消息广播给所有人
实现用户私聊功能 "私聊@用户ID:消息"

要求： 设计消息协议格式，处理用户标识

练习 3: 房间管理系统

扩展聊天室功能：

支持创建多个聊天房间
用户可以选择加入特定房间
房间内消息只在同一房间内广播
实现房间人数统计功能

提示： 使用 map 存储多个 ChatRoom 实例

常见错误

1. 并发读写问题

// 错误示例：在多个协程中同时读写同一个连接
go func() {
    conn.WriteMessage(websocket.TextMessage, []byte("msg1"))
}()
go func() {
    conn.WriteMessage(websocket.TextMessage, []byte("msg2")) // 并发写入可能导致 panic
}()

// 正确做法：使用单个协程处理写入，通过 channel 传递消息
type safeWriter struct {
    conn *websocket.Conn
    send chan []byte
}

func (sw *safeWriter) run() {
    for msg := range sw.send {
        if err := sw.conn.WriteMessage(websocket.TextMessage, msg); err != nil {
            break
        }
    }
}

2. 未处理连接关闭

// 错误示例：忽略连接关闭错误
for {
    _, msg, err := conn.ReadMessage()
    if err != nil {
        log.Printf("读取错误: %v", err)
        continue // 继续读取可能导致无限循环
    }
    // 处理消息...
}

// 正确做法：区分正常关闭和错误
for {
    _, msg, err := conn.ReadMessage()
    if err != nil {
        if websocket.IsCloseError(err, 
            websocket.CloseNormalClosure, 
            websocket.CloseGoingAway) {
            log.Println("客户端正常关闭")
        } else {
            log.Printf("读取错误: %v", err)
        }
        break // 退出循环
    }
    // 处理消息...
}

3. 内存泄漏

// 错误示例：客户端断开但未从列表中移除
type ChatRoom struct {
    clients map[*Client]bool
}

func (cr *ChatRoom) handleClient(conn *websocket.Conn) {
    client := &Client{conn: conn}
    cr.clients[client] = true
    
    // 如果客户端异常断开，不会从 map 中移除
    // 导致内存泄漏
}

// 正确做法：使用 defer 和清理函数
func (cr *ChatRoom) handleClient(conn *websocket.Conn) {
    client := &Client{conn: conn}
    cr.clients[client] = true
    
    defer func() {
        delete(cr.clients, client)
        conn.Close()
    }()
    
    // 处理连接...
}

4. 忽略消息缓冲区大小

// 潜在问题：客户端发送超大消息导致内存问题
conn.SetReadLimit(4096) // 设置最大消息大小为 4KB

// 添加读取超时
conn.SetReadDeadline(time.Now().Add(60 * time.Second))

小结

关键知识点

WebSocket 优势：全双工通信、低延迟、适合实时应用
gorilla/websocket：Go 生态中成熟的 WebSocket 实现库
连接管理：需要妥善处理连接的建立、维护和关闭
并发安全：WebSocket 连接不是并发安全的，需要正确同步
资源清理：及时移除断开连接的客户端，避免内存泄漏

设计建议

消息协议：设计清晰的消息格式（JSON 或自定义协议）
心跳机制：定期发送 ping/pong 保持连接活跃
错误处理：区分不同类型的错误并采取相应措施
限流控制：设置消息大小限制和频率限制

下一步学习

学习使用 WebSocket 实现文件传输
了解 WebSocket 安全最佳实践（WSS、认证、授权）
探索生产环境部署方案（Nginx 代理、负载均衡）
研究其他实时通信技术：SSE、gRPC 流式传输

// 完整项目结构建议
/*
websocket-chat/
├── main.go            # 主程序入口
├── handler.go         # HTTP 处理器
├── client.go          # 客户端结构体
├── chatroom.go        # 聊天室逻辑
├── public/
│   └── index.html     # 前端页面
└── go.mod             # 模块定义
*/

通过本课学习，你已经掌握了 Go 语言中 WebSocket 编程的核心技术。WebSocket 是构建实时 Web 应用的重要基础，在聊天室、实时数据看板、协同编辑、在线游戏等场景中有广泛应用。

WebSocket 实时通信