fix: adding monitoring

v1/client/client.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
+	"sync"
 	"time"
 
 	customrwmutex "git.neurocipta.com/rogerferdinan/custom-rwmutex"
@@ -137,9 +138,11 @@ func (b *SafeWebsocketClientBuilder) Build(ctx context.Context) (*SafeWebsocketC
 		isConnected: false,
 		doneMap:     safemap.NewSafeMap[string, chan struct{}](),
 		writeChan:   make(chan Message, *b.writeChannelSize),
+		monitor:     internal.NewMemoryMonitor(),
 	}
 
 	go wsClient.reconnectHandler()
+	go wsClient.monitor.Run(ctx, internal.DefaultLogFn("client-monitor"))
 
 	if err := wsClient.connect(); err != nil {
 		return nil, fmt.Errorf("failed to establish initial connection: %v", err)
@@ -158,11 +161,12 @@ type SafeWebsocketClient struct {
 	path     *string
 	rawQuery *string
 
-	mu          *customrwmutex.CustomRwMutex
-	conn        *websocket.Conn
-	ctx         context.Context
-	cancelFuncs []context.CancelFunc
-	dataChannel chan []byte
+	mu              *customrwmutex.CustomRwMutex
+	conn            *websocket.Conn
+	ctx             context.Context
+	cancelFuncs     []context.CancelFunc
+	dataChannel     chan []byte
+	dataChannelOnce sync.Once
 
 	reconnectCh    chan struct{}
 	reconnectChans []chan struct{}
@@ -170,8 +174,10 @@ type SafeWebsocketClient struct {
 	doneMap        *safemap.SafeMap[string, chan struct{}]
 
 	writeChan chan Message
+	monitor   *internal.MemoryMonitor
 }
 
+
 func (wsClient *SafeWebsocketClient) connect() error {
 	var scheme string
 	if wsClient.useTLS {
@@ -190,7 +196,7 @@ func (wsClient *SafeWebsocketClient) connect() error {
 		newURL.RawQuery = *wsClient.rawQuery
 	}
 
-	header := make(http.Header)
+	header := make(http.Header, 0)
 	if wsClient.headers != nil {
 		for k, v := range *wsClient.headers {
 			header.Set(k, v)
@@ -226,9 +232,13 @@ func (wsClient *SafeWebsocketClient) connect() error {
 			return err
 		}
 
-		wsClient.writeChan <- Message{
+		select {
+		case wsClient.writeChan <- Message{
 			MessageType: MessageTypePong,
 			Data:        []byte(pingData),
+		}:
+		default:
+			log.Println("writeChan full, dropping pong")
 		}
 		return nil
 	})
@@ -291,14 +301,10 @@ func (wsClient *SafeWebsocketClient) reconnectHandler() {
 				for _, reconnectCh := range wsClient.reconnectChans {
 					select {
 					case reconnectCh <- struct{}{}:
-					default: // prevent blocking if chan is full
+					default:
 					}
 				}
-				if len(wsClient.reconnectChans) > 1 {
-					wsClient.reconnectChans = wsClient.reconnectChans[1:]
-				} else {
-					wsClient.reconnectChans = nil
-				}
+				wsClient.reconnectChans = nil
 			}
 		case <-wsClient.ctx.Done():
 			log.Println("reconnect handler stopped due to client shutdown")
@@ -334,8 +340,11 @@ func (wsClient *SafeWebsocketClient) writePump(ctx context.Context, c *websocket
 }
 
 func (wsClient *SafeWebsocketClient) readPump(ctx context.Context, c *websocket.Conn) {
+	canceled := false
 	defer func() {
-		wsClient.triggerReconnect()
+		if !canceled {
+			wsClient.triggerReconnect()
+		}
 		c.Close()
 	}()
 
@@ -343,6 +352,7 @@ func (wsClient *SafeWebsocketClient) readPump(ctx context.Context, c *websocket.
 		select {
 		case <-ctx.Done():
 			log.Println("Reader canceled by context")
+			canceled = true
 			return
 		default:
 			// Set read deadline
@@ -363,6 +373,7 @@ func (wsClient *SafeWebsocketClient) readPump(ctx context.Context, c *websocket.
 			select {
 			case wsClient.dataChannel <- data:
 			case <-ctx.Done():
+				canceled = true
 				return
 			default:
 				if wsClient.isDrop {
@@ -383,9 +394,14 @@ func (wsClient *SafeWebsocketClient) startPingTicker(ctx context.Context) {
 			log.Println("ping ticker canceled by context")
 			return
 		case <-ticker.C:
-			wsClient.writeChan <- Message{
+			select {
+			case wsClient.writeChan <- Message{
 				MessageType: websocket.PingMessage,
 				Data:        []byte{},
+			}:
+			case <-ctx.Done():
+				log.Println("ping ticker canceled by context")
+				return
 			}
 		}
 	}
@@ -413,7 +429,9 @@ func (wsClient *SafeWebsocketClient) DataChannel() <-chan []byte {
 }
 
 func (wsClient *SafeWebsocketClient) CloseDataChannel() {
-	close(wsClient.dataChannel)
+	wsClient.dataChannelOnce.Do(func() {
+		close(wsClient.dataChannel)
+	})
 }
 
 func (wsClient *SafeWebsocketClient) Write(data []byte) error {
@@ -425,7 +443,7 @@ func (wsClient *SafeWebsocketClient) Write(data []byte) error {
 }
 
 func (wsClient *SafeWebsocketClient) Close() error {
-	wsClient.mu.ReadHandler(func() error {
+	wsClient.mu.WriteHandler(func() error {
 		if wsClient.cancelFuncs != nil {
 			for _, cancel := range wsClient.cancelFuncs {
 				cancel()
@@ -443,7 +461,10 @@ func (wsClient *SafeWebsocketClient) Close() error {
 		wsClient.conn.Close()
 	}
 	wsClient.isConnected = false
-	close(wsClient.dataChannel)
+
+	wsClient.dataChannelOnce.Do(func() {
+		close(wsClient.dataChannel)
+	})
 
 	return nil
 }