第一次提交

service/worker.go

@@ -0,0 +1,176 @@
+package service
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"time"
+
+	"model-asynch/dao"
+	"model-asynch/model/entity"
+
+	"github.com/gogf/gf/v2/frame/g"
+	"github.com/gogf/gf/v2/os/grpool"
+)
+
+var AsyncWorker = &asyncWorker{}
+
+type asyncWorker struct {
+}
+
+// RunOnce 由上层定时任务触发：一次性抢占并处理一批任务
+// - batchSize: 本次抢占数量
+// - goroutines: 本次并发数（协程池大小）
+func (w *asyncWorker) RunOnce(ctx context.Context, batchSize, goroutines int) (claimed int, err error) {
+	if batchSize <= 0 {
+		batchSize = 10
+	}
+	if goroutines <= 0 {
+		goroutines = 1
+	}
+	tasks, err := dao.Task.ClaimPendingGlobal(ctx, batchSize)
+	if err != nil {
+		return 0, err
+	}
+	if len(tasks) == 0 {
+		return 0, nil
+	}
+	pool := grpool.New(goroutines)
+	defer pool.Close()
+
+	claimed = len(tasks)
+	done := make(chan struct{}, claimed)
+	for _, t := range tasks {
+		task := t
+		_ = pool.AddWithRecover(ctx, func(ctx context.Context) {
+			w.handleOne(ctx, task)
+			done <- struct{}{}
+		}, func(ctx context.Context, e error) {
+			if e != nil {
+				_ = dao.Task.UpdateFailedGlobal(ctx, task.Id, fmt.Sprintf("worker panic: %v", e))
+				ReleaseQueueSlot(ctx, task.ModelName, task.TaskID)
+			}
+			done <- struct{}{}
+		})
+	}
+	for i := 0; i < claimed; i++ {
+		<-done
+	}
+	return claimed, nil
+}
+
+func (w *asyncWorker) handleOne(ctx context.Context, t *entity.AsynchTask) {
+	// 从任务入库的 request_payload 里恢复 payload + headers，给 OSS 上传透传鉴权用
+	payload, headers := parseStoredPayload(t.RequestPayload)
+	if len(headers) > 0 {
+		ctx = setTaskHeadersToCtx(ctx, headers)
+	}
+
+	// 1) 拉取模型配置
+	m, err := dao.Model.GetByModelNameForTenant(ctx, t.TenantId, t.ModelName)
+	if err != nil {
+		_ = dao.Task.UpdateFailedGlobal(ctx, t.Id, err.Error())
+		ReleaseQueueSlot(ctx, t.ModelName, t.TaskID)
+		return
+	}
+	if m == nil || m.Enabled != 1 {
+		_ = dao.Task.UpdateFailedGlobal(ctx, t.Id, "模型不存在或未启用")
+		ReleaseQueueSlot(ctx, t.ModelName, t.TaskID)
+		return
+	}
+
+	// 2) 分布式并发限制（按 model_name 全局维度）
+	semKey := fmt.Sprintf("asynch:sem:%s", t.ModelName)
+	leaseSeconds := int64(3600) // 兜底1小时
+	maxC := GetRuntimeMaxConcurrency(ctx, t.ModelName, m.MaxConcurrency)
+	acquired, err := acquireSemaphore(ctx, semKey, maxC, leaseSeconds)
+	if err != nil {
+		_ = dao.Task.UpdateFailedGlobal(ctx, t.Id, err.Error())
+		ReleaseQueueSlot(ctx, t.ModelName, t.TaskID)
+		return
+	}
+	if !acquired {
+		// 并发满了：放回排队（重新置回 state=0），下一轮再抢占
+		_ = w.rollbackToPending(ctx, t.Id)
+		return
+	}
+	defer func() {
+		_ = releaseSemaphore(ctx, semKey)
+	}()
+
+	// 3) 调用模型服务
+	if payload == nil {
+		payload = map[string]any{
+			"taskId":   t.TaskID,
+			"inputRef": t.InputRef,
+		}
+	}
+	var (
+		data        []byte
+		contentType string
+		ext         string
+	)
+
+	// phase=1 表示模型已成功但 OSS 上传失败：优先从临时文件加载，避免重复跑模型
+	if t.Phase == 1 && strings.TrimSpace(t.TmpFile) != "" {
+		data, err = loadTmpResult(t.TmpFile)
+		if err == nil && len(data) > 0 {
+			contentType, ext = DetectFileType(data)
+		} else {
+			// 临时文件不可用：回退重新调用模型
+			data = nil
+		}
+	}
+	if data == nil {
+		// 统计：仅在真正请求模型时 +1（OSS 重试不计入）
+		_ = dao.Stat.IncRequestCount(ctx, time.Now(), int64(t.TenantId), t.Creator, t.ModelName)
+
+		data, err = InvokeModel(ctx, m, payload, t.ModelKey)
+		if err != nil {
+			_ = dao.Task.UpdateFailedGlobal(ctx, t.Id, err.Error())
+			ReleaseQueueSlot(ctx, t.ModelName, t.TaskID)
+			return
+		}
+		contentType, ext = DetectFileType(data)
+		// 将模型输出写入临时文件，后续若 OSS 失败可只重试 OSS
+		tmpPath, err := saveTmpResult(t.TaskID, data, ext)
+		if err == nil && tmpPath != "" {
+			t.TmpFile = tmpPath
+			t.Phase = 1
+			_ = dao.Task.UpdateTmpAfterModelGlobal(ctx, t.Id, tmpPath)
+		}
+	}
+
+	// 4) 存储 OSS
+	ossURL, err := Storage.UploadByTask(ctx, t, data, ext, contentType)
+	if err != nil {
+		// OSS 阶段失败：保留临时文件，下一轮仅重试 OSS
+		_ = dao.Task.UpdateFailedKeepTmpGlobal(ctx, t.Id, err.Error())
+		ReleaseQueueSlot(ctx, t.ModelName, t.TaskID)
+		return
+	}
+
+	// 5) 更新任务状态成功
+	// 注意：expire_at 的计算改为“已下载(state=4)后开始计时”，因此成功(state=2)不写 expire_at。
+	fileType := strings.TrimPrefix(ext, ".")
+	if fileType == "" {
+		fileType = contentType
+	}
+	if err := dao.Task.UpdateSuccessGlobal(ctx, t.Id, ossURL, fileType, int64(len(data)), nil); err != nil {
+		g.Log().Errorf(ctx, "[worker] update success failed: %v", err)
+		return
+	}
+	// 成功/失败均不再占用 queue_limit（state=0/1 才占用）
+	ReleaseQueueSlot(ctx, t.ModelName, t.TaskID)
+	// 6) 成功回调（不影响主流程）
+	t.State = 2
+	t.OssFile = ossURL
+	t.FileType = fileType
+	go triggerSuccessCallback(context.WithoutCancel(ctx), t)
+	// 成功后清理临时文件
+	deleteTmpResult(t.TmpFile)
+}
+
+func (w *asyncWorker) rollbackToPending(ctx context.Context, id int64) error {
+	return dao.Task.RollbackToPendingGlobal(ctx, id)
+}