pkg/exchange: add custom heart beat func to StandardStream

This commit is contained in:
Edwin 2023-08-02 17:47:18 +08:00
parent fe0a630fa8
commit 5064615df8

View File

@ -41,6 +41,9 @@ type Parser func(message []byte) (interface{}, error)
type Dispatcher func(e interface{}) type Dispatcher func(e interface{})
// HeartBeat keeps connection alive by sending the heartbeat packet.
type HeartBeat func(ctxConn context.Context, conn *websocket.Conn, cancelConn context.CancelFunc)
//go:generate callbackgen -type StandardStream -interface //go:generate callbackgen -type StandardStream -interface
type StandardStream struct { type StandardStream struct {
parser Parser parser Parser
@ -106,6 +109,8 @@ type StandardStream struct {
FuturesPositionUpdateCallbacks []func(futuresPositions FuturesPositionMap) FuturesPositionUpdateCallbacks []func(futuresPositions FuturesPositionMap)
FuturesPositionSnapshotCallbacks []func(futuresPositions FuturesPositionMap) FuturesPositionSnapshotCallbacks []func(futuresPositions FuturesPositionMap)
heartBeat HeartBeat
} }
type StandardStreamEmitter interface { type StandardStreamEmitter interface {
@ -350,6 +355,9 @@ func (s *StandardStream) DialAndConnect(ctx context.Context) error {
go s.Read(connCtx, conn, connCancel) go s.Read(connCtx, conn, connCancel)
go s.ping(connCtx, conn, connCancel, pingInterval) go s.ping(connCtx, conn, connCancel, pingInterval)
if s.heartBeat != nil {
go s.heartBeat(connCtx, conn, connCancel)
}
return nil return nil
} }
@ -419,6 +427,11 @@ func (s *StandardStream) Close() error {
return nil return nil
} }
// SetHeartBeat sets the custom heart beat implementation if needed
func (s *StandardStream) SetHeartBeat(fn HeartBeat) {
s.heartBeat = fn
}
type Depth string type Depth string
const ( const (