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

package okx

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"sig-pub/pkg/utils/collect"
	"sig-pub/pkg/zlog"
	"sync"
	"time"

	"github.com/bytedance/sonic"
	"github.com/gorilla/websocket"
	"golang.org/x/net/http/httpproxy"
)

// https://my.okx.com/docs-v5/zh/#overview-websocket-overview
// 连接限制：3 次/秒 (基于IP)
// 当订阅公有频道时，使用公有服务的地址；当订阅私有频道时，使用私有服务的地址
// 请求限制：每个连接 对于 订阅/取消订阅/登录 请求的总次数限制为 480 次/小时

// 公共频道无需登录，包括行情频道，K线频道，交易数据频道，资金费率频道，限价范围频道，深度数据频道，标记价格频道等。
// 私有频道需登录，包括用户账户频道，用户交易频道，用户持仓频道等。

const (
	WsBaseUrl = "wss://ws.okx.com:8443"
	// WsBaseUrl = "wss://wseeapap.okx.com:8443" // 模拟盘
)

// SubscribeStatus 订阅状态
type SubscribeStatus int32

const (
	StatusUnsubscribe SubscribeStatus = 0 // 未订阅
	StatusSubscribing SubscribeStatus = 1 // 订阅中
	StatusSubscribed  SubscribeStatus = 2 // 已订阅
)

// WebSocketEvent okx推送事件
type WebSocketEvent struct {
	// 订阅数据
	Event  ChannelEventType `json:"event"`
	Code   string           `json:"code,omitempty"`
	Msg    string           `json:"msg,omitempty"`
	ConnId string           `json:"connId,omitempty"`
	// 推送数据
	Arg    WebSocketEventArg `json:"arg,omitempty"`
	Action ActionType        `json:"action"` // books交易深度订阅,推送数据动作类型
	Data   json.RawMessage   `json:"data"`
}

type WebSocketEventArg struct {
	Channel string `json:"channel"` // 订阅的频道 candle1s/candle1m/candle3M
	InstId  string `json:"instId"`  // 产品ID
}

// 订阅配置
type wsChannelConfig[T, R any] struct {
	channelId         string                           // ws日志id
	httpProxy         string                           // http 代理
	wsUrl             string                           // 服务地址, K线-> /ws/v5/business
	subscribeChannels []string                         // 订阅频道列表
	dataInstanceFunc  func() T                         // 生成一个data实例用以反序列化
	dataMappingFunc   func(*ChannelData[T]) (R, error) // 数据映射函数
}

// okx websocket 频道订阅
// 行情订阅 todo 单连接上限, 上百种币多个ws连接订阅
type wsChannel[T, R any] struct {
	sync.Mutex
	channelId      string
	cfg            wsChannelConfig[T, R]
	conn           *websocket.Conn
	subscribeInsts *collect.ConcurrentMap[string, SubscribeStatus] // <instId, status>
	dataC          chan R                                          // data channel

	prevLogFullMs int64
	// todo subscribe logs...
}

func newWsChannel[T, R any](cfg wsChannelConfig[T, R]) *wsChannel[T, R] {
	subInsts := collect.NewConcurrentMap[string, SubscribeStatus](8, func(k string) string { return k })
	ms := &wsChannel[T, R]{
		channelId:      cfg.channelId,
		cfg:            cfg,
		subscribeInsts: subInsts,
		dataC:          make(chan R, 4*1024), // 4k
	}
	return ms
}

func (c *wsChannel[T, R]) Init() (err error) {
	if c.cfg.dataInstanceFunc == nil {
		return errors.New("ws channel payloadInstanceFunc can't be nil")
	}
	err = c.connect()
	return
}

func (c *wsChannel[T, R]) connect() (err error) {
	var proxyFunc = http.ProxyFromEnvironment
	if c.cfg.httpProxy != "" {
		// 自定义 http proxy
		proxyFunc = func(req *http.Request) (*url.URL, error) {
			return (&httpproxy.Config{
				HTTPProxy:  c.cfg.httpProxy,
				HTTPSProxy: c.cfg.httpProxy,
			}).ProxyFunc()(req.URL)
		}
	}
	wsDialer := websocket.Dialer{
		Proxy:            proxyFunc,
		HandshakeTimeout: 10 * time.Second,
	}

	url := fmt.Sprintf("%s%s", WsBaseUrl, c.cfg.wsUrl)
	conn, _, err := wsDialer.Dial(url, nil)
	if err != nil {
		zlog.Error("Error connecting to websocket channel:", url, err)
		go c.reconnect()
		return
	}
	c.conn = conn
	zlog.Infof("websocket channel %s connected: %s", c.channelId, url)

	go c.readPump()
	return
}

