sig-pub/internal/exchange/exchange_grpc_server.go

package exchange

import (
	"context"
	"fmt"
	"io"
	"math"
	"sig-pub/api/pb"
	"sig-pub/pkg/publish"
	"sig-pub/pkg/types"
	"sig-pub/pkg/types/decimals"
	"sig-pub/pkg/zlog"
	"sync/atomic"

	"google.golang.org/grpc"
)

type ExchangeGrpcServer struct {
	pb.UnimplementedExchangeServiceServer
	exchangeService *ExchangeService
	klineStreamId   int64

	klineSubscriber *publish.Publisher[int64, grpc.BidiStreamingServer[pb.ReqStreamSubscribeKline, pb.RspStreamSubscribeKline]]
}

// exchanges: 支持的数据源交易所
func NewExchangeGrpcServer(exchangeService *ExchangeService) *ExchangeGrpcServer {
	return &ExchangeGrpcServer{
		exchangeService: exchangeService,
	}
}

func (svr *ExchangeGrpcServer) Init() (err error) {
	svr.klineSubscriber = svr.exchangeService.GetKlineSubscriber()
	return
}

// SubscribeKline 订阅k线stream
func (svr *ExchangeGrpcServer) SubscribeKline(stream grpc.BidiStreamingServer[pb.ReqStreamSubscribeKline, pb.RspStreamSubscribeKline]) (err0 error) {
	streamId := atomic.AddInt64(&svr.klineStreamId, 1)
	// subKey = /kline/exchange/instId/interval/confirm -> /kline/OKX/DOGE-USDT-SWAP/1s/1

	// 接收消息的goroutine
	recvChan := make(chan *pb.ReqStreamSubscribeKline)
	go func() {
		for {
			msg, err := stream.Recv()
			if err == io.EOF {
				zlog.Infof("recv close EOF")
				close(recvChan)
				return
			}
			if err != nil {
				zlog.Infof("recv error: %v", err)
				close(recvChan)
				return
			}
			zlog.Infof("recv stream msg: %#v", msg)
			recvChan <- msg
		}
	}()

	// 发送和处理消息
	for {
		select {
		case <-stream.Context().Done():
			// 客户端断开连接
			svr.klineSubscriber.UnsubscribeAll(streamId)
			return stream.Context().Err()
		case msg, ok := <-recvChan:
			if !ok {
				// 接收通道关闭，结束流
				svr.klineSubscriber.UnsubscribeAll(streamId)
				return
			}

			if msg.SubType == pb.SubscribeType_UNSUB_ALL {
				svr.klineSubscriber.UnsubscribeAll(streamId)
				continue
			}
			for _, exchange := range msg.Exchanges {
				for _, instId := range msg.InstIds {
					for _, interval := range msg.Intervals {
						confirms := []int{1}
						if !msg.OnlyConfirm {
							confirms = append(confirms, 0)
						}
						for _, confirm := range confirms {
							subKey := fmt.Sprintf("/kline/%s/%s/%s/%d", exchange.String(), instId, interval, confirm)
							zlog.Debugf("stream: id=%d, sub %s", streamId, subKey)
							switch msg.SubType {
							case pb.SubscribeType_SUB:
								svr.klineSubscriber.Subscribe(subKey, streamId, stream)
							case pb.SubscribeType_UNSUB:
								svr.klineSubscriber.Unsubscribe(subKey, streamId)
							}
						}
					}
				}
			}
		}
	}
}

// Exchanges 获取支持的交易所列表
func (svr *ExchangeGrpcServer) Exchanges(ctx context.Context, req *pb.ReqExchanges) (rsp *pb.RspExchanges, err error) {
	exchanges, err := svr.exchangeService.Exchanges()
	if err != nil {
		return
	}
	return &pb.RspExchanges{Exchanges: exchanges}, nil
}

