sig-pub/internal/trading/kline_store.go

package trading

import (
	"context"
	"io"
	"math"
	"sig-pub/api/pb"
	"sig-pub/pkg/data"
	"sig-pub/pkg/mq"
	"sig-pub/pkg/types"
	"sig-pub/pkg/utils/collect"
	"sig-pub/pkg/utils/retry"
	"sig-pub/pkg/zlog"
	"time"

	"google.golang.org/grpc"
)

type KlineStore struct {
	exchangeClient pb.ExchangeServiceClient

	store             *types.ExchangeState[*collect.ConcurrentMap[string, *TradeInstanceKlineSeries]]  // K线列表: []exchange<instId, interval, klines>
	subKlineIntervals []string                                                                         // 订阅的k线的周期列表
	subKlineInsts     *types.ExchangeState[*collect.SyncMap[string, bool]]                             // 订阅k线中的交易产品列表
	subKlineStream    grpc.BidiStreamingClient[pb.ReqStreamSubscribeKline, pb.RspStreamSubscribeKline] // 订阅k线的stream
}

func NewKlineSeriesStore(exchangeClient pb.ExchangeServiceClient) (kss *KlineStore) {
	kss = &KlineStore{exchangeClient: exchangeClient}

	// 周期列表
	kss.subKlineIntervals = collect.Map2Slice(types.SupportedIntervals, func(interval types.Interval, _ types.IntervalAdder) string {
		return string(interval)
	})
	// 产品列表
	kss.subKlineInsts = types.NewExchangeStateInit(func() *collect.SyncMap[string, bool] {
		return collect.NewSyncMap[string, bool]()
	})
	// 各交易所 store 初始化
	kss.store = types.NewExchangeStateInit(func() *collect.ConcurrentMap[string, *TradeInstanceKlineSeries] {
		return collect.NewConcurrentMap[string, *TradeInstanceKlineSeries](64, func(s string) string {
			return s
		})
	})
	return
}

func (s *KlineStore) Init() (err error) {
	// 连接 exchange kline stream
	go s.connectSubscribeKline(false)

	// 订阅交易产品初始化完成事件
	mq.NatsCreateConsumer("trading", mq.StreamExchange, mq.TopicExchangeTradeInstanceInited, func() *mq.PublishExchangeTradeInstanceInited { return new(mq.PublishExchangeTradeInstanceInited) },
		func(msg *mq.PublishExchangeTradeInstanceInited) (err error) {
			// 初始化k线, 开始订阅k线
			zlog.Infof("subscribed TopicExchangeTradeInstanceInited: %#v", msg)
			go s.inititalKlineSeries(msg.Exchange, msg.InstId)
			return
		})

	go func() {
		//  拉取已初始化完成交易产品, 初始化k线, 开始订阅k线
		rsp, _ := retry.DoWithFixDelay(math.MaxInt32, time.Second, func(retryTimes uint32) (rsp *pb.RspExchangeInstanceState, err error) {
			rsp, err = s.exchangeClient.ExchangeInstanceState(context.Background(), &pb.ReqExchangeInstanceState{
				AllExchange: true,
				AllInsts:    true,
				Status:      []int32{int32(data.StatusOk)},
			})
			if err != nil {
				zlog.Errorf("ExchangeInstanceState error: retry=%d, %v", retryTimes, err)
			}
			return
		})
		for _, inst := range rsp.InstsState {
			s.inititalKlineSeries(inst.Exchange, inst.InstId)
		}
	}()
	return
}

