接入层实现: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 章接着定义这条连接上跑的字节到底是什么。

笔记加载中…