lone-services/services/percutaneous/internal/logic/wsAddUserStreamLogic.go

package logic

import (
	"context"
	"lone-services/pkg/utils"
	user "lone-services/rpc/user/pb"
	"strconv"
	"time"

	"lone-services/rpc/percutaneous/pb"
	"lone-services/services/percutaneous/internal/svc"

	"github.com/zeromicro/go-zero/core/logx"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

type grpcStreamWrapper struct {
	stream percutaneous.Percutaneous_WsAddUserStreamServer
}

func (w *grpcStreamWrapper) Send(payload []byte) error {
	frame := &percutaneous.WsStreamRequest{
		RawBody: payload,
	}
	return w.stream.Send(frame)
}

type WsAddUserStreamLogic struct {
	ctx    context.Context
	svcCtx *svc.ServiceContext
	logx.Logger
}

type userInfo struct {
	Id     int64  `json:"id"`
	Name   string `json:"name"`
	Avatar string `json:"avatar"`
}

func NewWsAddUserStreamLogic(ctx context.Context, svcCtx *svc.ServiceContext) *WsAddUserStreamLogic {
	return &WsAddUserStreamLogic{
		ctx:    ctx,
		svcCtx: svcCtx,
		Logger: logx.WithContext(ctx),
	}
}

// 添加用户
func (l *WsAddUserStreamLogic) WsAddUserStream(stream percutaneous.Percutaneous_WsAddUserStreamServer) error {
	ctx := stream.Context()
	timeout := time.Duration(utils.GetConfigInt("ws.out_time")) * time.Second
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer func() {
		cancel()
	}()

	// BFF自动发来第一条消息，读取host_code
	recvFrame, err := stream.Recv()
	if err != nil {
		l.Logger.Errorf("recv host_code failed: %v", err)
		return err
	}
	var hostCode struct {
		Code string `json:"host_code"`
	}
	// 保留你确认的 !err 判断
	if ok := utils.JsonStringToStruct(string(recvFrame.RawBody), &hostCode); !ok {
		return status.Errorf(codes.InvalidArgument, "device code parse error")
	}

	msgCh := make(chan []byte, 1)
	wrapper := &grpcStreamWrapper{stream: stream}
	sess := &utils.WaitingSession{
		Stream: wrapper,
		Ctx:    ctx,
		Cancel: cancel,
		MsgCh:  msgCh,
	}
	l.svcCtx.SessionMgr.Register(hostCode.Code, sess)
	defer l.svcCtx.SessionMgr.Remove(hostCode.Code)

	select {
	case <-ctx.Done():
		if ctx.Err() == context.DeadlineExceeded {
			return status.Errorf(codes.DeadlineExceeded, "扫码等待超时")
		}
		return ctx.Err()
	case payload := <-msgCh:
		id, _ := strconv.Atoi(string(payload))
		info := l.getUserInfo(int64(id))
		jsonData := utils.StructToJson(info)
		if err := sess.Stream.Send([]byte(jsonData)); err != nil {
			l.Logger.Errorf("send to device err: %v", err)
			return err
		}
		return nil
	}
}

func (l *WsAddUserStreamLogic) getUserInfo(id int64) *userInfo {
	if id == utils.NumberZero {
		return nil
	}
	logx.Errorf("get user id: %v", id)
	cli, err := svc.GetRpcClient(l.svcCtx.UserName)
	if err != nil {
		logx.Errorf("get user err: %v", err)
		return nil
	}

	client := user.NewUserClient(cli.Conn())

	data := user.InternalGetUserInfoReq{
		UserId: id,
	}
	res, err := client.InternalGetUserInfo(context.Background(), &data)
	if err != nil {
		logx.Errorf("get user err: %v", err)
		return nil
	}

	return &userInfo{
		Id:     res.Id,
		Name:   res.Name,
		Avatar: res.Avatar,
	}
}