aiops/backend/controller/ws_pipeline_log.go

package controller

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	ezredis "ezops/pkg/redis"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"gorm.io/gorm"
)

var wsUpgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

// WSPipelineLogHandler WebSocket 流水线日志推流
type WSPipelineLogHandler struct {
	db    *gorm.DB
	cache ezredis.Cache
}

// NewWSPipelineLogHandler 创建 WebSocket 日志处理器
func NewWSPipelineLogHandler(db *gorm.DB, cache ezredis.Cache) *WSPipelineLogHandler {
	return &WSPipelineLogHandler{db: db, cache: cache}
}

// StreamLog WebSocket 端点：/ws/pipeline/:runId/log
func (h *WSPipelineLogHandler) StreamLog(c *gin.Context) {
	runIDStr := c.Param("runId")
	runID, err := strconv.ParseUint(runIDStr, 10, 32)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "无效的 runId"})
		return
	}

	ws, err := wsUpgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		return
	}
	defer ws.Close()

	logKey := fmt.Sprintf("pipeline:run:%d:logs", uint(runID))
	ctx := c.Request.Context()

	// 首次发送：已有日志
	raw, _ := h.cache.Get(ctx, logKey)
	if raw != "" {
		ws.WriteJSON(map[string]interface{}{"type": "init", "logs": raw})
	}

	// 轮询新日志（500ms 间隔）
	lastLen := 0
	ticker := time.NewTicker(500 * time.Millisecond)
	defer ticker.Stop()

	done := make(chan struct{})
	go func() {
		defer close(done)
		for {
			_, _, err := ws.ReadMessage()
			if err != nil {
				return
			}
		}
	}()

	for {
		select {
		case <-done:
			return
		case <-ticker.C:
			raw, err := h.cache.Get(ctx, logKey)
			if err != nil || raw == "" {
				continue
			}
			var entries []LogEntry
			if err := json.Unmarshal([]byte(raw), &entries); err != nil {
				continue
			}
			if len(entries) > lastLen {
				// 增量推送
				newEntries := entries[lastLen:]
				lastLen = len(entries)
				for _, e := range newEntries {
					buf, _ := json.Marshal(e)
					if err := ws.WriteMessage(websocket.TextMessage, buf); err != nil {
						return
					}
				}
			}

			// 检查流水线是否已结束（日志 key 被删除 = 已写入 DB）
			if lastLen > 0 {
				exists, _ := h.cache.Exists(ctx, logKey)
				if exists == 0 {
					// 发送完成信号
					ws.WriteJSON(map[string]interface{}{"type": "done"})
					return
				}
			}
		case <-ctx.Done():
			return
		}
	}
}

// GetLogHistory HTTP 日志历史（已完成的任务从 DB 读取）
func (h *WSPipelineLogHandler) GetLogHistory(c *gin.Context) {
	runIDStr := c.Param("runId")
	runID, err := strconv.ParseUint(runIDStr, 10, 32)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "无效的 runId"})
		return
	}

	// 先查 Redis（执行中）
	logKey := fmt.Sprintf("pipeline:run:%d:logs", uint(runID))
	ctx := c.Request.Context()
	
	raw, _ := h.cache.Get(ctx, logKey)
	if raw != "" && !strings.HasPrefix(raw, "[") {
		raw = "[]"
	}
	if raw != "" {
		c.JSON(http.StatusOK, gin.H{"code": 200, "data": gin.H{"source": "redis", "logs": raw}})
		return
	}

	// 查 DB（已完成）
	var logs string
	h.db.Table("ops_pipeline_runs").Where("id = ?", uint(runID)).Select("logs").Scan(&logs)
	if logs == "" {
		logs = "[]"
	}
	c.JSON(http.StatusOK, gin.H{"code": 200, "data": gin.H{"source": "db", "logs": logs}})
}