go语言使用gin库实现SSE长连接

项目需求：在go项目中实现SSE长连接，耗时操作完成后，后端主动通知前端消息。

1.建立长连接

最主要的操作是修改请求中的Content-Type类型为"text/event-stream"。

需要注意几点，后端首先不能让请求处理代码跑完，如果跑完这个请求就会断掉，保存的*gin.Context信息也就没用了，因此要新建一个range chan维持处理状态。

另外这里如果用for循环代替chan，会导致前端持续发送请求到后端。

import (
    "encoding/json"
    "fmt"
    "github.com/gin-gonic/gin"
    "github.com/labstack/gommon/log"
    "net/http"
    "strings"
    "time"
)

var ifChannelsMapInit = false

var channelsMap = map[string]chan string{}

func initChannelsMap() {
    channelsMap = make(map[string]chan string)
}

func AddChannel(userEmail string, traceId string) {
    if !ifChannelsMapInit {
        initChannelsMap()
        ifChannelsMapInit = true
    }
    var newChannel = make(chan string)
    channelsMap[userEmail+traceId] = newChannel
    log.Infof("Build SSE connection for user = " + userEmail + ", trace id = " + traceId)
}

func BuildNotificationChannel(userEmail string, traceId string, c *gin.Context) {
    AddChannel(userEmail, traceId)
    c.Writer.Header().Set("Content-Type", "text/event-stream")
    c.Writer.Header().Set("Cache-Control", "no-cache")
    c.Writer.Header().Set("Connection", "keep-alive")
    w := c.Writer
    flusher, _ := w.(http.Flusher)
    closeNotify := c.Request.Context().Done()
    go func() {
        <-closeNotify
        delete(channelsMap, userEmail+traceId)
        log.Infof("SSE close for user = " + userEmail + ", trace id = " + traceId)
        return
    }()
    fmt.Fprintf(w, "data: %s\n\n", "--ping--")
    flusher.Flush()
    for msg := range channelsMap[userEmail+traceId] {
        fmt.Fprintf(w, "data: %s\n\n", msg)
        flusher.Flush()
    }
}

2.后端主动通知前端消息

当耗时操作处理完成后，调用该方法，前端会收到通知。

func SendNotification(userEmail string, messageBody string, actionType string) {
    log.Infof("Send notification to user = " + userEmail)
    var msg = dao.NotificationLog{}
    msg.MessageBody = messageBody
    msg.UserEmail = userEmail
    msg.Type = actionType
    msg.Status = UNREAD
    msg.CreateTime = time.Now()
    msg.Create()
    msgBytes, _ := json.Marshal(msg)
    for key := range channelsMap {
        if strings.Contains(key, userEmail) {
            channel := channelsMap[key]
            channel <- string(msgBytes)
        }
    }
}

3.调试

准备两个接口，分别是建立SSE和触发耗时操作

GroupV1Rest.GET("/notification/socket-connection", controllers.SocketConnection)
GroupV1Rest.GET("/notification/export-excel", controllers.ExportExcel)

打开浏览器，进入调试模式，在console页输入

e = new EventSource('/business/v1/notification/socket-connection');
e.onmessage = function(event) {
    console.log(event.data);
};

看到日志打印‘--ping--’，长连接已建立

此时发送第二个请求，调试模式看到通知被chan处理

返回浏览器，可以看到已经收到通知

4.关闭长连接

前端关闭页面后，自动触发监听事件，后端清理连接信息