139 lines
2.8 KiB
Go
139 lines
2.8 KiB
Go
/*
|
|
* Copyright (c) 2025 Beijing Hardman Automation Equipment Co., LTD. All rights reserved.
|
|
*
|
|
* Project: adapter
|
|
* File: sse.go
|
|
* Last: 2025-05-07 14:12:54
|
|
* Author: wangcheng@bj-hardman.com
|
|
*/
|
|
|
|
package sse
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"log/slog"
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/zeromicro/go-zero/core/logx"
|
|
)
|
|
|
|
type Handler struct {
|
|
mu sync.RWMutex
|
|
clients map[chan Event]bool
|
|
}
|
|
|
|
type Event struct {
|
|
Event string `json:"event"`
|
|
Data string `json:"data"`
|
|
}
|
|
|
|
func NewHandler() *Handler {
|
|
return &Handler{
|
|
clients: make(map[chan Event]bool),
|
|
}
|
|
}
|
|
|
|
// Serve 处理 SSE 连接
|
|
func (h *Handler) Serve(w http.ResponseWriter, r *http.Request) {
|
|
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.Header().Set("Cache-Control", "no-cache")
|
|
w.Header().Set("Connection", "keep-alive")
|
|
|
|
// 为每个客户端创建一个 channel
|
|
clientChan := make(chan Event)
|
|
h.mu.Lock()
|
|
h.clients[clientChan] = true
|
|
h.mu.Unlock()
|
|
|
|
// 客户端断开时清理
|
|
defer func() {
|
|
h.mu.Lock()
|
|
delete(h.clients, clientChan)
|
|
h.mu.Unlock()
|
|
close(clientChan)
|
|
}()
|
|
|
|
w.(http.Flusher).Flush() // 刷新缓冲区
|
|
|
|
heartbeat := time.NewTicker(30 * time.Second)
|
|
defer heartbeat.Stop()
|
|
|
|
// 持续监听并推送事件
|
|
for {
|
|
select {
|
|
case msg := <-clientChan:
|
|
// 发送事件数据
|
|
var err error
|
|
if msg.Data != "" {
|
|
_, err = fmt.Fprintf(w, "event: %s\ndata: %s\n\n", msg.Event, msg.Data)
|
|
} else {
|
|
_, err = fmt.Fprintf(w, "event: %s\n\n", msg.Event)
|
|
}
|
|
if err != nil {
|
|
logx.Errorf("sse sent %+v failed: %s", msg, err)
|
|
return
|
|
}
|
|
w.(http.Flusher).Flush()
|
|
slog.InfoContext(r.Context(), fmt.Sprintf("sse sent %+v", msg))
|
|
case <-heartbeat.C:
|
|
_, err := fmt.Fprintf(w, ": heartbeat\n\n")
|
|
if err != nil {
|
|
return
|
|
}
|
|
w.(http.Flusher).Flush()
|
|
case <-r.Context().Done():
|
|
//客户端断开连接
|
|
logx.Error(r.Context().Err())
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (h *Handler) Emit(event string, data ...any) {
|
|
// 广播给所有客户端
|
|
|
|
var js []byte
|
|
if len(data) > 0 {
|
|
if s, ok := data[0].(string); ok {
|
|
js = []byte(s)
|
|
} else {
|
|
js, _ = json.Marshal(data[0])
|
|
}
|
|
}
|
|
|
|
h.mu.RLock()
|
|
defer h.mu.RUnlock()
|
|
for msgChan := range h.clients {
|
|
go func(ch chan Event) {
|
|
defer func() { recover() }() // 客户端断连后 channel 已关闭，忽略 panic
|
|
ch <- Event{
|
|
Event: event,
|
|
Data: string(js),
|
|
}
|
|
}(msgChan)
|
|
}
|
|
}
|
|
|
|
func (h *Handler) Toast(msg ToastMessage) {
|
|
h.Emit("toast", msg)
|
|
}
|
|
|
|
func (h *Handler) Notify(msg NoticeMessage) {
|
|
h.Emit("notify", msg)
|
|
}
|
|
|
|
// EmitProgressEvent sends a structured progress event to all SSE clients.
|
|
func (h *Handler) EmitProgressEvent(eventType string, taskID int, stepName string, data map[string]any) {
|
|
payload := map[string]any{
|
|
"event": eventType,
|
|
"task_id": taskID,
|
|
"step_name": stepName,
|
|
"data": data,
|
|
}
|
|
h.Emit("workflow", payload)
|
|
}
|