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

feat: implement agent custom scheduler config and queue model #101

Browse Source
This commit is contained in:
duorameng
2026-05-19 20:23:05 +08:00
parent b7fd473bc9
commit 0e1a95a3f8
9 changed files with 332 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/controllers/agent_controller.go
+63 -26
View File
@@ -45,6 +45,29 @@ func (c *AgentController) List(ctx *gin.Context) {
utils.Success(ctx, vo.ToAgentVOListFromModels(agents))
}
// getActiveSchedulerConfig 获取 Agent 的实际调度配置（若为空或零值，则使用系统默认的 settings）
func (c *AgentController) getActiveSchedulerConfig(agent *models.Agent) map[string]interface{} {
workerCount := agent.SchedulerConfig.WorkerCount
queueSize := agent.SchedulerConfig.QueueSize
rateInterval := int(agent.SchedulerConfig.RateInterval / time.Millisecond)
strictQueue := agent.SchedulerConfig.StrictQueue
// 如果未配置（WorkerCount <= 0），则使用全局系统设置
if workerCount <= 0 {
workerCount = getIntSetting(c.settingsService, constant.SectionScheduler, constant.KeyWorkerCount, 4)
queueSize = getIntSetting(c.settingsService, constant.SectionScheduler, constant.KeyQueueSize, 100)
rateInterval = getIntSetting(c.settingsService, constant.SectionScheduler, constant.KeyRateInterval, 200)
strictQueue = false
}
return map[string]interface{}{
"worker_count": workerCount,
"queue_size": queueSize,
"rate_interval": rateInterval,
"strict_queue": strictQueue,
}
}
// Update 更新 Agent
func (c *AgentController) Update(ctx *gin.Context) {
id := ctx.Param("id")
@@ -54,16 +77,17 @@ func (c *AgentController) Update(ctx *gin.Context) {
}
var req struct {
Name string `json:"name" binding:"required"`
Description string `json:"description"`
Enabled bool `json:"enabled"`
Name string `json:"name" binding:"required"`
Description string `json:"description"`
Enabled bool `json:"enabled"`
SchedulerConfig *vo.AgentSchedulerConfigVO `json:"scheduler_config"`
}
if err := ctx.ShouldBindJSON(&req); err != nil {
utils.BadRequest(ctx, "参数错误")
return
}
// 获取旧状态
oldAgent := c.agentService.GetByID(id)
if oldAgent == nil {
@@ -71,12 +95,21 @@ func (c *AgentController) Update(ctx *gin.Context) {
return
}
wasEnabled := utils.DerefBool(oldAgent.Enabled, true)
var schedulerConfig models.AgentSchedulerConfig
if req.SchedulerConfig != nil {
schedulerConfig.WorkerCount = req.SchedulerConfig.WorkerCount
schedulerConfig.QueueSize = req.SchedulerConfig.QueueSize
schedulerConfig.RateInterval = time.Duration(req.SchedulerConfig.RateInterval) * time.Millisecond
schedulerConfig.Verbose = req.SchedulerConfig.Verbose
schedulerConfig.StrictQueue = req.SchedulerConfig.StrictQueue
}
if err := c.agentService.Update(id, req.Name, req.Description, req.Enabled); err != nil {
if err := c.agentService.Update(id, req.Name, req.Description, req.Enabled, schedulerConfig); err != nil {
utils.ServerError(ctx, err.Error())
return
}
// 如果启用状态发生变化，通知 Agent
if wasEnabled != req.Enabled {
if req.Enabled {
@@ -93,7 +126,20 @@ func (c *AgentController) Update(ctx *gin.Context) {
})
}
}
// 推送最新的调度配置给 Agent (如果 Agent 在线)
if req.Enabled {
// 重新加载已更新的 Agent 信息以获取正确的 SchedulerConfig
updatedAgent := c.agentService.GetByID(id)
if updatedAgent != nil {
c.wsManager.SendToAgent(id, services.WSTypeConnected, map[string]interface{}{
"agent_id": id,
"name": req.Name,
"scheduler_config": c.getActiveSchedulerConfig(updatedAgent),
})
}
}
utils.SuccessMsg(ctx, "更新成功")
}
@@ -415,26 +461,17 @@ func (c *AgentController) WSConnect(ctx *gin.Context) {
// 更新 Agent 状态
c.agentService.Heartbeat(token, ip, "", "", "", "", "")
// 获取调度配置
workerCount := getIntSetting(c.settingsService, constant.SectionScheduler, constant.KeyWorkerCount, 4)
queueSize := getIntSetting(c.settingsService, constant.SectionScheduler, constant.KeyQueueSize, 100)
rateInterval := getIntSetting(c.settingsService, constant.SectionScheduler, constant.KeyRateInterval, 200)
// 发送连接成功消息（包含注册状态和调度配置）
// 获取调度配置并发送连接成功消息（包含注册状态和调度配置）
schedCfg := c.getActiveSchedulerConfig(agent)
c.wsManager.SendToAgent(agent.ID, services.WSTypeConnected, map[string]interface{}{
"agent_id": agent.ID,
"name": agent.Name,
"is_new_agent": isNewAgent,
"machine_id": machineID,
"scheduler_config": map[string]interface{}{
"worker_count": workerCount,
"queue_size": queueSize,
"rate_interval": rateInterval,
},
"agent_id": agent.ID,
"name": agent.Name,
"is_new_agent": isNewAgent,
"machine_id": machineID,
"scheduler_config": schedCfg,
})
logger.Infof("[AgentWS] Agent #%s 连接成功 (配置: workers=%d, queue=%d, rate=%d)",
agent.ID, workerCount, queueSize, rateInterval)
logger.Infof("[AgentWS] Agent #%s 连接成功 (配置: %v)", agent.ID, schedCfg)
// 启动读写协程
go c.wsWritePump(ac)
internal/executor/scheduler.go
+13 -5
View File
@@ -37,6 +37,7 @@ type SchedulerConfig struct {
QueueSize int // 队列大小
RateInterval time.Duration // 速率限制间隔
Verbose bool // 是否开启详细日志
StrictQueue bool // 是否开启严格排队（满时拒绝执行，不降级直接执行）
}
// TaskType 任务类型
@@ -281,9 +282,16 @@ func (s *Scheduler) EnqueueOrExecute(req *ExecutionRequest) {
s.handler.OnTaskScheduled(req)
}
default:
// 队列满，直接执行（降级处理）
s.logger.Warnf("[Scheduler] 任务队列已满，直接执行任务 %s", req.TaskID)
go s.executeTask(req)
if s.config.StrictQueue {
s.logger.Errorf("[Scheduler] 任务队列已满，拒绝执行任务 %s", req.TaskID)
if s.handler != nil {
s.handler.OnTaskFailed(req, fmt.Errorf("任务队列已满，拒绝执行"))
}
} else {
// 队列满，直接执行（降级处理）
s.logger.Warnf("[Scheduler] 任务队列已满，直接执行任务 %s", req.TaskID)
go s.executeTask(req)
}
}
}
@@ -592,8 +600,8 @@ func (s *Scheduler) Reload(config SchedulerConfig) {
// 重启 workers
s.Start()
s.logger.Infof("[Scheduler] 配置已重载: workers=%d, queue=%d, rate=%v",
config.WorkerCount, config.QueueSize, config.RateInterval)
s.logger.Infof("[Scheduler] 配置已重载: workers=%d, queue=%d, rate=%v, strict=%t",
config.WorkerCount, config.QueueSize, config.RateInterval, config.StrictQueue)
}
// GetQueueSize 获取当前队列大小
internal/models/agent.go
+57 -17
View File
@@ -1,28 +1,68 @@
package models
import (
"database/sql/driver"
"encoding/json"
"errors"
"time"
"github.com/engigu/baihu-panel/internal/constant"
)
// AgentSchedulerConfig Agent 调度器配置
type AgentSchedulerConfig struct {
WorkerCount int `json:"worker_count"`
QueueSize int `json:"queue_size"`
RateInterval time.Duration `json:"rate_interval"`
Verbose bool `json:"verbose"`
StrictQueue bool `json:"strict_queue"`
}
// Value 序列化为数据库字符串
func (c AgentSchedulerConfig) Value() (driver.Value, error) {
bytes, err := json.Marshal(c)
if err != nil {
return nil, err
}
return string(bytes), nil
}
// Scan 反序列化数据库字符串为结构体
func (c *AgentSchedulerConfig) Scan(value interface{}) error {
if value == nil {
return nil
}
bytes, ok := value.([]byte)
if !ok {
str, ok := value.(string)
if !ok {
return errors.New("invalid type for AgentSchedulerConfig")
}
bytes = []byte(str)
}
return json.Unmarshal(bytes, c)
}
// Agent 远程执行代理
type Agent struct {
ID string `json:"id" gorm:"primaryKey;size:20"`
Name string `json:"name" gorm:"size:100;not null"` // Agent 名称
Token string `json:"token" gorm:"size:64;index"` // 认证 Token（可重复使用）
MachineID string `json:"machine_id" gorm:"size:64;uniqueIndex"` // 机器识别码（唯一）
Description string `json:"description" gorm:"size:255"` // 描述
Status string `json:"status" gorm:"size:20;default:'pending';index"` // 状态: constant.AgentStatusOnline, constant.AgentStatusOffline
LastSeen *LocalTime `json:"last_seen"` // 最后心跳时间
IP string `json:"ip" gorm:"size:45"` // Agent IP 地址
Version string `json:"version" gorm:"size:50"` // Agent 版本
BuildTime string `json:"build_time" gorm:"size:30"` // Agent 构建时间
Hostname string `json:"hostname" gorm:"size:100"` // Agent 主机名
OS string `json:"os" gorm:"size:20"` // 操作系统
Arch string `json:"arch" gorm:"size:20"` // 架构
ForceUpdate bool `json:"force_update" gorm:"default:false"` // 强制更新标志
Enabled *bool `json:"enabled" gorm:"default:true"` // 是否启用
CreatedAt LocalTime `json:"created_at"`
UpdatedAt LocalTime `json:"updated_at"`
ID string `json:"id" gorm:"primaryKey;size:20"`
Name string `json:"name" gorm:"size:100;not null"` // Agent 名称
Token string `json:"token" gorm:"size:64;index"` // 认证 Token（可重复使用）
MachineID string `json:"machine_id" gorm:"size:64;uniqueIndex"` // 机器识别码（唯一）
Description string `json:"description" gorm:"size:255"` // 描述
Status string `json:"status" gorm:"size:20;default:'pending';index"` // 状态: constant.AgentStatusOnline, constant.AgentStatusOffline
LastSeen *LocalTime `json:"last_seen"` // 最后心跳时间
IP string `json:"ip" gorm:"size:45"` // Agent IP 地址
Version string `json:"version" gorm:"size:50"` // Agent 版本
BuildTime string `json:"build_time" gorm:"size:30"` // Agent 构建时间
Hostname string `json:"hostname" gorm:"size:100"` // Agent 主机名
OS string `json:"os" gorm:"size:20"` // 操作系统
Arch string `json:"arch" gorm:"size:20"` // 架构
ForceUpdate bool `json:"force_update" gorm:"default:false"` // 强制更新标志
Enabled *bool `json:"enabled" gorm:"default:true"` // 是否启用
SchedulerConfig AgentSchedulerConfig `json:"scheduler_config" gorm:"type:text"` // 调度配置，以 JSON 字符串形式存储在 Text 类型字段中
CreatedAt LocalTime `json:"created_at"`
UpdatedAt LocalTime `json:"updated_at"`
}
func (Agent) TableName() string {
internal/models/vo/agent_vo.go
+31 -8
View File
@@ -1,6 +1,8 @@
package vo
import (
"time"
"github.com/engigu/baihu-panel/internal/models"
"github.com/engigu/baihu-panel/internal/utils"
)
@@ -18,10 +20,11 @@ type AgentVO struct {
Hostname string `json:"hostname"`
OS string `json:"os"`
Arch string `json:"arch"`
ForceUpdate bool `json:"force_update"`
Enabled bool `json:"enabled"`
CreatedAt models.LocalTime `json:"created_at"`
UpdatedAt models.LocalTime `json:"updated_at"`
ForceUpdate bool `json:"force_update"`
Enabled bool `json:"enabled"`
SchedulerConfig *AgentSchedulerConfigVO `json:"scheduler_config"`
CreatedAt models.LocalTime `json:"created_at"`
UpdatedAt models.LocalTime `json:"updated_at"`
// 隐藏 Token 和 MachineID
}
@@ -30,6 +33,16 @@ func ToAgentVO(agent *models.Agent) *AgentVO {
if agent == nil {
return nil
}
var schedulerConfigVO *AgentSchedulerConfigVO
if agent.SchedulerConfig.WorkerCount > 0 {
schedulerConfigVO = &AgentSchedulerConfigVO{
WorkerCount: agent.SchedulerConfig.WorkerCount,
QueueSize: agent.SchedulerConfig.QueueSize,
RateInterval: int(agent.SchedulerConfig.RateInterval / time.Millisecond),
Verbose: agent.SchedulerConfig.Verbose,
StrictQueue: agent.SchedulerConfig.StrictQueue,
}
}
return &AgentVO{
ID: agent.ID,
Name: agent.Name,
@@ -42,10 +55,11 @@ func ToAgentVO(agent *models.Agent) *AgentVO {
Hostname: agent.Hostname,
OS: agent.OS,
Arch: agent.Arch,
ForceUpdate: agent.ForceUpdate,
Enabled: utils.DerefBool(agent.Enabled, true),
CreatedAt: agent.CreatedAt,
UpdatedAt: agent.UpdatedAt,
ForceUpdate: agent.ForceUpdate,
Enabled: utils.DerefBool(agent.Enabled, true),
SchedulerConfig: schedulerConfigVO,
CreatedAt: agent.CreatedAt,
UpdatedAt: agent.UpdatedAt,
}
}
@@ -119,3 +133,12 @@ func ToAgentTokenVOListFromModels(tokens []models.AgentToken) []*AgentTokenVO {
}
return vos
}
// AgentSchedulerConfigVO 调度配置视图对象
type AgentSchedulerConfigVO struct {
WorkerCount int `json:"worker_count"`
QueueSize int `json:"queue_size"`
RateInterval int `json:"rate_interval"` // 毫秒
Verbose bool `json:"verbose"`
StrictQueue bool `json:"strict_queue"`
}
internal/services/agent_service.go
+5 -4
View File
@@ -203,11 +203,12 @@ func (s *AgentService) Register(req *models.AgentRegisterRequest, ip string) (*m
}
// Update 更新 Agent
func (s *AgentService) Update(id string, name, description string, enabled bool) error {
func (s *AgentService) Update(id string, name, description string, enabled bool, schedulerConfig models.AgentSchedulerConfig) error {
return database.DB.Model(&models.Agent{}).Where("id = ?", id).Updates(map[string]interface{}{
"name": name,
"description": description,
"enabled": &enabled,
"name": name,
"description": description,
"enabled": &enabled,
"scheduler_config": schedulerConfig,
}).Error
}
internal/services/tasks/executor_service.go
+22 -2
View File
@@ -937,7 +937,17 @@ func (es *ExecutorService) CheckConcurrency(taskID string) error {
}
if config.Concurrency == 0 && len(goids) > 0 {
return fmt.Errorf("任务正在运行中，拒绝并行执行，请前往日志查看")
// 检查目标 Agent 是否开启了排队机制
var isAgentQueueing bool
if task.AgentID != nil && *task.AgentID != "" {
var agent models.Agent
if err := database.DB.Select("scheduler_config").Where("id = ?", *task.AgentID).First(&agent).Error; err == nil {
isAgentQueueing = agent.SchedulerConfig.StrictQueue
}
}
if !isAgentQueueing {
return fmt.Errorf("任务正在运行中，拒绝并行执行，请前往日志查看")
}
}
return nil
}
@@ -969,7 +979,17 @@ func (es *ExecutorService) AddRunningGo(taskID string) (int64, error) {
// 如果并发为0(禁用)且已有执行中的任务，返回错误
if config.Concurrency == 0 && len(goids) > 0 {
return fmt.Errorf("task is running")
// 检查目标 Agent 是否开启了排队机制
var isAgentQueueing bool
if task.AgentID != nil && *task.AgentID != "" {
var agent models.Agent
if err := tx.Select("scheduler_config").Where("id = ?", *task.AgentID).First(&agent).Error; err == nil {
isAgentQueueing = agent.SchedulerConfig.StrictQueue
}
}
if !isAgentQueueing {
return fmt.Errorf("task is running")
}
}
goids = append(goids, goid)