102 lines
1.9 KiB
Go
102 lines
1.9 KiB
Go
package chat
|
|
|
|
import (
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/assert"
|
|
"go.uber.org/zap/zaptest"
|
|
)
|
|
|
|
func TestHub_ConcurrentConnections(t *testing.T) {
|
|
logger := zaptest.NewLogger(t)
|
|
hub := NewHub(logger, nil)
|
|
go hub.Run()
|
|
time.Sleep(10 * time.Millisecond)
|
|
defer hub.Shutdown()
|
|
|
|
const numUsers = 100
|
|
const messagesPerUser = 10
|
|
roomID := uuid.New()
|
|
|
|
clients := make([]*Client, numUsers)
|
|
for i := 0; i < numUsers; i++ {
|
|
clients[i] = &Client{
|
|
Hub: hub,
|
|
UserID: uuid.New(),
|
|
send: make(chan []byte, 4096),
|
|
}
|
|
hub.Register(clients[i])
|
|
}
|
|
time.Sleep(100 * time.Millisecond)
|
|
|
|
for _, c := range clients {
|
|
hub.JoinRoom(c, roomID)
|
|
}
|
|
|
|
// Drain send channels to prevent blocking
|
|
done := make(chan struct{})
|
|
go func() {
|
|
for {
|
|
select {
|
|
case <-done:
|
|
return
|
|
default:
|
|
for _, c := range clients {
|
|
select {
|
|
case <-c.send:
|
|
default:
|
|
}
|
|
}
|
|
time.Sleep(time.Millisecond)
|
|
}
|
|
}
|
|
}()
|
|
|
|
var errors atomic.Int64
|
|
var wg sync.WaitGroup
|
|
|
|
for i := 0; i < numUsers; i++ {
|
|
wg.Add(1)
|
|
go func(idx int) {
|
|
defer wg.Done()
|
|
for j := 0; j < messagesPerUser; j++ {
|
|
msg := []byte(`{"type":"NewMessage","content":"bench"}`)
|
|
func() {
|
|
defer func() {
|
|
if r := recover(); r != nil {
|
|
errors.Add(1)
|
|
}
|
|
}()
|
|
hub.BroadcastToRoom(roomID, msg, clients[idx])
|
|
}()
|
|
}
|
|
}(i)
|
|
}
|
|
|
|
wg.Wait()
|
|
close(done)
|
|
time.Sleep(200 * time.Millisecond)
|
|
|
|
assert.Equal(t, int64(0), errors.Load(), "expected 0 errors during concurrent broadcast")
|
|
assert.Equal(t, numUsers, hub.GetConnectedUsersCount())
|
|
|
|
for _, c := range clients {
|
|
hub.Unregister(c)
|
|
}
|
|
time.Sleep(100 * time.Millisecond)
|
|
assert.Equal(t, 0, hub.GetConnectedUsersCount())
|
|
}
|
|
|
|
func BenchmarkRateLimiter_InMemory(b *testing.B) {
|
|
rl := NewRateLimiter(nil, nil)
|
|
userID := uuid.New()
|
|
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
rl.Allow(userID, "send_message")
|
|
}
|
|
}
|