diff --git a/api/stream/client.go b/api/stream/client.go index 74ad0467..dcd6d6b7 100644 --- a/api/stream/client.go +++ b/api/stream/client.go @@ -24,6 +24,7 @@ type client struct { conn *websocket.Conn onClose func(*client) write chan *model.MessageExternal + closed chan struct{} userID uint token string once once @@ -33,6 +34,7 @@ func newClient(conn *websocket.Conn, userID uint, token string, onClose func(*cl return &client{ conn: conn, write: make(chan *model.MessageExternal, 1), + closed: make(chan struct{}), userID: userID, token: token, onClose: onClose, @@ -43,7 +45,7 @@ func newClient(conn *websocket.Conn, userID uint, token string, onClose func(*cl func (c *client) Close() { c.once.Do(func() { c.conn.Close() - close(c.write) + close(c.closed) }) } @@ -51,7 +53,7 @@ func (c *client) Close() { func (c *client) NotifyClose() { c.once.Do(func() { c.conn.Close() - close(c.write) + close(c.closed) c.onClose(c) }) } @@ -87,11 +89,9 @@ func (c *client) startWriteHandler(pingPeriod time.Duration) { for { select { - case message, ok := <-c.write: - if !ok { - return - } - + case <-c.closed: + return + case message := <-c.write: c.conn.SetWriteDeadline(time.Now().Add(writeWait)) if err := writeJSON(c.conn, message); err != nil { printWebSocketError("WriteError", err) diff --git a/api/stream/stream.go b/api/stream/stream.go index 30074715..96877983 100644 --- a/api/stream/stream.go +++ b/api/stream/stream.go @@ -84,7 +84,10 @@ func (a *API) Notify(userID uint, msg *model.MessageExternal) { defer a.lock.RUnlock() if clients, ok := a.clients[userID]; ok { for _, c := range clients { - c.write <- msg + select { + case c.write <- msg: + case <-c.closed: + } } } } diff --git a/api/stream/stream_test.go b/api/stream/stream_test.go index 0d1e5873..0995d66c 100644 --- a/api/stream/stream_test.go +++ b/api/stream/stream_test.go @@ -99,6 +99,24 @@ func TestWritePingFails(t *testing.T) { user.expectNoMessage() } +func TestNotifyDoesNotPanicWhenClientIsClosed(t *testing.T) { + mode.Set(mode.TestDev) + defer leaktest.Check(t)() + + server, api := bootTestServer(staticUserID()) + defer server.Close() + defer api.Close() + + ws, _, err := websocket.DefaultDialer.Dial(wsURL(server.URL), nil) + assert.Nil(t, err) + defer ws.Close() + + waitForConnectedClients(api, 1) + + clients(api, 1)[0].Close() + api.Notify(1, &model.MessageExternal{Message: "after close"}) +} + func TestPing(t *testing.T) { mode.Set(mode.TestDev)