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)