sonet/internal/gateway_ws/server/ws_server.go

package server

import (
	"context"
	"github.com/gorilla/websocket"
	"runtime/debug"
	"sonet/internal/gateway_ws/session"
	"sonet/pkg/grpc/generic"
	"sonet/pkg/protocol"
	"sonet/pkg/utils/logger"
)

type ConnHandler struct {
	grpcFactory *generic.GrpcGenericClientFactory
}

func NewConnHandler(grpcFactory *generic.GrpcGenericClientFactory) *ConnHandler {
	return &ConnHandler{
		grpcFactory: grpcFactory,
	}
}

func (c *ConnHandler) handleConn(conn *websocket.Conn) {
	client := session.NewNetClient(conn, ReadDeadline, WriteDeadline)

	defer func() {
		// TODO 连接关闭，mq发送关闭事件

		// 捕获其他错误
		if r := recover(); r != nil {
			logger.Error("NetClient recover error: ", r)
			// 输出堆栈信息
			logger.Error("NetClient recover error stack: ", string(debug.Stack()))
		}
	}()

	defer client.Close()

	for {
		ignore, message, err := client.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
				logger.Error("unexpected close error: ", err)
			} else {
				// 读失败
				logger.Errorf("NetClient ReadMessage error: %T, %v", err, err)
			}
			return
		}
		if ignore {
			continue
		}

		payload, err := protocol.DecodeSo(message)
		if err != nil {
			logger.Errorf("decode message error: len=%d", len(message), err)
			return
		}
		header := payload.Header

		// grpc generic call
		ctx := context.Background()
		grpcClient, err := c.grpcFactory.GetClient(ctx, header.Svc)
		if err != nil {
			logger.Error("get grpc generic client error: ", err)
			continue
		}
		resp, err := grpcClient.InvokeUnary(ctx, header.Target, payload.Body)
		if err != nil {
			// todo 提取 err: fmt.Sprintf("rpc error: code = %s desc = %s", s.Code(), s.Message())
			writeError(client, header.SeqId, err.Error())
			logger.Error("grpc generic call error: ", err)
			continue
		}

		// write response
		header.Type = protocol.TypeResponse
		header.Target = resp.XXX_MessageName()
		payload.Body, err = resp.Marshal()
		if err != nil {
			logger.Error("generic call response marshal error: ", err)
			continue
		}
		resMessage, err := protocol.EncodeSo(payload)
		if err != nil {
			logger.Error("grpc generic call error: ", err)
			continue
		}
		client.MustWrite(resMessage)
	}

}

func writeError(client *session.NetClient, seqId int32, errMsg string) {
	header := &protocol.Header{}
	payload := &protocol.Payload{Header: header}

	header.Magic = protocol.Magic
	header.Type = protocol.TypeError
	header.Status = 50
	header.SeqId = seqId
	payload.Body = []byte(errMsg)
	message, err := protocol.EncodeSo(payload)
	if err != nil {
		logger.Error("NetClient writeError encode error:", err)
		return
	}
	client.MustWrite(message)
	return
}