Golang+RabbitMQ路由模式+websocket实时显示日志

一、环境介绍

Go版本：1.13.1
开发工具：IntelliJ IDEA 2019.2.3 x64
开发环境：windows10 64位
部署环境：centos7

业务场景：
       AI算法训练过程中的日志，实时在Web UI界面显示
       每次训练的日志可以在多个Web UI展示，同一个训练的同一条消息可以被多个消费端消费
       不同训练的消息不能混合显示，使用RabbitMQ的路由匹配模式
RountingKey=rtk.train.id，queue=que.logque.id，每个训练根据id匹配自己的日志。Id=训练id
Exchange持久化， queue不持久化
       消息流转过程、RabbitMQ队列设计如下

代码已上传github：https://github.com/reachyu/realtimelog

二、RabbitMQ代码

路由模式生产者和消费者代码如下

package msgmq

import (
	"aisvc/common/vo"
	"encoding/json"
	"fmt"
	"github.com/streadway/amqp"
	"log"
	"strconv"
	"unsafe"
)

// rabbitmq 路由routing模式

// 发送消息(生产者)
func PublishMsgRout(exName string,rtKey string, msg string) {

	ch, err := MQInstance().GetMQChannel()
	failOnError(err, "Failed to open a channel")
	defer ch.Close()

	err = ch.ExchangeDeclare(
		exName, // name
		"direct",      // type
		true,         // durable
		false,        // auto-deleted
		false,        // internal
		false,        // no-wait
		nil,          // arguments
	)
	failOnError(err, "Failed to declare an exchange")
	err = ch.Publish(
		exName,          // exchange
		rtKey, // routing key
		//如果为true，根据exchange类型和routekey类型，如果无法找到符合条件的队列，name会把发送的信息返回给发送者
		true, // mandatory
		false, // immediate
		amqp.Publishing{
			// 消息持久化
			DeliveryMode: amqp.Persistent,
			ContentType: "text/plain",
			Body:        []byte(msg),
		})
	failOnError(err, "Failed to publish a message")
}


type Callback func(trainId string,msg string)

// 消费消息(消费者)
func ConsumeMsgRout(exName string,queName string,rtKey string,callback Callback) {
	ch, err := MQInstance().GetMQChannel()
	failOnError(err, "Failed to receive a message")
	defer ch.Close()

	err = ch.ExchangeDeclare(
		exName,   // name
		"direct", // type
		true,     // durable
		false,    // auto-deleted
		false,    // internal
		false,    // no-wait
		nil,      // arguments
	)
	failOnError(err, "Failed to receive a message")

	_, err = ch.QueueDeclare(
		queName,    // name
		false, // durable
		false, // delete when usused
		false,  // exclusive
		false, // no-wait
		nil,   // arguments
	)
	failOnError(err, "Failed to receive a message")

	err = ch.QueueBind(
		queName, // queue name
		rtKey,     // routing key
		exName, // exchange
		false,
		nil,
	)
	failOnError(err, "Failed to receive a message")

	msgs, err := ch.Consume(
		queName, // queue
		"",     // consumer
		true,   // auto-ack
		false,  // exclusive
		false,  // no-local
		false,  // no-wait
		nil,    // args
	)

	forever := make(chan bool)
	go func() {
		for d := range msgs {
			var msgTrainLog vo.TrainLog
			_ = json.Unmarshal(d.Body, &msgTrainLog)
			trainId := msgTrainLog.TrainId
			id64 := strconv.FormatInt(trainId,10)

			// 避免循环引用  callback = ws.SendLogsToWeb(trainId string,msg string)
			strAddress := &callback
			strPointer := fmt.Sprintf("%d", unsafe.Pointer(strAddress))
			intPointer, _ := strconv.ParseInt(strPointer, 10, 0)
			var pointer *Callback
			pointer = *(**Callback)(unsafe.Pointer(&intPointer))
			(Callback)(*pointer)(id64,msgTrainLog.TrainLog)

			log.Printf("路由队列接收到消息======== [x] %s", d.Body)
		}
	}()
	<-forever
}

func failOnError(err error, msg string) {
	if err != nil {
		log.Fatalf("%s: %s", msg, err)
	}
}

三、websocket代码

websocket代码参考了
https://github.com/516134941/websocket-gin-demo/tree/master/message-chat

wsclient.go

package ws

import (
	_const "aisvc/common/const"
	mq "aisvc/msgmq"
	"bytes"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"

	"github.com/gorilla/websocket"
)

