19904408334
6166c24412
重新启用任务并发限制逻辑，使用 Redis 计数器控制模型并发数； 引入分布式锁防止重复创建，并优化优雅退出顺序确保请求完成； 移除 video_duration 计费计算，支持外部传入 taskId，为轮询查询添加独立超时上下文。
413 lines
12 KiB
Go
413 lines
12 KiB
Go
package task
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"model-gateway/common/util"
|
|
"model-gateway/consts/public"
|
|
"time"
|
|
|
|
"model-gateway/dao"
|
|
"model-gateway/model/dto"
|
|
"model-gateway/model/entity"
|
|
|
|
"gitea.redpowerfuture.com/red-future/common/beans"
|
|
"gitea.redpowerfuture.com/red-future/common/utils"
|
|
"github.com/gogf/gf/v2/database/gdb"
|
|
"github.com/gogf/gf/v2/database/gredis"
|
|
"github.com/gogf/gf/v2/errors/gerror"
|
|
"github.com/gogf/gf/v2/frame/g"
|
|
"github.com/gogf/gf/v2/os/glog"
|
|
"github.com/gogf/gf/v2/util/gconv"
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
var ModelGatewayTask = &taskService{}
|
|
|
|
type taskService struct{}
|
|
|
|
// Create 创建任务
|
|
func (s *taskService) Create(ctx context.Context, req *dto.CreateTaskReq) (res *dto.CreateTaskRes, err error) {
|
|
taskID := req.TaskId
|
|
if taskID == "" {
|
|
taskID = uuid.NewString()
|
|
}
|
|
startAt := time.Now()
|
|
|
|
// 1) 获取用户信息
|
|
userInfo, err := utils.GetUserInfo(ctx)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// 2) 检查模型配置
|
|
model, err := dao.ModelGatewayModels.Get(ctx, &entity.ModelGatewayModel{
|
|
SQLBaseDO: beans.SQLBaseDO{
|
|
TenantId: userInfo.TenantId,
|
|
Creator: userInfo.UserName,
|
|
},
|
|
ModelName: req.ModelName,
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if model == nil || (model.Enabled != nil && *model.Enabled != 1) {
|
|
return nil, errors.New("模型不存在或未启用")
|
|
}
|
|
lockKey := fmt.Sprintf("lock:tenantId-%s:model-%s", gconv.String(userInfo.TenantId), req.ModelName)
|
|
success, e := Lock(ctx, lockKey, -1, int64(time.Minute.Seconds()*5), func(ctx context.Context) error {
|
|
const (
|
|
keyExpireSec = 600 // 计数Key兜底过期时间 10min
|
|
waitInterval = 10 * time.Second // 轮询等待间隔
|
|
)
|
|
// Redis 操作统一使用独立上下文，避免外部 ctx canceled
|
|
redisCtx := context.WithoutCancel(ctx)
|
|
// 模型并发计数Key
|
|
concurrencyKey := fmt.Sprintf("model:concurrency:%s", req.ModelName)
|
|
maxCon := gconv.Int64(model.MaxConcurrency)
|
|
|
|
// 循环尝试获取并发名额，超限则等待重试
|
|
var held bool // 标记当前是否持有未释放的计数
|
|
for {
|
|
// 检测全局上下文取消
|
|
if ctx.Err() != nil {
|
|
if held {
|
|
g.Redis().Decr(redisCtx, concurrencyKey)
|
|
}
|
|
return ctx.Err()
|
|
}
|
|
// 计数自增
|
|
currentCon, e := g.Redis().Incr(redisCtx, concurrencyKey)
|
|
if e != nil {
|
|
if held {
|
|
g.Redis().Decr(redisCtx, concurrencyKey)
|
|
}
|
|
glog.Errorf(ctx, "redis incr concurrency key err: %v", e)
|
|
return e
|
|
}
|
|
held = true
|
|
// 首次创建Key时设置过期时间（避免重复执行EXPIRE）
|
|
exists, errr := g.Redis().Exists(redisCtx, concurrencyKey)
|
|
if errr == nil && exists == 1 {
|
|
g.Redis().Expire(redisCtx, concurrencyKey, keyExpireSec)
|
|
}
|
|
|
|
// 未超限：跳出循环，执行业务
|
|
if currentCon <= maxCon {
|
|
glog.Infof(ctx, "并发数: %s %d/%d", concurrencyKey, currentCon, maxCon)
|
|
break
|
|
}
|
|
// 超限立刻回减，撤销本次计数
|
|
g.Redis().Decr(redisCtx, concurrencyKey)
|
|
held = false
|
|
glog.Infof(ctx, "并发超限等待: %s %d/%d", concurrencyKey, currentCon, maxCon)
|
|
time.Sleep(waitInterval)
|
|
}
|
|
|
|
// 3) 构建任务实体
|
|
task := &entity.ModelGatewayTask{
|
|
ModelName: model.ModelName,
|
|
TaskID: taskID,
|
|
State: public.TaskStatusRunning,
|
|
BizName: req.BizName,
|
|
CallbackURL: req.CallbackUrl,
|
|
RequestPayload: &entity.RequestPayload{
|
|
Body: req.RequestPayload,
|
|
Headers: util.ParseHeadMsgHeaders(model.HeadMsg),
|
|
},
|
|
EpicycleId: req.EpicycleId,
|
|
BuildModelName: req.BuildModelName,
|
|
}
|
|
|
|
// 4) 插入任务记录
|
|
id, errr := dao.ModelGatewayTask.Insert(ctx, task)
|
|
if errr != nil {
|
|
g.Redis().Decr(redisCtx, concurrencyKey)
|
|
// TODO: 恢复排队逻辑后，此处需要回滚排队占位
|
|
//queue.ReleaseQueueSlot(ctx, req.ModelName, taskID)
|
|
return errr
|
|
}
|
|
task.Id = id
|
|
|
|
// 5) 记录操作日志（非关键路径，失败不影响主流程）
|
|
ip, ua := "", ""
|
|
if r := g.RequestFromCtx(ctx); r != nil {
|
|
ip = utils.GetLocalIP()
|
|
ua = r.UserAgent()
|
|
}
|
|
_, _ = dao.ModelGatewayLogsOp.Insert(ctx, &entity.ModelGatewayLogsOp{
|
|
IP: ip,
|
|
UserAgent: ua,
|
|
APIPath: "/task/createTask",
|
|
HttpMethod: "POST",
|
|
BizName: req.BizName,
|
|
ModelName: req.ModelName,
|
|
TaskID: taskID,
|
|
OpType: "createTask",
|
|
Success: 1,
|
|
CostMs: time.Since(startAt).Milliseconds(),
|
|
RequestPayload: task.RequestPayload,
|
|
ResponsePayload: gdb.Map{"taskId": taskID},
|
|
})
|
|
|
|
// 6) 模型计费
|
|
if len(model.BillingConfig) > 0 {
|
|
requestData := util.ExtractRequestBilling(ctx, model.BillingConfig, req.RequestPayload)
|
|
// 请求数据作为计费记录的基础字段，先存入数组
|
|
task.BillingData = append(task.BillingData, requestData)
|
|
_, _ = dao.ModelGatewayTask.Update(ctx, &entity.ModelGatewayTask{
|
|
SQLBaseDO: beans.SQLBaseDO{Id: task.Id},
|
|
BillingData: task.BillingData,
|
|
})
|
|
}
|
|
|
|
// 7) 异步执行任务
|
|
go AsyncWorker.handleOne(util.AsyncCtx(ctx), task, model, req)
|
|
|
|
return nil
|
|
})
|
|
if e != nil {
|
|
err = e
|
|
return
|
|
}
|
|
if !success {
|
|
err = gerror.New("任务排队已满，请稍后再试")
|
|
return
|
|
}
|
|
|
|
return &dto.CreateTaskRes{TaskID: taskID}, nil
|
|
}
|
|
|
|
// Lock 分布式锁 纯原生命令、无Lua、隔离上下文防 context canceled
|
|
func Lock(ctx context.Context, key string, limit, expireSeconds int64, fn func(ctx context.Context) error) (success bool, err error) {
|
|
if limit <= 0 {
|
|
limit = -1
|
|
}
|
|
|
|
// 过期时间合法校验（单位：秒）
|
|
const maxExpireSec = 86400 * 7
|
|
if expireSeconds < 1 || expireSeconds > maxExpireSec {
|
|
glog.Warningf(ctx, "锁过期时间非法，原值:%d，兜底为60秒", expireSeconds)
|
|
expireSeconds = 60
|
|
}
|
|
|
|
lockVal := "1"
|
|
|
|
LOOP:
|
|
// 检测父级上下文取消，防止无限重试阻塞 goroutine
|
|
if ctx.Err() != nil {
|
|
return false, ctx.Err()
|
|
}
|
|
if limit != -1 {
|
|
if limit < 0 {
|
|
return false, errors.New("锁重试次数耗尽，获取锁失败")
|
|
}
|
|
limit--
|
|
}
|
|
|
|
// 核心：创建独立上下文，不受外部 ctx 取消影响
|
|
redisCtx := context.WithoutCancel(ctx)
|
|
|
|
// 加锁
|
|
val, err := g.Redis().Set(redisCtx, key, lockVal, gredis.SetOption{
|
|
TTLOption: gredis.TTLOption{
|
|
EX: &expireSeconds,
|
|
},
|
|
NX: true,
|
|
})
|
|
if err != nil {
|
|
glog.Errorf(ctx, "redis set lock failed: %v", err)
|
|
time.Sleep(time.Second)
|
|
goto LOOP
|
|
}
|
|
|
|
if val.Bool() {
|
|
// 执行业务逻辑（使用原上下文）
|
|
runErr := fn(ctx)
|
|
|
|
// 释放锁：同样使用独立上下文 + 先GET再DEL防误删
|
|
getRes, err := g.Redis().Get(redisCtx, key)
|
|
if err != nil {
|
|
glog.Errorf(ctx, "redis get lock value failed: %v", err)
|
|
} else if getRes.String() == lockVal {
|
|
_, delErr := g.Redis().Del(redisCtx, key)
|
|
if delErr != nil {
|
|
glog.Errorf(ctx, "redis del lock failed: %v", delErr)
|
|
}
|
|
}
|
|
|
|
return true, runErr
|
|
}
|
|
|
|
// 抢锁失败，休眠重试
|
|
time.Sleep(time.Second)
|
|
goto LOOP
|
|
}
|
|
|
|
// GetResult 获取任务结果
|
|
func (s *taskService) GetResult(ctx context.Context, taskID string) (res *dto.GetTaskResultRes, err error) {
|
|
t, err := dao.ModelGatewayTask.Get(ctx, &entity.ModelGatewayTask{
|
|
TaskID: taskID,
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if t == nil {
|
|
return nil, errors.New("任务不存在")
|
|
}
|
|
return &dto.GetTaskResultRes{
|
|
OssFile: t.ResultFile.OssFile,
|
|
State: t.State,
|
|
}, nil
|
|
}
|
|
|
|
// GetBatch 批量查询任务；将成功(state=2)的任务更新为已下载(state=4)，并写入过期时间
|
|
func (s *taskService) GetBatch(ctx context.Context, req *dto.GetTaskBatchReq) (res *dto.GetTaskBatchRes, err error) {
|
|
if req == nil || len(req.TaskIDs) == 0 {
|
|
return &dto.GetTaskBatchRes{List: []dto.GetTaskBatchItem{}}, nil
|
|
}
|
|
// 1) 先查当前租户下的任务列表
|
|
list, err := dao.ModelGatewayTask.ListByTaskIDs(ctx, req.TaskIDs)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// 2) 对成功(state=2)的任务：标记为已下载(state=4)
|
|
for _, t := range list {
|
|
if t == nil {
|
|
continue
|
|
}
|
|
if t.State != public.BuildTypeNode {
|
|
continue
|
|
}
|
|
_ = dao.ModelGatewayTask.MarkDownloadedByID(ctx, t.Id)
|
|
|
|
// 为了本次返回一致性，内存里也更新
|
|
t.State = public.TaskStatusDownloaded
|
|
}
|
|
|
|
// 3) 组装返回
|
|
items := make([]dto.GetTaskBatchItem, 0, len(list))
|
|
for _, t := range list {
|
|
if t == nil {
|
|
continue
|
|
}
|
|
items = append(items, dto.GetTaskBatchItem{
|
|
TaskID: t.TaskID,
|
|
State: t.State,
|
|
OssFile: t.ResultFile.OssFile,
|
|
TextResult: t.TextResult,
|
|
})
|
|
}
|
|
return &dto.GetTaskBatchRes{List: items}, nil
|
|
}
|
|
|
|
// List 获取任务列表
|
|
func (s *taskService) List(ctx context.Context, req *dto.ListTaskReq) (*dto.ListTaskRes, error) {
|
|
if req.PageNum <= 0 {
|
|
req.PageNum = 1
|
|
}
|
|
if req.PageSize <= 0 {
|
|
req.PageSize = 10
|
|
}
|
|
user, err := utils.GetUserInfo(ctx)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
list, total, err := dao.ModelGatewayTask.List(ctx, req.PageNum, req.PageSize, &entity.ModelGatewayTask{
|
|
SQLBaseDO: beans.SQLBaseDO{
|
|
Creator: user.UserName,
|
|
},
|
|
ModelName: req.ModelName,
|
|
BizName: req.BizName,
|
|
State: req.State,
|
|
TaskID: req.TaskID,
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &dto.ListTaskRes{List: list, Total: total}, nil
|
|
}
|
|
|
|
// ModelTaskCallback 模型异步任务的回调通知
|
|
func (s *taskService) ModelTaskCallback(ctx context.Context, req *dto.ModelTaskCallbackReq) (*dto.ModelTaskCallbackRes, error) {
|
|
g.Log().Infof(ctx, "[模型回调] 收到通知 taskID=%s status=%s", req.TaskID, req.Status)
|
|
// 1. 查本地任务
|
|
task, err := dao.ModelGatewayTask.Get(ctx, &entity.ModelGatewayTask{
|
|
TaskID: req.TaskID,
|
|
})
|
|
if err != nil || task == nil {
|
|
return nil, fmt.Errorf("任务不存在: %s", req.TaskID)
|
|
}
|
|
|
|
// 2. 成功：取 video_url 和 usage
|
|
if req.Status == "succeeded" {
|
|
result := map[string]any{
|
|
"video_url": req.Content["video_url"],
|
|
"usage": req.Usage,
|
|
}
|
|
NotifyAsyncResult(req.TaskID, result, nil)
|
|
return &dto.ModelTaskCallbackRes{Success: true}, nil
|
|
}
|
|
|
|
// 3. 失败/过期
|
|
if req.Status == "failed" || req.Status == "expired" {
|
|
NotifyAsyncResult(req.TaskID, nil, fmt.Errorf(req.Status))
|
|
return &dto.ModelTaskCallbackRes{Success: true}, nil
|
|
}
|
|
|
|
return &dto.ModelTaskCallbackRes{Success: true}, nil
|
|
}
|
|
|
|
// QueryPendingTasks 批量轮询进行中的异步任务
|
|
func (s *taskService) QueryPendingTasks(ctx context.Context, req *dto.QueryPendingTasksReq) (*dto.QueryPendingTasksRes, error) {
|
|
limit := req.Limit
|
|
if limit <= 0 {
|
|
limit = g.Cfg().MustGet(ctx, "asynch.queryPending.limit", 10).Int()
|
|
}
|
|
|
|
// 1. 查 state=1（执行中）的异步任务
|
|
tasks, err := dao.ModelGatewayTask.GetPendingAsyncTasks(ctx, limit)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// 2. 逐个查询
|
|
var results []dto.QueryTaskItem
|
|
for _, t := range tasks {
|
|
// 拿到模型配置
|
|
model, err := dao.ModelGatewayModels.GetByModelNameForTenant(ctx, t.TenantId, t.ModelName)
|
|
if err != nil || model == nil || model.QueryConfig == nil {
|
|
continue
|
|
}
|
|
// 每个任务使用独立的超时上下文，防止单个任务阻塞整个轮询
|
|
pullCtx, pullCancel := context.WithTimeout(ctx, 30*time.Second)
|
|
result, err := util.PullTaskResult(pullCtx, nil, model.QueryConfig, model.HeadMsg)
|
|
pullCancel()
|
|
if err != nil {
|
|
g.Log().Warningf(ctx, "[轮询] 查询失败 taskID=%s err=%v", t.TaskID, err)
|
|
continue
|
|
}
|
|
|
|
status := gconv.String(result["status"])
|
|
item := dto.QueryTaskItem{
|
|
TaskID: t.TaskID,
|
|
Status: status,
|
|
Content: result["content"].(map[string]any),
|
|
Usage: result["usage"].(map[string]any),
|
|
}
|
|
results = append(results, item)
|
|
|
|
// 如果任务完成，通知等待通道
|
|
if status == "succeeded" || status == "failed" || status == "expired" {
|
|
NotifyAsyncResult(t.TaskID, result["content"].(map[string]any), nil)
|
|
}
|
|
}
|
|
|
|
return &dto.QueryPendingTasksRes{
|
|
Total: len(results),
|
|
Results: results,
|
|
}, nil
|
|
}
|