From fef9f8d3908c142e7df72572beff51f5b1748a27 Mon Sep 17 00:00:00 2001 From: laodaming <11058467+laudamine@user.noreply.gitee.com> Date: Fri, 28 Jul 2023 17:15:37 +0800 Subject: [PATCH] fix --- goctl_template/api/context.tpl | 2 + initalize/rabbitmq.go | 93 ++++++++++++++++--- .../internal/logic/datatransferlogic.go | 19 ++-- .../internal/logic/ws_render_image_logic.go | 6 +- .../websocket/internal/svc/servicecontext.go | 2 + 5 files changed, 101 insertions(+), 21 deletions(-) diff --git a/goctl_template/api/context.tpl b/goctl_template/api/context.tpl index ca4b66ff..e5884b34 100644 --- a/goctl_template/api/context.tpl +++ b/goctl_template/api/context.tpl @@ -18,6 +18,7 @@ type ServiceContext struct { {{.middleware}} MysqlConn *gorm.DB AllModels *gmodel.AllModelsGen + RabbitMq *initalize.RabbitMqHandle } func NewServiceContext(c {{.config}}) *ServiceContext { @@ -26,6 +27,7 @@ func NewServiceContext(c {{.config}}) *ServiceContext { Config: c, MysqlConn: initalize.InitMysql(c.SourceMysql), AllModels: gmodel.NewAllModels(initalize.InitMysql(c.SourceMysql)), + RabbitMq:initalize.InitRabbitMq(c.SourceRabbitMq, nil), {{.middlewareAssignment}} } } diff --git a/initalize/rabbitmq.go b/initalize/rabbitmq.go index f3a70475..5bfffa9e 100644 --- a/initalize/rabbitmq.go +++ b/initalize/rabbitmq.go @@ -2,22 +2,28 @@ package initalize import ( "crypto/tls" + "errors" "fusenapi/constants" "github.com/streadway/amqp" + "github.com/zeromicro/go-zero/core/logx" "log" ) -// handle -type QueueItem struct { - Ch *amqp.Channel - Queue amqp.Queue +type RabbitMqHandle struct { } -var mapMq = make(map[string]*QueueItem) +// 连接属性 +type queueItem struct { + ch *amqp.Channel + queue amqp.Queue +} -func InitRabbitMq(url string, config *tls.Config) { +// 存储连接 +var mapMq = make(map[constants.RABBIT_MQ]queueItem) + +func InitRabbitMq(url string, config *tls.Config) *RabbitMqHandle { if url == "" { - return + return nil } conn, err := amqp.DialTLS(url, config) if err != nil { @@ -39,13 +45,76 @@ func InitRabbitMq(url string, config *tls.Config) { nil, // 其他参数 ) if err != nil { - conn.Close() - ch.Close() log.Fatalf("Failed to declare a queue: %v", err) } - mapMq[string(queueName)] = &QueueItem{ - Ch: ch, - Queue: q, + mapMq[queueName] = queueItem{ + ch: ch, + queue: q, } } + return &RabbitMqHandle{} +} + +// 发送消息 +func (h *RabbitMqHandle) SendMsg(queueName constants.RABBIT_MQ, message []byte) error { + object, ok := mapMq[queueName] + if !ok { + return errors.New("unknown queue") + } + // 发送消息到队列 + return object.ch.Publish( + "", // exchange，如果为空，则使用默认交换机 + object.queue.Name, // routing key，将消息发送到指定队列 + false, // 是否等待服务器响应 + false, // 是否立即发送 + amqp.Publishing{ + ContentType: "text/plain", //普通文本 + Body: message, + }, // 消息内容 + ) +} + +// 消费消息 +func (h *RabbitMqHandle) Consume(queueName constants.RABBIT_MQ, handleFunc func(data []byte) error) error { + object, ok := mapMq[queueName] + if !ok { + return errors.New("unknown queue") + } + msgs, err := object.ch.Consume( + object.queue.Name, // 队列名 + "", // 消费者名，如果为空，则是随机生成一个 + false, // 自动应答 + false, // 是否排他 + false, // 是否阻塞 + false, // 是否等待服务器响应 + nil, // 其他参数 + ) + if err != nil { + log.Fatalf("Failed to register a consumer: %v", err) + } + //允许20的并发 + limit := make(chan struct{}, 20) + defer close(limit) + // 消费消息 + for msg := range msgs { + limit <- struct{}{} + go func(m amqp.Delivery) { + if err := recover(); err != nil { + logx.Error(err) + } + defer func() { + <-limit + }() + if err = handleFunc(m.Body); err != nil { + logx.Error("failed to deal with MQ message:", string(m.Body)) + return + } + if err = object.ch.Ack(m.DeliveryTag, false); err != nil { + logx.Error("failed to ack mq to delete") + } else { + log.Printf("Consume Mq message success: %s", m.Body) + } + }(msg) + } + return nil } diff --git a/server/websocket/internal/logic/datatransferlogic.go b/server/websocket/internal/logic/datatransferlogic.go index 44ba0abb..25bf50e7 100644 --- a/server/websocket/internal/logic/datatransferlogic.go +++ b/server/websocket/internal/logic/datatransferlogic.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "fusenapi/constants" + "fusenapi/initalize" "fusenapi/server/websocket/internal/types" "fusenapi/utils/auth" "fusenapi/utils/id_generator" @@ -64,13 +65,14 @@ var ( // 每个连接的连接基本属性 type wsConnectItem struct { conn *websocket.Conn //websocket的连接 - closeChan chan struct{} //ws连接关闭chan - isClose bool //是否已经关闭 - uniqueId uint64 //ws连接唯一标识 - inChan chan []byte //接受消息缓冲通道 - outChan chan []byte //发送回客户端的消息 - mutex sync.Mutex //互斥锁 - renderProperty renderProperty //扩展云渲染属性 + rabbitMq *initalize.RabbitMqHandle + closeChan chan struct{} //ws连接关闭chan + isClose bool //是否已经关闭 + uniqueId uint64 //ws连接唯一标识 + inChan chan []byte //接受消息缓冲通道 + outChan chan []byte //发送回客户端的消息 + mutex sync.Mutex //互斥锁 + renderProperty renderProperty //扩展云渲染属性 } func (l *DataTransferLogic) DataTransfer(svcCtx *svc.ServiceContext, w http.ResponseWriter, r *http.Request) { @@ -100,6 +102,7 @@ func (l *DataTransferLogic) DataTransfer(svcCtx *svc.ServiceContext, w http.Resp uniqueId := websocketIdGenerator.Get() ws := wsConnectItem{ conn: conn, + rabbitMq: l.svcCtx.RabbitMq, uniqueId: uniqueId, closeChan: make(chan struct{}, 1), inChan: make(chan []byte, 1000), @@ -272,7 +275,7 @@ func (w *wsConnectItem) dealwithReciveData(data []byte) { switch parseInfo.T { //图片渲染 case constants.WEBSOCKET_RENDER_IMAGE: - go w.SendToCloudRender(d) + w.SendToCloudRender(d) default: } diff --git a/server/websocket/internal/logic/ws_render_image_logic.go b/server/websocket/internal/logic/ws_render_image_logic.go index 9a5f7d11..901a01d8 100644 --- a/server/websocket/internal/logic/ws_render_image_logic.go +++ b/server/websocket/internal/logic/ws_render_image_logic.go @@ -59,8 +59,12 @@ func (w *wsConnectItem) SendToCloudRender(data []byte) { Option: 1, //0删除 1添加 Key: key, } - // TODO 数据发送给云渲染服务器 + if err := w.rabbitMq.SendMsg(constants.RABBIT_MQ_ASSEMBLE_RENDER_DATA, data); err != nil { + logx.Error(err) + continue + } + logx.Info("发送渲染数据到rabbitmq成功") } } } diff --git a/server/websocket/internal/svc/servicecontext.go b/server/websocket/internal/svc/servicecontext.go index 83c5d7a3..3ca743ec 100644 --- a/server/websocket/internal/svc/servicecontext.go +++ b/server/websocket/internal/svc/servicecontext.go @@ -18,6 +18,7 @@ type ServiceContext struct { MysqlConn *gorm.DB AllModels *gmodel.AllModelsGen + RabbitMq *initalize.RabbitMqHandle } func NewServiceContext(c config.Config) *ServiceContext { @@ -26,6 +27,7 @@ func NewServiceContext(c config.Config) *ServiceContext { Config: c, MysqlConn: initalize.InitMysql(c.SourceMysql), AllModels: gmodel.NewAllModels(initalize.InitMysql(c.SourceMysql)), + RabbitMq: initalize.InitRabbitMq(c.SourceRabbitMq, nil), } }