在Golang中使用通道同步WebSocket循环

Got it, let's figure out how to sync multiple WebSocket connections for a specific user using Go channels. Looking at your existing code structure, here's a practical approach that aligns with your setup:

The key is to use a shared broadcast channel per user. When a message needs to be sent to the user, we send it to this channel, and each of the user's WebSocket connection goroutines listens to this channel to receive and forward the message to their respective sockets. This ensures all connections get the same message in sync.

First, let's tweak your existing structs to support per-user broadcast and connection management:

Your original structure:

type msg struct { Key string Value string }
type connStruct struct {
 //...
 ConnRoutineChans []*chan string
 LoggedIn bool
 Login string
 //...
 Sockets []*websocket.Conn
}
var (
 //...
 /* 已连接用户及其IP地址列表 */
 guestMap sync.Map
)

Modified Structure #

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

type msg struct { 
    Key   string `json:"key"`
    Value string `json:"value"`
}

// wsConn wraps a WebSocket connection with a quit signal for graceful shutdown
type wsConn struct {
    conn *websocket.Conn
    quit chan struct{} // Signals the goroutine to exit
}

type connStruct struct {
    //...
    broadcastChan chan msg       // Shared channel for all user connections
    LoggedIn      bool
    Login         string
    //...
    connections []wsConn         // Track all active WebSocket connections for the user
    mu          sync.Mutex       // Protect concurrent access to connections slice
}

var (
    //...
    /* 已连接用户及其IP地址列表 */
    guestMap sync.Map
)

Let's walk through how to manage new connections, message broadcasting, and cleanup:

1. Handle New WebSocket Connections #

When a user opens a new WebSocket, we either retrieve their existing connStruct or create a new one, then start a goroutine to listen for broadcast messages:

func wsHandler(w http.ResponseWriter, r *http.Request) {
    // Upgrade HTTP to WebSocket
    upgrader := websocket.Upgrader{
        CheckOrigin: func(r *http.Request) bool {
            // Adjust origin check based on your needs
            return true
        },
    }
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println("WebSocket upgrade error:", err)
        return
    }
    defer conn.Close()

    // Get user's unique identifier (adjust this to your auth method: cookie, token, etc.)
    login := r.URL.Query().Get("login")
    if login == "" {
        _ = conn.WriteMessage(websocket.TextMessage, []byte("Missing user identifier"))
        return
    }

    // Retrieve or create user's connStruct from guestMap
    userVal, ok := guestMap.Load(login)
    var user *connStruct
    if !ok {
        user = &connStruct{
            broadcastChan: make(chan msg, 10), // Buffered channel to prevent blocking
            Login:         login,
            connections:   []wsConn{},
            mu:            sync.Mutex{},
        }
        guestMap.Store(login, user)
        // Start a dedicated broadcast loop for the user (optional but cleaner)
        go user.broadcastLoop()
    } else {
        user = userVal.(*connStruct)
    }

    // Initialize connection with quit signal
    quit := make(chan struct{})
    newWsConn := wsConn{conn: conn, quit: quit}

    // Add new connection to user's list (thread-safe)
    user.mu.Lock()
    user.connections = append(user.connections, newWsConn)
    user.mu.Unlock()

    // Cleanup routine when connection closes
    defer func() {
        user.mu.Lock()
        // Filter out the closed connection
        updatedConns := []wsConn{}
        for _, c := range user.connections {
            if c.quit != quit {
                updatedConns = append(updatedConns, c)
            }
        }
        user.connections = updatedConns
        user.mu.Unlock()

        close(quit)
        // Remove user from guestMap if no active connections left
        if len(user.connections) == 0 {
            guestMap.Delete(login)
            close(user.broadcastChan)
        }
    }()

    // Handle incoming messages from the client (if needed)
    for {
        _, msgData, err := conn.ReadMessage()
        if err != nil {
            log.Printf("Read message error for user %s: %v", login, err)
            break
        }

        // Parse incoming message
        var receivedMsg msg
        if err := json.Unmarshal(msgData, &receivedMsg); err != nil {
            log.Printf("Unmarshal message error: %v", err)
            continue
        }

        // Broadcast the message to all of the user's connections
        user.broadcastChan <- receivedMsg
    }
}

2. Dedicated Broadcast Loop #

Instead of having each connection goroutine listen to the broadcast channel, you can use a single goroutine per user to distribute messages to all connections. This reduces redundant channel listeners:

func (u *connStruct) broadcastLoop() {
    for msg := range u.broadcastChan {
        // Serialize message to JSON
        msgBytes, err := json.Marshal(msg)
        if err != nil {
            log.Printf("Marshal message error for user %s: %v", u.Login, err)
            continue
        }

        // Send message to all active connections (thread-safe)
        u.mu.Lock()
        for _, conn := range u.connections {
            select {
            case <-conn.quit:
                // Skip closed connections
                continue
            default:
                if err := conn.conn.WriteMessage(websocket.TextMessage, msgBytes); err != nil {
                    log.Printf("Write message error for user %s: %v", u.Login, err)
                    close(conn.quit) // Mark connection for cleanup
                }
            }
        }
        u.mu.Unlock()
    }
}

• Concurrent Safety: Always use a mutex (sync.Mutex) when modifying the connections slice, since multiple goroutines may add/remove connections at the same time.
• Graceful Shutdown: The quit channel ensures connection goroutines exit cleanly when a socket closes, preventing goroutine leaks.
• Channel Buffering: A buffered broadcastChan prevents blocking when sending messages, especially if one of the connections is slow to receive.
• Resource Cleanup: When a user's last connection closes, delete their entry from guestMap and close the broadcast channel to free up resources.

内容的提问来源于stack exchange，提问作者hewiefreeman