多协议同时监听教程

// internal/ws/hub.go
package ws

import (
    "sync"
    "github.com/gorilla/websocket"
)

type Hub struct {
    mu      sync.RWMutex
    clients map[string]map[*websocket.Conn]bool  // room_id → connections
}

func NewHub() *Hub {
    return &Hub{
        clients: make(map[string]map[*websocket.Conn]bool),
    }
}

func (h *Hub) Register(roomID string, conn *websocket.Conn) {
    h.mu.Lock()
    defer h.mu.Unlock()
    if h.clients[roomID] == nil {
        h.clients[roomID] = make(map[*websocket.Conn]bool)
    }
    h.clients[roomID][conn] = true
}

func (h *Hub) Unregister(roomID string, conn *websocket.Conn) {
    h.mu.Lock()
    defer h.mu.Unlock()
    if clients, ok := h.clients[roomID]; ok {
        delete(clients, conn)
        if len(clients) == 0 {
            delete(h.clients, roomID)
        }
    }
}

func (h *Hub) Broadcast(roomID string, message []byte) {
    h.mu.RLock()
    defer h.mu.RUnlock()
    for conn := range h.clients[roomID] {
        conn.WriteMessage(websocket.TextMessage, message)
    }
}

// internal/service/chat_service.go
package service

type ChatService struct {
    db    *ent.Client
    hub   *ws.Hub
    broker broker.Broker
}

func (s *ChatService) SendMessage(ctx context.Context, roomID, userID, content string) (*Message, error) {
    msg := &Message{
        ID:        uuid.New().String(),
        RoomID:    roomID,
        UserID:    userID,
        Content:   content,
        Timestamp: time.Now().UnixMilli(),
    }

    // 存储到数据库
    s.db.Message.Create().
        SetID(msg.ID).
        SetRoomID(roomID).
        SetUserID(userID).
        SetContent(content).
        Save(ctx)

    // 通过 WebSocket 推送给当前房间的在线用户
    data, _ := json.Marshal(msg)
    s.hub.Broadcast(roomID, data)

    // 同时发布到 Broker（跨实例推送）
    s.broker.Publish(ctx, "chat.messages", &broker.Message{Body: data})

    return msg, nil
}

// internal/handler/http_handler.go
package handler

func (h *ChatHandler) Routes() map[string]http.HandlerFunc {
    return map[string]http.HandlerFunc{
        "POST /api/rooms/{room_id}/messages": h.sendMessage,
        "GET  /api/rooms/{room_id}/messages": h.getHistory,
        "POST /api/rooms":                    h.createRoom,
        "GET  /api/health":                   h.health,
    }
}

// internal/handler/grpc_handler.go
package handler

type GRPCChatHandler struct {
    pb.UnimplementedChatServiceServer
    svc *service.ChatService
}

func (h *GRPCChatHandler) SendMessage(ctx context.Context, req *pb.SendMessageRequest) (*pb.SendMessageResponse, error) {
    userID := wind.UserID(ctx)
    msg, err := h.svc.SendMessage(ctx, req.RoomId, userID, req.Content)
    if err != nil {
        return nil, err
    }
    return &pb.SendMessageResponse{
        Message: toProtoMessage(msg),
    }, nil
}

func (h *GRPCChatHandler) StreamMessages(req *pb.StreamMessagesRequest, stream pb.ChatService_StreamMessagesServer) error {
    // 服务端流式推送
    ctx := stream.Context()
    sub, _ := h.svc.broker.Subscribe(ctx, "chat.messages", func(ctx context.Context, msg *broker.Message) error {
        var m Message
        json.Unmarshal(msg.Body, &m)
        if m.RoomID == req.RoomId {
            stream.Send(toProtoMessage(&m))
        }
        return nil
    })
    defer sub.Unsubscribe()

    <-ctx.Done()
    return nil
}

// internal/app.go
package internal

func init() {
    bootstrap.OnReady(func(app *bootstrap.App) {
        hub := ws.NewHub()
        chatSvc := service.NewChatService(app.DB(), hub, app.Broker())

        // HTTP 路由
        mux := http.NewServeMux()
        httpHandler := handler.NewHTTPHandler(chatSvc)
        for pattern, h := range httpHandler.Routes() {
            mux.HandleFunc(pattern, h)
        }
        app.SetHTTPHandler(mux)

        // gRPC 服务注册
        grpcHandler := handler.NewGRPCChatHandler(chatSvc)
        app.RegisterGRPC(func(s *grpc.Server) {
            pb.RegisterChatServiceServer(s, grpcHandler)
        })

        // WebSocket 处理器
        app.SetWebSocketHandler(func(conn *websocket.Conn) {
            roomID := conn.URL.Query().Get("room")
            hub.Register(roomID, conn)
            defer hub.Unregister(roomID, conn)

            for {
                _, msg, err := conn.ReadMessage()
                if err != nil { break }
                chatSvc.SendMessage(context.Background(), roomID, "ws-user", string(msg))
            }
        })
    })
}

// cmd/server/main.go
package main

import (
    _ "github.com/tx7do/go-wind-plugins/transport/http"
    _ "github.com/tx7do/go-wind-plugins/transport/grpc"
    _ "github.com/tx7do/go-wind-plugins/transport/websocket"
    _ "github.com/tx7do/go-wind-plugins/log/zap"
    _ "github.com/tx7do/go-wind-plugins/database/postgres"
    _ "github.com/tx7do/go-wind-plugins/cache/redis"
    _ "github.com/tx7do/go-wind-plugins/broker/redis"
    _ "github.com/tx7do/go-wind-plugins/registry/consul"

    "github.com/tx7do/go-wind-bootstrap"
    _ "chat-service/internal"
)

func main() {
    app := bootstrap.New("configs/config.yaml")
    app.Run()
}