feat: adding channelSize to safe websocket client

v1/client/client.go

@@ -101,7 +101,7 @@ func (b *SafeWebsocketClientBuilder) RawQuery(rawQuery string) *SafeWebsocketCli
 	return b
 }
 
-func (b *SafeWebsocketClientBuilder) Build() (*SafeWebsocketClient, error) {
+func (b *SafeWebsocketClientBuilder) Build(channelSize int) (*SafeWebsocketClient, error) {
 	if err := internal.NilChecker(b); err != nil {
 		return nil, err
 	}
@@ -117,7 +117,7 @@ func (b *SafeWebsocketClientBuilder) Build() (*SafeWebsocketClient, error) {
 		useTLS:      useTLS,
 		path:        b.path,
 		rawQuery:    b.rawQuery,
-		dataChannel: make(chan []byte, 1),
+		dataChannel: make(chan []byte, channelSize),
 		mu:          custom_rwmutex.NewCustomRwMutex(),
 		reconnectCh: make(chan struct{}, 1),
 		isConnected: false,
@@ -275,6 +275,9 @@ func (wsClient *SafeWebsocketClient) reconnectHandler() {
 	backoff := 1 * time.Second
 	maxBackoff := 30 * time.Second
 	for {
+		if wsClient.ctx == nil {
+			continue
+		}
 		select {
 		case <-wsClient.reconnectCh:
 			log.Println("Reconnect triggered")
@@ -288,12 +291,13 @@ func (wsClient *SafeWebsocketClient) reconnectHandler() {
 
 			time.Sleep(100 * time.Millisecond)
 
-			for {
-				log.Println("Attempting reconnect in %v...", backoff)
+			isInnerLoop := true
+			for isInnerLoop {
+				log.Printf("Attempting reconnect in %v...", backoff)
 				select {
 				case <-time.After(backoff):
 					if err := wsClient.connect(); err != nil {
-						log.Println("Reconnect failed: %v", err)
+						log.Printf("Reconnect failed: %v", err)
 						if backoff < maxBackoff {
 							backoff *= 2
 						}
@@ -301,7 +305,10 @@ func (wsClient *SafeWebsocketClient) reconnectHandler() {
 					}
 					log.Println("Reconnected successfully")
 					backoff = 1 * time.Second
-					break
+					isInnerLoop = false
+					continue
+				case <-wsClient.ctx.Done():
+					return
 				}
 			}
 		case <-wsClient.ctx.Done():
@@ -314,3 +321,19 @@ func (wsClient *SafeWebsocketClient) reconnectHandler() {
 func (wsClient *SafeWebsocketClient) DataChannel() <-chan []byte {
 	return wsClient.dataChannel
 }
+
+func (wsClient *SafeWebsocketClient) Close() error {
+	wsClient.mu.WriteHandler(func() error {
+		if wsClient.cancel != nil {
+			wsClient.cancel()
+		}
+		if wsClient.conn != nil {
+			wsClient.conn.Close()
+		}
+		wsClient.isConnected = false
+		return nil
+	})
+	close(wsClient.dataChannel)
+
+	return nil
+}