sig-pub/internal/trading/kline_store.go

package trading

import (
	"context"
	"fmt"
	"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
	subscribeKlineIntervals []string
	store                   [3]*collect.ConcurrentMap[string, *KlineStoreInstance] // K线列表: []exchange<instId, interval, klines>
}

func NewKlineSeriesStore(exchangeClient pb.ExchangeServiceClient) (kss *KlineStore) {
	// 订阅实时k线周期列表
	subscribeKlineIntervals := collect.Map2Slice(types.SupportedIntervals, func(interval types.Interval, _ types.IntervalAdder) string {
		return string(interval)
	})
	kss = &KlineStore{
		subscribeKlineIntervals: subscribeKlineIntervals,
		exchangeClient:          exchangeClient,
	}
	kss.store[pb.ExchangeType_OKX] = collect.NewConcurrentMap[string, *KlineStoreInstance](64, func(s string) string { return s })
	// kss.klines[pb.ExchangeType_BINANCE] =
	return
}

func (s *KlineStore) Init() (err error) {
	// 拉取已初始化完成交易产品, 初始化k线, 开始订阅k线

	// 订阅交易产品初始化完成事件
	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.subscribeKlines(msg.Exchange, msg.InstId)
			return
		})
	return
}

func (s *KlineStore) subscribeKlines(exchange pb.ExchangeType, instId string) {
	storeInst := s.store[exchange].ComputeIfAbsent(instId, func(k string) *KlineStoreInstance {
		return NewKlineStoreInstance(exchange, k)
	})

	// 初始化最新的 klineSeries
	for _, interval := range s.subscribeKlineIntervals {
		for {
			after, before := int64(0), 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:    after,
					Before:   before,
					Live:     false,
					Asc:      true,
				}, 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 {
				zlog.Errorf("trade instance initial failed: %s(%s), %v", instId, exchange, err)
				return
			}
			for _, kline := range rsp.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)
				after = k.Ts
			}

		}
	}

	// 开始订阅k线
}

// Update
// kline klineStore -> klineSeries -> strategy -> indicator -> klineSeries.Series
func (s *KlineStore) Update(exchange pb.ExchangeType, instId string, kline *types.Kline) {
	storeInst := s.store[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
}