WebSocket #

一、WebSocket概述 #

WebSocket是一种在单个TCP连接上进行全双工通信的协议。

1.1 WebSocket特点 #

  • 全双工通信
  • 低延迟
  • 服务端主动推送
  • 保持连接

1.2 应用场景 #

  • 实时聊天
  • 在线协作
  • 实时数据推送
  • 游戏
  • 股票行情

二、基本使用 #

2.1 安装gorilla/websocket #

bash 
go get github.com/gorilla/websocket

2.2 WebSocket升级 #

go 
package main

import (
    "net/http"
    "github.com/gorilla/websocket"
    "github.com/labstack/echo/v4"
)

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

func main() {
    e := echo.New()
    
    e.GET("/ws", func(c echo.Context) error {
        ws, err := upgrader.Upgrade(c.Response(), c.Request(), nil)
        if err != nil {
            return err
        }
        defer ws.Close()
        
        for {
            messageType, message, err := ws.ReadMessage()
            if err != nil {
                break
            }
            
            err = ws.WriteMessage(messageType, message)
            if err != nil {
                break
            }
        }
        
        return nil
    })
    
    e.Logger.Fatal(e.Start(":8080"))
}

2.3 客户端示例 #

html 
<!DOCTYPE html>
<html>
<head>
    <title>WebSocket Demo</title>
</head>
<body>
    <script>
        const ws = new WebSocket('ws://localhost:8080/ws');
        
        ws.onopen = function() {
            console.log('Connected');
            ws.send('Hello, Server!');
        };
        
        ws.onmessage = function(event) {
            console.log('Received:', event.data);
        };
        
        ws.onclose = function() {
            console.log('Disconnected');
        };
        
        ws.onerror = function(error) {
            console.error('Error:', error);
        };
    </script>
</body>
</html>

三、消息类型 #

3.1 文本消息 #

go 
ws.WriteMessage(websocket.TextMessage, []byte("Hello"))

3.2 二进制消息 #

go 
ws.WriteMessage(websocket.BinaryMessage, []byte{0x01, 0x02, 0x03})

3.3 JSON消息 #

go 
type Message struct {
    Type    string      `json:"type"`
    Content interface{} `json:"content"`
}

func sendJSON(ws *websocket.Conn, msg Message) error {
    data, err := json.Marshal(msg)
    if err != nil {
        return err
    }
    return ws.WriteMessage(websocket.TextMessage, data)
}

func readJSON(ws *websocket.Conn) (*Message, error) {
    _, data, err := ws.ReadMessage()
    if err != nil {
        return nil, err
    }
    
    var msg Message
    if err := json.Unmarshal(data, &msg); err != nil {
        return nil, err
    }
    
    return &msg, nil
}

四、连接管理 #

4.1 客户端结构 #

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

func NewClient(conn *websocket.Conn) *Client {
    return &Client{
        ID:   uuid.New().String(),
        Conn: conn,
        Send: make(chan []byte, 256),
    }
}

func (c *Client) Read() {
    defer c.Conn.Close()
    
    for {
        _, message, err := c.Conn.ReadMessage()
        if err != nil {
            break
        }
        
        handleMessage(c, message)
    }
}

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

4.2 Hub中心 #

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

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

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)
                }
            }
        }
    }
}

五、聊天室实现 #

5.1 完整示例 #

go 
package main

import (
    "encoding/json"
    "log"
    "net/http"
    "sync"
    "time"
    "github.com/gorilla/websocket"
    "github.com/labstack/echo/v4"
    "github.com/labstack/echo/v4/middleware"
)

type Message struct {
    Type      string `json:"type"`
    Sender    string `json:"sender"`
    Content   string `json:"content"`
    Timestamp int64  `json:"timestamp"`
}

type Client struct {
    ID     string
    Name   string
    Conn   *websocket.Conn
    Room   *Room
    Send   chan []byte
}

type Room struct {
    ID      string
    Name    string
    Clients map[*Client]bool
    Mutex   sync.RWMutex
}

type ChatServer struct {
    Rooms map[string]*Room
    Mutex sync.RWMutex
}

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

func NewChatServer() *ChatServer {
    return &ChatServer{
        Rooms: make(map[string]*Room),
    }
}

