Compare commits

...

3 Commits

2 changed files with 68 additions and 30 deletions

View File

@@ -71,6 +71,7 @@ type SafeWebsocketClientBuilder struct {
rawQuery *string
useTLS *bool
channelSize *int64
authenticateFn func() error
}
func NewSafeWebsocketClientBuilder() *SafeWebsocketClientBuilder {
@@ -92,6 +93,11 @@ func (b *SafeWebsocketClientBuilder) UseTLS(useTLS bool) *SafeWebsocketClientBui
return b
}
func (b *SafeWebsocketClientBuilder) AuthenticateFn(authenticateFn func() 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 {
@@ -157,9 +167,10 @@ type SafeWebsocketClient struct {
cancelFuncs []context.CancelFunc
ctx context.Context
reconnectCh chan struct{}
reconnectChans []chan struct{}
isConnected bool
doneMap *SafeMap[string, chan struct{}]
authenticateFn func() error
}
func (wsClient *SafeWebsocketClient) connect() error {
@@ -200,15 +211,18 @@ func (wsClient *SafeWebsocketClient) connect() error {
return nil
})
wsClient.mu.WriteHandler(func() error {
if wsClient.conn != nil {
wsClient.conn.Close()
}
wsClient.conn = conn
wsClient.isConnected = true
if err := wsClient.authenticateFn(); err != nil {
return err
}
wsClient.mu.WriteHandler(func() error {
pingCtx, pingCancel := context.WithCancel(context.Background())
go wsClient.startPingTicker(pingCtx)
wsClient.cancelFuncs = append(wsClient.cancelFuncs, pingCancel)
@@ -269,7 +283,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 +314,6 @@ func (wsClient *SafeWebsocketClient) reconnectHandler() {
}
}
wsClient.isConnected = false
time.Sleep(100 * time.Millisecond)
isInnerLoop := true
@@ -324,7 +336,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,6 +349,16 @@ 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
}
@@ -350,6 +376,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 {