fix: adding SafeWebSocketClient parameter

v1/client/client.go

@@ -65,12 +65,13 @@ func (sm *SafeMap[K, V]) Len() int {
 }
 
 type SafeWebsocketClientBuilder struct {
-	baseHost    *string `nil_checker:"required"`
-	basePort    *uint16 `nil_checker:"required"`
-	path        *string
-	rawQuery    *string
-	useTLS      *bool
-	channelSize *int64
+	baseHost       *string `nil_checker:"required"`
+	basePort       *uint16 `nil_checker:"required"`
+	path           *string
+	rawQuery       *string
+	useTLS         *bool
+	channelSize    *int64
+	authenticateFn func(*SafeWebsocketClient) error
 }
 
 func NewSafeWebsocketClientBuilder() *SafeWebsocketClientBuilder {
@@ -92,6 +93,11 @@ func (b *SafeWebsocketClientBuilder) UseTLS(useTLS bool) *SafeWebsocketClientBui
 	return b
 }
 
+func (b *SafeWebsocketClientBuilder) AuthenticateFn(authenticateFn func(*SafeWebsocketClient) error) *SafeWebsocketClientBuilder {
+	b.authenticateFn = authenticateFn
+	return b
+}
+
 func (b *SafeWebsocketClientBuilder) Path(path string) *SafeWebsocketClientBuilder {
 	b.path = &path
 	return b
@@ -136,6 +142,10 @@ func (b *SafeWebsocketClientBuilder) Build(ctx context.Context) (*SafeWebsocketC
 		doneMap:     NewSafeMap[string, chan struct{}](),
 	}
 
+	if b.authenticateFn != nil {
+		wsClient.authenticateFn = b.authenticateFn
+	}
+
 	go wsClient.reconnectHandler()
 
 	if err := wsClient.connect(); err != nil {
@@ -146,20 +156,21 @@ func (b *SafeWebsocketClientBuilder) Build(ctx context.Context) (*SafeWebsocketC
 }
 
 type SafeWebsocketClient struct {
-	baseHost    string
-	basePort    uint16
-	useTLS      bool
-	path        *string
-	rawQuery    *string
-	dataChannel chan []byte
-	mu          *custom_rwmutex.CustomRwMutex
-	conn        *websocket.Conn
-	cancelFuncs []context.CancelFunc
-	ctx         context.Context
-	reconnectCh chan struct{}
-	isConnected bool
-
-	doneMap *SafeMap[string, chan struct{}]
+	baseHost       string
+	basePort       uint16
+	useTLS         bool
+	path           *string
+	rawQuery       *string
+	dataChannel    chan []byte
+	mu             *custom_rwmutex.CustomRwMutex
+	conn           *websocket.Conn
+	cancelFuncs    []context.CancelFunc
+	ctx            context.Context
+	reconnectCh    chan struct{}
+	reconnectChans []chan struct{}
+	isConnected    bool
+	doneMap        *SafeMap[string, chan struct{}]
+	authenticateFn func(*SafeWebsocketClient) error
 }
 
 func (wsClient *SafeWebsocketClient) connect() error {
@@ -200,15 +211,19 @@ func (wsClient *SafeWebsocketClient) connect() error {
 		return nil
 	})
 
-	wsClient.mu.WriteHandler(func() error {
-		if wsClient.conn != nil {
-			wsClient.conn.Close()
+	if wsClient.conn != nil {
+		wsClient.conn.Close()
+	}
+	wsClient.conn = conn
+	wsClient.isConnected = true
+
+	if wsClient.authenticateFn != nil {
+		if err := wsClient.authenticateFn(wsClient); err != nil {
+			return err
 		}
+	}
 
-		wsClient.conn = conn
-
-		wsClient.isConnected = true
-
+	wsClient.mu.WriteHandler(func() error {
 		pingCtx, pingCancel := context.WithCancel(context.Background())
 		go wsClient.startPingTicker(pingCtx)
 		wsClient.cancelFuncs = append(wsClient.cancelFuncs, pingCancel)
@@ -269,7 +284,6 @@ func (wsClient *SafeWebsocketClient) startReceiveHandler(ctx context.Context) {
 				case wsClient.dataChannel <- message:
 				case <-ctx.Done():
 					log.Println("Reconnect handler stopped")
-					return nil
 				default:
 					log.Println("Data channel full, dropping message")
 				}
@@ -301,7 +315,6 @@ func (wsClient *SafeWebsocketClient) reconnectHandler() {
 				}
 			}
 			wsClient.isConnected = false
-
 			time.Sleep(100 * time.Millisecond)
 
 			isInnerLoop := true
@@ -324,7 +337,11 @@ func (wsClient *SafeWebsocketClient) reconnectHandler() {
 					return
 				}
 			}
-
+			if wsClient.reconnectChans != nil {
+				for _, reconnectCh := range wsClient.reconnectChans {
+					reconnectCh <- struct{}{}
+				}
+			}
 		case <-wsClient.ctx.Done():
 			log.Println("reconnect handler stopped due to client shutdown")
 			wsClient.Close()
@@ -333,10 +350,26 @@ func (wsClient *SafeWebsocketClient) reconnectHandler() {
 	}
 }
 
+func (wsClient *SafeWebsocketClient) ReconnectChannel() <-chan struct{} {
+	reconnectCh := make(chan struct{})
+	wsClient.mu.WriteHandler(func() error {
+		wsClient.reconnectChans = append(wsClient.reconnectChans, reconnectCh)
+		return nil
+	})
+
+	return reconnectCh
+}
+
 func (wsClient *SafeWebsocketClient) DataChannel() <-chan []byte {
 	return wsClient.dataChannel
 }
 
+func (wsClient *SafeWebsocketClient) WriteJSON(message any) error {
+	return wsClient.mu.WriteHandler(func() error {
+		return wsClient.conn.WriteJSON(message)
+	})
+}
+
 func (wsClient *SafeWebsocketClient) Close() error {
 	wsClient.mu.WriteHandler(func() error {
 		if wsClient.cancelFuncs != nil {
@@ -344,6 +377,12 @@ func (wsClient *SafeWebsocketClient) Close() error {
 				cancel()
 			}
 		}
+
+		if wsClient.reconnectChans != nil {
+			for _, reconnectChan := range wsClient.reconnectChans {
+				close(reconnectChan)
+			}
+		}
 		if wsClient.conn != nil {
 			wsClient.conn.Close()
 		}

v1/examples/client/main.go

@@ -34,6 +34,12 @@ func main() {
 		log.Fatal(err)
 	}
 
+	go func() {
+		for range wsClient.ReconnectChannel() {
+			fmt.Println("Reconnection Success")
+		}
+	}()
+
 	dataChannel := wsClient.DataChannel()
 
 	for data := range dataChannel {