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

feat: 添加模型并发控制与分布式锁

Browse Source 
重新启用任务并发限制逻辑，使用 Redis 计数器控制模型并发数；
引入分布式锁防止重复创建，并优化优雅退出顺序确保请求完成；
移除 video_duration 计费计算，支持外部传入 taskId，为轮询查询添加独立超时上下文。
This commit is contained in:
19904408334
2026-07-09 13:47:08 +08:00
parent a7f30d1640
commit 6166c24412
5 changed files with 222 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/util/billing.go
+22 -23
View File
@@ -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
}
main.go
+3 -2
View File
@@ -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) {
model/dto/model_gateway_task_dto.go
+1
View File
@@ -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 {
service/task/task_service.go
+190 -66
View File
@@ -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
service/task/worker.go
+6 -3
View File
@@ -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) // 触发回调
}