sig-pub/internal/exchange/okx/ws_channel_subscribe.go

package okx

import (
	"sig-pub/pkg/zlog"

	"github.com/bytedance/sonic"
	"github.com/gorilla/websocket"
)

func (c *wsChannel[T, R]) Consumer() <-chan R {
	return c.dataC
}

func (c *wsChannel[T, R]) Subscribe(instIds ...string) (err error) {
	c.Lock()
	defer c.Unlock()

	var args []WebSocketEventArg
	for _, instId := range instIds {
		if status, ok := c.subscribeInsts.Load(instId); !ok || status == StatusUnsubscribe {
			for _, channel := range c.cfg.subscribeChannels {
				args = append(args, WebSocketEventArg{
					Channel: channel,
					InstId:  instId,
				})
			}
			c.subscribeInsts.Store(instId, StatusSubscribing) // 订阅中
		}
	}
	if len(args) == 0 {
		return
	}

	// 发送订阅消息
	subscribeMsg := map[string]any{
		"op":   "subscribe",
		"args": args,
	}
	msg, err := sonic.Marshal(subscribeMsg)
	if err != nil {
		return
	}
	err = c.conn.WriteMessage(websocket.TextMessage, msg)
	if err != nil {
		zlog.Errorf("channel %s write subscribe msg error: %v", c.channelId, err)
		return
	}
	return
}

// Unsubscribe todo 无订阅时关闭 channel
func (c *wsChannel[T, R]) Unsubscribe(instIds ...string) (err error) {
	c.Lock()
	defer c.Unlock()

	var args []WebSocketEventArg
	for _, instId := range instIds {
		if status, ok := c.subscribeInsts.Load(instId); ok && status != StatusUnsubscribe {
			for _, channel := range c.cfg.subscribeChannels {
				args = append(args, WebSocketEventArg{
					Channel: channel,
					InstId:  instId,
				})
			}
			c.subscribeInsts.Store(instId, StatusUnsubscribe) // todo 取消订阅中
		}
	}
	if len(args) == 0 {
		return
	}

	// send subscribe message
	unsubscribeMsg := map[string]any{
		"op":   ChannelEventTypeUnsubscribe,
		"args": args,
	}
	msg, err := sonic.Marshal(unsubscribeMsg)
	if err != nil {
		return
	}
	err = c.conn.WriteMessage(websocket.TextMessage, msg)
	if err != nil {
		zlog.Error("Error write connecting to server:", err)
		return
	}
	return
}

// 产品是否已订阅
func (c *wsChannel[T, R]) IsSubscribed(instId string) bool {
	status, ok := c.subscribeInsts.Load(instId)
	return ok && (status == StatusSubscribed)
}

// 已订阅产品列表
func (c *wsChannel[T, R]) GetSubscribes() (instIds []string) {
	c.subscribeInsts.Range(func(instId string, status SubscribeStatus) bool {
		if status == StatusSubscribed {
			instIds = append(instIds, instId)
		}
		return true
	})
	return
}