215 lines
8.1 KiB
Go
215 lines
8.1 KiB
Go
package dao
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
|
|
"model-gateway/consts/public"
|
|
"model-gateway/model/entity"
|
|
|
|
"gitea.com/red-future/common/db/gfdb"
|
|
"github.com/gogf/gf/v2/database/gdb"
|
|
"github.com/gogf/gf/v2/os/gtime"
|
|
)
|
|
|
|
// ======================== 查询辅助 ========================
|
|
|
|
// taskColumns 查询用的公共字段
|
|
const taskColumns = `id, tenant_id, creator, model_name, task_id, biz_name, callback_url, model_key, retry_count, input_ref, request_payload, phase, tmp_file`
|
|
|
|
// ======================== 事务抢占 ========================
|
|
|
|
// claimTasks 事务内抢占任务并更新 state=1
|
|
func claimTasks(ctx context.Context, where string, args ...any) ([]*entity.AsynchTask, error) {
|
|
var tasks []*entity.AsynchTask
|
|
err := gfdb.DB(ctx).Transaction(ctx, func(ctx context.Context, tx gdb.TX) error {
|
|
sql := fmt.Sprintf(`SELECT %s FROM %s WHERE deleted_at IS NULL AND state = 0 %s LIMIT 1 FOR UPDATE SKIP LOCKED`, taskColumns, public.TableNameTask, where)
|
|
r, err := tx.GetOne(sql, args...)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if r.IsEmpty() {
|
|
return nil
|
|
}
|
|
var task entity.AsynchTask
|
|
if err := r.Struct(&task); err != nil {
|
|
return err
|
|
}
|
|
now := time.Now()
|
|
_, err = tx.Exec(fmt.Sprintf(`UPDATE %s SET state=1, started_at=?, updated_at=? WHERE id=?`, public.TableNameTask), now, now, task.Id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
tasks = []*entity.AsynchTask{&task}
|
|
return nil
|
|
})
|
|
return tasks, err
|
|
}
|
|
|
|
// ClaimPendingGlobal 批量抢占 pending 任务
|
|
func (d *taskDao) ClaimPendingGlobal(ctx context.Context, batchSize int) ([]*entity.AsynchTask, error) {
|
|
if batchSize <= 0 {
|
|
batchSize = 1
|
|
}
|
|
var tasks []*entity.AsynchTask
|
|
err := gfdb.DB(ctx).Transaction(ctx, func(ctx context.Context, tx gdb.TX) error {
|
|
sql := fmt.Sprintf(`SELECT %s FROM %s WHERE deleted_at IS NULL AND state = 0 ORDER BY enqueue_at ASC LIMIT %d FOR UPDATE SKIP LOCKED`, taskColumns, public.TableNameTask, batchSize)
|
|
r, err := tx.GetAll(sql)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if r.IsEmpty() {
|
|
return nil
|
|
}
|
|
if err := r.Structs(&tasks); err != nil {
|
|
return err
|
|
}
|
|
now := time.Now()
|
|
for _, t := range tasks {
|
|
_, err = tx.Exec(fmt.Sprintf(`UPDATE %s SET state=1, started_at=?, updated_at=? WHERE id=?`, public.TableNameTask), now, now, t.Id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
})
|
|
return tasks, err
|
|
}
|
|
|
|
// ClaimPendingByTaskIDGlobal 按 task_id 抢占
|
|
func (d *taskDao) ClaimPendingByTaskIDGlobal(ctx context.Context, taskID string) (*entity.AsynchTask, error) {
|
|
if taskID == "" {
|
|
return nil, nil
|
|
}
|
|
tasks, err := claimTasks(ctx, "AND task_id = ?", taskID)
|
|
if err != nil || len(tasks) == 0 {
|
|
return nil, err
|
|
}
|
|
return tasks[0], nil
|
|
}
|
|
|
|
// ======================== 更新辅助 ========================
|
|
|
|
func execSQL(ctx context.Context, sql string, args ...any) error {
|
|
_, err := gfdb.DB(ctx).Exec(ctx, sql, args...)
|
|
return err
|
|
}
|
|
|
|
// updateTask 通用更新
|
|
func updateTask(ctx context.Context, id int64, data entity.AsynchTask) error {
|
|
_, err := gfdb.DB(ctx).Model(ctx, public.TableNameTask).OmitEmpty().
|
|
Where(entity.AsynchTaskCol.Id, id).Data(data).Update()
|
|
return err
|
|
}
|
|
|
|
// UpdateSuccessGlobal 更新任务成功
|
|
func (d *taskDao) UpdateSuccessGlobal(ctx context.Context, t *entity.AsynchTask) error {
|
|
return updateTask(ctx, t.Id, entity.AsynchTask{
|
|
State: 2,
|
|
OssFile: t.OssFile,
|
|
FileType: t.FileType,
|
|
TextResult: t.TextResult,
|
|
FileSize: t.FileSize,
|
|
ErrorMsg: "",
|
|
FinishedAt: gtime.Now(),
|
|
Phase: 0,
|
|
TmpFile: "",
|
|
ExpendTokens: t.ExpendTokens,
|
|
DurationSeconds: t.DurationSeconds,
|
|
})
|
|
}
|
|
|
|
// UpdateFailedGlobal 模型调用失败
|
|
func (d *taskDao) UpdateFailedGlobal(ctx context.Context, t *entity.AsynchTask) error {
|
|
return updateTask(ctx, t.Id, entity.AsynchTask{
|
|
State: 3,
|
|
ErrorMsg: t.ErrorMsg,
|
|
FinishedAt: gtime.Now(),
|
|
Phase: 0,
|
|
TmpFile: "",
|
|
TextResult: t.TextResult,
|
|
DurationSeconds: t.DurationSeconds,
|
|
})
|
|
}
|
|
|
|
// UpdateFailedKeepTmpGlobal OSS 上传失败
|
|
func (d *taskDao) UpdateFailedKeepTmpGlobal(ctx context.Context, id int64, errorMsg string) error {
|
|
return execSQL(ctx, fmt.Sprintf(`UPDATE %s SET state=3, error_msg=?, finished_at=?, phase=1, updated_at=? WHERE id=?`, public.TableNameTask), errorMsg, gtime.Now(), gtime.Now(), id)
|
|
}
|
|
|
|
// UpdateTmpAfterModelGlobal 写临时文件
|
|
func (d *taskDao) UpdateTmpAfterModelGlobal(ctx context.Context, id int64, tmpFile string) error {
|
|
return execSQL(ctx, fmt.Sprintf(`UPDATE %s SET phase=1, tmp_file=?, updated_at=NOW() WHERE id=?`, public.TableNameTask), tmpFile, id)
|
|
}
|
|
|
|
// RollbackToPendingGlobal 回滚
|
|
func (d *taskDao) RollbackToPendingGlobal(ctx context.Context, id int64) error {
|
|
return execSQL(ctx, fmt.Sprintf(`UPDATE %s SET state=0, enqueue_at=NOW(), updated_at=NOW() WHERE id=? AND state=1`, public.TableNameTask), id)
|
|
}
|
|
|
|
// IncRetryCountGlobal 重试计数+1
|
|
func (d *taskDao) IncRetryCountGlobal(ctx context.Context, id int64) error {
|
|
return execSQL(ctx, fmt.Sprintf(`UPDATE %s SET retry_count=retry_count+1, updated_at=NOW() WHERE id=?`, public.TableNameTask), id)
|
|
}
|
|
|
|
// RequeueForRetryGlobal 重新入队
|
|
func (d *taskDao) RequeueForRetryGlobal(ctx context.Context, id int64, enqueueAt time.Time) error {
|
|
return execSQL(ctx, fmt.Sprintf(`UPDATE %s SET state=0, retry_count=retry_count+1, enqueue_at=?, updated_at=NOW() WHERE id=? AND state=3 AND deleted_at IS NULL`, public.TableNameTask), enqueueAt, id)
|
|
}
|
|
|
|
// ======================== 列表查询 ========================
|
|
|
|
// ListExpiredDownloadedGlobal
|
|
func (d *taskDao) ListExpiredDownloadedGlobal(ctx context.Context, limit int) ([]*entity.AsynchTask, error) {
|
|
return queryTasks(ctx, fmt.Sprintf(`SELECT * FROM %s WHERE deleted_at IS NULL AND state=4 AND expire_at IS NOT NULL AND expire_at < ? LIMIT ?`, public.TableNameTask), gtime.Now(), clampLimit(limit, 200))
|
|
}
|
|
|
|
// ListFailedRetryableGlobal
|
|
func (d *taskDao) ListFailedRetryableGlobal(ctx context.Context, limit int) ([]*entity.AsynchTask, error) {
|
|
return queryTasks(ctx, fmt.Sprintf(`SELECT t.*, m.retry_queue_max_seconds FROM %s t JOIN %s m ON t.tenant_id=m.tenant_id AND t.model_name=m.model_name WHERE t.deleted_at IS NULL AND t.state=3 AND t.retry_count < m.retry_times ORDER BY t.updated_at ASC LIMIT ?`, public.TableNameTask, public.TableNameModel), clampLimit(limit, 200))
|
|
}
|
|
|
|
// ListFailedExhaustedGlobal
|
|
func (d *taskDao) ListFailedExhaustedGlobal(ctx context.Context, limit int) ([]*entity.AsynchTask, error) {
|
|
return queryTasks(ctx, fmt.Sprintf(`SELECT t.* FROM %s t JOIN %s m ON t.tenant_id=m.tenant_id AND t.model_name=m.model_name WHERE t.deleted_at IS NULL AND t.state=3 AND t.retry_count >= m.retry_times ORDER BY t.updated_at ASC LIMIT ?`, public.TableNameTask, public.TableNameModel), clampLimit(limit, 200))
|
|
}
|
|
|
|
// ListTimeoutTasksGlobal
|
|
func (d *taskDao) ListTimeoutTasksGlobal(ctx context.Context, limit int) ([]*entity.AsynchTask, error) {
|
|
return queryTasks(ctx, fmt.Sprintf(`SELECT t.* FROM %s t JOIN %s m ON t.tenant_id=m.tenant_id AND t.model_name=m.model_name WHERE t.deleted_at IS NULL AND t.state IN (0,1) AND m.expected_seconds > 0 AND t.created_at < (NOW() - (m.expected_seconds || ' seconds')::interval) LIMIT ?`, public.TableNameTask, public.TableNameModel), clampLimit(limit, 200))
|
|
}
|
|
|
|
// HardDeleteByIDGlobal
|
|
func (d *taskDao) HardDeleteByIDGlobal(ctx context.Context, id int64) error {
|
|
return execSQL(ctx, fmt.Sprintf(`DELETE FROM %s WHERE id=?`, public.TableNameTask), id)
|
|
}
|
|
|
|
// ======================== 内部辅助 ========================
|
|
|
|
func queryTasks(ctx context.Context, sql string, args ...any) ([]*entity.AsynchTask, error) {
|
|
r, err := gfdb.DB(ctx).GetAll(ctx, sql, args...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var list []*entity.AsynchTask
|
|
err = r.Structs(&list)
|
|
return list, err
|
|
}
|
|
|
|
func clampLimit(limit, defaultVal int) int {
|
|
if limit <= 0 {
|
|
return defaultVal
|
|
}
|
|
return limit
|
|
}
|
|
|
|
// UpdateColumns 更新指定字段（结构体版）
|
|
func (d *taskDao) UpdateColumns(ctx context.Context, id int64, data entity.AsynchTask) error {
|
|
_, err := gfdb.DB(ctx).Model(ctx, public.TableNameTask).OmitEmpty().
|
|
Where(entity.AsynchTaskCol.Id, id).
|
|
Data(data).
|
|
Update()
|
|
return err
|
|
}
|