func (s *ChatServer) GetOrCreateRoom(roomID string) *Room {
    s.Mutex.Lock()
    defer s.Mutex.Unlock()
    
    if room, ok := s.Rooms[roomID]; ok {
        return room
    }
    
    room := &Room{
        ID:      roomID,
        Name:    "Room " + roomID,
        Clients: make(map[*Client]bool),
    }
    s.Rooms[roomID] = room
    
    return room
}

func (r *Room) AddClient(client *Client) {
    r.Mutex.Lock()
    r.Clients[client] = true
    r.Mutex.Unlock()
}

func (r *Room) RemoveClient(client *Client) {
    r.Mutex.Lock()
    delete(r.Clients, client)
    r.Mutex.Unlock()
}

func (r *Room) Broadcast(message []byte, sender *Client) {
    r.Mutex.RLock()
    defer r.Mutex.RUnlock()
    
    for client := range r.Clients {
        if client != sender {
            select {
            case client.Send <- message:
            default:
                close(client.Send)
                delete(r.Clients, client)
            }
        }
    }
}

func (c *Client) ReadPump() {
    defer func() {
        c.Room.RemoveClient(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
        }
        
        var msg Message
        if err := json.Unmarshal(message, &msg); err != nil {
            continue
        }
        
        msg.Sender = c.Name
        msg.Timestamp = time.Now().Unix()
        
        data, _ := json.Marshal(msg)
        c.Room.Broadcast(data, c)
    }
}

func (c *Client) WritePump() {
    ticker := time.NewTicker(30 * 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
            }
            
            c.Conn.WriteMessage(websocket.TextMessage, message)
            
        case <-ticker.C:
            c.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
            if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
                return
            }
        }
    }
}

func main() {
    e := echo.New()
    
    e.Use(middleware.Logger())
    e.Use(middleware.Recover())
    
    server := NewChatServer()
    
    e.GET("/ws/:room", func(c echo.Context) error {
        roomID := c.Param("room")
        name := c.QueryParam("name")
        
        if name == "" {
            name = "Anonymous"
        }
        
        conn, err := upgrader.Upgrade(c.Response(), c.Request(), nil)
        if err != nil {
            return err
        }
        
        room := server.GetOrCreateRoom(roomID)
        
        client := &Client{
            ID:   time.Now().String(),
            Name: name,
            Conn: conn,
            Room: room,
            Send: make(chan []byte, 256),
        }
        
        room.AddClient(client)
        
        joinMsg := Message{
            Type:      "join",
            Sender:    "System",
            Content:   name + " joined the room",
            Timestamp: time.Now().Unix(),
        }
        data, _ := json.Marshal(joinMsg)
        room.Broadcast(data, nil)
        
        go client.WritePump()
        go client.ReadPump()
        
        return nil
    })
    
    e.Static("/", "static")
    
    e.Logger.Fatal(e.Start(":8080"))
}

5.2 客户端页面 #

html 
<!DOCTYPE html>
<html>
<head>
    <title>Chat Room</title>
    <style>
        #messages { height: 400px; overflow-y: scroll; border: 1px solid #ccc; }
        #message { width: 80%; }
    </style>
</head>
<body>
    <div id="messages"></div>
    <input type="text" id="message" placeholder="Type a message...">
    <button onclick="send()">Send</button>
    
    <script>
        const room = 'default';
        const name = prompt('Enter your name:') || 'Anonymous';
        const ws = new WebSocket(`ws://localhost:8080/ws/${room}?name=${name}`);
        
        ws.onmessage = function(event) {
            const msg = JSON.parse(event.data);
            const div = document.createElement('div');
            div.textContent = `${msg.sender}: ${msg.content}`;
            document.getElementById('messages').appendChild(div);
            document.getElementById('messages').scrollTop = 999999;
        };
        
        function send() {
            const input = document.getElementById('message');
            const msg = {
                type: 'message',
                content: input.value
            };
            ws.send(JSON.stringify(msg));
            input.value = '';
        }
        
        document.getElementById('message').addEventListener('keypress', function(e) {
            if (e.key === 'Enter') send();
        });
    </script>
</body>
</html>

六、总结 #

WebSocket要点：

要点 说明
Upgrader HTTP升级到WebSocket
ReadMessage 读取消息
WriteMessage 发送消息
Ping/Pong 心跳检测
Hub 连接管理中心
Room 房间管理

准备好学习性能优化了吗？让我们进入下一章！

最后更新：2026-03-28