Go语言WebSocket实时通信实战：从基础到广播机制-深圳市維司達科技有限公司

Go语言WebSocket实时通信实战：从基础到广播机制

引言

WebSocket是一种全双工通信协议，允许服务器主动向客户端推送数据。Go语言的gorilla/websocket库提供了优秀的WebSocket支持。本文将深入探讨Go语言的WebSocket编程实践，帮助您构建实时通信应用。

一、WebSocket基础

1.1 WebSocket协议

HTTP握手请求: GET /ws HTTP/1.1 Host: example.com Upgrade: websocket Connection: Upgrade Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ== Sec-WebSocket-Version: 13 HTTP握手响应: HTTP/1.1 101 Switching Protocols Upgrade: websocket Connection: Upgrade Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=

1.2 安装gorilla/websocket

go get github.com/gorilla/websocket

二、基础WebSocket服务器

2.1 简单服务器

import ( "log" "net/http" "github.com/gorilla/websocket" ) var upgrader = websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 1024, } func wsHandler(w http.ResponseWriter, r *http.Request) { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { log.Println("Upgrade error:", err) return } defer conn.Close() for { messageType, p, err := conn.ReadMessage() if err != nil { log.Println("Read error:", err) break } log.Printf("Received: %s", p) err = conn.WriteMessage(messageType, p) if err != nil { log.Println("Write error:", err) break } } } func main() { http.HandleFunc("/ws", wsHandler) log.Fatal(http.ListenAndServe(":8080", nil)) }

2.2 配置选项

var upgrader = websocket.Upgrader{ ReadBufferSize: 4096, WriteBufferSize: 4096, CheckOrigin: func(r *http.Request) bool { // 允许所有来源 return true }, Subprotocols: []string{"chat", "binary"}, }

三、广播机制

3.1 Hub-Client模型

type Hub struct { clients map[*Client]bool broadcast chan []byte register chan *Client unregister chan *Client } func NewHub() *Hub { return &Hub{ broadcast: make(chan []byte), register: make(chan *Client), unregister: make(chan *Client), clients: make(map[*Client]bool), } } func (h *Hub) Run() { for { select { case client := <-h.register: h.clients[client] = true case client := <-h.unregister: if _, ok := h.clients[client]; ok { delete(h.clients, client) close(client.send) } case message := <-h.broadcast: for client := range h.clients { select { case client.send <- message: default: close(client.send) delete(h.clients, client) } } } } }

3.2 Client结构体

type Client struct { hub *Hub conn *websocket.Conn send chan []byte } func (c *Client) readPump() { defer func() { c.hub.unregister <- c c.conn.Close() }() c.conn.SetReadLimit(512) c.conn.SetReadDeadline(time.Now().Add(60 * time.Second)) for { _, message, err := c.conn.ReadMessage() if err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) { log.Printf("error: %v", err) } break } c.hub.broadcast <- message } } func (c *Client) writePump() { defer func() { c.conn.Close() }() for message := range c.send { c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) w, err := c.conn.NextWriter(websocket.TextMessage) if err != nil { return } w.Write(message) if err := w.Close(); err != nil { return } } }

3.3 完整服务器

func serveWs(hub *Hub, w http.ResponseWriter, r *http.Request) { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { log.Println(err) return } client := &Client{ hub: hub, conn: conn, send: make(chan []byte, 256), } client.hub.register <- client go client.writePump() go client.readPump() } func main() { hub := NewHub() go hub.Run() http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) { serveWs(hub, w, r) }) log.Fatal(http.ListenAndServe(":8080", nil)) }

四、心跳检测

4.1 服务器端心跳

func (c *Client) readPump() { defer func() { c.hub.unregister <- c c.conn.Close() }() c.conn.SetReadLimit(512) c.conn.SetReadDeadline(time.Now().Add(60 * time.Second)) c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(60 * time.Second)) return nil }) for { _, message, err := c.conn.ReadMessage() if err != nil { break } c.hub.broadcast <- message } } func (c *Client) writePump() { ticker := time.NewTicker(60 * time.Second) 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 { c.conn.WriteMessage(websocket.CloseMessage, []byte{}) return } w, err := c.conn.NextWriter(websocket.TextMessage) if err != nil { return } w.Write(message) if err := w.Close(); err != nil { return } case <-ticker.C: c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil { return } } } }

五、消息压缩

5.1 启用压缩

var upgrader = websocket.Upgrader{ ReadBufferSize: 4096, WriteBufferSize: 4096, EnableCompression: true, }

5.2 自定义压缩配置

import "github.com/gorilla/websocket/compression" func init() { compression.SetLevel(compression.BestCompression) }

六、WSS配置

6.1 自签名证书

openssl req -x509 -newkey rsa:4096 -keyout key.pem -out cert.pem -days 365

6.2 HTTPS服务器

func main() { hub := NewHub() go hub.Run() http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) { serveWs(hub, w, r) }) log.Fatal(http.ListenAndServeTLS(":8443", "cert.pem", "key.pem", nil)) }

七、性能优化

7.1 连接池

type ConnectionPool struct { pool sync.Pool } func NewConnectionPool() *ConnectionPool { return &ConnectionPool{ pool: sync.Pool{ New: func() interface{} { return make([]byte, 4096) }, }, } } func (p *ConnectionPool) Get() []byte { return p.pool.Get().([]byte) } func (p *ConnectionPool) Put(buf []byte) { p.pool.Put(buf[:0]) }

7.2 批量发送

func (c *Client) batchSend(messages [][]byte) error { w, err := c.conn.NextWriter(websocket.TextMessage) if err != nil { return err } for _, msg := range messages { w.Write(msg) w.Write([]byte("\n")) } return w.Close() }

八、实战：聊天室应用

type ChatServer struct { hub *Hub messages []string mu sync.RWMutex } func NewChatServer() *ChatServer { cs := &ChatServer{ hub: NewHub(), messages: make([]string, 0, 100), } go cs.hub.Run() return cs } func (cs *ChatServer) broadcastMessage(msg string) { cs.mu.Lock() cs.messages = append(cs.messages, msg) if len(cs.messages) > 100 { cs.messages = cs.messages[1:] } cs.mu.Unlock() cs.hub.broadcast <- []byte(msg) } func (cs *ChatServer) getHistory() []string { cs.mu.RLock() defer cs.mu.RUnlock() return append([]string{}, cs.messages...) }