diff --git a/pkg/events/websocket/websocket_producer.go b/pkg/events/websocket/websocket_producer.go index f28ff9aa..336c2fe6 100644 --- a/pkg/events/websocket/websocket_producer.go +++ b/pkg/events/websocket/websocket_producer.go @@ -19,17 +19,29 @@ var upgrader = websocket.Upgrader{ }, } +// wsConn serializa writes: gorilla/websocket permite apenas um writer concorrente. +type wsConn struct { + conn *websocket.Conn + writeMu sync.Mutex +} + +func (w *wsConn) writeJSON(v any) error { + w.writeMu.Lock() + defer w.writeMu.Unlock() + return w.conn.WriteJSON(v) +} + type websocketProducer struct { - clients map[string]*websocket.Conn // conexões específicas por instância - broadcast []*websocket.Conn // conexões que recebem todos os eventos + clients map[string]*wsConn // conexões específicas por instância + broadcast []*wsConn // conexões que recebem todos os eventos clientsMux sync.RWMutex loggerWrapper *logger_wrapper.LoggerManager } func NewWebsocketProducer(loggerWrapper *logger_wrapper.LoggerManager) *websocketProducer { return &websocketProducer{ - clients: make(map[string]*websocket.Conn), - broadcast: make([]*websocket.Conn, 0), + clients: make(map[string]*wsConn), + broadcast: make([]*wsConn, 0), clientsMux: sync.RWMutex{}, loggerWrapper: loggerWrapper, } @@ -72,7 +84,7 @@ func ServeWs(w http.ResponseWriter, r *http.Request, instanceId string, producer func (p *websocketProducer) AddBroadcastClient(conn *websocket.Conn) { p.clientsMux.Lock() defer p.clientsMux.Unlock() - p.broadcast = append(p.broadcast, conn) + p.broadcast = append(p.broadcast, &wsConn{conn: conn}) logger.LogInfo("Cliente broadcast websocket adicionado") } @@ -80,7 +92,7 @@ func (p *websocketProducer) RemoveBroadcastClient(conn *websocket.Conn) { p.clientsMux.Lock() defer p.clientsMux.Unlock() for i, c := range p.broadcast { - if c == conn { + if c.conn == conn { p.broadcast = append(p.broadcast[:i], p.broadcast[i+1:]...) break } @@ -91,7 +103,7 @@ func (p *websocketProducer) RemoveBroadcastClient(conn *websocket.Conn) { func (p *websocketProducer) AddClient(instanceID string, conn *websocket.Conn) { p.clientsMux.Lock() defer p.clientsMux.Unlock() - p.clients[instanceID] = conn + p.clients[instanceID] = &wsConn{conn: conn} p.loggerWrapper.GetLogger(instanceID).LogInfo("Cliente websocket adicionado para instância: %s", instanceID) } @@ -113,7 +125,7 @@ func (p *websocketProducer) Produce(queueName string, payload []byte, instanceID // Envia para cliente específico da instância if client, exists := p.clients[instanceID]; exists { - err := client.WriteJSON(message) + err := client.writeJSON(message) if err != nil { p.loggerWrapper.GetLogger(instanceID).LogError("Erro ao enviar mensagem websocket para %s: %v", instanceID, err) // Não remove o cliente aqui pois estamos com o RLock @@ -124,7 +136,7 @@ func (p *websocketProducer) Produce(queueName string, payload []byte, instanceID // Envia para todos os clientes broadcast for _, conn := range p.broadcast { - err := conn.WriteJSON(message) + err := conn.writeJSON(message) if err != nil { p.loggerWrapper.GetLogger(instanceID).LogError("Erro ao enviar mensagem broadcast websocket: %v", err) continue diff --git a/pkg/events/websocket/websocket_producer_test.go b/pkg/events/websocket/websocket_producer_test.go new file mode 100644 index 00000000..34b3e925 --- /dev/null +++ b/pkg/events/websocket/websocket_producer_test.go @@ -0,0 +1,154 @@ +package websocket_producer + +import ( + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/evolution-foundation/evolution-go/pkg/config" + logger_wrapper "github.com/evolution-foundation/evolution-go/pkg/logger" + "github.com/gorilla/websocket" +) + +func newTestLoggerManager(t *testing.T) *logger_wrapper.LoggerManager { + t.Helper() + return logger_wrapper.NewLoggerManager(&config.Config{ + LogDirectory: t.TempDir(), + LogMaxSize: 1, + LogMaxBackups: 1, + LogMaxAge: 1, + LogCompress: false, + }) +} + +// dialTestWS starts a local websocket server, dials a client, and returns the +// server-side connection registered for Produce writes. +func dialTestWS(t *testing.T) (serverConn *websocket.Conn, cleanup func()) { + t.Helper() + + upgraded := make(chan *websocket.Conn, 1) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade: %v", err) + return + } + upgraded <- conn + // Keep the connection open until the test finishes. + for { + if _, _, err := conn.ReadMessage(); err != nil { + return + } + } + })) + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + client, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + srv.Close() + t.Fatalf("dial: %v", err) + } + + var serverConnResult *websocket.Conn + select { + case serverConnResult = <-upgraded: + case <-time.After(2 * time.Second): + client.Close() + srv.Close() + t.Fatal("timed out waiting for server upgrade") + } + + // Drain client reads so WriteJSON on the server does not block on buffer full. + go func() { + for { + if _, _, err := client.ReadMessage(); err != nil { + return + } + } + }() + + cleanup = func() { + _ = client.Close() + if serverConnResult != nil { + _ = serverConnResult.Close() + } + srv.Close() + } + return serverConnResult, cleanup +} + +func TestProduceConcurrentWrites(t *testing.T) { + const ( + goroutines = 50 + messages = 20 + ) + + t.Run("concurrent instance writes", func(t *testing.T) { + serverConn, cleanup := dialTestWS(t) + defer cleanup() + + producer := NewWebsocketProducer(newTestLoggerManager(t)) + const instanceID = "instance-race-test" + producer.AddClient(instanceID, serverConn) + + var wg sync.WaitGroup + errCh := make(chan error, goroutines*messages) + + for g := 0; g < goroutines; g++ { + wg.Add(1) + go func(n int) { + defer wg.Done() + for i := 0; i < messages; i++ { + payload := []byte(fmt.Sprintf(`{"event":"Receipt","n":%d}`, n)) + if err := producer.Produce("instance-race-test.receipt", payload, instanceID, ""); err != nil { + errCh <- err + return + } + } + }(g) + } + + wg.Wait() + close(errCh) + + for err := range errCh { + t.Fatalf("Produce failed under concurrency: %v", err) + } + }) + + t.Run("concurrent broadcast writes", func(t *testing.T) { + serverConn, cleanup := dialTestWS(t) + defer cleanup() + + producer := NewWebsocketProducer(newTestLoggerManager(t)) + producer.AddBroadcastClient(serverConn) + + var wg sync.WaitGroup + errCh := make(chan error, goroutines*messages) + + for g := 0; g < goroutines; g++ { + wg.Add(1) + go func(n int) { + defer wg.Done() + for i := 0; i < messages; i++ { + payload := []byte(fmt.Sprintf(`{"event":"Message","n":%d}`, n)) + if err := producer.Produce("broadcast.message", payload, "any-instance", ""); err != nil { + errCh <- err + return + } + } + }(g) + } + + wg.Wait() + close(errCh) + + for err := range errCh { + t.Fatalf("Produce failed under concurrency: %v", err) + } + }) +}