Compare commits

...

3 Commits

2 changed files with 66 additions and 37 deletions

View File

@@ -146,20 +146,20 @@ 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{}]
}
func (wsClient *SafeWebsocketClient) connect() error {
@@ -253,30 +253,28 @@ func (wsClient *SafeWebsocketClient) startReceiveHandler(ctx context.Context) {
log.Println("receive handler stopped")
return
default:
if err := wsClient.mu.ReadHandler(func() error {
conn := wsClient.conn
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")
return nil
default:
log.Println("Data channel full, dropping message")
}
return nil
}); err != nil {
if conn == nil {
wsClient.triggerReconnect()
return
// return fmt.Errorf("connection closed")
}
_, message, err := conn.ReadMessage()
if err != nil {
wsClient.triggerReconnect()
// return fmt.Errorf("failed to read message: %v", err)
return
}
select {
case wsClient.dataChannel <- message:
case <-ctx.Done():
log.Println("Reconnect handler stopped")
return
default:
log.Println("Data channel full, dropping message")
}
return
}
}
}
@@ -301,7 +299,6 @@ func (wsClient *SafeWebsocketClient) reconnectHandler() {
}
}
wsClient.isConnected = false
time.Sleep(100 * time.Millisecond)
isInnerLoop := true
@@ -324,7 +321,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 +334,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 +361,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()
}

View File

@@ -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 {