127 lines
3.3 KiB
Go
127 lines
3.3 KiB
Go
package grpcclient
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"git.hlsq.asia/alex/proto/proto/rs/rspb"
|
|
"git.hlsq.asia/alex/service-common/discover"
|
|
"git.hlsq.asia/alex/service-common/discover/common"
|
|
"git.hlsq.asia/alex/service-common/log"
|
|
"git.hlsq.asia/alex/service-common/utils/locator"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/protobuf/proto"
|
|
"sync"
|
|
)
|
|
|
|
type GatewayFun int
|
|
|
|
const (
|
|
FunToClient GatewayFun = iota
|
|
FunBroadcastClient
|
|
FunBroadcastChannel
|
|
)
|
|
|
|
var gatewayServer sync.Map // map[string]*gatewayStream
|
|
|
|
type gatewayStream struct {
|
|
mu sync.Mutex
|
|
stream grpc.ClientStream
|
|
}
|
|
|
|
func findGatewayBySID(sid int64, fun GatewayFun) (*gatewayStream, error) {
|
|
key := gatewayKey(sid, fun)
|
|
|
|
if v, ok := gatewayServer.Load(key); ok {
|
|
return v.(*gatewayStream), nil
|
|
}
|
|
|
|
client, err := GatewayNewClient(sid)
|
|
if err != nil {
|
|
log.Errorf("findGatewayBySID cannot find client: %v", err)
|
|
return nil, err
|
|
}
|
|
var stream grpc.ClientStream
|
|
switch fun {
|
|
case FunToClient:
|
|
stream, err = client.InnerToClient(context.Background())
|
|
case FunBroadcastClient:
|
|
stream, err = client.InnerBroadcastClient(context.Background())
|
|
case FunBroadcastChannel:
|
|
stream, err = client.InnerBroadcastChannel(context.Background())
|
|
}
|
|
if err != nil {
|
|
log.Errorf("findGatewayBySID %v err: %v, sid: %v", fun, err, sid)
|
|
return nil, err
|
|
}
|
|
|
|
ss := &gatewayStream{stream: stream}
|
|
if actual, loaded := gatewayServer.LoadOrStore(key, ss); loaded {
|
|
go func() { _ = stream.CloseSend() }()
|
|
return actual.(*gatewayStream), nil
|
|
}
|
|
|
|
return ss, nil
|
|
}
|
|
|
|
func sendMessageToGateway(sid int64, fun GatewayFun, msg proto.Message, re ...bool) error {
|
|
ss, err := findGatewayBySID(sid, fun)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
ss.mu.Lock()
|
|
err = ss.stream.SendMsg(msg)
|
|
ss.mu.Unlock()
|
|
|
|
if err != nil {
|
|
key := gatewayKey(sid, fun)
|
|
if v, ok := gatewayServer.Load(key); ok && v == ss {
|
|
gatewayServer.Delete(key)
|
|
_ = ss.stream.CloseSend()
|
|
}
|
|
// 如果没有标识本次是重试的，就重试一次（默认重试）
|
|
if re == nil || !re[0] {
|
|
return sendMessageToGateway(sid, fun, msg, true)
|
|
}
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// SendMessageToClient 单播消息给客户端
|
|
func SendMessageToClient(ctx context.Context, usn int64, msg proto.Message, platform ...rspb.PlatformType) error {
|
|
if platform == nil || len(platform) == 0 {
|
|
return nil
|
|
}
|
|
for _, pf := range platform {
|
|
if gatewaySID := locator.NewLocator(ctx, usn, int32(pf)).LocateGateway(); gatewaySID > 0 {
|
|
return sendMessageToGateway(gatewaySID, FunToClient, msg)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// BroadcastMessageToClient 广播消息给客户端
|
|
func BroadcastMessageToClient(msg *rspb.InnerBroadcastClientReq) error {
|
|
return broadcastMessageToGateway(FunBroadcastClient, msg)
|
|
}
|
|
|
|
// BroadcastMessageToChannel 广播消息到频道（加入频道的能收到，不区分平台）
|
|
func BroadcastMessageToChannel(msg *rspb.InnerBroadcastChannelReq) error {
|
|
return broadcastMessageToGateway(FunBroadcastChannel, msg)
|
|
}
|
|
|
|
func broadcastMessageToGateway(fun GatewayFun, msg proto.Message) error {
|
|
sidArr := discover.FindServerAllSid(common.KeyDiscoverGateway)
|
|
for _, sid := range sidArr {
|
|
if err := sendMessageToGateway(sid, fun, msg); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func gatewayKey(sid int64, fun GatewayFun) string {
|
|
return fmt.Sprintf("%v-%v", sid, fun)
|
|
}
|