构建WebSocket服务器

WebSocket协议已经成为现代Web应用中实时通信的基石，它提供了全双工通信通道，使服务器能够主动向客户端推送数据。在本文中，我们将探索如何使用Go语言轻松构建一个高性能的WebSocket服务器。

go get github.com/gorilla/websocket

基本实现

package main

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

var upgrader = websocket.Upgrader{
    ReadBufferSize:  1024,
    WriteBufferSize: 1024,
    CheckOrigin: func(r *http.Request) bool {
        return true // 允许所有来源（生产环境应限制）
    },
}

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println("升级到WebSocket失败:", err)
        return
    }
    defer conn.Close()
    
    log.Println("客户端连接成功:", conn.RemoteAddr())
    
    for {
        // 读取客户端消息
        messageType, p, err := conn.ReadMessage()
        if err != nil {
            log.Println("读取消息失败:", err)
            return
        }
        
        log.Printf("收到消息: %s\n", p)
        
        // 原样返回消息（echo）
        if err := conn.WriteMessage(messageType, p); err != nil {
            log.Println("发送消息失败:", err)
            return
        }
    }
}

func main() {
    http.HandleFunc("/ws", handleWebSocket)
    log.Println("WebSocket服务器启动 :8080")
    log.Fatal(http.ListenAndServe(":8080", nil))
}

聊天室

创建一个简单的聊天室，支持多个客户端

package main

import (
    "log"
    "net/http"
    "sync"
    
    "github.com/gorilla/websocket"
)

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

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

var (
    clients    = make(map[*Client]bool)
    clientsMux sync.Mutex
    broadcast  = make(chan []byte)
)
var connectionLimit = make(chan struct{}, 100) // 限制100个连接

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println("升级失败:", err)
        return
    }
    conn.SetReadLimit(1024 * 1024) // 限制1MB消息
    // 链接限制
    connectionLimit <- struct{}{}
    defer func() { <-connectionLimit }()

    client := &Client{
        conn: conn,
        send: make(chan []byte, 256),
    }
    
    // 注册客户端
    clientsMux.Lock()
    clients[client] = true
    clientsMux.Unlock()
    
    log.Printf("新客户端连接: %s (当前客户端数: %d)", 
        conn.RemoteAddr(), len(clients))
    
    // 启动读写goroutine
    go client.writePump()
    go client.readPump()
}

func (c *Client) readPump() {
    defer func() {
        c.conn.Close()
        clientsMux.Lock()
        delete(clients, c)
        clientsMux.Unlock()
        log.Printf("客户端断开连接 (剩余: %d)", len(clients))
    }()

    go c.heartbeat()

    for {
        _, message, err := c.conn.ReadMessage()
        if err != nil {
            if websocket.IsUnexpectedCloseError(err, 
                websocket.CloseGoingAway, 
                websocket.CloseAbnormalClosure) {
                log.Printf("读取错误: %v", err)
            }
            break
        }
        
        // 广播消息给所有客户端
        broadcast <- message
    }
}

func (c *Client) writePump() {
    defer c.conn.Close()
    
    for {
        select {
        case message, ok := <-c.send:
            if !ok {
                // 通道关闭
                c.conn.WriteMessage(websocket.CloseMessage, []byte{})
                return
            }
            
            if err := c.conn.WriteMessage(websocket.TextMessage, message); err != nil {
                log.Println("发送失败:", err)
                return
            }
        }
    }
}
func (c *Client) heartbeat() {
    ticker := time.NewTicker(30 * time.Second)
    defer ticker.Stop()
    
    for {
        select {
        case <-ticker.C:
            if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
                log.Println("心跳失败:", err)
                return
            }
        }
    }
}


func broadcastMessages() {
    for {
        msg := <-broadcast
        
        clientsMux.Lock()
        for client := range clients {
            select {
            case client.send <- msg:
            default:
                // 发送失败则关闭连接
                close(client.send)
                delete(clients, client)
            }
        }
        clientsMux.Unlock()
    }
}

func main() {
    go broadcastMessages()
    
    http.HandleFunc("/ws", handleWebSocket)
    http.Handle("/", http.FileServer(http.Dir("./static")))
    
    log.Println("聊天室服务器启动 :8080")
    log.Fatal(http.ListenAndServe(":8080", nil))
}

客户端测试页面

<!DOCTYPE html>
<html>
<head>
    <title>WebSocket聊天室</title>
    <style>
        body { font-family: Arial, sans-serif; max-width: 600px; margin: 0 auto; }
        #messages { height: 300px; border: 1px solid #ccc; padding: 10px; overflow-y: auto; }
        #messageInput { width: 75%; padding: 8px; }
        #sendBtn { padding: 8px 16px; }
    </style>
</head>
<body>
    <h1>Go WebSocket聊天室</h1>
    <div id="messages"></div>
    <div>
        <input type="text" id="messageInput" placeholder="输入消息...">
        <button id="sendBtn">发送</button>
    </div>

    <script>
        const messages = document.getElementById('messages');
        const messageInput = document.getElementById('messageInput');
        const sendBtn = document.getElementById('sendBtn');
        
        // 创建WebSocket连接
        const socket = new WebSocket('ws://' + '127.0.0.1:8080' + '/ws');
        
        socket.onopen = () => {
            addMessage('系统: 已连接到服务器');
        };
        
        socket.onmessage = (event) => {
            addMessage(event.data);
        };
        
        socket.onclose = () => {
            addMessage('系统: 连接已断开');
        };
        
        function addMessage(msg) {
            const msgElement = document.createElement('div');
            msgElement.textContent = msg;
            messages.appendChild(msgElement);
            messages.scrollTop = messages.scrollHeight;
        }
        
        function sendMessage() {
            const message = messageInput.value.trim();
            if (message) {
                socket.send(message);
                messageInput.value = '';
            }
        }
        
        sendBtn.addEventListener('click', sendMessage);
        messageInput.addEventListener('keypress', (e) => {
            if (e.key === 'Enter') sendMessage();
        });
    </script>
</body>
</html>