sig-pub/internal/exchange/exchange_grpc_server.go

package exchange

import (
	"context"
	"fmt"
	"io"
	"sig-pub/api/pb"
	"sig-pub/pkg/aside"
	"sig-pub/pkg/data/entity"
	"sig-pub/pkg/types"
	"sig-pub/pkg/zlog"
	"sync"
	"sync/atomic"

	"google.golang.org/grpc"
)

type Exchange struct {
	Type       pb.Exchange
	Subscriber ExchangeSubscriber
	Insts      map[string]*entity.TradeInstanceExchange
	sync.RWMutex
}

type ExchangeGrpcServer struct {
	pb.UnimplementedExchangeServiceServer
	exchangeMap         map[pb.Exchange]*Exchange
	tradeInstanceAside  *aside.TradeInstanceAside
	exchangeDataService *ExchangeDataService

	klineStreamId  int64
	klinePublisher *Publisher[int64, grpc.BidiStreamingServer[pb.ReqStreamSubscribeKline, pb.RspStreamSubscribeKline]]
}

// exchanges: 支持的数据源交易所
func NewExchangeGrpcServer(tradeInstanceAside *aside.TradeInstanceAside, exchangeDataService *ExchangeDataService, exchanges ...ExchangeSubscriber) *ExchangeGrpcServer {
	exchangeMap := make(map[pb.Exchange]*Exchange)
	for _, exchange := range exchanges {
		exchangeType := exchange.ExhcangeType()
		pbExchangeType, ok := exchangeType.Exchange2PB()
		if !ok {
			panic(fmt.Errorf("unknown exchange: %#v", exchangeType))
		}
		exchangeMap[pbExchangeType] = &Exchange{
			Type:       pbExchangeType,
			Subscriber: exchange,
			Insts:      make(map[string]*entity.TradeInstanceExchange),
		}
	}

	return &ExchangeGrpcServer{
		exchangeMap:         exchangeMap,
		tradeInstanceAside:  tradeInstanceAside,
		exchangeDataService: exchangeDataService,
		klinePublisher:      NewPublisher[int64, grpc.BidiStreamingServer[pb.ReqStreamSubscribeKline, pb.RspStreamSubscribeKline]](16),
	}
}

func (svc *ExchangeGrpcServer) Init() (err error) {
	svc.subscribeExchanges()
	return
}

// 订阅交易所推送行情
func (svc *ExchangeGrpcServer) subscribeExchanges() {
	// consumerKline
	// 交易所订阅交易产品
	for _, exchange := range svc.exchangeMap {
		go func(exchange *Exchange) {
			// get exchange trade instances
			insts, err := svc.tradeInstanceAside.ListExchangeTradeInstance(context.Background(), exchange.Type)
			if err != nil {
				zlog.Error(err)
				return
			}
			var exchangeInstIds []string
			exchange.Lock()
			for _, inst := range insts {
				exchangeInstIds = append(exchangeInstIds, inst.ExchangeInstId)
				exchange.Insts[inst.ExchangeInstId] = inst
			}
			exchange.Unlock()

			// instIds := []string{"BTC-USDT", "DOGE-USDT-SWAP"}
			err = exchange.Subscriber.SubscribeKline(exchangeInstIds...)
			if err != nil {
				zlog.Error(err)
				return
			}
			c := exchange.Subscriber.ConsumerKline()
			svc.consumerKline(exchange, c)
			// todo subscribe books 订单簿
			zlog.Infof("unsubscribe exchange: %s", exchange.Type.String())
		}(exchange)
	}
}

// consumerKline 消费交易所k线数据
func (svc *ExchangeGrpcServer) consumerKline(exchange *Exchange, c <-chan *types.ChannelKline) {
	for {
		channelK, ok := <-c
		if !ok {
			return
		}

		// 交易所 instid 转 sig-instid
		var exInst *entity.TradeInstanceExchange
		exchange.RLock()
		if inst, ok := exchange.Insts[channelK.InstId]; ok && inst != nil {
			exchange.RUnlock()
			exInst = inst
		} else {
			exchange.RUnlock()
			zlog.Errorf("unknown exchange instId: %v, %s", channelK.Exchange, channelK.InstId)
			continue
		}

		// tsdb storage
		typeInst := types.TradeInstance{
			InstId: exInst.InstId, // channelK.InstId
			TickSz: 0,
			MinSz:  0,
		}
		// todo 异步处理
		err := svc.exchangeDataService.SaveKlines(typeInst, channelK.Klines)
		if err != nil {
			zlog.Errorf("kline save to tsdb error: ", err)
		}

		// publish to subscribers
		pubMsgMap := make(map[string]*pb.StreamKline)

		// instId := channelK.InstId
		exchange, ok := channelK.Exchange.Exchange2PB()
		if !ok {
			zlog.Errorf("unknown exchange kline: %v", channelK.Exchange)
			continue
		}
		exchangeName := exchange.String()

		for _, kline := range channelK.Klines {
			// zlog.Infof("recv kline: %#v", kline)
			confirm := 0
			if kline.Confirm {
				confirm = 1
			}
			pubKey := fmt.Sprintf("/kline/%s/%s/%s/%d", exchangeName, exInst.InstId, kline.Interval, confirm)
			// todo 优化没有订阅者就跳过
			msg, ok := pubMsgMap[pubKey]
			if !ok {
				msg = new(pb.StreamKline)
				msg.InstId = exInst.InstId
				msg.Exchange = exchange
				pubMsgMap[pubKey] = msg
			}

			pbk := kline.ToPBKline()
			msg.Klines = append(msg.Klines, pbk)
		}

		for pubKey, msg := range pubMsgMap {
			if len(msg.Klines) == 0 {
				continue
			}
			subs := svc.klinePublisher.Publisher(pubKey)
			for _, sub := range subs {
				if err := sub.Send(&pb.RspStreamSubscribeKline{Kline: msg}); err != nil {
					zlog.Error(err)
				}
			}
		}
	}
}

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

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

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

			if msg.SubType == pb.SubscribeType_UnsubscribeAll {
				svc.klinePublisher.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.Infof("stream: %d sub: %s", streamId, subKey)
							switch msg.SubType {
							case pb.SubscribeType_Subscribe:
								svc.klinePublisher.Subscribe(subKey, streamId, stream)
							case pb.SubscribeType_Unsubscribe:
								svc.klinePublisher.Unsubscribe(subKey, streamId)
							}
						}
					}
				}
			}
		}
	}

	// for {
	// 	select {
	// 	case <-ctx.Done():
	// 		// 客户端断开连接
	// 		// log.Printf("Client disconnected from topic: %s", topic)
	// 		return
	// 	default:
	// 		// 模拟事件生成
	// 		event := &pb.SubscribeKlineStreamRsp{
	// 			Kline: &pb.StreamKline{
	// 				InstId:   "DOGE/USDT",
	// 				Exchange: pb.Exchange_OKX,
	// 				Klines: []*pb.Kline{
	// 					{
	// 						Ts: time.Now().Unix(),
	// 					},
	// 				},
	// 			},
	// 		}

	// 		// 推送事件
	// 		if err := streamRsp.Send(event); err != nil {
	// 			log.Printf("Failed to send event to client: %v", err)
	// 			return err
	// 		}

	// 		// 模拟事件间隔
	// 		time.Sleep(2 * time.Second)
	// 	}
	// }
}