const (
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second
	// Time allowed to read the next pong message from the peer.
	pongWait = 60 * time.Second
	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10
	// Maximum message size allowed from peer.
	maxMessageSize = 512
)

var clientMaps map[string]*WSClient
var once sync.Once

func GetClientMaps() map[string]*WSClient {
	once.Do(func() {
		clientMaps = make(map[string]*WSClient)
	})
	return clientMaps
}

var (
	newline = []byte{'\n'}
	space   = []byte{' '}
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

// Client is a middleman between the websocket connection and the hub.
type WSClient struct {
	hub *WSHub
	// The websocket connection.
	conn *websocket.Conn
	// Buffered channel of outbound messages.
	send chan []byte
	trainId []byte
}

func SendLogsToWeb(trainId string,msg string) {
	clientMaps := GetClientMaps()
	c := clientMaps[trainId]

	if c == nil{
		return
	}
	c.conn.SetReadLimit(maxMessageSize)
	c.conn.SetReadDeadline(time.Now().Add(pongWait))
	c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })

	message := bytes.TrimSpace(bytes.Replace([]byte(msg), newline, space, -1))
	message = []byte(trainId + "&" + msg)
	fmt.Println("websocket读取到的消息====="+string(message))
	c.hub.broadcast <- []byte(message)

}

func (c *WSClient) readPump() {
	defer func() {
		c.hub.unregister <- c
		c.conn.Close()
	}()
	c.conn.SetReadLimit(maxMessageSize)
	c.conn.SetReadDeadline(time.Now().Add(pongWait))
	c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		_, message, err := c.conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("error: %v", err)
			}
			break
		}
		message = bytes.TrimSpace(bytes.Replace(message, newline, space, -1))
		message = []byte(string(c.trainId) + "&" + string(message))
		fmt.Println("websocket读取到的消息====="+string(message))
		c.hub.broadcast <- []byte(message)
	}
}

