red-future/model-gateway
4
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

refactor(task): 重构任务调度系统,新增 JobTask 定时任务功能,并优化配置结构

Browse Source
This commit is contained in:
WangLiZhao
2026-06-22 16:13:39 +08:00
parent 3ccce74465
commit 47b38f221c
7 changed files with 140 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
config.yml
+8 -13
View File
@@ -60,16 +60,11 @@ jaeger:
addr: 192.168.3.30:4318
# 本地调试用：可选自动执行 worker/cleaner（默认关闭）
asynch:
queryPending:
enabled: false
intervalSeconds: 10 # 每10秒轮询一次
limit: 10 # 每次查10条
worker:
enabled: false
intervalSeconds: 5
batchSize: 10
goroutines: 1
cleaner:
enabled: false
intervalSeconds: 30
queryPending:
enabled: false
intervalSeconds: 10 # 每10秒轮询一次
limit: 10 # 每次查10条
jobTask:
intervalSeconds: 10 # 轮询间隔(秒)
batchSize: 10 # 每批处理条数
poolSize: 5 # 协程池大小
dao/model_gateway_task_dao.go
+17
View File
@@ -103,6 +103,23 @@ func (d *modelGatewayTaskDao) ListByTaskIDs(ctx context.Context, taskIDs []strin
return
}
// ListPending 查询待处理任务
func (d *modelGatewayTaskDao) ListPending(ctx context.Context, limit int) (list []*entity.ModelGatewayTask, err error) {
if limit <= 0 {
limit = 10
}
model := gfdb.DB(ctx, public.DbNameModelGateway).Model(ctx, public.TableNameTask).
Where(entity.ModelGatewayTaskCol.State, 0).
OrderAsc(entity.ModelGatewayTaskCol.CreatedAt).
Limit(limit)
r, err := model.All()
if err != nil {
return nil, err
}
err = r.Structs(&list)
return
}
// MarkDownloadedByID 标记已下载
func (d *modelGatewayTaskDao) MarkDownloadedByID(ctx context.Context, id int64) error {
_, err := gfdb.DB(ctx, public.DbNameModelGateway).Model(ctx, public.TableNameTask).
main.go
+11 -7
View File
@@ -17,6 +17,7 @@ import (
_ "github.com/gogf/gf/contrib/drivers/pgsql/v2"
_ "github.com/gogf/gf/contrib/nosql/redis/v2"
"github.com/gogf/gf/v2/frame/g"
"github.com/gogf/gf/v2/os/grpool"
)
func main() {
@@ -24,6 +25,11 @@ func main() {
defer cancel()
defer jaeger.ShutDown(ctx)
// 初始化全局协程池
poolSize := g.Cfg().MustGet(ctx, "jobTask.poolSize", 5).Int()
task.JobPool = grpool.New(poolSize)
defer task.JobPool.Close()
// 注册路由
http.RouteRegister([]interface{}{
controller.ModelGatewayModels,
@@ -31,25 +37,23 @@ func main() {
controller.ModelGatewayLogsStat,
})
// 本地调试：可选自动触发 worker/cleaner（由配置文件控制）
// 本地调试：可选自动触发 worker/cleaner
startAutoRunner(ctx)
// 监听退出信号，确保 Ctrl+C 能完整退出（停止 worker/cleaner 并关闭 gateway server）
// 监听退出信号
quit := make(chan os.Signal, 1)
signal.Notify(quit, os.Interrupt, syscall.SIGTERM)
<-quit
g.Log().Infof(ctx, "[main] 收到退出信号，开始优雅退出...")
cancel()
// 关闭 gateway server（RouteRegister 内部是 go Httpserver.Run() 启动的）
_ = http.Httpserver.Shutdown()
}
func startAutoRunner(ctx context.Context) {
// queryPending
if g.Cfg().MustGet(ctx, "asynch.queryPending.enabled").Bool() {
interval := g.Cfg().MustGet(ctx, "asynch.queryPending.intervalSeconds", 10).Int()
limit := g.Cfg().MustGet(ctx, "asynch.queryPending.limit", 10).Int()
if g.Cfg().MustGet(ctx, "queryPending.enabled").Bool() {
interval := g.Cfg().MustGet(ctx, "queryPending.intervalSeconds", 10).Int()
limit := g.Cfg().MustGet(ctx, "queryPending.limit", 10).Int()
ticker := time.NewTicker(time.Duration(interval) * time.Second)
go func() {
defer ticker.Stop()
model/dto/model_gateway_task_dto.go
+1
View File
@@ -22,6 +22,7 @@ type JobTaskReq struct {
g.Meta `path:"/jobTask" method:"post" tags:"任务管理" summary:"定时任务" dc:"循环执行待处理任务，按间隔时间和批次大小处理"`
Interval int `json:"interval" dc:"循环间隔(秒)"`
BatchSize int `json:"batchSize" dc:"每批执行条数"`
PoolSize int `json:"poolSize" dc:"协程池大小"`
}
type JobTaskRes struct {
model/entity/model_gateway_task.go
+3
View File
@@ -16,6 +16,7 @@ type modelGatewayTaskCol struct {
ExpendTokens string
DurationSeconds string
RetryCount string
BuildType string
RequestPayload string
EpicycleId string
}
@@ -32,6 +33,7 @@ var ModelGatewayTaskCol = modelGatewayTaskCol{
ExpendTokens: "expend_tokens",
DurationSeconds: "duration_seconds",
RetryCount: "retry_count",
BuildType: "build_type",
RequestPayload: "request_payload",
EpicycleId: "epicycle_id",
}
@@ -44,6 +46,7 @@ type ModelGatewayTask struct {
BizName string `orm:"biz_name" json:"bizName"`
CallbackURL string `orm:"callback_url" json:"callbackUrl"`
State int `orm:"state" json:"state"`
BuildType int64 `orm:"build_type" json:"buildType"`
ErrorMsg string `orm:"error_msg" json:"errorMsg"`
ResultFile *ResultFile `orm:"result_file" json:"resultFile"`
ExpendTokens int64 `orm:"expend_tokens" json:"expendTokens"`
service/task/task_service.go
+92 -2
View File
@@ -6,6 +6,7 @@ import (
"fmt"
"model-gateway/common/util"
"model-gateway/consts/public"
"sync"
"time"
"model-gateway/dao"
@@ -16,6 +17,7 @@ import (
"gitea.redpowerfuture.com/red-future/common/utils"
"github.com/gogf/gf/v2/database/gdb"
"github.com/gogf/gf/v2/frame/g"
"github.com/gogf/gf/v2/os/grpool"
"github.com/gogf/gf/v2/util/gconv"
"github.com/google/uuid"
)
@@ -71,6 +73,7 @@ func (s *taskService) Create(ctx context.Context, req *dto.CreateTaskReq) (res *
CallbackURL: req.CallbackUrl,
RequestPayload: req.RequestPayload,
EpicycleId: req.EpicycleId,
BuildType: req.BuildType,
}
// 4) 插入任务记录
@@ -104,15 +107,102 @@ func (s *taskService) Create(ctx context.Context, req *dto.CreateTaskReq) (res *
})
// 6) 异步执行任务
go AsyncWorker.handleOne(util.AsyncCtx(ctx), task, model, req)
go AsyncWorker.handleOne(util.AsyncCtx(ctx), task, model)
return &dto.CreateTaskRes{TaskID: taskID}, nil
}
var JobPool *grpool.Pool
// JobTask 定时任务：循环执行待处理任务
func (s *taskService) JobTask(ctx context.Context, req *dto.JobTaskReq) (res *dto.JobTaskRes, err error) {
// 1) 参数默认值从配置取
if req.Interval <= 0 {
req.Interval = g.Cfg().MustGet(ctx, "jobTask.intervalSeconds", 5).Int()
}
if req.BatchSize <= 0 {
req.BatchSize = g.Cfg().MustGet(ctx, "jobTask.batchSize", 10).Int()
}
return nil, err
var (
totalProcessed int
successCount int
failCount int
mu sync.Mutex
wg sync.WaitGroup
)
// 2) 循环查询待处理任务
for {
select {
case <-ctx.Done():
wg.Wait()
return &dto.JobTaskRes{
TotalProcessed: totalProcessed,
SuccessCount: successCount,
FailCount: failCount,
}, nil
default:
}
// 3) 查询 state=0 的任务列表
tasks, err := dao.ModelGatewayTask.ListPending(ctx, req.BatchSize)
if err != nil {
g.Log().Warningf(ctx, "[定时任务] 查询任务失败: %v", err)
time.Sleep(time.Second)
continue
}
if len(tasks) == 0 {
time.Sleep(time.Duration(req.Interval) * time.Second)
continue
}
// 4) 提交到全局协程池执行
for _, task := range tasks {
wg.Add(1)
t := task
err = JobPool.Add(ctx, func(ctx context.Context) {
defer wg.Done()
mu.Lock()
totalProcessed++
mu.Unlock()
if execErr := s.executeTask(ctx, t); execErr != nil {
mu.Lock()
failCount++
mu.Unlock()
g.Log().Errorf(ctx, "[定时任务] 执行失败 taskId=%s err=%v", t.TaskID, execErr)
} else {
mu.Lock()
successCount++
mu.Unlock()
}
})
if err != nil {
return nil, err
}
}
}
}
// executeTask 执行单个任务
func (s *taskService) executeTask(ctx context.Context, task *entity.ModelGatewayTask) error {
// 1) 查询模型配置
model, err := dao.ModelGatewayModels.Get(ctx, &entity.ModelGatewayModel{
ModelName: task.ModelName,
})
if err != nil {
return fmt.Errorf("查询模型配置失败: %w", err)
}
if model == nil || (model.Enabled != nil && *model.Enabled != 1) {
return fmt.Errorf("模型不存在或未启用: %s", task.ModelName)
}
// 3) 调用 handleOne
AsyncWorker.handleOne(ctx, task, model)
return nil
}
// GetResult 获取任务结果
service/task/worker.go
+8 -9
View File
@@ -6,7 +6,6 @@ import (
"encoding/json"
"fmt"
"io"
"model-gateway/model/dto"
"net/http"
"strings"
"sync"
@@ -31,7 +30,7 @@ type asyncWorker struct {
}
// handleOne 执行一次完整的任务
func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTask, model *entity.ModelGatewayModel, req *dto.CreateTaskReq) {
func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTask, model *entity.ModelGatewayModel) {
var (
body = task.RequestPayload
maxRetry = model.RetryTimes
@@ -59,12 +58,12 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa
result, err = util.ParseStreamResponse(rawBytes, model.StreamConfig)
}
case model.CallMode != nil && *model.CallMode == public.CallModeAsync:
result, err = w.callModel(ctx, task, model, body)
result, err = w.callModel(ctx, model, body)
if err == nil {
result, err = util.PullTaskResult(ctx, result, model.QueryConfig, model.HeadMsg)
}
default:
result, err = w.callModel(ctx, task, model, body)
result, err = w.callModel(ctx, model, body)
}
if err == nil {
@@ -83,7 +82,7 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa
// ============================================
// 2) 解析校验 + 响应映射（可重试）
// ============================================
result, err = w.parseAndRetry(ctx, result, task, model, maxRetry, req)
result, err = w.parseAndRetry(ctx, result, task, model, maxRetry)
if err != nil {
w.failTask(ctx, task, startTime, err.Error())
return
@@ -144,7 +143,7 @@ var asyncTaskChan = sync.Map{} // taskID → chan asyncResult
func (w *asyncWorker) callModelAsync(ctx context.Context, task *entity.ModelGatewayTask, model *entity.ModelGatewayModel, body map[string]any) (map[string]any, error) {
// 1. 提交异步任务
body, err := w.callModel(ctx, task, model, body)
body, err := w.callModel(ctx, model, body)
if err != nil {
return nil, err
}
@@ -186,7 +185,7 @@ func NotifyAsyncResult(taskID string, result map[string]any, err error) {
}
// callModel 调用模型 + 提取文本结果
func (w *asyncWorker) callModel(ctx context.Context, task *entity.ModelGatewayTask, model *entity.ModelGatewayModel, body map[string]any) (map[string]any, error) {
func (w *asyncWorker) callModel(ctx context.Context, model *entity.ModelGatewayModel, body map[string]any) (map[string]any, error) {
data, err := InvokeModel(ctx, model, body)
if err != nil {
return nil, err
@@ -206,7 +205,7 @@ func (w *asyncWorker) callModel(ctx context.Context, task *entity.ModelGatewayTa
}
// parseAndRetry 解析模型返回结果，并重试
func (w *asyncWorker) parseAndRetry(ctx context.Context, body map[string]any, task *entity.ModelGatewayTask, model *entity.ModelGatewayModel, maxRetry int, req *dto.CreateTaskReq) (map[string]any, error) {
func (w *asyncWorker) parseAndRetry(ctx context.Context, body map[string]any, task *entity.ModelGatewayTask, model *entity.ModelGatewayModel, maxRetry int) (map[string]any, error) {
var lastErr error
for attempt := 0; attempt <= maxRetry; attempt++ {
if attempt > 0 {
@@ -234,7 +233,7 @@ func (w *asyncWorker) parseAndRetry(ctx context.Context, body map[string]any, ta
// 3) 解析 + 校验
var parsed map[string]any
switch req.BuildType {
switch task.BuildType {
case public.BuildTypePrompt, public.BuildTypeNode:
parsed, err = util.ParseAndValidate(mapped, model)
if err == nil {