Go WebSocket 连接管理与消息广播实践

func EchoServer(ws *websocket.Conn) {
    io.Copy(ws, ws)
}

func main() {
    http.Handle("/echo", websocket.Handler(EchoServer))
    http.ListenAndServe(":12345", nil)
}

package main

import (
    "fmt"
    "io"
    "log"
    "net/http"
    "sync" // For potential future use, though not strictly needed for the channel approach
    "golang.org/x/net/websocket" // Using the recommended package
)

// 定义通道
var connects = make(chan *websocket.Conn)    // 用于接收新连接
var broadcasts = make(chan []byte)           // 用于接收要广播的消息
var disconnects = make(chan *websocket.Conn) // 用于接收断开连接的通知

// 活跃连接存储
var activeConnections = make(map[*websocket.Conn]struct{}) // 使用 struct{} 节省内存

// connectionManager 负责管理所有连接和消息广播
func connectionManager() {
    for {
        select {
        case newConn := <-connects:
            // 添加新连接
            activeConnections[newConn] = struct{}{}
            log.Printf("New client connected: %s. Total connections: %d", newConn.RemoteAddr(), len(activeConnections))

        case msg := <-broadcasts:
            // 广播消息给所有活跃连接
            log.Printf("Broadcasting message: %s", string(msg))
            for conn := range activeConnections {
                if _, err := conn.Write(msg); err != nil {
                    // 如果写入失败，通常表示客户端已断开，将其标记为待移除
                    log.Printf("Failed to write to client %s, marking for removal: %v", conn.RemoteAddr(), err)
                    select {
                    case disconnects <- conn: // 尝试发送到断开连接通道
                    default: // 防止通道阻塞
                        log.Printf("Disconnects channel full, dropping disconnect notification for %s", conn.RemoteAddr())
                    }
                }
            }

        case disconnectedConn := <-disconnects:
            // 移除断开的连接
            if _, ok := activeConnections[disconnectedConn]; ok {
                delete(activeConnections, disconnectedConn)
                disconnectedConn.Close() // 确保连接被关闭
                log.Printf("Client disconnected: %s. Total connections: %d", disconnectedConn.RemoteAddr(), len(activeConnections))
            }
        }
    }
}

// EchoServer 处理单个 WebSocket 连接
func EchoServer(ws *websocket.Conn) {
    // 1. 将新连接发送到 connectionManager
    connects <- ws

    // 2. 循环读取客户端消息并发送到广播通道
    buff := make([]byte, 1024)
    for {
        n, err := ws.Read(buff)
        if err != nil {
            if err == io.EOF {
                log.Printf("Client %s disconnected normally.", ws.RemoteAddr())
            } else {
                log.Printf("Read error from client %s: %v", ws.RemoteAddr(), err)
            }
            // 客户端读取错误或断开，发送到断开连接通道
            disconnects <- ws
            break // 退出循环，结束当前 goroutine
        }
        // 将接收到的消息发送到广播通道
        broadcasts <- buff[:n]
    }
}

func main() {
    // 启动连接管理器 goroutine
    go connectionManager()

    // 设置 HTTP 路由处理 WebSocket 连接
    http.Handle("/echo", websocket.Handler(EchoServer))

    // 启动 HTTP 服务器
    log.Println("WebSocket server started on :12345")
    err := http.ListenAndServe(":12345", nil)
    if err != nil {
        log.Fatalf("ListenAndServe error: %v", err)
    }
}

package main

import (
    "fmt"
    "io"
    "log"
    "net/http"
    "sync"
    "golang.org/x/net/websocket"
)

// 全局连接存储，由互斥锁保护
var globalConnections = struct {
    sync.RWMutex
    m map[*websocket.Conn]struct{}
}{
    m: make(map[*websocket.Conn]struct{}),
}

// EchoServerMutex 版本，直接操作全局连接
func EchoServerMutex(ws *websocket.Conn) {
    // 添加新连接
    globalConnections.Lock()
    globalConnections.m[ws] = struct{}{}
    log.Printf("New client connected: %s. Total connections: %d", ws.RemoteAddr(), len(globalConnections.m))
    globalConnections.Unlock()

    buff := make([]byte, 1024)
    for {
        n, err := ws.Read(buff)
        if err != nil {
            if err == io.EOF {
                log.Printf("Client %s disconnected normally.", ws.RemoteAddr())
            } else {
                log.Printf("Read error from client %s: %v", ws.RemoteAddr(), err)
            }
            break // 退出循环
        }

        // 广播消息
        msg := buff[:n]
        globalConnections.RLock() // 读锁
        for conn := range globalConnections.m {
            if conn == ws { // 避免回显给自己
                continue
            }
            if _, writeErr := conn.Write(msg); writeErr != nil {
                log.Printf("Failed to write to client %s: %v", conn.RemoteAddr(), writeErr)
                // 注意：在读锁中删除元素会导致死锁或并发修改错误
                // 正确的做法是收集需要移除的连接，然后在释放读锁后，用写锁进行删除
            }
        }
        globalConnections.RUnlock() // 释放读锁
    }

    // 客户端断开，移除连接
    globalConnections.Lock() // 写锁
    delete(globalConnections.m, ws)
    log.Printf("Client disconnected: %s. Total connections: %d", ws.RemoteAddr(), len(globalConnections.m))
    globalConnections.Unlock()
    ws.Close()
}

func main() {
    // http.Handle("/echo", websocket.Handler(EchoServer)) // 使用通道版本
    http.Handle("/echo_mutex", websocket.Handler(EchoServerMutex)) // 使用互斥锁版本

    log.Println("WebSocket server started on :12345 (mutex version)")
    err := http.ListenAndServe(":12345", nil)
    if err != nil {
        log.Fatalf("ListenAndServe error: %v", err)
    }
}