sig-pub/internal/trading/kline_store.go

package trading

import (
	"context"
	"fmt"
	"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"
	"sync"
	"sync/atomic"
	"time"

	"google.golang.org/grpc"
)

type KlineStore struct {
	exchangeClient pb.ExchangeServiceClient

	store             *types.ExchangeState[*collect.ConcurrentMap[string, *KlineStoreInstance]]        // 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.NewExchangeState0(func() *collect.SyncMap[string, bool] {
		return collect.NewSyncMap[string, bool]()
	})
	// 各交易所 store 初始化
	kss.store = types.NewExchangeState0(func() *collect.ConcurrentMap[string, *KlineStoreInstance] {
		return collect.NewConcurrentMap[string, *KlineStoreInstance](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.initKlineSeries(msg.Exchange, msg.InstId)
			return
		})

	// todo 拉取已初始化完成交易产品, 初始化k线, 开始订阅k线
	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)
			zlog.Debugf("recv: streamId=%d, %v, %s, %#v", msg.Kline.StreamId, msg.Kline.Exchange, msg.Kline.InstId, kline)
			// 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("send stream subscribe kline msg: retry=%d, %#v", retry, subMsg)
		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) initKlineSeries(exchange pb.ExchangeType, instId string) {
	if !s.store.IsSupport(exchange) {
		return
	}

	storeInst := s.store.Get(exchange).ComputeIfAbsent(instId, func(k string) *KlineStoreInstance {
		return NewKlineStoreInstance(exchange, k)
	})

	// 初始化最新的 klineSeries
	for _, interval := range s.subKlineIntervals {
		for {
			before := int64(0)
			rsp, err := retry.DoWithFixDelay(math.MaxInt32, 2*time.Second, func(retryTimes uint32) (rsp *pb.RspHistoryKline, err error) {
				rsp, err = s.exchangeClient.HistoryKline(context.Background(), &pb.ReqHistoryKline{
					Exchange: exchange,
					InstId:   instId,
					Interval: string(interval),
					Count:    MaxSeriesKlines,
					After:    0,
					Before:   before,
					Live:     false,
					Asc:      true,
				}, grpc.UseCompressor("snappy"))
				if err != nil {
					zlog.Errorf("fetch missing klines error: instId=%s(%s) before=%d, err=%v", instId, exchange, before, err)
				}
				return
			})
			if err != nil {
				zlog.Errorf("trade instance initial failed: %s(%s), %v", instId, exchange, err)
				return
			}
			klines := rsp.Klines
			if len(klines) == 0 {
				break
			}
			for _, kline := range klines {
				k := new(types.Kline)
				k.ParsePBKline(exchange, kline)
				_, _, err = storeInst.Update(k)
				if err != nil {
					zlog.Error("update initial kline series error: instId=%s(%s)", instId, exchange, err)
					return
				}

				s.Update(exchange, instId, k)
			}

			before = klines[len(klines)-1].Ts
			if !rsp.Next {
				break
			}
		}
	}
	// 初始化历史k线完成, 开始订阅k线
	storeInst.Status.Store(int32(data.StatusOk))
	s.sendSubscribeKline(true, exchange, instId)
}

// Update
// kline klineStore -> klineSeries -> strategy -> indicator -> klineSeries.Series
func (s *KlineStore) Update(exchange pb.ExchangeType, instId string, kline *types.Kline) {
	storeInst := s.store.Get(exchange).ComputeIfAbsent(instId, func(k string) *KlineStoreInstance {
		return NewKlineStoreInstance(exchange, k)
	})

	// 只处理已初始化完成的交易产品k线
	if storeInst.Status.Load() != int32(data.StatusOk) {
		return
	}
	before, serial, err := storeInst.Update(kline)
	if err != nil {
		zlog.Error("update kline series error: instId=%s(%s)", instId, exchange, err)
		return
	}
	if !serial {
		// 拉取缺失的k线
		s.paddingMissKlines(storeInst, exchange, instId, kline.Interval, kline.Ts, before)
		// emit kline event
	}
}

// 拉取缺失的k线
func (s *KlineStore) paddingMissKlines(storeInst *KlineStoreInstance, exchange pb.ExchangeType, instId string, interval types.Interval, after, before int64) {
	// 拉取缺失的k线
	rsp, err := retry.DoWithFixDelay(math.MaxInt32, 2*time.Second, func(retryTimes uint32) (rsp *pb.RspHistoryKline, err error) {
		rsp, err = s.exchangeClient.HistoryKline(context.Background(), &pb.ReqHistoryKline{
			Exchange: exchange,
			InstId:   instId,
			Interval: string(interval),
			After:    after,
			Before:   before,
			Live:     false,
		}, grpc.UseCompressor("snappy"))
		if err != nil {
			zlog.Errorf("fetch missing klines error: instId=%s(%s) after=%d before=%d, err=%v", instId, exchange, after, before, err)
		}
		return
	})
	if err != nil {
		return
	}

	collect.Reverse(rsp.Klines)
	for _, kline := range rsp.Klines {
		k := new(types.Kline)
		k.ParsePBKline(exchange, kline)
		zlog.Infof("padding missing kline: instId=%s(%s) %s %#v", instId, exchange, interval, k)
		_, _, err := storeInst.Update(k)
		if err != nil {
			zlog.Error("update missing kline series error: instId=%s(%s)", instId, exchange, err)
		}
	}
}

// KlineStoreInstance 单个交易产品所有周期k线
type KlineStoreInstance struct {
	sync.RWMutex
	Exchange       pb.ExchangeType
	InstId         string
	intervalKlines *types.IntervalState[*KlineSeries]
	Status         atomic.Int32 // 交易产品状态
}

func NewKlineStoreInstance(exchange pb.ExchangeType, instId string) *KlineStoreInstance {
	si := &KlineStoreInstance{
		Exchange:       exchange,
		InstId:         instId,
		intervalKlines: types.NewIntervalState[*KlineSeries](),
	}
	si.Status.Store(int32(data.StatusProcessing))

	for interval := range types.SupportedIntervals {
		si.intervalKlines.Set(interval, NewKlineSeries(exchange, instId, interval))
	}
	return si
}

// Update 更新k线
// serial k线是否连续
func (si *KlineStoreInstance) Update(kline *types.Kline) (before int64, serial bool, err error) {
	if _, ok := types.SupportedIntervals[kline.Interval]; !ok {
		err = fmt.Errorf("unsupport interval: %s", kline.Interval)
		return
	}

	si.Lock()
	defer si.Unlock()

	ks := si.intervalKlines.Get(kline.Interval)
	before, serial = ks.Update(kline)
	if !serial {
		// si.status.Store(int32(data.StatusProcessing))
	}
	return
}