fix: read issues

v1/client/client.go

@@ -156,6 +156,7 @@ func (b *SafeWebsocketClientBuilder) Build(ctx context.Context) (*SafeWebsocketC
 		reconnectCh: make(chan struct{}, 1),
 		isConnected: false,
 		doneMap:     NewSafeMap[string, chan struct{}](),
+		writeChan:   make(chan Message),
 	}
 
 	if b.authenticateFn != nil {
@@ -258,6 +259,11 @@ func (wsClient *SafeWebsocketClient) connect() error {
 			wsClient.cancelFuncs = append(wsClient.cancelFuncs, cancel)
 			return nil
 		})
+
+		writer, err := conn.NextWriter(websocket.TextMessage)
+		if err != nil {
+			return
+		}
 		for {
 			select {
 			case <-ctx.Done():
@@ -268,21 +274,9 @@ func (wsClient *SafeWebsocketClient) connect() error {
 					wsClient.triggerReconnect()
 					return
 				}
-				messageType := websocket.TextMessage
-				switch data.MessageType {
-				case MessageTypePing:
-					messageType = websocket.PingMessage
-				case MessageTypePong:
-					messageType = websocket.PongMessage
-				case MessageTypeClose:
-					messageType = websocket.CloseMessage
-				}
-				writer, err := conn.NextWriter(messageType)
-				if err != nil {
-					return
-				}
 
 				if _, err := writer.Write(data.Data); err != nil {
+					wsClient.triggerReconnect()
 					return
 				}
 			}
@@ -295,6 +289,7 @@ func (wsClient *SafeWebsocketClient) connect() error {
 			wsClient.cancelFuncs = append(wsClient.cancelFuncs, cancel)
 			return nil
 		})
+
 		for {
 			select {
 			case <-ctx.Done():
@@ -308,13 +303,17 @@ func (wsClient *SafeWebsocketClient) connect() error {
 					fmt.Printf("error on read deadline: %v\n", err)
 					return
 				}
-				_, reader, err := conn.NextReader()
+
+				mt, reader, err := conn.NextReader()
 				if err != nil {
-					fmt.Printf("Next Reader Closed: %v\n", err)
 					wsClient.triggerReconnect()
 					return
 				}
 
+				if mt != websocket.TextMessage {
+					continue
+				}
+
 				readerBytes, err := io.ReadAll(reader)
 				if err != nil {
 					fmt.Printf("io reader failed: %v\n", err)
@@ -355,40 +354,6 @@ func (wsClient *SafeWebsocketClient) startPingTicker(ctx context.Context) {
 	}
 }
 
-// func (wsClient *SafeWebsocketClient) startReceiveHandler(ctx context.Context) {
-// 	for {
-// 		select {
-// 		case <-ctx.Done():
-// 			log.Println("receive handler stopped")
-// 			return
-// 		default:
-// 			if err := wsClient.mu.ReadHandler(func() error {
-// 				conn := wsClient.conn
-
-// 				if conn == nil {
-// 					wsClient.triggerReconnect()
-// 					return fmt.Errorf("connection closed")
-// 				}
-// 				_, message, err := conn.ReadMessage()
-// 				if err != nil {
-// 					wsClient.triggerReconnect()
-// 					return fmt.Errorf("failed to read message: %v", err)
-// 				}
-// 				select {
-// 				case wsClient.dataChannel <- message:
-// 				case <-ctx.Done():
-// 					log.Println("Reconnect handler stopped")
-// 				default:
-// 					log.Println("")
-// 				}
-// 				return nil
-// 			}); err != nil {
-// 				return
-// 			}
-// 		}
-// 	}
-// }
-
 func (wsClient *SafeWebsocketClient) triggerReconnect() {
 	select {
 	case wsClient.reconnectCh <- struct{}{}:
@@ -404,11 +369,15 @@ func (wsClient *SafeWebsocketClient) reconnectHandler() {
 		case <-wsClient.reconnectCh:
 			log.Println("Reconnect triggered")
 
-			if wsClient.cancelFuncs != nil {
-				for _, cancel := range wsClient.cancelFuncs {
-					cancel()
+			wsClient.mu.ReadHandler(func() error {
+				if wsClient.cancelFuncs != nil {
+					for _, cancel := range wsClient.cancelFuncs {
+						cancel()
+					}
 				}
-			}
+				return nil
+			})
+
 			wsClient.isConnected = false
 			time.Sleep(100 * time.Millisecond)
 
@@ -470,11 +439,14 @@ func (wsClient *SafeWebsocketClient) Write(data []byte) error {
 }
 
 func (wsClient *SafeWebsocketClient) Close() error {
-	if wsClient.cancelFuncs != nil {
-		for _, cancel := range wsClient.cancelFuncs {
-			cancel()
+	wsClient.mu.ReadHandler(func() error {
+		if wsClient.cancelFuncs != nil {
+			for _, cancel := range wsClient.cancelFuncs {
+				cancel()
+			}
 		}
-	}
+		return nil
+	})
 
 	if wsClient.reconnectChans != nil {
 		for _, reconnectChan := range wsClient.reconnectChans {

v1/examples/client/main.go

@@ -29,6 +29,7 @@ func main() {
 		BasePort(8080).
 		Path("/ws/test/data_1").
 		UseTLS(false).
+		ChannelSize(30).
 		Build(ctx)
 	if err != nil {
 		log.Fatal(err)