接入层实现:WebSocket 网关与连接管理
本章目标:在 gateway/ 里实现「升级 → 注册 → 一读一写 → 心跳 → 注销」,让连接稳定保活、下行永远只有一个协程在写。
server.go 管握手与托管,client.go 是单条连接(读协程 + 写协程 + 心跳),hub.go 是本机连接注册表。协议编解码交给第 03 章,令牌校验交给第 04 章,网关只留两个钩子。
1. Hub:本机连接注册表
// gateway/hub.go
type Hub struct {
mu sync.RWMutex
byUID map[uint64]map[uint64]*Client // uid → connID → 连接(支持多端)
}
func NewHub() *Hub { return &Hub{byUID: map[uint64]map[uint64]*Client{}} }
func (h *Hub) Register(c *Client) {
h.mu.Lock()
defer h.mu.Unlock()
if h.byUID[c.UID] == nil {
h.byUID[c.UID] = map[uint64]*Client{}
}
h.byUID[c.UID][c.ConnID] = c
}
func (h *Hub) Unregister(c *Client) {
h.mu.Lock()
defer h.mu.Unlock()
delete(h.byUID[c.UID], c.ConnID)
if len(h.byUID[c.UID]) == 0 {
delete(h.byUID, c.UID) // 该用户已无连接,键也删掉
}
}
上下行分离的要害在 Client.send:它是一个带缓冲的通道,下行只做非阻塞入队(select 带 default,队列满就丢,由 ACK 重发兜底,第 10 章),网络写全部收敛到 WritePump;若把网络写放进持锁区,一个慢客户端就会拖住整张注册表。
2. Client:一读一写加心跳
// gateway/client.go:Client 六个字段——ConnID、UID、conn、hub、send(下行队列)、closed(关闭信号)
const writeWait, pongWait, pingPeriod = 10 * time.Second, 70 * time.Second, 30 * time.Second
// ReadPump 是唯一读协程:读到的字节交给 handler,出错即注销连接。
func (c *Client) ReadPump(handler func(*Client, []byte)) {
defer func() { c.hub.Unregister(c); _ = c.conn.Close() }()
if err := c.conn.SetReadDeadline(time.Now().Add(pongWait)); err != nil {
return
}
c.conn.SetPongHandler(func(string) error { // 收到 Pong 就续命
return c.conn.SetReadDeadline(time.Now().Add(pongWait))
})
for {
_, data, err := c.conn.ReadMessage()
if err != nil {
return // 读超时、对端关闭或帧非法都会走到这里
}
handler(c, data)
}
}
// WritePump 是唯一写协程:所有下行都进 send 通道,写操作天然串行。
func (c *Client) WritePump() {
ticker := time.NewTicker(pingPeriod)
defer ticker.Stop()
for {
select {
case <-c.closed:
return
case payload := <-c.send:
if err := c.conn.SetWriteDeadline(time.Now().Add(writeWait)); err != nil {
return
}
if err := c.conn.WriteMessage(websocket.BinaryMessage, payload); err != nil {
return
}
case <-ticker.C:
if err := c.conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(writeWait)); err != nil {
return
}
}
}
}
关闭用幂等的 Close(code, reason):sync.Once 里发关闭帧、close(c.closed)、conn.Close() 三步。用 closed 通道而不是 close(send),因为并发下行无法保证「没有人在发」,往已关闭的通道发送会 panic。
3. Server:升级与托管
// gateway/server.go
func (s *Server) HandleWS(c *gin.Context) {
uid, deviceID, err := s.Authenticate(c.Request) // 第 04 章注入
if err != nil {
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "unauthorized"})
return
}
conn, err := s.upgrader.Upgrade(c.Writer, c.Request, nil)
if err != nil {
return // 库已经把响应写回去了,这里只记日志
}
client := NewClient(s.nextConnID.Add(1), uid, deviceID, conn, s.hub)
s.hub.Register(client)
go client.WritePump()
client.ReadPump(s.OnMessage) // 阻塞到断开,注销由读协程负责
}
upgrader 要配三处:读写缓冲各 4096、CheckOrigin 按白名单校验来源(别一律放行)、SetReadLimit 按业务最大消息的两倍设上限。
4. 验证方式
go build -o bin/gateway ./cmd/gateway && ./bin/gateway -conf configs/config.yaml
curl -i -N -H "Connection: Upgrade" -H "Upgrade: websocket" -H "Sec-WebSocket-Version: 13" -H "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==" http://127.0.0.1:8080/ws # 期望 401
websocat "ws://127.0.0.1:8080/ws?token=<令牌>&device_id=dev-1" # 挂 90 秒不掉线即心跳正常
ss -tnp '( sport = :8080 )' | grep -c ESTAB # 断开后应回到 0
常见坑
| 坑 | 后果 | 做法 |
|---|---|---|
多协程直接 WriteMessage | 库会 panic(并发写) | 下行只进 send 通道,由 WritePump 独占写 |
用 close(send) 表示关闭 | 并发发送时 panic | 另用 closed 通道做关闭信号 |
| 心跳周期大于代理空闲超时 | 被 Nginx 60 秒静默断开 | ping 周期取 30 秒以内(第 14 章) |
小结:接入层的骨架就是「一张注册表 + 每连接两个协程 + 一条写队列」——Hub 管谁在本机,ReadPump 只管读、WritePump 只管写,心跳靠 ping/pong 与读超时配合。第 03 章接着定义这条连接上跑的字节到底是什么。