Go sse推送任务进度

测试一个账户在多个设备登录、同时接收消息、模拟视频任务成功、发送推送

package main

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

	sse "github.com/r3labs/sse/v2"
)

var sseServer = newSSEServer()

func main() {
	mux := http.NewServeMux()
	mux.HandleFunc("/", indexHandler)
	mux.HandleFunc("/events", eventsHandler)
	mux.HandleFunc("/mock/task-done", taskDoneHandler)
	mux.HandleFunc("/mock/task-failed", taskFailedHandler)

	addr := ":3801"
	log.Printf("SSE demo listening on http://localhost%s", addr)
	log.Printf("Open the same URL in multiple browsers/tabs with the same userId, then click send.")

	if err := http.ListenAndServe(addr, mux); err != nil {
		log.Fatal(err)
	}
}

func newSSEServer() *sse.Server {
	server := sse.New()
	server.AutoReplay = false
	return server
}

func streamID(userID uint64) string {
	return fmt.Sprintf("user:%d", userID)
}

func ensureUserStream(userID uint64) string {
	id := streamID(userID)
	sseServer.CreateStream(id)
	return id
}

func publishToUser(userID uint64, eventName string, payload any) error {
	data, err := json.Marshal(payload)
	if err != nil {
		return err
	}

	sseServer.Publish(ensureUserStream(userID), &sse.Event{
		Event: []byte(eventName),
		Data:  data,
	})

	return nil
}

func indexHandler(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.NotFound(w, r)
		return
	}

	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	_, _ = w.Write([]byte(indexHTML))
}

func eventsHandler(w http.ResponseWriter, r *http.Request) {
	userID, ok := parseUserID(w, r)
	if !ok {
		return
	}

	query := r.URL.Query()
	query.Set("stream", ensureUserStream(userID))
	r.URL.RawQuery = query.Encode()

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("X-Accel-Buffering", "no")

	log.Printf("SSE connected: user=%d remote=%s", userID, r.RemoteAddr)
	sseServer.ServeHTTP(w, r)
	log.Printf("SSE disconnected: user=%d remote=%s", userID, r.RemoteAddr)
}

func taskDoneHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var req taskRequest
	if !decodeJSON(w, r, &req) {
		return
	}
	if req.TaskID == "" {
		req.TaskID = fmt.Sprintf("task_%d", time.Now().UnixMilli())
	}
	if req.VideoURL == "" {
		req.VideoURL = "https://example.com/demo-video.mp4"
	}

	payload := map[string]any{
		"taskId":   req.TaskID,
		"status":   "succeeded",
		"type":     "video",
		"videoUrl": req.VideoURL,
		"time":     time.Now().Format(time.RFC3339),
	}
	if err := publishToUser(req.UserID, "video_task_done", payload); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	writeJSON(w, map[string]any{"message": "published", "event": payload})
}

func taskFailedHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var req taskRequest
	if !decodeJSON(w, r, &req) {
		return
	}
	if req.TaskID == "" {
		req.TaskID = fmt.Sprintf("task_%d", time.Now().UnixMilli())
	}
	if req.Message == "" {
		req.Message = "mock generation failed"
	}

	payload := map[string]any{
		"taskId":  req.TaskID,
		"status":  "failed",
		"type":    "video",
		"message": req.Message,
		"time":    time.Now().Format(time.RFC3339),
	}
	if err := publishToUser(req.UserID, "video_task_failed", payload); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	writeJSON(w, map[string]any{"message": "published", "event": payload})
}

type taskRequest struct {
	UserID   uint64 `json:"userId"`
	TaskID   string `json:"taskId"`
	VideoURL string `json:"videoUrl"`
	Message  string `json:"message"`
}

func parseUserID(w http.ResponseWriter, r *http.Request) (uint64, bool) {
	value := r.URL.Query().Get("userId")
	if value == "" {
		http.Error(w, "missing userId", http.StatusBadRequest)
		return 0, false
	}

	userID, err := strconv.ParseUint(value, 10, 64)
	if err != nil || userID == 0 {
		http.Error(w, "invalid userId", http.StatusBadRequest)
		return 0, false
	}

	return userID, true
}

func decodeJSON(w http.ResponseWriter, r *http.Request, target any) bool {
	defer r.Body.Close()
	if err := json.NewDecoder(r.Body).Decode(target); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return false
	}
	return true
}

func writeJSON(w http.ResponseWriter, value any) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	_ = json.NewEncoder(w).Encode(value)
}

