Compare commits

...

2 Commits

View File

@@ -65,12 +65,13 @@ func (sm *SafeMap[K, V]) Len() int {
} }
type SafeWebsocketClientBuilder struct { type SafeWebsocketClientBuilder struct {
baseHost *string `nil_checker:"required"` baseHost *string `nil_checker:"required"`
basePort *uint16 `nil_checker:"required"` basePort *uint16 `nil_checker:"required"`
path *string path *string
rawQuery *string rawQuery *string
useTLS *bool useTLS *bool
channelSize *int64 channelSize *int64
authenticateFn func() error
} }
func NewSafeWebsocketClientBuilder() *SafeWebsocketClientBuilder { func NewSafeWebsocketClientBuilder() *SafeWebsocketClientBuilder {
@@ -92,6 +93,11 @@ func (b *SafeWebsocketClientBuilder) UseTLS(useTLS bool) *SafeWebsocketClientBui
return b return b
} }
func (b *SafeWebsocketClientBuilder) AuthenticateFn(authenticateFn func() error) *SafeWebsocketClientBuilder {
b.authenticateFn = authenticateFn
return b
}
func (b *SafeWebsocketClientBuilder) Path(path string) *SafeWebsocketClientBuilder { func (b *SafeWebsocketClientBuilder) Path(path string) *SafeWebsocketClientBuilder {
b.path = &path b.path = &path
return b return b
@@ -136,6 +142,10 @@ func (b *SafeWebsocketClientBuilder) Build(ctx context.Context) (*SafeWebsocketC
doneMap: NewSafeMap[string, chan struct{}](), doneMap: NewSafeMap[string, chan struct{}](),
} }
if b.authenticateFn != nil {
wsClient.authenticateFn = b.authenticateFn
}
go wsClient.reconnectHandler() go wsClient.reconnectHandler()
if err := wsClient.connect(); err != nil { if err := wsClient.connect(); err != nil {
@@ -160,6 +170,7 @@ type SafeWebsocketClient struct {
reconnectChans []chan struct{} reconnectChans []chan struct{}
isConnected bool isConnected bool
doneMap *SafeMap[string, chan struct{}] doneMap *SafeMap[string, chan struct{}]
authenticateFn func() error
} }
func (wsClient *SafeWebsocketClient) connect() error { func (wsClient *SafeWebsocketClient) connect() error {
@@ -200,14 +211,17 @@ func (wsClient *SafeWebsocketClient) connect() error {
return nil return nil
}) })
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 { wsClient.mu.WriteHandler(func() error {
if wsClient.conn != nil {
wsClient.conn.Close()
}
wsClient.conn = conn
wsClient.isConnected = true
pingCtx, pingCancel := context.WithCancel(context.Background()) pingCtx, pingCancel := context.WithCancel(context.Background())
go wsClient.startPingTicker(pingCtx) go wsClient.startPingTicker(pingCtx)
@@ -269,7 +283,6 @@ func (wsClient *SafeWebsocketClient) startReceiveHandler(ctx context.Context) {
case wsClient.dataChannel <- message: case wsClient.dataChannel <- message:
case <-ctx.Done(): case <-ctx.Done():
log.Println("Reconnect handler stopped") log.Println("Reconnect handler stopped")
return nil
default: default:
log.Println("Data channel full, dropping message") log.Println("Data channel full, dropping message")
} }