Golang Gin框架中实现Server Sent Events(SSE)

被这个功能困扰了好久，一直没实现，因为网上的Gin框架实现的SSE都非常复杂，写成了一整套复杂的中间件，经过研究，写出了一套精简的实现。

SSE接口代码, 发送信息修改sendEvents内代码即可：

func GetTaskResult(c *gin.Context) {
    tag := "compositioncontroller.GetTaskResult"
    ctx := xesgin.TransferToContext(c)
    var args proto.GetTaskResultReq
    err := c.ShouldBind(&args)
    if err != nil {
        logger.Ix(ctx, tag, "c.ShouldBind failed", "err:[%+v]", err)
        c.JSON(http.StatusOK, xesgin.Error(logger.NewError("", logger.PARAM_ERROR)))
        return
    }

    err = compositionSvc.NewCompositionSvc().TaskResultStream(ctx, args, c.Writer)
    if err != nil {
        logger.Ix(ctx, tag, "compositionSvc.GetTaskResultStreaming failed", "err:[%+v]", err)
        c.JSON(http.StatusOK, xesgin.Error(err))
        return
    }
}

// 流式响应任务结果
func (c *compositionSvc) TaskResultStream(ctx context.Context, args proto.GetTaskResultReq, w http.ResponseWriter) (err error) {
    w.Header().Set("Content-Type", "text/event-stream")
    w.Header().Set("Cache-Control", "no-cache")
    w.Header().Set("Connection", "keep-alive")
    w.Header().Set("Access-Control-Allow-Origin", "*")

    closeNotifier := w.(http.CloseNotifier).CloseNotify()
    stopCh := make(chan struct{})

    go c.sendEvents(ctx, &args, w, stopCh)

    <-closeNotifier
    if _, ok := <-stopCh; ok {
        close(stopCh)
    }
    return
}

func (c *compositionSvc) sendEvents(ctx context.Context, args *proto.GetTaskResultReq, w http.ResponseWriter, stopCh chan struct{}) {
    tag := "compositionSvc.sendEvents"

    var resp proto.GetTaskResultResp
    for {
        var data []string
        select {
        case <-time.After(time.Millisecond * 300):
            data, err := redisRepo.Get(ctx, args.TaskId)
            if err != nil {
                logger.Ix(ctx, tag, "redisRepo.Get failed", "err:[%+v]", err)
                return
            }
            if data.Content == "" {
                continue
            }

            // 响应数据
            respByte, err := json.Marshal(resp)
            if err != nil {
                logger.Ix(ctx, tag, "json.Marshal failed", "err:[%+v]", err)
                return
            }
            sseData := fmt.Sprintf("data: %s\n\n", string(respByte))
            _, err = io.WriteString(w, sseData)
            if err != nil {
                logger.Ix(ctx, tag, "io.WriteString failed", "err:[%+v]", err)
                return
            }
            if flusher, ok := w.(http.Flusher); ok {
                flusher.Flush()
            }
            if data.status == cast.ToString(aigc.TaskStatusSuccess) || data.status == cast.ToString(aigc.TaskStatusFail) {
                return
            }
        case <-stopCh:
            return
        }
    }
}

原创的版权均归本人所有，任何人或团体、机构全部转载或者部分转载、摘录，请保留本博客链接或标注来源。