flvx/go-backend/internal/http/handler/diagnosis_stream.go

package handler

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"strings"
	"time"

	"go-backend/internal/http/response"
)

type diagnosisStreamEvent struct {
	Type string      `json:"type"`
	Data interface{} `json:"data,omitempty"`
	TS   int64       `json:"ts"`
}

func prepareDiagnosisStreamResponse(w http.ResponseWriter) (http.Flusher, error) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		return nil, errors.New("当前服务不支持流式响应")
	}
	w.Header().Set("Content-Type", "application/x-ndjson; charset=utf-8")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("X-Accel-Buffering", "no")
	return flusher, nil
}

func writeDiagnosisStreamEvent(encoder *json.Encoder, flusher http.Flusher, eventType string, data interface{}) error {
	if encoder == nil || flusher == nil {
		return errors.New("流式响应写入器未初始化")
	}
	event := diagnosisStreamEvent{Type: eventType, Data: data, TS: time.Now().UnixMilli()}
	if err := encoder.Encode(event); err != nil {
		return err
	}
	flusher.Flush()
	return nil
}

func summarizeDiagnosisProgress(results []map[string]interface{}) diagnosisProgress {
	progress := diagnosisProgress{Total: len(results)}
	for _, item := range results {
		progress.Completed++
		if asBool(item["success"], false) {
			progress.Success++
		} else {
			progress.Failed++
		}
	}
	return progress
}

func shouldIgnoreDiagnosisStreamError(err error) bool {
	if err == nil {
		return false
	}
	if errors.Is(err, context.Canceled) {
		return true
	}
	msg := strings.ToLower(strings.TrimSpace(err.Error()))
	if strings.Contains(msg, "broken pipe") || strings.Contains(msg, "connection reset by peer") {
		return true
	}
	if strings.Contains(msg, "stream already closed") {
		return true
	}
	return false
}

func (h *Handler) streamDiagnosisRuntime(ctx context.Context, cancel context.CancelFunc, w http.ResponseWriter, startPayload map[string]interface{}, workItems []diagnosisWorkItem) error {
	flusher, err := prepareDiagnosisStreamResponse(w)
	if err != nil {
		return err
	}
	encoder := json.NewEncoder(w)

	payload := map[string]interface{}{
		"total":     len(workItems),
		"timestamp": time.Now().UnixMilli(),
	}
	for key, value := range startPayload {
		payload[key] = value
	}
	if err := writeDiagnosisStreamEvent(encoder, flusher, "start", payload); err != nil {
		return err
	}

	streamBroken := false
	emitter := func(index int, item map[string]interface{}, progress diagnosisProgress) {
		if streamBroken {
			return
		}
		itemPayload := map[string]interface{}{
			"index":    index,
			"result":   item,
			"progress": progress,
		}
		if err := writeDiagnosisStreamEvent(encoder, flusher, "item", itemPayload); err != nil {
			streamBroken = true
			if cancel != nil {
				cancel()
			}
		}
	}

	results := h.runDiagnosisWorkItems(ctx, workItems, emitter)
	if streamBroken {
		return context.Canceled
	}

	progress := summarizeDiagnosisProgress(results)
	donePayload := map[string]interface{}{
		"progress": progress,
		"timedOut": errors.Is(ctx.Err(), context.DeadlineExceeded),
	}
	return writeDiagnosisStreamEvent(encoder, flusher, "done", donePayload)
}

func (h *Handler) tunnelDiagnoseStream(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		response.WriteJSON(w, response.ErrDefault("请求失败"))
		return
	}
	id := asInt64FromBodyKey(r, w, "tunnelId")
	if id <= 0 {
		return
	}

	tunnelName, tunnelType, workItems, err := h.prepareTunnelDiagnosis(id)
	if err != nil {
		if strings.Contains(err.Error(), "不存在") || strings.Contains(err.Error(), "不完整") {
			response.WriteJSON(w, response.ErrDefault(err.Error()))
			return
		}
		response.WriteJSON(w, response.Err(-2, err.Error()))
		return
	}

	ctx, cancel := context.WithTimeout(r.Context(), diagnosisRequestTimeout)
	defer cancel()

	startPayload := map[string]interface{}{
		"tunnelName": tunnelName,
		"tunnelType": tunnelType,
	}
	if err := h.streamDiagnosisRuntime(ctx, cancel, w, startPayload, workItems); err != nil {
		if shouldIgnoreDiagnosisStreamError(err) {
			return
		}
		if strings.Contains(err.Error(), "不支持流式响应") {
			response.WriteJSON(w, response.ErrDefault(err.Error()))
			return
		}
		return
	}
}

func (h *Handler) forwardDiagnoseStream(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		response.WriteJSON(w, response.ErrDefault("请求失败"))
		return
	}
	id := asInt64FromBodyKey(r, w, "forwardId")
	if id <= 0 {
		return
	}

	forward, _, _, err := h.resolveForwardAccess(r, id)
	if err != nil {
		if errors.Is(err, errForwardNotFound) {
			response.WriteJSON(w, response.ErrDefault("转发不存在"))
			return
		}
		response.WriteJSON(w, response.Err(-2, err.Error()))
		return
	}

	forwardName, workItems, err := h.prepareForwardDiagnosis(forward)
	if err != nil {
		if strings.Contains(err.Error(), "不存在") || strings.Contains(err.Error(), "不能为空") || strings.Contains(err.Error(), "错误") {
			response.WriteJSON(w, response.ErrDefault(err.Error()))
			return
		}
		response.WriteJSON(w, response.Err(-2, err.Error()))
		return
	}

	ctx, cancel := context.WithTimeout(r.Context(), diagnosisRequestTimeout)
	defer cancel()

	startPayload := map[string]interface{}{
		"forwardName": forwardName,
	}
	if err := h.streamDiagnosisRuntime(ctx, cancel, w, startPayload, workItems); err != nil {
		if shouldIgnoreDiagnosisStreamError(err) {
			return
		}
		if strings.Contains(err.Error(), "不支持流式响应") {
			response.WriteJSON(w, response.ErrDefault(err.Error()))
			return
		}
		return
	}
}