func (c *wsChannel[T, R]) reconnect() {
	zlog.Infof("channel %s reconnecting", c.channelId)

	var err error
	func() {
		c.Lock()
		defer c.Unlock()

		c.close()
		// wait 1sec
		<-time.After(time.Second)
		err = c.connect()
	}()
	if err != nil {
		return
	}

	// 重新订阅
	var insts []string
	c.subscribeInsts.Range(func(instId string, status SubscribeStatus) bool {
		if status == StatusUnsubscribe {
			insts = append(insts, instId)
		}
		return true
	})

	if len(insts) > 0 {
		err = c.Subscribe(insts...)
		if err != nil {
			zlog.Error("channel %s subscribe %v error", c.channelId, insts, err)
			go c.reconnect()
		}
	}
}

func (c *wsChannel[T, R]) close() {
	if c.conn != nil {
		err := c.conn.Close()
		if err != nil {
			zlog.Infof("channel %s close error: %v", c.channelId, err)
		}
		c.conn = nil
	}

	// 重置到待订阅状态
	c.subscribeInsts.RangeUpdate(func(instId string, status SubscribeStatus) (bool, bool, SubscribeStatus) {
		return true, false, StatusUnsubscribe
	})
}

// 从ws读取数据并解析
func (c *wsChannel[T, R]) readPump() {
	for {
		t, bytes, err := c.conn.ReadMessage()
		if err != nil {
			go c.reconnect()
			zlog.Errorf("channel %s reading message error: %v", c.channelId, err)
			return
		}
		switch t {
		default:
			zlog.Errorf("channel %s unknown message type: %d", c.channelId, t)
			continue
		case websocket.PingMessage:
			c.conn.WriteMessage(websocket.PongMessage, []byte{})
			continue
		case websocket.CloseMessage:
			go c.reconnect()
			zlog.Infof("channel %s close message", c.channelId)
			return
		case websocket.TextMessage:
		}

		// fmt.Println(string(bytes))
		var event WebSocketEvent
		if err = sonic.Unmarshal(bytes, &event); err != nil {
			zlog.Errorf("parse websocket event error: ", err)
			continue
		}

		switch event.Event {
		case ChannelEventTypeSubscribe:
			c.subscribeInsts.Store(event.Arg.InstId, StatusSubscribed)
			zlog.Infof("channel %s subscribed %s %s", c.channelId, event.Arg.InstId, event.Arg.Channel)
			continue
		case ChannelEventTypeUnsubscribe:
			c.subscribeInsts.Delete(event.Arg.InstId)
			zlog.Infof("channel %s unsubscribed %s", c.channelId, event.Arg.InstId)
			continue
		case ChannelEventTypeError:
			zlog.Errorf("channel %s event error: %#v", c.channelId, event)
			continue
		case ChannelEventTypeConnectionInfo: // 新链接订阅频道时, 消息同步链接数量
			zlog.Infof("channel-conn-count event %s: %#v", c.channelId, event)
			continue
		case ChannelEventTypeConnectionError: // 当超出限制时
			zlog.Errorf("channel-conn-count-error event %s: %#v", c.channelId, event)
			continue
		case ChannelEventTypeNotice: // websocket 服务升级断线通知
			zlog.Errorf("channel event type notice %s: %#v", c.channelId, event)
			continue
		// case WsEventTypeLogin: // todo
		default:
			// 推送数据
		}
		if event.Event != "" {
			zlog.Infof("unhandle event %s: %#v", c.channelId, event)
			continue
		}

		// 解析数据
		data := c.cfg.dataInstanceFunc()
		if err := sonic.Unmarshal(event.Data, data); err != nil {
			zlog.Error("unmarshal event data error: ", err)
			continue
		}
		// zlog.Infof("read data: %v", data)

		// 数据打包
		channelData := &ChannelData[T]{
			Channel: event.Arg.Channel,
			InstId:  event.Arg.InstId,
			Action:  event.Action,
			Data:    data,
		}

		// 数据映射
		dataR, err := c.cfg.dataMappingFunc(channelData)
		if err != nil {
			zlog.Error("okx channel data mapping error: ", err)
			continue
		}

		select {
		case c.dataC <- dataR:
		default:
			// channel 满了
			now := time.Now().UnixMilli()
			if now-c.prevLogFullMs > 5000 { // 每5秒打印不要太频繁
				c.prevLogFullMs = now
				zlog.Warningf("DataC %s full, sub %d insts", c.channelId, c.subscribeInsts.Size())
			}
		}
	}
}