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 { type SafeWebsocketClient struct {
baseHost string baseHost string
basePort uint16 basePort uint16
useTLS bool useTLS bool
path *string path *string
rawQuery *string rawQuery *string
dataChannel chan []byte dataChannel chan []byte
mu *custom_rwmutex.CustomRwMutex mu *custom_rwmutex.CustomRwMutex
conn *websocket.Conn conn *websocket.Conn
cancelFuncs []context.CancelFunc cancelFuncs []context.CancelFunc
ctx context.Context ctx context.Context
reconnectCh chan struct{} reconnectCh chan struct{}
isConnected bool reconnectChans []chan struct{}
isConnected bool
doneMap *SafeMap[string, chan struct{}] doneMap *SafeMap[string, chan struct{}]
} }
func (wsClient *SafeWebsocketClient) connect() error { func (wsClient *SafeWebsocketClient) connect() error {
@@ -253,30 +253,28 @@ func (wsClient *SafeWebsocketClient) startReceiveHandler(ctx context.Context) {
log.Println("receive handler stopped") log.Println("receive handler stopped")
return return
default: default:
if err := wsClient.mu.ReadHandler(func() error { conn := wsClient.conn
conn := wsClient.conn
if conn == nil { if conn == nil {
wsClient.triggerReconnect() wsClient.triggerReconnect()
return fmt.Errorf("connection closed") return
} // return fmt.Errorf("connection closed")
_, message, err := conn.ReadMessage() }
if err != nil { _, message, err := conn.ReadMessage()
wsClient.triggerReconnect() if err != nil {
return fmt.Errorf("failed to read message: %v", err) 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 {
return 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 wsClient.isConnected = false
time.Sleep(100 * time.Millisecond) time.Sleep(100 * time.Millisecond)
isInnerLoop := true isInnerLoop := true
@@ -324,7 +321,11 @@ func (wsClient *SafeWebsocketClient) reconnectHandler() {
return return
} }
} }
if wsClient.reconnectChans != nil {
for _, reconnectCh := range wsClient.reconnectChans {
reconnectCh <- struct{}{}
}
}
case <-wsClient.ctx.Done(): case <-wsClient.ctx.Done():
log.Println("reconnect handler stopped due to client shutdown") log.Println("reconnect handler stopped due to client shutdown")
wsClient.Close() 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 { func (wsClient *SafeWebsocketClient) DataChannel() <-chan []byte {
return wsClient.dataChannel 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 { func (wsClient *SafeWebsocketClient) Close() error {
wsClient.mu.WriteHandler(func() error { wsClient.mu.WriteHandler(func() error {
if wsClient.cancelFuncs != nil { if wsClient.cancelFuncs != nil {
@@ -344,6 +361,12 @@ func (wsClient *SafeWebsocketClient) Close() error {
cancel() cancel()
} }
} }
if wsClient.reconnectChans != nil {
for _, reconnectChan := range wsClient.reconnectChans {
close(reconnectChan)
}
}
if wsClient.conn != nil { if wsClient.conn != nil {
wsClient.conn.Close() wsClient.conn.Close()
} }

View File

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