diff --git a/util/ws/websocket.go b/util/ws/websocket.go index f36c5d3..e9aed61 100644 --- a/util/ws/websocket.go +++ b/util/ws/websocket.go @@ -3,6 +3,7 @@ package ws import ( "context" "github.com/gorilla/websocket" + "log" ) type Packet struct { @@ -27,16 +28,21 @@ func Binary(buf []byte) *Packet { type Conn struct { WriteChan chan *Packet ReadChan chan []byte - ctx context.Context + Context context.Context + cancel context.CancelFunc conn *websocket.Conn close bool } func Wrap(conn *websocket.Conn) *Conn { + ctx, cancel := context.WithCancel(context.Background()) return &Conn{ WriteChan: make(chan *Packet), ReadChan: make(chan []byte), conn: conn, + Context: ctx, + cancel: cancel, + close: false, } } @@ -44,6 +50,9 @@ func (conn *Conn) ReadLoop() { for !conn.close { _, message, err := conn.conn.ReadMessage() if err != nil { + if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseAbnormalClosure) { + log.Println("websocket read error:", err) + } conn.Close() break } @@ -63,16 +72,14 @@ func (conn *Conn) WriteLoop() { if err != nil { break } - default: - if conn.close { - return - } + case <-conn.Context.Done(): + return } } } func (conn *Conn) Close() error { - conn.ReadChan <- nil + conn.cancel() conn.close = true return conn.conn.Close() }