WebSocket 支持

WebSocket 提供全双工通信能力，适合实时应用。Gin 可以配合 gorilla/websocket 实现 WebSocket 功能。

基本配置

package main

import (
    "log"
    "net/http"
    
    "github.com/gin-gonic/gin"
    "github.com/gorilla/websocket"
)

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

func main() {
    r := gin.Default()
    
    r.GET("/ws", func(c *gin.Context) {
        conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
        if err != nil {
            log.Printf("WebSocket upgrade error: %v", err)
            return
        }
        defer conn.Close()
        
        for {
            messageType, message, err := conn.ReadMessage()
            if err != nil {
                log.Printf("Read error: %v", err)
                break
            }
            
            log.Printf("Received: %s", message)
            
            err = conn.WriteMessage(messageType, message)
            if err != nil {
                log.Printf("Write error: %v", err)
                break
            }
        }
    })
    
    r.Run(":8080")
}

聊天室示例

type Client struct {
    Conn *websocket.Conn
    Send chan []byte
}

type ChatRoom struct {
    Clients    map[*Client]bool
    Register   chan *Client
    Unregister chan *Client
    Broadcast  chan []byte
}

func NewChatRoom() *ChatRoom {
    return &ChatRoom{
        Clients:    make(map[*Client]bool),
        Register:   make(chan *Client),
        Unregister: make(chan *Client),
        Broadcast:  make(chan []byte),
    }
}

func (r *ChatRoom) Run() {
    for {
        select {
        case client := <-r.Register:
            r.Clients[client] = true
            
        case client := <-r.Unregister:
            if _, ok := r.Clients[client]; ok {
                delete(r.Clients, client)
                close(client.Send)
            }
            
        case message := <-r.Broadcast:
            for client := range r.Clients {
                select {
                case client.Send <- message:
                default:
                    close(client.Send)
                    delete(r.Clients, client)
                }
            }
        }
    }
}

func (c *Client) ReadPump(room *ChatRoom) {
    defer func() {
        room.Unregister <- c
        c.Conn.Close()
    }()
    
    for {
        _, message, err := c.Conn.ReadMessage()
        if err != nil {
            break
        }
        room.Broadcast <- message
    }
}

func (c *Client) WritePump() {
    defer c.Conn.Close()
    
    for message := range c.Send {
        if err := c.Conn.WriteMessage(websocket.TextMessage, message); err != nil {
            break
        }
    }
}

var chatRoom = NewChatRoom()

func main() {
    go chatRoom.Run()
    
    r := gin.Default()
    
    r.GET("/ws/chat", func(c *gin.Context) {
        conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
        if err != nil {
            return
        }
        
        client := &Client{
            Conn: conn,
            Send: make(chan []byte, 256),
        }
        
        chatRoom.Register <- client
        
        go client.WritePump()
        client.ReadPump(chatRoom)
    })
    
    r.Run(":8080")
}

心跳检测

const (
    writeWait      = 10 * time.Second
    pongWait       = 60 * time.Second
    pingPeriod     = (pongWait * 9) / 10
    maxMessageSize = 512
)

func (c *Client) ReadPumpWithPing(room *ChatRoom) {
    defer func() {
        room.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 {
            break
        }
        room.Broadcast <- message
    }
}

func (c *Client) WritePumpWithPing() {
    ticker := time.NewTicker(pingPeriod)
    defer func() {
        ticker.Stop()
        c.Conn.Close()
    }()
    
    for {
        select {
        case message, ok := <-c.Send:
            c.Conn.SetWriteDeadline(time.Now().Add(writeWait))
            if !ok {
                c.Conn.WriteMessage(websocket.CloseMessage, []byte{})
                return
            }
            
            c.Conn.WriteMessage(websocket.TextMessage, message)
            
        case <-ticker.C:
            c.Conn.SetWriteDeadline(time.Now().Add(writeWait))
            if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
                return
            }
        }
    }
}

前端连接示例

const socket = new WebSocket('ws://localhost:8080/ws');

socket.onopen = function(e) {
    console.log('Connected');
    socket.send('Hello Server');
};

socket.onmessage = function(event) {
    console.log('Received:', event.data);
};

socket.onclose = function(event) {
    console.log('Disconnected');
};

socket.onerror = function(error) {
    console.log('Error:', error);
};

小结

WebSocket 适合实时通信场景。使用 gorilla/websocket 可以轻松在 Gin 中实现 WebSocket 功能。注意处理连接关闭、心跳检测和并发安全。