fix: write & reader posiiton

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 {
@@ -235,7 +236,10 @@ func (wsClient *SafeWebsocketClient) connect() error {
 
 	pingCtx, pingCancel := context.WithCancel(context.Background())
 	go wsClient.startPingTicker(pingCtx)
+	wsClient.mu.WriteHandler(func() error {
 		wsClient.cancelFuncs = append(wsClient.cancelFuncs, pingCancel)
+		return nil
+	})
 
 	if wsClient.conn != nil {
 		wsClient.conn.Close()
@@ -251,7 +255,15 @@ func (wsClient *SafeWebsocketClient) connect() error {
 
 	go func() {
 		ctx, cancel := context.WithCancel(context.Background())
+		wsClient.mu.WriteHandler(func() error {
 			wsClient.cancelFuncs = append(wsClient.cancelFuncs, cancel)
+			return nil
+		})
+
+		writer, err := conn.NextWriter(websocket.TextMessage)
+		if err != nil {
+			return
+		}
 		for {
 			select {
 			case <-ctx.Done():
@@ -262,19 +274,6 @@ 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 {
 					return
@@ -285,7 +284,14 @@ func (wsClient *SafeWebsocketClient) connect() error {
 
 	go func() {
 		ctx, cancel := context.WithCancel(context.Background())
+		wsClient.mu.WriteHandler(func() error {
 			wsClient.cancelFuncs = append(wsClient.cancelFuncs, cancel)
+			return nil
+		})
+		_, reader, err := conn.NextReader()
+		if err != nil {
+			return
+		}
 		for {
 			select {
 			case <-ctx.Done():
@@ -299,12 +305,6 @@ func (wsClient *SafeWebsocketClient) connect() error {
 					fmt.Printf("error on read deadline: %v\n", err)
 					return
 				}
-				_, reader, err := conn.NextReader()
-				if err != nil {
-					fmt.Printf("Next Reader Closed: %v\n", err)
-					wsClient.triggerReconnect()
-					return
-				}
 
 				readerBytes, err := io.ReadAll(reader)
 				if err != nil {
@@ -395,11 +395,15 @@ func (wsClient *SafeWebsocketClient) reconnectHandler() {
 		case <-wsClient.reconnectCh:
 			log.Println("Reconnect triggered")
 
+			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)
 
@@ -461,11 +465,14 @@ func (wsClient *SafeWebsocketClient) Write(data []byte) error {
 }
 
 func (wsClient *SafeWebsocketClient) Close() error {
+	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 {