78 lines
1.6 KiB
Go
78 lines
1.6 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"ky-go-kratos/app/websocket/internal/biz"
|
|
"log"
|
|
"sync"
|
|
)
|
|
|
|
// WebSocketService is a WebSocket service.
|
|
type WebSocketService struct {
|
|
uc *biz.WebSocketUsecase
|
|
clients map[string]*Client
|
|
clientsMux sync.RWMutex
|
|
}
|
|
|
|
type Client struct {
|
|
ID string
|
|
Send chan []byte
|
|
Topics map[string]bool
|
|
}
|
|
|
|
// NewWebSocketService 创建 WebSocket 服务
|
|
func NewWebSocketService(uc *biz.WebSocketUsecase) *WebSocketService {
|
|
return &WebSocketService{
|
|
uc: uc,
|
|
clients: make(map[string]*Client),
|
|
}
|
|
}
|
|
|
|
// HandleMessage 处理 WebSocket 消息
|
|
func (s *WebSocketService) HandleMessage(ctx context.Context, message []byte) error {
|
|
return s.uc.HandleMessage(ctx, message)
|
|
}
|
|
|
|
// Broadcast 广播消息给所有客户端
|
|
func (s *WebSocketService) Broadcast(message json.RawMessage) error {
|
|
s.clientsMux.RLock()
|
|
defer s.clientsMux.RUnlock()
|
|
|
|
for _, client := range s.clients {
|
|
select {
|
|
case client.Send <- message:
|
|
default:
|
|
log.Printf("Client %s buffer is full, message dropped", client.ID)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// SendPrivateMessage 发送私聊消息
|
|
func (s *WebSocketService) SendPrivateMessage(message json.RawMessage) error {
|
|
var msg struct {
|
|
To string `json:"to"`
|
|
Content json.RawMessage `json:"content"`
|
|
}
|
|
|
|
if err := json.Unmarshal(message, &msg); err != nil {
|
|
return err
|
|
}
|
|
|
|
s.clientsMux.RLock()
|
|
client, exists := s.clients[msg.To]
|
|
s.clientsMux.RUnlock()
|
|
|
|
if !exists {
|
|
return nil
|
|
}
|
|
|
|
select {
|
|
case client.Send <- msg.Content:
|
|
default:
|
|
log.Printf("Client %s buffer is full, message dropped", client.ID)
|
|
}
|
|
return nil
|
|
}
|