const indexHTML = `<!doctype html>
<html lang="zh-CN">
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1" />
  <title>SSE 多客户端测试</title>
  <style>
    body { font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; margin: 32px; color: #111827; }
    .card { max-width: 900px; border: 1px solid #e5e7eb; border-radius: 14px; padding: 20px; box-shadow: 0 10px 30px rgba(15, 23, 42, .06); }
    label { display: block; margin: 12px 0 6px; font-weight: 600; }
    input { width: 100%; max-width: 420px; box-sizing: border-box; padding: 10px 12px; border: 1px solid #d1d5db; border-radius: 10px; }
    button { margin: 12px 8px 0 0; padding: 10px 14px; border: 0; border-radius: 10px; color: white; background: #2563eb; cursor: pointer; }
    button.secondary { background: #4b5563; }
    button.danger { background: #dc2626; }
    button:disabled { background: #9ca3af; cursor: not-allowed; }
    pre { min-height: 280px; max-height: 520px; overflow: auto; padding: 14px; background: #0f172a; color: #d1fae5; border-radius: 12px; white-space: pre-wrap; }
    .hint { color: #6b7280; line-height: 1.7; }
    .status { display: inline-block; margin-top: 12px; padding: 4px 10px; border-radius: 999px; background: #f3f4f6; }
  </style>
</head>
<body>
  <div class="card">
    <h1>SSE 多客户端测试</h1>
    <p class="hint">
      打开多个浏览器或标签页，填同一个 userId 后点击“连接 SSE”。任意一个页面点击“发送成功事件”，所有连接到同一 userId 的页面都会收到消息。
    </p>

    <label for="userId">userId</label>
    <input id="userId" value="123" />

    <label for="taskId">taskId</label>
    <input id="taskId" value="task_demo_001" />

    <label for="videoUrl">videoUrl</label>
    <input id="videoUrl" value="https://example.com/demo-video.mp4" />

    <div>
      <button id="connectBtn">连接 SSE</button>
      <button id="closeBtn" class="secondary" disabled>断开</button>
      <button id="doneBtn">发送成功事件</button>
      <button id="failedBtn" class="danger">发送失败事件</button>
      <button id="clearBtn" class="secondary">清空日志</button>
    </div>

    <div class="status" id="status">未连接</div>
    <h3>事件日志</h3>
    <pre id="log"></pre>
  </div>

  <script>
    let es = null
    const userIdInput = document.getElementById('userId')
    const taskIdInput = document.getElementById('taskId')
    const videoUrlInput = document.getElementById('videoUrl')
    const connectBtn = document.getElementById('connectBtn')
    const closeBtn = document.getElementById('closeBtn')
    const doneBtn = document.getElementById('doneBtn')
    const failedBtn = document.getElementById('failedBtn')
    const clearBtn = document.getElementById('clearBtn')
    const statusEl = document.getElementById('status')
    const logEl = document.getElementById('log')

    function log(message, data) {
      const line = '[' + new Date().toLocaleTimeString() + '] ' + message + (data ? '\n' + JSON.stringify(data, null, 2) : '')
      logEl.textContent = line + '\n\n' + logEl.textContent
    }

    function setConnected(connected) {
      connectBtn.disabled = connected
      closeBtn.disabled = !connected
      statusEl.textContent = connected ? '已连接 userId=' + userIdInput.value : '未连接'
    }

    connectBtn.onclick = () => {
      if (es) es.close()
      const userId = encodeURIComponent(userIdInput.value.trim())
      es = new EventSource('/events?userId=' + userId)

      es.onopen = () => {
        setConnected(true)
        log('SSE connected')
      }

      es.onerror = () => {
        log('SSE error or reconnecting')
      }

      es.addEventListener('video_task_done', (event) => {
        log('video_task_done', JSON.parse(event.data))
      })

      es.addEventListener('video_task_failed', (event) => {
        log('video_task_failed', JSON.parse(event.data))
      })
    }

    closeBtn.onclick = () => {
      if (es) {
        es.close()
        es = null
      }
      setConnected(false)
      log('SSE closed')
    }

    doneBtn.onclick = async () => {
      const payload = {
        userId: Number(userIdInput.value),
        taskId: taskIdInput.value,
        videoUrl: videoUrlInput.value,
      }
      const res = await fetch('/mock/task-done', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify(payload),
      })
      log('POST /mock/task-done ' + res.status, await res.json())
    }

    failedBtn.onclick = async () => {
      const payload = {
        userId: Number(userIdInput.value),
        taskId: taskIdInput.value,
        message: '模拟生成失败',
      }
      const res = await fetch('/mock/task-failed', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify(payload),
      })
      log('POST /mock/task-failed ' + res.status, await res.json())
    }

    clearBtn.onclick = () => {
      logEl.textContent = ''
    }
  </script>
</body>
</html>`