pkg/types: add BeforeConnect hook function

This commit is contained in:
Edwin 2023-08-10 15:08:31 +08:00
parent 0b03336fb0
commit 509f9ac8ca

View File

@ -44,6 +44,8 @@ type Dispatcher func(e interface{})
// HeartBeat keeps connection alive by sending the heartbeat packet. // HeartBeat keeps connection alive by sending the heartbeat packet.
type HeartBeat func(ctxConn context.Context, conn *websocket.Conn, cancelConn context.CancelFunc) type HeartBeat func(ctxConn context.Context, conn *websocket.Conn, cancelConn context.CancelFunc)
type BeforeConnect func(ctx context.Context) error
//go:generate callbackgen -type StandardStream -interface //go:generate callbackgen -type StandardStream -interface
type StandardStream struct { type StandardStream struct {
parser Parser parser Parser
@ -111,6 +113,8 @@ type StandardStream struct {
FuturesPositionSnapshotCallbacks []func(futuresPositions FuturesPositionMap) FuturesPositionSnapshotCallbacks []func(futuresPositions FuturesPositionMap)
heartBeat HeartBeat heartBeat HeartBeat
beforeConnect BeforeConnect
} }
type StandardStreamEmitter interface { type StandardStreamEmitter interface {
@ -306,6 +310,11 @@ func (s *StandardStream) Reconnect() {
// Connect starts the stream and create the websocket connection // Connect starts the stream and create the websocket connection
func (s *StandardStream) Connect(ctx context.Context) error { func (s *StandardStream) Connect(ctx context.Context) error {
if s.beforeConnect != nil {
if err := s.beforeConnect(ctx); err != nil {
return err
}
}
err := s.DialAndConnect(ctx) err := s.DialAndConnect(ctx)
if err != nil { if err != nil {
return err return err
@ -432,6 +441,11 @@ func (s *StandardStream) SetHeartBeat(fn HeartBeat) {
s.heartBeat = fn s.heartBeat = fn
} }
// SetBeforeConnect sets the custom hook function before connect
func (s *StandardStream) SetBeforeConnect(fn BeforeConnect) {
s.beforeConnect = fn
}
type Depth string type Depth string
const ( const (