// ExchangeInstanceState 获取交易产品系统状态
func (svr *ExchangeGrpcServer) ExchangeInstanceState(ctx context.Context, req *pb.ReqExchangeInstanceState) (rsp *pb.RspExchangeInstanceState, err error) {
	states, err := svr.exchangeService.ExchangeInstanceState(req)
	if err != nil {
		return
	}
	return &pb.RspExchangeInstanceState{InstsState: states}, nil
}

// 查询SeriesRange时间范围
func (svr *ExchangeGrpcServer) SeriesRange(ctx context.Context, req *pb.ReqSeriesRange) (rsp *pb.RspSeriesRange, err error) {
	rsp = new(pb.RspSeriesRange)
	rsp.After, rsp.Before, rsp.Total, err = svr.exchangeService.CalcSeriesRange(req.Series)
	return
}

// HistoryKline 获取交易产品历史k线 (before < klines... < after)
func (svr *ExchangeGrpcServer) HistoryKline(ctx context.Context, req *pb.ReqHistoryKline) (rsp *pb.RspHistoryKline, err error) {
	if req.Series == nil {
		err = fmt.Errorf("series arg is required")
		return
	}
	// 查询范围数据条数检查
	_, _, total, err := svr.exchangeService.CalcSeriesRange(req.Series)
	if err != nil {
		return
	}
	if total > MaxHistoryKlines {
		err = fmt.Errorf("time range too large max %d", MaxHistoryKlines)
		return
	}

	rsp = new(pb.RspHistoryKline)
	rsp = &pb.RspHistoryKline{
		Exchange: req.Series.Exchange,
		InstId:   req.Series.InstId,
		Interval: req.Series.Interval,
	}

	kSeries := &pb.KlineSeries{
		Interval: req.Series.Interval,
		Time:     make([]int64, 0, total),
		Open:     make([]float64, 0, total),
		High:     make([]float64, 0, total),
		Low:      make([]float64, 0, total),
		Close:    make([]float64, 0, total),
		Vol:      make([]float64, 0, total),
		VolQuote: make([]float64, 0, total),
	}
	live, err := svr.exchangeService.HistoryKline(req.Series, math.MaxInt, func(liveK *types.Kline, klines []*types.Kline) error {
		for _, kline := range klines {
			kSeries.Time = append(kSeries.Time, kline.Ts)
			kSeries.Open = append(kSeries.Open, decimals.MustToFloat64(kline.Open))
			kSeries.High = append(kSeries.High, decimals.MustToFloat64(kline.High))
			kSeries.Low = append(kSeries.Low, decimals.MustToFloat64(kline.Low))
			kSeries.Close = append(kSeries.Close, decimals.MustToFloat64(kline.Close))
			kSeries.Vol = append(kSeries.Vol, decimals.MustToFloat64(kline.Vol))
			kSeries.VolQuote = append(kSeries.VolQuote, decimals.MustToFloat64(kline.VolQuote))
			if liveK != nil {
				kSeries.Live = liveK.ToPBKline()
			}
		}
		return nil
	})
	if err != nil {
		return
	}
	rsp.Live = live
	rsp.Klines = kSeries
	return
}

// HistoryKlineStream 获取交易产品历史k线(流式返回)
func (svr *ExchangeGrpcServer) HistoryKlineStream(req *pb.ReqHistoryKlineStream, stream grpc.ServerStreamingServer[pb.RspHistoryKlineStream]) (err error) {
	if req.Series == nil {
		err = fmt.Errorf("series arg is required")
		return
	}
	_, err = svr.exchangeService.HistoryKline(req.Series, 200, func(_ *types.Kline, klines []*types.Kline) error {
		rsp := &pb.RspHistoryKlineStream{}
		rsp.Klines = make([]*pb.Kline, 0, len(klines))
		for _, kline := range klines {
			rsp.Klines = append(rsp.Klines, kline.ToPBKline())
		}
		return stream.Send(rsp)
	})
	if err != nil {
		return
	}
	return
}