diff --git a/common/util/billing.go b/common/util/billing.go index 85e8eed..4d56344 100644 --- a/common/util/billing.go +++ b/common/util/billing.go @@ -4,7 +4,6 @@ import ( "context" "fmt" "math" - "model-gateway/service/gateway" "strings" "github.com/gogf/gf/v2/encoding/gjson" @@ -159,28 +158,28 @@ func ExtractRequestBilling(ctx context.Context, config map[string]any, requestPa } } - if compute, ok := config["compute"].(map[string]any); ok { - for targetField, rule := range compute { - r := rule.(map[string]any) - dependsOn := gconv.String(r["depends_on"]) - dependsValue := gconv.Bool(r["depends_value"]) - - if gconv.Bool(data[dependsOn]) != dependsValue { - continue - } - - switch r["service"] { - case "video_duration": - urls := extractVideoUrls(requestPayload) - if len(urls) > 0 { - resp, err := gateway.GetVideoDuration(ctx, urls) - if err == nil { - data[targetField] = resp.TotalDuration - } - } - } - } - } + //if compute, ok := config["compute"].(map[string]any); ok { + // for targetField, rule := range compute { + // r := rule.(map[string]any) + // dependsOn := gconv.String(r["depends_on"]) + // dependsValue := gconv.Bool(r["depends_value"]) + // + // if gconv.Bool(data[dependsOn]) != dependsValue { + // continue + // } + // + // switch r["service"] { + // case "video_duration": + // urls := extractVideoUrls(requestPayload) + // if len(urls) > 0 { + // resp, err := gateway.GetVideoDuration(ctx, urls) + // if err == nil { + // data[targetField] = resp.TotalDuration + // } + // } + // } + // } + //} return data } diff --git a/main.go b/main.go index d00e8fe..ff2821a 100644 --- a/main.go +++ b/main.go @@ -40,9 +40,10 @@ func main() { <-quit g.Log().Infof(ctx, "[main] 收到退出信号，开始优雅退出...") - cancel() - // 关闭 gateway server（RouteRegister 内部是 go Httpserver.Run() 启动的） + // 先关闭 gateway server，等待 in-flight 请求处理完成 _ = http.Httpserver.Shutdown() + // 再取消上下文，避免活跃请求被中断 + cancel() } func startAutoRunner(ctx context.Context) { diff --git a/model/dto/model_gateway_task_dto.go b/model/dto/model_gateway_task_dto.go index 710e39f..863aa03 100644 --- a/model/dto/model_gateway_task_dto.go +++ b/model/dto/model_gateway_task_dto.go @@ -14,6 +14,7 @@ type CreateTaskReq struct { EpicycleId int64 `json:"epicycleId" dc:"轮次ID"` BuildType int64 `json:"buildType" dc:"构建类型：1-提示词构建 2-节点构建"` BuildModelName string `json:"buildModelName" json:"buildModelName" dc:"构建模型名称"` + TaskId string `json:"taskId" dc:"任务ID"` } type CreateTaskRes struct { diff --git a/service/task/task_service.go b/service/task/task_service.go index 9b89032..a30b0e9 100644 --- a/service/task/task_service.go +++ b/service/task/task_service.go @@ -15,7 +15,10 @@ import ( "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" ) @@ -26,7 +29,10 @@ type taskService struct{} // Create 创建任务 func (s *taskService) Create(ctx context.Context, req *dto.CreateTaskReq) (res *dto.CreateTaskRes, err error) { - taskID := uuid.NewString() + taskID := req.TaskId + if taskID == "" { + taskID = uuid.NewString() + } startAt := time.Now() // 1) 获取用户信息 @@ -49,81 +55,196 @@ func (s *taskService) Create(ctx context.Context, req *dto.CreateTaskReq) (res * 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) - // TODO: 排队控制暂时关闭，后续需要时取消注释 - // limit := queue.GetRuntimeQueueLimit(ctx, req.ModelName, model.MaxConcurrency*2) - // if limit > 0 { - // ok, err := queue.AcquireQueueSlot(ctx, req.ModelName, taskID, limit, model.TimeoutSeconds) - // if err != nil { - // return nil, err - // } - // if !ok { - // return nil, errors.New("任务排队已满，请稍后再试") - // } - // } + // 循环尝试获取并发名额，超限则等待重试 + 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) + } - // 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, - } + // 未超限：跳出循环，执行业务 + 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) + } - // 4) 插入任务记录 - id, err := dao.ModelGatewayTask.Insert(ctx, task) - if err != nil { - // TODO: 恢复排队逻辑后，此处需要回滚排队占位 - // queue.ReleaseQueueSlot(ctx, req.ModelName, taskID) - return nil, err - } - task.Id = id + // 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, + } - // 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}, - }) + // 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 - // 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, + // 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}, }) - } - // 7) 异步执行任务 - go AsyncWorker.handleOne(util.AsyncCtx(ctx), task, model, req) + // 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{ @@ -260,7 +381,10 @@ func (s *taskService) QueryPendingTasks(ctx context.Context, req *dto.QueryPendi if err != nil || model == nil || model.QueryConfig == nil { continue } - result, err := util.PullTaskResult(ctx, nil, model.QueryConfig, model.HeadMsg) + // 每个任务使用独立的超时上下文，防止单个任务阻塞整个轮询 + 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 diff --git a/service/task/worker.go b/service/task/worker.go index 357c8ce..e27d8c6 100644 --- a/service/task/worker.go +++ b/service/task/worker.go @@ -189,10 +189,11 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa g.Log().Errorf(ctx, "[handleOne] 更新DB失败 taskId=%s err=%v", task.TaskID, err) return } - - go gateway.TriggerCallback(util.AsyncCtx(ctx), task) + concurrencyKey := fmt.Sprintf("model:concurrency:%s", req.ModelName) + g.Redis().Decr(ctx, concurrencyKey) + gateway.TriggerCallback(ctx, task) if req.EpicycleId != 0 { - go gateway.TriggerPromptsCallback(util.AsyncCtx(ctx), task, req.EpicycleId) + gateway.TriggerPromptsCallback(ctx, task, req.EpicycleId) } g.Log().Infof(ctx, "[handleOne] 成功 taskId=%s duration=%ds fileType=%s", @@ -550,6 +551,8 @@ func (w *asyncWorker) failTask(ctx context.Context, t *entity.ModelGatewayTask, t.State = 3 t.ErrorMsg = errMsg t.DurationSeconds = int64(time.Since(startTime).Seconds()) + concurrencyKey := fmt.Sprintf("model:concurrency:%s", t.ModelName) + g.Redis().Decr(ctx, concurrencyKey) _, _ = dao.ModelGatewayTask.Update(ctx, t) // 更新任务状态 go gateway.TriggerCallback(util.AsyncCtx(ctx), t) // 触发回调 }