A:增加SSE
This commit is contained in:
1 parent
8d51fc17b5
commit
3d63040636
4 files changed
+1231
No files matched your search
@@ -0,0 +1,282 @@
|
|||||||
|
# SSE (Server-Sent Events) 管理器
|
||||||
|
|
||||||
|
一个功能完整、生产就绪的 Go SSE 管理器,用于实现服务端到客户端的实时数据推送。
|
||||||
|
|
||||||
|
## 功能特性
|
||||||
|
|
||||||
|
- ✅ **动态客户端管理**: 自动添加/删除客户端连接,支持断线检测
|
||||||
|
- ✅ **多种推送模式**: 广播消息、按客户端ID定向推送
|
||||||
|
- ✅ **业务解耦设计**: 通过 `chan string` 接收业务层事件
|
||||||
|
- ✅ **断线重连支持**: 客户端自动重连,服务端自动清理失效连接
|
||||||
|
- ✅ **高可用设计**: 内置心跳机制、超时处理、panic恢复
|
||||||
|
- ✅ **配置灵活**: 支持自定义选项配置
|
||||||
|
- ✅ **CORS支持**: 跨域访问支持
|
||||||
|
- ✅ **并发安全**: 使用读写锁保护并发访问
|
||||||
|
|
||||||
|
## 快速开始
|
||||||
|
|
||||||
|
### 基本使用
|
||||||
|
|
||||||
|
```go
|
||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"rear/pkg/sse"
|
||||||
|
"github.com/gin-gonic/gin"
|
||||||
|
)
|
||||||
|
|
||||||
|
func main() {
|
||||||
|
r := gin.Default()
|
||||||
|
|
||||||
|
// 创建 SSE 管理器
|
||||||
|
manager := sse.NewManager(nil) // 使用默认配置
|
||||||
|
|
||||||
|
// 设置 SSE 连接端点
|
||||||
|
r.GET("/events", func(c *gin.Context) {
|
||||||
|
manager.RegisterClient(c)
|
||||||
|
})
|
||||||
|
|
||||||
|
// 广播消息端点
|
||||||
|
r.POST("/broadcast", func(c *gin.Context) {
|
||||||
|
var req struct {
|
||||||
|
Message string `json:"message"`
|
||||||
|
}
|
||||||
|
c.ShouldBindJSON(&req)
|
||||||
|
|
||||||
|
manager.Broadcast(req.Message)
|
||||||
|
c.JSON(200, gin.H{"status": "sent"})
|
||||||
|
})
|
||||||
|
|
||||||
|
r.Run(":8080")
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
### 集成到现有项目
|
||||||
|
|
||||||
|
1. **创建 SSE 处理器**
|
||||||
|
|
||||||
|
```go
|
||||||
|
// internal/handler/sse_handler.go
|
||||||
|
sseHandler := handler.NewSSEHandler()
|
||||||
|
```
|
||||||
|
|
||||||
|
2. **添加路由**
|
||||||
|
|
||||||
|
```go
|
||||||
|
// internal/router/router.go
|
||||||
|
func SetupRoutes(r *gin.Engine, contain *container.DbContainer, imgContain *container.TaskContainer) {
|
||||||
|
// 现有路由...
|
||||||
|
|
||||||
|
// SSE 相关路由
|
||||||
|
api := r.Group("/api/v1")
|
||||||
|
{
|
||||||
|
api.GET("/events", sseHandler.HandleSSEConnection)
|
||||||
|
api.POST("/events/broadcast", sseHandler.BroadcastMessage)
|
||||||
|
api.POST("/events/send/:clientId", sseHandler.SendToClient)
|
||||||
|
api.GET("/events/clients", sseHandler.GetClients)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
3. **业务层集成**
|
||||||
|
|
||||||
|
```go
|
||||||
|
// 获取事件通道
|
||||||
|
eventChannel := sseHandler.GetEventChannel()
|
||||||
|
|
||||||
|
// 在业务逻辑中推送事件
|
||||||
|
func ProcessPhoto(photoID string) {
|
||||||
|
// 处理照片...
|
||||||
|
|
||||||
|
// 发送处理完成事件
|
||||||
|
eventData := map[string]interface{}{
|
||||||
|
"type": "photo_processed",
|
||||||
|
"photo_id": photoID,
|
||||||
|
"status": "completed",
|
||||||
|
}
|
||||||
|
eventJSON, _ := json.Marshal(eventData)
|
||||||
|
|
||||||
|
select {
|
||||||
|
case eventChannel <- string(eventJSON):
|
||||||
|
// 事件发送成功
|
||||||
|
case <-time.After(1 * time.Second):
|
||||||
|
// 超时处理
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
## 配置选项
|
||||||
|
|
||||||
|
```go
|
||||||
|
opts := &sse.Options{
|
||||||
|
PingInterval: 30 * time.Second, // 心跳间隔
|
||||||
|
MaxClients: 1000, // 最大客户端数
|
||||||
|
AllowedOrigins: []string{"*"}, // 允许的来源
|
||||||
|
EnableCORS: true, // 启用CORS
|
||||||
|
EventBufferSize: 100, // 事件缓冲区大小
|
||||||
|
ClientTimeout: 5 * time.Minute, // 客户端超时时间
|
||||||
|
}
|
||||||
|
|
||||||
|
manager := sse.NewManager(opts)
|
||||||
|
```
|
||||||
|
|
||||||
|
## API 接口
|
||||||
|
|
||||||
|
### 客户端连接
|
||||||
|
```
|
||||||
|
GET /api/v1/events
|
||||||
|
```
|
||||||
|
|
||||||
|
### 广播消息
|
||||||
|
```
|
||||||
|
POST /api/v1/events/broadcast
|
||||||
|
Content-Type: application/json
|
||||||
|
|
||||||
|
{
|
||||||
|
"message": "Hello World",
|
||||||
|
"event_type": "notification" // 可选
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
### 定向发送
|
||||||
|
```
|
||||||
|
POST /api/v1/events/send/{clientId}
|
||||||
|
Content-Type: application/json
|
||||||
|
|
||||||
|
{
|
||||||
|
"message": "Private message"
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
### 获取客户端列表
|
||||||
|
```
|
||||||
|
GET /api/v1/events/clients
|
||||||
|
```
|
||||||
|
|
||||||
|
### 获取统计信息
|
||||||
|
```
|
||||||
|
GET /api/v1/events/stats
|
||||||
|
```
|
||||||
|
|
||||||
|
## 客户端连接示例
|
||||||
|
|
||||||
|
### JavaScript
|
||||||
|
```javascript
|
||||||
|
const eventSource = new EventSource('/api/v1/events');
|
||||||
|
|
||||||
|
eventSource.onopen = function(event) {
|
||||||
|
console.log('Connected to SSE');
|
||||||
|
};
|
||||||
|
|
||||||
|
eventSource.onmessage = function(event) {
|
||||||
|
const data = JSON.parse(event.data);
|
||||||
|
console.log('Received:', data);
|
||||||
|
};
|
||||||
|
|
||||||
|
// 监听特定事件类型
|
||||||
|
eventSource.addEventListener('photo_processed', function(event) {
|
||||||
|
const data = JSON.parse(event.data);
|
||||||
|
console.log('Photo processed:', data);
|
||||||
|
});
|
||||||
|
|
||||||
|
eventSource.onerror = function(event) {
|
||||||
|
console.log('SSE error:', event);
|
||||||
|
};
|
||||||
|
```
|
||||||
|
|
||||||
|
### cURL 测试
|
||||||
|
```bash
|
||||||
|
# 连接 SSE
|
||||||
|
curl -N -H "Accept: text/event-stream" http://localhost:8080/api/v1/events
|
||||||
|
|
||||||
|
# 发送广播消息
|
||||||
|
curl -X POST http://localhost:8080/api/v1/events/broadcast \
|
||||||
|
-H "Content-Type: application/json" \
|
||||||
|
-d '{"message": "Hello World", "event_type": "notification"}'
|
||||||
|
```
|
||||||
|
|
||||||
|
## 架构设计
|
||||||
|
|
||||||
|
```
|
||||||
|
业务层 (Business Logic)
|
||||||
|
↓ (通过 chan string)
|
||||||
|
SSE 管理器 (Manager)
|
||||||
|
↓ (WebSocket/HTTP连接)
|
||||||
|
客户端 (Frontend/Mobile App)
|
||||||
|
```
|
||||||
|
|
||||||
|
### 核心组件
|
||||||
|
|
||||||
|
1. **Manager**: 核心管理器,负责客户端生命周期管理
|
||||||
|
2. **Client**: 客户端连接抽象,包含连接信息和通信通道
|
||||||
|
3. **Event**: 事件数据结构,支持 SSE 标准格式
|
||||||
|
4. **Options**: 配置选项,支持自定义行为
|
||||||
|
|
||||||
|
### 线程安全
|
||||||
|
|
||||||
|
- 使用 `sync.RWMutex` 保护客户端集合
|
||||||
|
- 通过 channel 实现 goroutine 间通信
|
||||||
|
- 自动清理失效连接和资源
|
||||||
|
|
||||||
|
## 生产环境考虑
|
||||||
|
|
||||||
|
### 性能优化
|
||||||
|
- 事件缓冲区避免阻塞
|
||||||
|
- 读写锁优化并发访问
|
||||||
|
- 超时机制防止资源泄露
|
||||||
|
|
||||||
|
### 监控指标
|
||||||
|
- 客户端连接数
|
||||||
|
- 事件发送成功率
|
||||||
|
- 连接异常次数
|
||||||
|
|
||||||
|
### 扩展性
|
||||||
|
- 支持集群部署(需要外部消息队列)
|
||||||
|
- 支持客户端分组
|
||||||
|
- 支持事件过滤
|
||||||
|
|
||||||
|
## 故障排除
|
||||||
|
|
||||||
|
### 常见问题
|
||||||
|
|
||||||
|
1. **客户端连接失败**
|
||||||
|
- 检查 CORS 设置
|
||||||
|
- 确认路由配置正确
|
||||||
|
- 检查防火墙设置
|
||||||
|
|
||||||
|
2. **消息丢失**
|
||||||
|
- 增大事件缓冲区
|
||||||
|
- 检查网络稳定性
|
||||||
|
- 实现客户端重连机制
|
||||||
|
|
||||||
|
3. **内存占用过高**
|
||||||
|
- 限制最大客户端数
|
||||||
|
- 减少事件缓冲区大小
|
||||||
|
- 启用客户端超时清理
|
||||||
|
|
||||||
|
### 调试技巧
|
||||||
|
|
||||||
|
```go
|
||||||
|
// 启用详细日志
|
||||||
|
manager.SetDebugMode(true)
|
||||||
|
|
||||||
|
// 获取连接统计
|
||||||
|
stats := manager.GetStats()
|
||||||
|
fmt.Printf("Active clients: %d\n", stats.ClientCount)
|
||||||
|
```
|
||||||
|
|
||||||
|
## 示例项目
|
||||||
|
|
||||||
|
查看 `examples/sse_integration_example.go` 获取完整的集成示例。
|
||||||
|
|
||||||
|
运行示例:
|
||||||
|
```bash
|
||||||
|
cd examples
|
||||||
|
go run sse_integration_example.go
|
||||||
|
```
|
||||||
|
|
||||||
|
访问管理面板:http://localhost:8080/admin/sse/
|
||||||
|
|
||||||
|
## 许可证
|
||||||
|
|
||||||
|
本项目采用 MIT 许可证。
|
||||||
@@ -0,0 +1,97 @@
|
|||||||
|
package sse
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/gin-gonic/gin"
|
||||||
|
)
|
||||||
|
|
||||||
|
func ExampleUsage() {
|
||||||
|
r := gin.Default()
|
||||||
|
|
||||||
|
opts := DefaultOptions()
|
||||||
|
opts.PingInterval = 30 * time.Second
|
||||||
|
opts.MaxClients = 100
|
||||||
|
|
||||||
|
manager := NewManager(opts)
|
||||||
|
|
||||||
|
r.GET("/events", func(c *gin.Context) {
|
||||||
|
manager.RegisterClient(c)
|
||||||
|
})
|
||||||
|
|
||||||
|
r.POST("/broadcast", func(c *gin.Context) {
|
||||||
|
var req struct {
|
||||||
|
Message string `json:"message"`
|
||||||
|
}
|
||||||
|
if err := c.ShouldBindJSON(&req); err != nil {
|
||||||
|
c.JSON(400, gin.H{"error": err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
manager.Broadcast(req.Message)
|
||||||
|
c.JSON(200, gin.H{"status": "sent"})
|
||||||
|
})
|
||||||
|
|
||||||
|
eventChan := manager.GetEventChannel()
|
||||||
|
|
||||||
|
go func() {
|
||||||
|
ticker := time.NewTicker(10 * time.Second)
|
||||||
|
defer ticker.Stop()
|
||||||
|
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-ticker.C:
|
||||||
|
data := map[string]interface{}{
|
||||||
|
"timestamp": time.Now().Format(time.RFC3339),
|
||||||
|
"message": "Heartbeat from business logic",
|
||||||
|
"clients": manager.GetClientCount(),
|
||||||
|
}
|
||||||
|
jsonData, _ := json.Marshal(data)
|
||||||
|
eventChan <- string(jsonData)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
|
r.Run(":8080")
|
||||||
|
}
|
||||||
|
|
||||||
|
type BusinessService struct {
|
||||||
|
eventChannel chan<- string
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewBusinessService(eventChannel chan<- string) *BusinessService {
|
||||||
|
return &BusinessService{
|
||||||
|
eventChannel: eventChannel,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (bs *BusinessService) ProcessData(data interface{}) {
|
||||||
|
message := fmt.Sprintf("Data processed: %v at %s", data, time.Now().Format(time.RFC3339))
|
||||||
|
|
||||||
|
select {
|
||||||
|
case bs.eventChannel <- message:
|
||||||
|
fmt.Println("Event sent to SSE manager")
|
||||||
|
case <-time.After(1 * time.Second):
|
||||||
|
fmt.Println("Failed to send event: channel full or closed")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (bs *BusinessService) NotifyUsers(userIDs []string, message string) {
|
||||||
|
notification := map[string]interface{}{
|
||||||
|
"type": "notification",
|
||||||
|
"users": userIDs,
|
||||||
|
"message": message,
|
||||||
|
"timestamp": time.Now().Format(time.RFC3339),
|
||||||
|
}
|
||||||
|
|
||||||
|
jsonData, _ := json.Marshal(notification)
|
||||||
|
|
||||||
|
select {
|
||||||
|
case bs.eventChannel <- string(jsonData):
|
||||||
|
fmt.Printf("Notification sent to %d users\n", len(userIDs))
|
||||||
|
case <-time.After(1 * time.Second):
|
||||||
|
fmt.Println("Failed to send notification: timeout")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,404 @@
|
|||||||
|
package sse
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/gin-gonic/gin"
|
||||||
|
"github.com/google/uuid"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Client struct {
|
||||||
|
ID string
|
||||||
|
Writer gin.ResponseWriter
|
||||||
|
Request *http.Request
|
||||||
|
Context *gin.Context
|
||||||
|
EventChan chan *Event
|
||||||
|
CloseChan chan struct{}
|
||||||
|
LastPing time.Time
|
||||||
|
UserAgent string
|
||||||
|
RemoteAddr string
|
||||||
|
ConnectedAt time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
type Event struct {
|
||||||
|
ID string `json:"id,omitempty"`
|
||||||
|
Event string `json:"event,omitempty"`
|
||||||
|
Data string `json:"data"`
|
||||||
|
Retry int `json:"retry,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type Manager struct {
|
||||||
|
clients map[string]*Client
|
||||||
|
clientsMutex sync.RWMutex
|
||||||
|
broadcastChan chan *Event
|
||||||
|
registerChan chan *Client
|
||||||
|
unregisterChan chan string
|
||||||
|
eventChan chan string
|
||||||
|
ctx context.Context
|
||||||
|
cancel context.CancelFunc
|
||||||
|
pingInterval time.Duration
|
||||||
|
options *Options
|
||||||
|
}
|
||||||
|
|
||||||
|
type Options struct {
|
||||||
|
PingInterval time.Duration
|
||||||
|
MaxClients int
|
||||||
|
AllowedOrigins []string
|
||||||
|
EnableCORS bool
|
||||||
|
EventBufferSize int
|
||||||
|
ClientTimeout time.Duration
|
||||||
|
}
|
||||||
|
|
||||||
|
func DefaultOptions() *Options {
|
||||||
|
return &Options{
|
||||||
|
PingInterval: 30 * time.Second,
|
||||||
|
MaxClients: 1000,
|
||||||
|
AllowedOrigins: []string{"*"},
|
||||||
|
EnableCORS: true,
|
||||||
|
EventBufferSize: 100,
|
||||||
|
ClientTimeout: 5 * time.Minute,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewManager(opts *Options) *Manager {
|
||||||
|
if opts == nil {
|
||||||
|
opts = DefaultOptions()
|
||||||
|
}
|
||||||
|
|
||||||
|
// 确保必要的字段有有效值
|
||||||
|
if opts.PingInterval <= 0 {
|
||||||
|
opts.PingInterval = 30 * time.Second
|
||||||
|
}
|
||||||
|
if opts.EventBufferSize <= 0 {
|
||||||
|
opts.EventBufferSize = 100
|
||||||
|
}
|
||||||
|
if opts.MaxClients <= 0 {
|
||||||
|
opts.MaxClients = 1000
|
||||||
|
}
|
||||||
|
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
|
||||||
|
manager := &Manager{
|
||||||
|
clients: make(map[string]*Client),
|
||||||
|
broadcastChan: make(chan *Event, opts.EventBufferSize),
|
||||||
|
registerChan: make(chan *Client),
|
||||||
|
unregisterChan: make(chan string),
|
||||||
|
eventChan: make(chan string, opts.EventBufferSize),
|
||||||
|
ctx: ctx,
|
||||||
|
cancel: cancel,
|
||||||
|
pingInterval: opts.PingInterval,
|
||||||
|
options: opts,
|
||||||
|
}
|
||||||
|
|
||||||
|
go manager.run()
|
||||||
|
go manager.handleExternalEvents()
|
||||||
|
if opts.PingInterval > 0 {
|
||||||
|
go manager.periodicPing()
|
||||||
|
}
|
||||||
|
|
||||||
|
return manager
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) run() {
|
||||||
|
defer func() {
|
||||||
|
if r := recover(); r != nil {
|
||||||
|
fmt.Printf("SSE Manager panic recovered: %v\n", r)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-m.ctx.Done():
|
||||||
|
return
|
||||||
|
case client := <-m.registerChan:
|
||||||
|
m.addClient(client)
|
||||||
|
case clientID := <-m.unregisterChan:
|
||||||
|
m.removeClient(clientID)
|
||||||
|
case event := <-m.broadcastChan:
|
||||||
|
m.broadcastToAllClients(event)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) handleExternalEvents() {
|
||||||
|
defer func() {
|
||||||
|
if r := recover(); r != nil {
|
||||||
|
fmt.Printf("SSE Event Handler panic recovered: %v\n", r)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-m.ctx.Done():
|
||||||
|
return
|
||||||
|
case eventData := <-m.eventChan:
|
||||||
|
event := &Event{
|
||||||
|
ID: uuid.New().String(),
|
||||||
|
Data: eventData,
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case m.broadcastChan <- event:
|
||||||
|
case <-time.After(1 * time.Second):
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) periodicPing() {
|
||||||
|
ticker := time.NewTicker(m.pingInterval)
|
||||||
|
defer ticker.Stop()
|
||||||
|
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-m.ctx.Done():
|
||||||
|
return
|
||||||
|
case <-ticker.C:
|
||||||
|
m.pingAllClients()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) addClient(client *Client) {
|
||||||
|
if m.options.MaxClients > 0 && len(m.clients) >= m.options.MaxClients {
|
||||||
|
close(client.CloseChan)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
m.clientsMutex.Lock()
|
||||||
|
m.clients[client.ID] = client
|
||||||
|
m.clientsMutex.Unlock()
|
||||||
|
|
||||||
|
fmt.Printf("Client connected: %s (Total: %d)\n", client.ID, len(m.clients))
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) removeClient(clientID string) {
|
||||||
|
m.clientsMutex.Lock()
|
||||||
|
client, exists := m.clients[clientID]
|
||||||
|
if exists {
|
||||||
|
delete(m.clients, clientID)
|
||||||
|
close(client.CloseChan)
|
||||||
|
}
|
||||||
|
m.clientsMutex.Unlock()
|
||||||
|
|
||||||
|
if exists {
|
||||||
|
fmt.Printf("Client disconnected: %s (Total: %d)\n", clientID, len(m.clients))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) broadcastToAllClients(event *Event) {
|
||||||
|
m.clientsMutex.RLock()
|
||||||
|
clients := make([]*Client, 0, len(m.clients))
|
||||||
|
for _, client := range m.clients {
|
||||||
|
clients = append(clients, client)
|
||||||
|
}
|
||||||
|
m.clientsMutex.RUnlock()
|
||||||
|
|
||||||
|
for _, client := range clients {
|
||||||
|
select {
|
||||||
|
case client.EventChan <- event:
|
||||||
|
case <-time.After(100 * time.Millisecond):
|
||||||
|
go m.UnregisterClient(client.ID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) pingAllClients() {
|
||||||
|
pingEvent := &Event{
|
||||||
|
Event: "ping",
|
||||||
|
Data: fmt.Sprintf(`{"time":"%s"}`, time.Now().Format(time.RFC3339)),
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case m.broadcastChan <- pingEvent:
|
||||||
|
case <-time.After(1 * time.Second):
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) RegisterClient(c *gin.Context) {
|
||||||
|
if m.options.EnableCORS {
|
||||||
|
m.setupCORS(c)
|
||||||
|
}
|
||||||
|
|
||||||
|
c.Header("Content-Type", "text/event-stream")
|
||||||
|
c.Header("Cache-Control", "no-cache")
|
||||||
|
c.Header("Connection", "keep-alive")
|
||||||
|
c.Header("Access-Control-Allow-Origin", "*")
|
||||||
|
|
||||||
|
clientID := uuid.New().String()
|
||||||
|
client := &Client{
|
||||||
|
ID: clientID,
|
||||||
|
Writer: c.Writer,
|
||||||
|
Request: c.Request,
|
||||||
|
Context: c,
|
||||||
|
EventChan: make(chan *Event, 50),
|
||||||
|
CloseChan: make(chan struct{}),
|
||||||
|
LastPing: time.Now(),
|
||||||
|
UserAgent: c.GetHeader("User-Agent"),
|
||||||
|
RemoteAddr: c.ClientIP(),
|
||||||
|
ConnectedAt: time.Now(),
|
||||||
|
}
|
||||||
|
|
||||||
|
m.registerChan <- client
|
||||||
|
|
||||||
|
c.Stream(func(w io.Writer) bool {
|
||||||
|
select {
|
||||||
|
case <-client.CloseChan:
|
||||||
|
return false
|
||||||
|
case event := <-client.EventChan:
|
||||||
|
m.writeEvent(w, event)
|
||||||
|
client.LastPing = time.Now()
|
||||||
|
return true
|
||||||
|
case <-c.Request.Context().Done():
|
||||||
|
go m.UnregisterClient(clientID)
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
go m.UnregisterClient(clientID)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) setupCORS(c *gin.Context) {
|
||||||
|
origin := c.GetHeader("Origin")
|
||||||
|
if m.isAllowedOrigin(origin) {
|
||||||
|
c.Header("Access-Control-Allow-Origin", origin)
|
||||||
|
}
|
||||||
|
c.Header("Access-Control-Allow-Methods", "GET, OPTIONS")
|
||||||
|
c.Header("Access-Control-Allow-Headers", "Accept, Cache-Control")
|
||||||
|
c.Header("Access-Control-Allow-Credentials", "true")
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) isAllowedOrigin(origin string) bool {
|
||||||
|
for _, allowed := range m.options.AllowedOrigins {
|
||||||
|
if allowed == "*" || allowed == origin {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) writeEvent(w io.Writer, event *Event) {
|
||||||
|
if event.ID != "" {
|
||||||
|
fmt.Fprintf(w, "id: %s\n", event.ID)
|
||||||
|
}
|
||||||
|
if event.Event != "" {
|
||||||
|
fmt.Fprintf(w, "event: %s\n", event.Event)
|
||||||
|
}
|
||||||
|
if event.Retry > 0 {
|
||||||
|
fmt.Fprintf(w, "retry: %d\n", event.Retry)
|
||||||
|
}
|
||||||
|
fmt.Fprintf(w, "data: %s\n\n", event.Data)
|
||||||
|
|
||||||
|
if flusher, ok := w.(http.Flusher); ok {
|
||||||
|
flusher.Flush()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) UnregisterClient(clientID string) {
|
||||||
|
select {
|
||||||
|
case m.unregisterChan <- clientID:
|
||||||
|
case <-time.After(1 * time.Second):
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) Broadcast(message string) {
|
||||||
|
event := &Event{
|
||||||
|
ID: uuid.New().String(),
|
||||||
|
Data: message,
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case m.broadcastChan <- event:
|
||||||
|
case <-time.After(1 * time.Second):
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) BroadcastEvent(eventType, message string) {
|
||||||
|
event := &Event{
|
||||||
|
ID: uuid.New().String(),
|
||||||
|
Event: eventType,
|
||||||
|
Data: message,
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case m.broadcastChan <- event:
|
||||||
|
case <-time.After(1 * time.Second):
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) SendToClient(clientID, message string) bool {
|
||||||
|
m.clientsMutex.RLock()
|
||||||
|
client, exists := m.clients[clientID]
|
||||||
|
m.clientsMutex.RUnlock()
|
||||||
|
|
||||||
|
if !exists {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
event := &Event{
|
||||||
|
ID: uuid.New().String(),
|
||||||
|
Data: message,
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case client.EventChan <- event:
|
||||||
|
return true
|
||||||
|
case <-time.After(100 * time.Millisecond):
|
||||||
|
go m.UnregisterClient(clientID)
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) GetEventChannel() chan<- string {
|
||||||
|
return m.eventChan
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) GetClientCount() int {
|
||||||
|
m.clientsMutex.RLock()
|
||||||
|
defer m.clientsMutex.RUnlock()
|
||||||
|
return len(m.clients)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) GetClientIDs() []string {
|
||||||
|
m.clientsMutex.RLock()
|
||||||
|
defer m.clientsMutex.RUnlock()
|
||||||
|
|
||||||
|
ids := make([]string, 0, len(m.clients))
|
||||||
|
for id := range m.clients {
|
||||||
|
ids = append(ids, id)
|
||||||
|
}
|
||||||
|
return ids
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) GetClientInfo(clientID string) (*Client, bool) {
|
||||||
|
m.clientsMutex.RLock()
|
||||||
|
defer m.clientsMutex.RUnlock()
|
||||||
|
|
||||||
|
client, exists := m.clients[clientID]
|
||||||
|
if !exists {
|
||||||
|
return nil, false
|
||||||
|
}
|
||||||
|
|
||||||
|
clientCopy := *client
|
||||||
|
return &clientCopy, true
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *Manager) Close() {
|
||||||
|
m.cancel()
|
||||||
|
|
||||||
|
m.clientsMutex.Lock()
|
||||||
|
for _, client := range m.clients {
|
||||||
|
close(client.CloseChan)
|
||||||
|
}
|
||||||
|
m.clients = make(map[string]*Client)
|
||||||
|
m.clientsMutex.Unlock()
|
||||||
|
|
||||||
|
close(m.broadcastChan)
|
||||||
|
close(m.registerChan)
|
||||||
|
close(m.unregisterChan)
|
||||||
|
close(m.eventChan)
|
||||||
|
}
|
||||||
@@ -0,0 +1,448 @@
|
|||||||
|
package sse
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestNewManager(t *testing.T) {
|
||||||
|
manager := NewManager(nil)
|
||||||
|
if manager == nil {
|
||||||
|
t.Fatal("Manager should not be nil")
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(manager.clients) != 0 {
|
||||||
|
t.Errorf("Expected 0 clients, got %d", len(manager.clients))
|
||||||
|
}
|
||||||
|
|
||||||
|
manager.Close()
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestManagerWithCustomOptions(t *testing.T) {
|
||||||
|
opts := &Options{
|
||||||
|
PingInterval: 10 * time.Second,
|
||||||
|
MaxClients: 50,
|
||||||
|
AllowedOrigins: []string{"http://localhost:3000"},
|
||||||
|
EnableCORS: true,
|
||||||
|
EventBufferSize: 10,
|
||||||
|
ClientTimeout: 2 * time.Minute,
|
||||||
|
}
|
||||||
|
|
||||||
|
manager := NewManager(opts)
|
||||||
|
if manager == nil {
|
||||||
|
t.Fatal("Manager should not be nil")
|
||||||
|
}
|
||||||
|
|
||||||
|
if manager.options.MaxClients != 50 {
|
||||||
|
t.Errorf("Expected MaxClients 50, got %d", manager.options.MaxClients)
|
||||||
|
}
|
||||||
|
|
||||||
|
if manager.options.PingInterval != 10*time.Second {
|
||||||
|
t.Errorf("Expected PingInterval 10s, got %v", manager.options.PingInterval)
|
||||||
|
}
|
||||||
|
|
||||||
|
manager.Close()
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestEventChannelBroadcast(t *testing.T) {
|
||||||
|
manager := NewManager(nil)
|
||||||
|
defer manager.Close()
|
||||||
|
|
||||||
|
eventChannel := manager.GetEventChannel()
|
||||||
|
|
||||||
|
// 创建一个测试客户端
|
||||||
|
testClient := &Client{
|
||||||
|
ID: "test-client-1",
|
||||||
|
EventChan: make(chan *Event, 10),
|
||||||
|
CloseChan: make(chan struct{}),
|
||||||
|
LastPing: time.Now(),
|
||||||
|
}
|
||||||
|
|
||||||
|
// 手动添加客户端到管理器
|
||||||
|
manager.clientsMutex.Lock()
|
||||||
|
manager.clients[testClient.ID] = testClient
|
||||||
|
manager.clientsMutex.Unlock()
|
||||||
|
|
||||||
|
// 通过事件通道发送消息
|
||||||
|
testMessage := "test broadcast message"
|
||||||
|
|
||||||
|
go func() {
|
||||||
|
eventChannel <- testMessage
|
||||||
|
}()
|
||||||
|
|
||||||
|
// 等待事件传播
|
||||||
|
select {
|
||||||
|
case event := <-testClient.EventChan:
|
||||||
|
if event.Data != testMessage {
|
||||||
|
t.Errorf("Expected message '%s', got '%s'", testMessage, event.Data)
|
||||||
|
}
|
||||||
|
case <-time.After(2 * time.Second):
|
||||||
|
t.Error("Timeout waiting for broadcast event")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBroadcastMethod(t *testing.T) {
|
||||||
|
manager := NewManager(nil)
|
||||||
|
defer manager.Close()
|
||||||
|
|
||||||
|
// 创建测试客户端
|
||||||
|
testClient := &Client{
|
||||||
|
ID: "test-client-2",
|
||||||
|
EventChan: make(chan *Event, 10),
|
||||||
|
CloseChan: make(chan struct{}),
|
||||||
|
LastPing: time.Now(),
|
||||||
|
}
|
||||||
|
|
||||||
|
// 手动添加客户端
|
||||||
|
manager.clientsMutex.Lock()
|
||||||
|
manager.clients[testClient.ID] = testClient
|
||||||
|
manager.clientsMutex.Unlock()
|
||||||
|
|
||||||
|
testMessage := "direct broadcast test"
|
||||||
|
manager.Broadcast(testMessage)
|
||||||
|
|
||||||
|
select {
|
||||||
|
case event := <-testClient.EventChan:
|
||||||
|
if event.Data != testMessage {
|
||||||
|
t.Errorf("Expected message '%s', got '%s'", testMessage, event.Data)
|
||||||
|
}
|
||||||
|
if event.ID == "" {
|
||||||
|
t.Error("Event ID should not be empty")
|
||||||
|
}
|
||||||
|
case <-time.After(2 * time.Second):
|
||||||
|
t.Error("Timeout waiting for direct broadcast")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBroadcastEventMethod(t *testing.T) {
|
||||||
|
manager := NewManager(nil)
|
||||||
|
defer manager.Close()
|
||||||
|
|
||||||
|
testClient := &Client{
|
||||||
|
ID: "test-client-3",
|
||||||
|
EventChan: make(chan *Event, 10),
|
||||||
|
CloseChan: make(chan struct{}),
|
||||||
|
LastPing: time.Now(),
|
||||||
|
}
|
||||||
|
|
||||||
|
manager.clientsMutex.Lock()
|
||||||
|
manager.clients[testClient.ID] = testClient
|
||||||
|
manager.clientsMutex.Unlock()
|
||||||
|
|
||||||
|
testEventType := "test-event"
|
||||||
|
testMessage := "test event message"
|
||||||
|
|
||||||
|
manager.BroadcastEvent(testEventType, testMessage)
|
||||||
|
|
||||||
|
select {
|
||||||
|
case event := <-testClient.EventChan:
|
||||||
|
if event.Data != testMessage {
|
||||||
|
t.Errorf("Expected message '%s', got '%s'", testMessage, event.Data)
|
||||||
|
}
|
||||||
|
if event.Event != testEventType {
|
||||||
|
t.Errorf("Expected event type '%s', got '%s'", testEventType, event.Event)
|
||||||
|
}
|
||||||
|
if event.ID == "" {
|
||||||
|
t.Error("Event ID should not be empty")
|
||||||
|
}
|
||||||
|
case <-time.After(2 * time.Second):
|
||||||
|
t.Error("Timeout waiting for event broadcast")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSendToClient(t *testing.T) {
|
||||||
|
manager := NewManager(nil)
|
||||||
|
defer manager.Close()
|
||||||
|
|
||||||
|
// 创建两个测试客户端
|
||||||
|
client1 := &Client{
|
||||||
|
ID: "test-client-1",
|
||||||
|
EventChan: make(chan *Event, 10),
|
||||||
|
CloseChan: make(chan struct{}),
|
||||||
|
LastPing: time.Now(),
|
||||||
|
}
|
||||||
|
|
||||||
|
client2 := &Client{
|
||||||
|
ID: "test-client-2",
|
||||||
|
EventChan: make(chan *Event, 10),
|
||||||
|
CloseChan: make(chan struct{}),
|
||||||
|
LastPing: time.Now(),
|
||||||
|
}
|
||||||
|
|
||||||
|
manager.clientsMutex.Lock()
|
||||||
|
manager.clients[client1.ID] = client1
|
||||||
|
manager.clients[client2.ID] = client2
|
||||||
|
manager.clientsMutex.Unlock()
|
||||||
|
|
||||||
|
testMessage := "private message to client 1"
|
||||||
|
success := manager.SendToClient(client1.ID, testMessage)
|
||||||
|
|
||||||
|
if !success {
|
||||||
|
t.Error("SendToClient should return true for existing client")
|
||||||
|
}
|
||||||
|
|
||||||
|
// 检查client1收到消息
|
||||||
|
select {
|
||||||
|
case event := <-client1.EventChan:
|
||||||
|
if event.Data != testMessage {
|
||||||
|
t.Errorf("Client1 expected message '%s', got '%s'", testMessage, event.Data)
|
||||||
|
}
|
||||||
|
case <-time.After(2 * time.Second):
|
||||||
|
t.Error("Client1 timeout waiting for private message")
|
||||||
|
}
|
||||||
|
|
||||||
|
// 检查client2没有收到消息
|
||||||
|
select {
|
||||||
|
case <-client2.EventChan:
|
||||||
|
t.Error("Client2 should not receive private message meant for Client1")
|
||||||
|
case <-time.After(100 * time.Millisecond):
|
||||||
|
// 这是预期行为
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSendToNonExistentClient(t *testing.T) {
|
||||||
|
manager := NewManager(nil)
|
||||||
|
defer manager.Close()
|
||||||
|
|
||||||
|
success := manager.SendToClient("non-existent-client", "test message")
|
||||||
|
if success {
|
||||||
|
t.Error("SendToClient should return false for non-existent client")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGetClientCount(t *testing.T) {
|
||||||
|
manager := NewManager(nil)
|
||||||
|
defer manager.Close()
|
||||||
|
|
||||||
|
if manager.GetClientCount() != 0 {
|
||||||
|
t.Errorf("Expected client count 0, got %d", manager.GetClientCount())
|
||||||
|
}
|
||||||
|
|
||||||
|
// 添加测试客户端
|
||||||
|
testClient := &Client{
|
||||||
|
ID: "test-client",
|
||||||
|
EventChan: make(chan *Event, 10),
|
||||||
|
CloseChan: make(chan struct{}),
|
||||||
|
LastPing: time.Now(),
|
||||||
|
}
|
||||||
|
|
||||||
|
manager.clientsMutex.Lock()
|
||||||
|
manager.clients[testClient.ID] = testClient
|
||||||
|
manager.clientsMutex.Unlock()
|
||||||
|
|
||||||
|
if manager.GetClientCount() != 1 {
|
||||||
|
t.Errorf("Expected client count 1, got %d", manager.GetClientCount())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGetClientIDs(t *testing.T) {
|
||||||
|
manager := NewManager(nil)
|
||||||
|
defer manager.Close()
|
||||||
|
|
||||||
|
ids := manager.GetClientIDs()
|
||||||
|
if len(ids) != 0 {
|
||||||
|
t.Errorf("Expected 0 client IDs, got %d", len(ids))
|
||||||
|
}
|
||||||
|
|
||||||
|
// 添加测试客户端
|
||||||
|
testClients := map[string]*Client{
|
||||||
|
"client-1": {
|
||||||
|
ID: "client-1",
|
||||||
|
EventChan: make(chan *Event, 10),
|
||||||
|
CloseChan: make(chan struct{}),
|
||||||
|
LastPing: time.Now(),
|
||||||
|
},
|
||||||
|
"client-2": {
|
||||||
|
ID: "client-2",
|
||||||
|
EventChan: make(chan *Event, 10),
|
||||||
|
CloseChan: make(chan struct{}),
|
||||||
|
LastPing: time.Now(),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
manager.clientsMutex.Lock()
|
||||||
|
for id, client := range testClients {
|
||||||
|
manager.clients[id] = client
|
||||||
|
}
|
||||||
|
manager.clientsMutex.Unlock()
|
||||||
|
|
||||||
|
ids = manager.GetClientIDs()
|
||||||
|
if len(ids) != 2 {
|
||||||
|
t.Errorf("Expected 2 client IDs, got %d", len(ids))
|
||||||
|
}
|
||||||
|
|
||||||
|
// 检查是否包含正确的ID
|
||||||
|
idMap := make(map[string]bool)
|
||||||
|
for _, id := range ids {
|
||||||
|
idMap[id] = true
|
||||||
|
}
|
||||||
|
|
||||||
|
if !idMap["client-1"] || !idMap["client-2"] {
|
||||||
|
t.Error("Client IDs not returned correctly")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGetClientInfo(t *testing.T) {
|
||||||
|
manager := NewManager(nil)
|
||||||
|
defer manager.Close()
|
||||||
|
|
||||||
|
// 测试不存在的客户端
|
||||||
|
_, exists := manager.GetClientInfo("non-existent")
|
||||||
|
if exists {
|
||||||
|
t.Error("GetClientInfo should return false for non-existent client")
|
||||||
|
}
|
||||||
|
|
||||||
|
// 添加测试客户端
|
||||||
|
originalClient := &Client{
|
||||||
|
ID: "test-client",
|
||||||
|
EventChan: make(chan *Event, 10),
|
||||||
|
CloseChan: make(chan struct{}),
|
||||||
|
LastPing: time.Now(),
|
||||||
|
RemoteAddr: "127.0.0.1",
|
||||||
|
UserAgent: "test-agent",
|
||||||
|
ConnectedAt: time.Now(),
|
||||||
|
}
|
||||||
|
|
||||||
|
manager.clientsMutex.Lock()
|
||||||
|
manager.clients[originalClient.ID] = originalClient
|
||||||
|
manager.clientsMutex.Unlock()
|
||||||
|
|
||||||
|
clientInfo, exists := manager.GetClientInfo("test-client")
|
||||||
|
if !exists {
|
||||||
|
t.Error("GetClientInfo should return true for existing client")
|
||||||
|
}
|
||||||
|
|
||||||
|
if clientInfo.ID != originalClient.ID {
|
||||||
|
t.Errorf("Expected client ID '%s', got '%s'", originalClient.ID, clientInfo.ID)
|
||||||
|
}
|
||||||
|
|
||||||
|
if clientInfo.RemoteAddr != originalClient.RemoteAddr {
|
||||||
|
t.Errorf("Expected remote addr '%s', got '%s'", originalClient.RemoteAddr, clientInfo.RemoteAddr)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestManagerClose(t *testing.T) {
|
||||||
|
manager := NewManager(nil)
|
||||||
|
|
||||||
|
// 添加一个测试客户端
|
||||||
|
testClient := &Client{
|
||||||
|
ID: "test-client",
|
||||||
|
EventChan: make(chan *Event, 10),
|
||||||
|
CloseChan: make(chan struct{}),
|
||||||
|
LastPing: time.Now(),
|
||||||
|
}
|
||||||
|
|
||||||
|
manager.clientsMutex.Lock()
|
||||||
|
manager.clients[testClient.ID] = testClient
|
||||||
|
manager.clientsMutex.Unlock()
|
||||||
|
|
||||||
|
// 关闭管理器
|
||||||
|
manager.Close()
|
||||||
|
|
||||||
|
// 验证客户端被清理
|
||||||
|
if len(manager.clients) != 0 {
|
||||||
|
t.Errorf("Expected 0 clients after close, got %d", len(manager.clients))
|
||||||
|
}
|
||||||
|
|
||||||
|
// 验证context被取消
|
||||||
|
select {
|
||||||
|
case <-manager.ctx.Done():
|
||||||
|
// 这是预期行为
|
||||||
|
case <-time.After(100 * time.Millisecond):
|
||||||
|
t.Error("Manager context should be cancelled after Close()")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMaxClientsLimit(t *testing.T) {
|
||||||
|
opts := &Options{
|
||||||
|
MaxClients: 2,
|
||||||
|
EventBufferSize: 10,
|
||||||
|
PingInterval: 30 * time.Second,
|
||||||
|
}
|
||||||
|
|
||||||
|
manager := NewManager(opts)
|
||||||
|
defer manager.Close()
|
||||||
|
|
||||||
|
// 创建3个测试客户端,但只有2个应该被接受
|
||||||
|
clients := []*Client{
|
||||||
|
{
|
||||||
|
ID: "client-1",
|
||||||
|
EventChan: make(chan *Event, 10),
|
||||||
|
CloseChan: make(chan struct{}),
|
||||||
|
LastPing: time.Now(),
|
||||||
|
},
|
||||||
|
{
|
||||||
|
ID: "client-2",
|
||||||
|
EventChan: make(chan *Event, 10),
|
||||||
|
CloseChan: make(chan struct{}),
|
||||||
|
LastPing: time.Now(),
|
||||||
|
},
|
||||||
|
{
|
||||||
|
ID: "client-3",
|
||||||
|
EventChan: make(chan *Event, 10),
|
||||||
|
CloseChan: make(chan struct{}),
|
||||||
|
LastPing: time.Now(),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
// 添加客户端
|
||||||
|
for _, client := range clients {
|
||||||
|
manager.addClient(client)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 验证只有2个客户端被接受
|
||||||
|
if manager.GetClientCount() != 2 {
|
||||||
|
t.Errorf("Expected max 2 clients, got %d", manager.GetClientCount())
|
||||||
|
}
|
||||||
|
|
||||||
|
// 验证第三个客户端的CloseChan被关闭
|
||||||
|
select {
|
||||||
|
case <-clients[2].CloseChan:
|
||||||
|
// 这是预期行为,第三个客户端应该被拒绝
|
||||||
|
case <-time.After(100 * time.Millisecond):
|
||||||
|
t.Error("Third client should be rejected and CloseChan should be closed")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func BenchmarkBroadcast(b *testing.B) {
|
||||||
|
manager := NewManager(&Options{
|
||||||
|
EventBufferSize: 1000,
|
||||||
|
MaxClients: 100,
|
||||||
|
PingInterval: 30 * time.Second,
|
||||||
|
AllowedOrigins: []string{"*"},
|
||||||
|
EnableCORS: true,
|
||||||
|
ClientTimeout: 5 * time.Minute,
|
||||||
|
})
|
||||||
|
defer manager.Close()
|
||||||
|
|
||||||
|
// 添加一些测试客户端
|
||||||
|
for i := 0; i < 10; i++ {
|
||||||
|
client := &Client{
|
||||||
|
ID: string(rune('A' + i)),
|
||||||
|
EventChan: make(chan *Event, 100),
|
||||||
|
CloseChan: make(chan struct{}),
|
||||||
|
LastPing: time.Now(),
|
||||||
|
}
|
||||||
|
manager.clientsMutex.Lock()
|
||||||
|
manager.clients[client.ID] = client
|
||||||
|
manager.clientsMutex.Unlock()
|
||||||
|
|
||||||
|
// 启动goroutine来消费事件,避免channel阻塞
|
||||||
|
go func(c *Client) {
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-c.EventChan:
|
||||||
|
// 消费事件
|
||||||
|
case <-c.CloseChan:
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}(client)
|
||||||
|
}
|
||||||
|
|
||||||
|
b.ResetTimer()
|
||||||
|
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
manager.Broadcast("benchmark test message")
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in new issue
Block a user