// connectSubscribeKline 连接exchange订阅实时k线
func (s *KlineStore) connectSubscribeKline(reconnect bool) {
	defer func() {
		if s.subKlineStream != nil {
			s.subKlineStream.CloseSend()
			s.subKlineStream = nil
		}
		go s.connectSubscribeKline(true)
	}()

	if reconnect {
		zlog.Infof("subscribeKlines will reconnect after 5s")
		time.Sleep(5 * time.Second)
	}

	stream, err := s.exchangeClient.SubscribeKline(context.Background())
	if err != nil {
		zlog.Error("subscribeKlines reqeust error: ", err)
		return
	}
	s.subKlineStream = stream

	// 发送所有交易产品订阅消息
	go func() {
		s.subKlineInsts.Range(func(exchange pb.ExchangeType, m *collect.SyncMap[string, bool]) {
			// todo 分批订阅
			var instIds []string
			m.Range(func(instId string, _ bool) bool {
				instIds = append(instIds, instId)
				return true
			})
			s.sendSubscribeKline(false, exchange, instIds...)
		})
	}()

	// 接收订阅k线消息
	for {
		msg, err := stream.Recv()
		if err == io.EOF {
			zlog.Debugf("subscribeKlines connection server closeed")
			return
		}
		if err != nil {
			zlog.Error("subscribeKlines recv error: ", err)
			return
		}

		for _, k := range msg.Kline.Klines {
			kline := new(types.Kline)
			kline.ParsePBKline(msg.Kline.Exchange, k)
			if kms, ok := kline.Interval.AddMul(kline.Ts, 1); ok {
				delay := time.Now().UnixMilli() - kms
				zlog.Debugf("recv kline: streamId=%d, delay=%dms, inst=%s(%v), interval=%s, close=%s(%v)", msg.Kline.StreamId, delay, msg.Kline.InstId, msg.Kline.Exchange, kline.Interval, kline.Close.String(), kline.Confirm)
			}
			// kline klineStore -> klineSeries -> strategy -> indicator -> klineSeries.Series
			s.Update(msg.Kline.Exchange, msg.Kline.InstId, kline)
		}
	}
}

// subscribeKline 发送订阅消息
func (s *KlineStore) sendSubscribeKline(save bool, exchange pb.ExchangeType, instIds ...string) {
	if len(instIds) == 0 {
		return
	}
	// 交易产品订阅记录
	if save {
		for _, instId := range instIds {
			s.subKlineInsts.Get(exchange).Store(instId, true)
		}
	}

	// 发送订阅消息
	subMsg := &pb.ReqStreamSubscribeKline{
		SubType:     pb.SubscribeType_Subscribe,
		Exchanges:   []pb.ExchangeType{exchange},
		InstIds:     instIds,
		Intervals:   s.subKlineIntervals,
		OnlyConfirm: true,
	}
	doSend := func(retry uint32) (_ int, err error) {
		if s.subKlineStream == nil {
			return
		}
		zlog.Debugf("sending stream subscribe kline: retry=%d, instId=%s%v", retry, exchange, instIds)
		if err = s.subKlineStream.Send(subMsg); err != nil {
			zlog.Errorf("send stream subscribe kline msg error: %#v", subMsg, err)
			return
		}
		return
	}
	if _, err := doSend(0); err == nil {
		return
	}

	go retry.DoWithFixDelay(math.MaxInt32, time.Second, doSend)
}

func (s *KlineStore) inititalKlineSeries(exchange pb.ExchangeType, instId string) {
	if !s.store.IsSupport(exchange) {
		zlog.Errorf("unsupport exchange %s", exchange)
		return
	}
	storeInst := s.store.Get(exchange).ComputeIfAbsent(instId, func(k string) *TradeInstanceKlineSeries {
		return NewTradeInstanceKlineSeries(exchange, k)
	})
	// 交易产品已初始化过
	if !storeInst.Status.CompareAndSwap(int32(data.StatusNone), int32(data.StatusProcessing)) {
		return
	}

	zlog.Infof("initial kline series starting: %s(%s),", instId, exchange)
	// 初始化最新的 klineSeries
	for _, interval := range s.subKlineIntervals {
		retry.DoWithFixDelay(math.MaxInt32, 2*time.Second, func(retryTimes uint32) (_ struct{}, err error) {
			_, err = s.fetchHistoryKlineToSeries(exchange, instId, interval, 0, 0, MaxSeriesKlines)
			return
		})
	}
	// 初始化历史k线完成, 开始订阅k线
	storeInst.Status.Store(int32(data.StatusOk))
	s.sendSubscribeKline(true, exchange, instId)
	zlog.Infof("initial kline series success: %s(%s)", instId, exchange)
}

