From 1efdfa8aff304a451c4e7ea7b2ae9c2c4db1f081 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=A1=D1=8B=D1=81?=
 =?UTF-8?q?=D0=BE=D0=B5=D0=B2?= <k.sysoev@me.com>
Date: Wed, 8 Nov 2023 22:39:53 +0800
Subject: [PATCH] Ignoring errors if connection is closed from our side

---
 pkg/ws/ws.go | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/pkg/ws/ws.go b/pkg/ws/ws.go
index 2fbd712..52e62c8 100644
--- a/pkg/ws/ws.go
+++ b/pkg/ws/ws.go
@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"strings"
 	"sync"
+	"sync/atomic"
 
 	"github.com/fatih/color"
 	"golang.org/x/net/websocket"
@@ -44,6 +45,7 @@ type Connection struct {
 	ws        *websocket.Conn
 	Messages  chan Message
 	waitGroup *sync.WaitGroup
+	isClosed  atomic.Bool
 }
 
 type Options struct {
@@ -92,6 +94,8 @@ func NewWS(url string, opts Options) (*Connection, error) {
 
 	messages := make(chan Message, WSMessageBufferSize)
 
+	wsInsp := &Connection{ws: ws, Messages: messages, waitGroup: &waitGroup}
+
 	go func() {
 		defer func() {
 			waitGroup.Wait()
@@ -103,6 +107,10 @@ func NewWS(url string, opts Options) (*Connection, error) {
 
 			err = websocket.Message.Receive(ws, &msg)
 			if err != nil {
+				if wsInsp.isClosed.Load() {
+					return
+				}
+
 				if err.Error() == "EOF" {
 					color.New(color.FgRed).Println("Connection closed by the server")
 				} else {
@@ -116,7 +124,7 @@ func NewWS(url string, opts Options) (*Connection, error) {
 		}
 	}()
 
-	return &Connection{ws: ws, Messages: messages, waitGroup: &waitGroup}, nil
+	return wsInsp, nil
 }
 
 func (wsInsp *Connection) Send(msg string) (*Message, error) {
@@ -133,5 +141,11 @@ func (wsInsp *Connection) Send(msg string) (*Message, error) {
 }
 
 func (wsInsp *Connection) Close() {
+	if wsInsp.isClosed.Load() {
+		return
+	}
+
+	wsInsp.isClosed.Store(true)
+
 	wsInsp.ws.Close()
 }