func (c *WSClient) writePump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		c.conn.Close()
	}()
	for {
		select {
		case message, ok := <-c.send:
			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if !ok {
				// The hub closed the channel.
				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			w, err := c.conn.NextWriter(websocket.TextMessage)
			if err != nil {
				return
			}
			w.Write(message)
			if err := w.Close(); err != nil {
				log.Printf("error: %v", err)
				return
			}
		case <-ticker.C:
			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

// ServeWs handles websocket requests from the peer.
func ServeWs(hub *WSHub, c *gin.Context) {
	GetClientMaps()
	trainId := c.Param("trainId")
	// 将网络请求变为websocket
	var upgrader = websocket.Upgrader{
		// 解决跨域问题
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Println(err)
		return
	}
	fmt.Println("websocket接收到前端的trainId======"+trainId)
	client := &WSClient{hub: hub, conn: conn, send: make(chan []byte, 256), trainId: []byte(trainId)}
	client.hub.register <- client

	clientMaps[trainId] = client

	// 监听日志消息
	go mq.ConsumeMsgRout(_const.RABBITMQ_ROUT_EXCHANGE_NAME,_const.RABBITMQ_ROUT_QUEUE + trainId,_const.RABBITMQ_ROUT_ROUTING_KEY + trainId,SendLogsToWeb)

	// Allow collection of memory referenced by the caller by doing all work in
	// new goroutines.
	go client.writePump()
	//go client.readPump()
}

wshub.go

package ws

import "strings"

// Hub maintains the set of active clients and broadcasts messages to the
// clients.
type WSHub struct {
	// Registered clients.
	clients map[*WSClient]bool
	// Inbound messages from the clients.
	broadcast chan []byte
	// Register requests from the clients.
	register chan *WSClient
	// Unregister requests from clients.
	unregister chan *WSClient
	// 唯一id key:client value:唯一id
	trainId map[*WSClient]string
}

// NewHub .
func NewHub() *WSHub {
	return &WSHub{
		broadcast:  make(chan []byte),
		register:   make(chan *WSClient),
		unregister: make(chan *WSClient),
		clients:    make(map[*WSClient]bool),
		trainId:    make(map[*WSClient]string),
	}
}

// Run .
func (h *WSHub) Run() {
	for {
		select {
		case client := <-h.register:
			h.clients[client] = true                 // 注册client端
			h.trainId[client] = string(client.trainId) // 给client端添加唯一id
		case client := <-h.unregister:
			if _, ok := h.clients[client]; ok {
				delete(h.clients, client)
				delete(h.trainId, client)
				close(client.send)
			}
		case message := <-h.broadcast:
			for client := range h.clients {
				// 使用“&”对message进行message切割 获取唯一id
				// 向信息所属的训练内的所有client 内添加send
				// msg[0]为唯一id msg[1]为打印内容
				msg := strings.Split(string(message), "&")
				if string(client.trainId) == msg[0] {
					select {
					case client.send <- []byte(msg[1]):
					default:
						close(client.send)
						delete(h.clients, client)
						delete(h.trainId, client)
					}
				}
			}
		}
	}
}

四、GIN路由

// 设置跨域
	router.Use(cors.New(cors.Config{
		AllowAllOrigins:  true,  // 这是允许访问所有域
		AllowMethods:     []string{"GET", "PUT", "POST", "DELETE", "OPTIONS"},   //服务器支持的所有跨域请求的方法,为了避免浏览次请求的多次'预检'请求
		AllowHeaders:     []string{"x-xq5-jwt", "Content-Type", "Origin", "Content-Length"},  // 允许跨域设置
		ExposeHeaders:    []string{"x-xq5-jwt"},  // 跨域关键设置 让浏览器可以解析
		AllowCredentials: true,  //  跨域请求是否需要带cookie信息 默认设置为true
		MaxAge:           12 * time.Hour,
	}))

hub := ws.NewHub()
	go hub.Run()
	router.GET("/ws/logs/:trainId", func(c *gin.Context) { ws.ServeWs(hub, c) })

五、HTML代码

<!DOCTYPE html>
<html lang="en">
<head>
<title>logs</title>
<script type="text/javascript">
window.onload = function () {
    var conn;
    var msg = document.getElementById("msg");
    var log = document.getElementById("log");
    function appendLog(item) {
        var doScroll = log.scrollTop > log.scrollHeight - log.clientHeight - 1;
        log.appendChild(item);
        if (doScroll) {
            log.scrollTop = log.scrollHeight - log.clientHeight;
        }
    }
    document.getElementById("form").onsubmit = function () {
        if (!conn) {
            return false;
        }
        if (!msg.value) {
            return false;
        }
        conn.send(msg.value);
        msg.value = "";
        return false;
    };
    if (window["WebSocket"]) {
        // 123456为trainId
        conn = new WebSocket("ws://" + "localhost:9090" + "/ws/logs/123456");
        conn.onclose = function (evt) {
            var item = document.createElement("div");
            item.innerHTML = "Connection closed.";
            appendLog(item);
        };
        conn.onmessage = function (evt) {
            var messages = evt.data.split('\n');
            for (var i = 0; i < messages.length; i++) {
                var item = document.createElement("div");
                item.innerText = messages[i];
                appendLog(item);
            }
        };
    } else {
        var item = document.createElement("div");
        item.innerHTML = "Your browser does not support WebSockets.";
        appendLog(item);
    }
};
</script>
<style type="text/css">
html {
    overflow: hidden;
}
body {
    overflow: hidden;
    padding: 0;
    margin: 0;
    width: 100%;
    height: 100%;
    background: gray;
}
#log {
    background: white;
    margin: 0;
    padding: 0.5em 0.5em 0.5em 0.5em;
    position: absolute;
    top: 0.5em;
    left: 0.5em;
    right: 0.5em;
    bottom: 3em;
    overflow: auto;
}
#form {
    padding: 0 0.5em 0 0.5em;
    margin: 0;
    position: absolute;
    bottom: 1em;
    left: 0px;
    width: 100%;
    overflow: hidden;
}
</style>
</head>
<body>
<div id="log"></div>
<form id="form">
    <input type="submit" value="Send" />
    <input type="text" id="msg" size="64"/>
</form>
</body>
</html>

六、测试代码

func main() {
	// 不调用的话，glog会报错----ERROR: logging before flag.Parse:
	flag.Parse()
	initenv.InitEnv()

	go httpserver.InitHttpServer()

msg := map[string]interface{}{
			"trainId":    123456,
			"trainLog":   "测试测试测试测试",
		}
msgJson, _ := json.Marshal(msg)

	forever := make(chan bool)
	go func() {
		for {

	mq.PublishMsgRout(_const.RABBITMQ_ROUT_EXCHANGE_NAME,_const.RABBITMQ_ROUT_ROUTING_KEY+"123456",string(msgJson))
		time.Sleep(1 * time.Second)
		}
	}()
	<-forever

}

实时日志效果