// fetchHistoryKlineToSeries 拉去历史k线数据更新series
func (s *KlineStore) fetchHistoryKlineToSeries(exchange pb.ExchangeType, instId, interval string, before, after int64, count uint32) (total int, err error) {
	// 拉取最新的1000条k线
	req := &pb.ReqHistoryKlineStream{
		Exchange: exchange,
		InstId:   instId,
		Interval: interval,
		Count:    count,
		Before:   before,
		After:    after,
	}
	stream, err := s.exchangeClient.HistoryKlineStream(context.Background(), req, grpc.UseCompressor("snappy"))
	if err != nil {
		zlog.Errorf("fetch history kline stream error: instId=%s(%s), interval=%s, %#v, err=%v", instId, exchange, interval, req, err)
		return
	}

	for {
		msg, err0 := stream.Recv()
		if err0 == io.EOF {
			// zlog.Debugf("fetch kline stream connection server closed")
			break
		}
		if err0 != nil {
			err = err0
			zlog.Error("fetch kline stream recv error: ", err0)
			return
		}
		total += len(msg.Klines)
		// zlog.Debugf("recv: %s(%s), %s, branch=%d, ts=%d~%d", instId, exchange, interval, len(msg.Klines), msg.Klines[0].Ts, msg.Klines[len(msg.Klines)-1].Ts)
		for _, k := range msg.Klines {
			kline := new(types.Kline)
			kline.ParsePBKline(exchange, k)
			s.Update(exchange, instId, kline)
		}
	}
	return
}

// Update
// kline klineStore -> klineSeries -> strategy -> indicator -> klineSeries.Series
func (s *KlineStore) Update(exchange pb.ExchangeType, instId string, kline *types.Kline) {
	if _, ok := types.SupportedIntervals[kline.Interval]; !ok {
		zlog.Warningf("unsupport interval: %s", kline.Interval)
		return
	}
	if !s.store.IsSupport(exchange) {
		zlog.Warningf("unsupport exchange: %v", exchange)
		return
	}

	instSeries := s.store.Get(exchange).ComputeIfAbsent(instId, func(k string) *TradeInstanceKlineSeries {
		return NewTradeInstanceKlineSeries(exchange, k)
	})

	before, serial := instSeries.IntervalKlines.Get(kline.Interval).Update(kline)
	if !serial {
		inprocessing := instSeries.Status.CompareAndSwap(int32(data.StatusOk), int32(data.StatusProcessing))
		if !inprocessing {
			return
		}
		// 拉取缺失的k线
		func() {
			defer instSeries.Status.Store(int32(data.StatusOk))
			after := kline.Ts
			total, err := s.fetchHistoryKlineToSeries(exchange, instId, string(kline.Interval), before, after, 0)
			if err != nil {
				zlog.Error("fetch padding kline series error: instId=%s(%s), interval=%s, ts=%d~%d, err=%v", instId, exchange, kline.Interval, before, after, err)
				return
			} else {
				zlog.Debugf("fetched padding kline series: instId=%s(%s), interval=%s, total=%d, ts=%d~%d", instId, exchange, kline.Interval, total, before, after)
			}
		}()
	}

	// 发布k线时间驱动策略执行
	if instSeries.Status.Load() != int32(data.StatusOk) {
		return
	}
	if nts, ok := kline.Interval.AddMul(kline.Ts, 2); ok {
		// k线已过期则不执行策略
		if nts < time.Now().UnixMilli() {
			return
		}
	}
	// 判断同一时刻k线
	var intervals []types.Interval
	endTs := kline.Interval.MustAddMul(kline.Ts, 1)
	instSeries.IntervalKlines.Range(func(interval types.Interval, v *KlineSeries) {
		if endTs == interval.MustAddMul(v.LastTs(), 1) {
			intervals = append(intervals, interval)
		}
	})
	zlog.Debugf("confirm kline intervals: instId=%s(%s), interval=%s, ts=%d, %v", instId, exchange, kline.Interval, kline.Ts, intervals)
	// todo emit kline update, calc indicator...
	// pubKey := fmt.Sprintf("/kline/%s/%s/%s/%d", exchangeType, tradeInst.InstId, kline.Interval, confirm)
	// interval/okx/BTC_USDT/1m,3m,5m
}