Go中实现RabbitMQ消息转WebSocket客户端且避免协程阻塞

核心思路 #

用线程安全的连接管理器统一管理所有活跃的WebSocket连接，RabbitMQ消费者通过用户ID从管理器中获取对应连接并转发消息，避免直接传递连接引用导致的阻塞问题。

具体实现步骤 #

1. 实现线程安全的连接管理器 #

创建一个ConnectionManager结构体，用sync.RWMutex保护内部连接映射，提供注册、注销、根据用户ID获取连接的方法：

package service

import (
    "sync"
)

// ConnectionManager 管理所有活跃的WebSocket连接
type ConnectionManager struct {
    mu      sync.RWMutex
    clients map[string]*ClientService // key: 用户ID
}

func NewConnectionManager() *ConnectionManager {
    return &ConnectionManager{
        clients: make(map[string]*ClientService),
    }
}

// Register 注册用户的WebSocket连接
func (cm *ConnectionManager) Register(userID string, client *ClientService) {
    cm.mu.Lock()
    defer cm.mu.Unlock()
    cm.clients[userID] = client
}

// Unregister 注销用户的WebSocket连接
func (cm *ConnectionManager) Unregister(userID string) {
    cm.mu.Lock()
    defer cm.mu.Unlock()
    delete(cm.clients, userID)
}

// GetClient 根据用户ID获取对应的ClientService
func (cm *ConnectionManager) GetClient(userID string) (*ClientService, bool) {
    cm.mu.RLock()
    defer cm.mu.RUnlock()
    client, exists := cm.clients[userID]
    return client, exists
}

2. 修改主函数集成连接管理器 #

初始化连接管理器，在WebSocket连接验证通过后注册连接，断开时自动注销：

func main() {
    loadEnvVariables()

    wg := new(sync.WaitGroup)
    wg.Add(2)

    // 初始化连接管理器
    connManager := service.NewConnectionManager()

    go func() {
        http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
            conn, err := upgrades.Upgrade(w, r, nil)
            if err != nil {
                w.WriteHeader(http.StatusInternalServerError)
                _, err := w.Write([]byte("failed upgrading connection"))
                if err != nil {
                    return
                }
                return
            }

            defer conn.Close()

            clientService := service.NewClientService(conn)
            done := make(chan struct{})
            clientID := uuid.New().String()
            var userID string

            if _, exists := r.Header["Authorization"]; exists {
                jwt := getJWT(r)
                grpc, err := grpcclient.GetGrpcClient(os.Getenv("AS"))
                if err != nil {
                    log.Fatalf("did not connect: %s", err)
                }
                defer grpc.Close()

                sessionClient := as.NewSessionServiceClient(grpc)
                authService := service.NewAuthService(sessionClient)
                res, err := authService.ValidateSession(jwt)
                if err != nil {
                    clientService.Send(fmt.Sprintf("Internal server error"))
                    return // 验证失败直接返回
                }

                userID = res.User.Id
                clientService.Send(fmt.Sprintf("Server: Welcome! Your ID is %s", clientID))
                clientService.SaveConnection(userID, clientID)

                // 注册连接到管理器
                connManager.Register(userID, clientService)
                // 连接断开时自动注销
                defer connManager.Unregister(userID)

                go clientService.WritePump(userID, done)
                clientService.ReadPump(userID, done)
                return // 验证成功跳过错误提示
            }
            clientService.Send(fmt.Sprintf("Authentication error"))
        })

        if err := http.ListenAndServe(":3071", nil); err != nil {
            panic(err)
        }
        wg.Done()
    }()

    go func() {
        rabbitService := service.NewRabbitService(&service.RabbitConfig{
            Schema:         model.SCHEMA,
            Username:       os.Getenv("RMQ_USERNAME"),
            Password:       os.Getenv("RMQ_PASSWORD"),
            Host:           os.Getenv("RMQ_HOST"),
            Port:           os.Getenv("RMQ_PORT"),
            VHost:          os.Getenv("RMQ_VHOST"),
            ConnectionName: model.CONNECTION_NAME,
            ExchangeName:   model.EXCHANGE_NAME,
            ExchangeType:   model.EXCHANGE_TOPIC,
        })
        if err := rabbitService.Connect(); err != nil {
            log.Fatalln("unable to connect to rabbit", err)
        }
        // 将连接管理器传给TestHandler
        testHandler := service.NewTestHandler(&service.TestHandlerConfig{
            RabbitService: rabbitService,
            ExchangeConfig: &service.ExchangeConfig{
                RoutingKey:   model.TEST_HANDLER_EVENT,
                QueueName:    model.QUE_NAME,
                ConsumerName: model.TEST_HANDLER_NAME,
            },
            ConnManager: connManager,
        })
        testHandler.Declare()
        testHandler.Handle()
        wg.Done()
    }()

    wg.Wait()
}

3. 修改TestHandler转发消息 #

更新TestHandler定义，加入连接管理器字段，消费RabbitMQ消息时根据用户ID获取连接并发送：

package service

import (
    "encoding/json"
    "log"
)

type TestHandlerConfig struct {
    RabbitService *RabbitService
    ExchangeConfig *ExchangeConfig
    ConnManager *ConnectionManager // 新增连接管理器字段
}

type TestHandler struct {
    rabbitService *RabbitService
    exchangeCfg *ExchangeConfig
    connManager *ConnectionManager
}

func NewTestHandler(cfg *TestHandlerConfig) *TestHandler {
    return &TestHandler{
        rabbitService: cfg.RabbitService,
        exchangeCfg: cfg.ExchangeConfig,
        connManager: cfg.ConnManager,
    }
}

// Handle 消费RabbitMQ消息并转发到WebSocket
func (th *TestHandler) Handle() {
    err := th.rabbitService.Consume(th.exchangeCfg.QueueName, th.exchangeCfg.ConsumerName, func(msg []byte) error {
        // 解析消息获取目标用户ID（根据实际消息格式调整）
        var msgData struct {
            UserID string `json:"user_id"`
            Content string `json:"content"`
        }
        if err := json.Unmarshal(msg, &msgData); err != nil {
            log.Printf("failed to parse message: %v", err)
            return nil
        }

        // 获取用户的WebSocket连接
        client, exists := th.connManager.GetClient(msgData.UserID)
        if !exists {
            log.Printf("no active connection for user %s", msgData.UserID)
            return nil
        }

        // 发送消息到客户端
        if err := client.Send(msgData.Content); err != nil {
            log.Printf("failed to send message to user %s: %v", msgData.UserID, err)
            th.connManager.Unregister(msgData.UserID)
        }
        return nil
    })

    if err != nil {
        log.Fatalln("failed to consume messages:", err)
    }
}

关键说明 #

• 线程安全：用sync.RWMutex保护连接映射，支持多goroutine并发读写，避免数据竞争。
• 连接生命周期管理：连接建立时注册，断开时自动注销，确保管理器只保留活跃连接。
• 组件解耦：RabbitMQ消费者与WebSocket处理器通过连接管理器交互，无需直接传递连接引用，避免阻塞新连接处理。

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