admin/TaskPool
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: add sharedworker for task stauts

Browse Source
This commit is contained in:
duorameng
2026-05-07 20:12:30 +08:00
parent e5f9081a64
commit 4bf9ea8802
12 changed files with 542 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/constant/constant.go
+2
View File
@@ -99,6 +99,8 @@ const (
EventTaskSuccess = "task_success"
EventTaskFailed = "task_failed"
EventTaskTimeout = "task_timeout"
EventTaskRunning = "task_running"
EventTaskQueued = "task_queued"
// 其他事件类型
EventSystemNotice = "system_notice"
internal/controllers/system_ws_controller.go
+90
View File
@@ -0,0 +1,90 @@
package controllers
import (
"time"
"github.com/engigu/baihu-panel/internal/constant"
"github.com/engigu/baihu-panel/internal/logger"
"github.com/engigu/baihu-panel/internal/services"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
)
type SystemWSController struct {
manager *services.SystemWSManager
}
func NewSystemWSController() *SystemWSController {
return &SystemWSController{
manager: services.GetSystemWSManager(),
}
}
func (sc *SystemWSController) HandleEvents(c *gin.Context) {
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
if err != nil {
logger.Errorf("[SystemWS] 升级 WebSocket 失败: %v", err)
return
}
client := sc.manager.Register(conn)
defer sc.manager.Unregister(client)
// 启动写循环
go sc.writeLoop(client)
// 启动读循环 (主要用于检测连接断开和维持心跳)
sc.readLoop(client)
}
func (sc *SystemWSController) readLoop(client *services.ClientConnection) {
defer client.Close()
client.Conn.SetReadLimit(constant.MaxMessageSize)
client.Conn.SetReadDeadline(time.Now().Add(constant.PongWait))
client.Conn.SetPongHandler(func(string) error {
client.Conn.SetReadDeadline(time.Now().Add(constant.PongWait))
return nil
})
for {
_, _, err := client.Conn.ReadMessage()
if err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
logger.Warnf("[SystemWS] 客户端异常断开: %v", err)
}
break
}
// 暂时不处理来自前端的消息，前端仅作为接收方
}
}
func (sc *SystemWSController) writeLoop(client *services.ClientConnection) {
ticker := time.NewTicker(constant.PingPeriod)
defer func() {
ticker.Stop()
client.Close()
}()
for {
select {
case message, ok := <-client.Send:
client.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
if !ok {
// 通道关闭
client.Conn.WriteMessage(websocket.CloseMessage, []byte{})
return
}
if err := client.Conn.WriteMessage(websocket.TextMessage, message); err != nil {
return
}
case <-ticker.C:
client.Conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
if err := client.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
return
}
}
}
}
internal/models/vo/ws_vo.go
+8
View File
@@ -0,0 +1,8 @@
package vo
// WSMessage 通用 WebSocket 消息结构
type WSMessage struct {
Type string `json:"type"` // 事件类型: task_status, notice, system_stats
Timestamp int64 `json:"timestamp"` // 毫秒时间戳
Payload interface{} `json:"payload"` // 负载数据
}
internal/router/api_routes.go
+5
View File
@@ -61,6 +61,7 @@ func initAuthorizedAPIRoutes(api *gin.RouterGroup, c *Controllers) {
registerMiseRoutes(adminOnly, c)
registerNotificationRoutes(adminOnly, c)
registerAppLogRoutes(adminOnly, c)
registerSystemWSRoutes(adminOnly, c)
}
}
@@ -256,6 +257,10 @@ func registerAppLogRoutes(g *gin.RouterGroup, c *Controllers) {
}
}
func registerSystemWSRoutes(g *gin.RouterGroup, c *Controllers) {
g.GET("/ws/events", c.SystemWS.HandleEvents)
}
func initAgentAPIRoutes(root *gin.RouterGroup, c *Controllers) {
// Agent API（供远程 Agent 调用，不使用 /v1 版本号）
agentAPI := root.Group("/api/agent")
internal/router/register.go
+3 -1
View File
@@ -23,6 +23,7 @@ func RegisterControllers() *Controllers {
scriptService := services.NewScriptService()
sendStatsService := services.NewSendStatsService()
agentWSManager := services.GetAgentWSManager()
systemWSManager := services.GetSystemWSManager()
taskLogService := tasks.NewTaskLogService(sendStatsService)
// 创建任务执行服务（需要依赖注入）
@@ -40,7 +41,7 @@ func RegisterControllers() *Controllers {
executorService.StartCron()
// 初始化所有关注系统总线的服务
setupEventHandlers(appLogService, notifyService, loginLogService)
setupEventHandlers(appLogService, notifyService, loginLogService, systemWSManager)
go startAppLogCleanup(appLogService)
// 初始化并返回控制器
@@ -61,6 +62,7 @@ func RegisterControllers() *Controllers {
Mise: controllers.NewMiseController(services.NewMiseService()),
Notification: controllers.NewNotificationController(),
AppLog: controllers.NewAppLogController(),
SystemWS: controllers.NewSystemWSController(),
}
}
internal/router/router.go
+1
View File
@@ -29,6 +29,7 @@ type Controllers struct {
Mise *controllers.MiseController
Notification *controllers.NotificationController
AppLog *controllers.AppLogController
SystemWS *controllers.SystemWSController
}
func Setup(c *Controllers) *gin.Engine {
internal/services/system_ws_service.go
+126
View File
@@ -0,0 +1,126 @@
package services
import (
"encoding/json"
"sync"
"time"
"github.com/engigu/baihu-panel/internal/constant"
"github.com/engigu/baihu-panel/internal/eventbus"
"github.com/engigu/baihu-panel/internal/logger"
"github.com/engigu/baihu-panel/internal/models/vo"
"github.com/gorilla/websocket"
)
// SystemWSManager 前端系统事件 WebSocket 管理器 (单例)
type SystemWSManager struct {
clients map[*ClientConnection]bool
mu sync.RWMutex
}
// ClientConnection 代表一个前端页面的 WebSocket 连接
type ClientConnection struct {
Conn *websocket.Conn
Send chan []byte
closed bool
mu sync.Mutex
}
var systemWSManager *SystemWSManager
var systemWSOnce sync.Once
// GetSystemWSManager 获取系统 WebSocket 管理器单例
func GetSystemWSManager() *SystemWSManager {
systemWSOnce.Do(func() {
systemWSManager = &SystemWSManager{
clients: make(map[*ClientConnection]bool),
}
})
return systemWSManager
}
// Register 注册一个新的前端连接
func (m *SystemWSManager) Register(conn *websocket.Conn) *ClientConnection {
m.mu.Lock()
defer m.mu.Unlock()
client := &ClientConnection{
Conn: conn,
Send: make(chan []byte, 256),
}
m.clients[client] = true
return client
}
// Unregister 注销一个前端连接
func (m *SystemWSManager) Unregister(client *ClientConnection) {
m.mu.Lock()
defer m.mu.Unlock()
if _, ok := m.clients[client]; ok {
delete(m.clients, client)
client.Close()
}
}
// Broadcast 广播消息给所有在线前端
func (m *SystemWSManager) Broadcast(msgType string, payload interface{}) {
msg := vo.WSMessage{
Type: msgType,
Timestamp: time.Now().UnixMilli(),
Payload: payload,
}
data, err := json.Marshal(msg)
if err != nil {
logger.Errorf("[SystemWS] 序列化消息失败: %v", err)
return
}
m.mu.RLock()
defer m.mu.RUnlock()
for client := range m.clients {
select {
case client.Send <- data:
default:
// 缓冲区满，可能该客户端连接已死
go m.Unregister(client)
}
}
}
// SubscribeEvents 订阅系统事件总线并分发给 WebSocket
func (m *SystemWSManager) SubscribeEvents(bus *eventbus.EventBus) {
// 任务相关事件
taskEvents := []string{
constant.EventTaskSuccess,
constant.EventTaskFailed,
constant.EventTaskTimeout,
constant.EventTaskRunning,
constant.EventTaskQueued,
}
for _, evt := range taskEvents {
bus.Subscribe(evt, func(e eventbus.Event) {
m.Broadcast(e.Type, e.Payload)
})
}
// 系统通知事件
bus.Subscribe(constant.EventSystemNotice, func(e eventbus.Event) {
m.Broadcast("notice", e.Payload)
})
}
func (c *ClientConnection) Close() {
c.mu.Lock()
defer c.mu.Unlock()
if c.closed {
return
}
c.closed = true
c.Conn.Close()
close(c.Send)
}
internal/services/tasks/executor_service.go
+19 -1
View File
@@ -116,7 +116,15 @@ type ServerSchedulerHandler struct {
}
func (h *ServerSchedulerHandler) OnTaskScheduled(req *executor.ExecutionRequest) {
// 任务入队事件，可以在此处更新数据库状态为 "pending"
if req.TaskID != "" {
eventbus.DefaultBus.Publish(eventbus.Event{
Type: constant.EventTaskQueued,
Payload: map[string]interface{}{
"task_id": req.TaskID,
"status": constant.TaskStatusQueued,
},
})
}
}
func (h *ServerSchedulerHandler) OnTaskExecuting(req *executor.ExecutionRequest) (io.Writer, io.Writer, error) {
@@ -167,6 +175,16 @@ func (h *ServerSchedulerHandler) OnTaskExecuting(req *executor.ExecutionRequest)
StartTime: time.Now(),
})
// 发布任务开始运行事件
eventbus.DefaultBus.Publish(eventbus.Event{
Type: constant.EventTaskRunning,
Payload: map[string]interface{}{
"task_id": req.TaskID,
"status": constant.TaskStatusRunning,
"log_id": req.LogID,
},
})
if req.Metadata.RetryIndex > 0 {
tl.Write([]byte(fmt.Sprintf("\n[System] 此为任务失败后的第 %d 次重试执行...\n\n", req.Metadata.RetryIndex)))
}