package task import ( "context" "encoding/json" "fmt" "runtime/debug" "sync" "time" "github.com/gogf/gf/v2/errors/gerror" "github.com/gogf/gf/v2/frame/g" "github.com/gogf/gf/v2/os/gctx" "github.com/gogf/gf/v2/os/gtime" "github.com/qianfree/team-api/internal/dao" do "github.com/qianfree/team-api/internal/model/do" ) // TaskStatus defines the status of an async task. type TaskStatus string const ( StatusPending TaskStatus = "pending" StatusRunning TaskStatus = "running" StatusSucceeded TaskStatus = "succeeded" StatusFailed TaskStatus = "failed" StatusCancelled TaskStatus = "cancelled" pendingTaskWorkerCount = 8 ) // Task defines an async task to be executed. type Task struct { Name string Handler string // handler function identifier Payload any MaxRetries int ScheduledAt *time.Time } // TaskResult represents the result of a task execution. type TaskResult struct { Success bool Data any Error string } // HandlerFunc is the function signature for task handlers. type HandlerFunc func(ctx context.Context, payload json.RawMessage) (any, error) var ( handlers = make(map[string]HandlerFunc) handlersMu sync.RWMutex ) // RegisterHandler registers a task handler function. func RegisterHandler(name string, handler HandlerFunc) { handlersMu.Lock() defer handlersMu.Unlock() handlers[name] = handler } // GetHandler retrieves a registered handler. func GetHandler(name string) (HandlerFunc, bool) { handlersMu.RLock() defer handlersMu.RUnlock() h, ok := handlers[name] return h, ok } // CreateTask creates a new task record in the database. func CreateTask(ctx context.Context, task *Task) (int64, error) { payload, err := json.Marshal(task.Payload) if err != nil { return 0, gerror.Wrapf(err, "marshal payload") } result, err := dao.TskTasks.Ctx(ctx).Data(do.TskTasks{ Name: task.Name, Handler: task.Handler, Payload: payload, Status: StatusPending, MaxRetries: task.MaxRetries, ScheduledAt: func() *gtime.Time { if task.ScheduledAt != nil { return gtime.NewFromTime(*task.ScheduledAt) } return nil }(), }).Insert() if err != nil { return 0, err } id, err := result.LastInsertId() if err != nil { return 0, err } return id, nil } // ExecuteTask executes a task by ID. // It loads the task from DB, updates status, calls the handler, and records the result. func ExecuteTask(ctx context.Context, taskID int64) (err error) { defer func() { if recovered := recover(); recovered != nil { message := fmt.Sprintf("panic: %v\n%s", recovered, debug.Stack()) err = failTask(ctx, taskID, message) } }() // Load task var task *struct { ID int64 `json:"id"` Name string `json:"name"` Handler string `json:"handler"` Payload json.RawMessage `json:"payload"` MaxRetries int `json:"max_retries"` RetryCount int `json:"retry_count"` } err = dao.TskTasks.Ctx(ctx). Where("id", taskID). Where("status", StatusPending). Scan(&task) if err != nil { return gerror.Wrapf(err, "load task") } if task == nil { return gerror.Newf("task %d not found or not pending", taskID) } // Mark as running —— 原子领取:仅当仍为 pending 时才能翻成 running。 // 用条件更新 + RowsAffected 做 CAS(对齐 async_provider.UpdateTaskCAS 的做法), // 避免「先 SELECT pending 再无条件 UPDATE」在多副本/并发调度下被两个 worker 同时领取、重复执行同一任务。 res, err := dao.TskTasks.Ctx(ctx). Where("id", taskID). Where("status", StatusPending). Data(do.TskTasks{ Status: StatusRunning, StartedAt: gtime.Now(), }).Update() if err != nil { return gerror.Wrapf(err, "update task status to running") } affected, err := res.RowsAffected() if err != nil { return gerror.Wrapf(err, "confirm task claim result") } if affected == 0 { // 另一 worker 已抢先领取(status 已非 pending),本次放弃,避免重复执行 return nil } // Get handler handler, ok := GetHandler(task.Handler) if !ok { return failTask(ctx, taskID, fmt.Sprintf("handler %q not registered", task.Handler)) } // Execute handler result, execErr := handler(ctx, task.Payload) // Record result var resultJSON []byte if result != nil { resultJSON, _ = json.Marshal(result) } if execErr != nil { // Check if should retry if task.RetryCount < task.MaxRetries { _, _ = dao.TskTasks.Ctx(ctx). Where("id", taskID). Data(do.TskTasks{ Status: StatusPending, RetryCount: task.RetryCount + 1, ErrorMessage: execErr.Error(), StartedAt: nil, }).Update() return execErr } return failTask(ctx, taskID, execErr.Error()) } // Success _, err = dao.TskTasks.Ctx(ctx). Where("id", taskID). Data(do.TskTasks{ Status: StatusSucceeded, Result: resultJSON, FinishedAt: gtime.Now(), }).Update() if err != nil { g.Log().Errorf(ctx, "update task %d to succeeded: %v", taskID, err) } return nil } // failTask marks a task as failed and logs the error. func failTask(ctx context.Context, taskID int64, errMsg string) error { _, err := dao.TskTasks.Ctx(ctx). Where("id", taskID). Data(do.TskTasks{ Status: StatusFailed, ErrorMessage: errMsg, FinishedAt: gtime.Now(), }).Update() if err != nil { g.Log().Errorf(ctx, "update task %d to failed: %v", taskID, err) } // Log error logTask(ctx, taskID, "error", errMsg) return gerror.Newf("task failed: %s", errMsg) } // logTask adds a log entry for a task. func logTask(ctx context.Context, taskID int64, level, message string) { _, err := dao.TskTaskLogs.Ctx(ctx).Data(do.TskTaskLogs{ TaskId: taskID, Level: level, Message: message, }).Insert() if err != nil { g.Log().Errorf(ctx, "insert task log: %v", err) } } // LogInfo adds an info log for a task. func LogInfo(ctx context.Context, taskID int64, message string) { logTask(ctx, taskID, "info", message) } // LogError adds an error log for a task. func LogError(ctx context.Context, taskID int64, message string) { logTask(ctx, taskID, "error", message) } // CancelTask cancels a pending task. func CancelTask(ctx context.Context, taskID int64) error { result, err := dao.TskTasks.Ctx(ctx). Where("id", taskID). WhereIn("status", []string{string(StatusPending), string(StatusRunning)}). Data(do.TskTasks{ Status: StatusCancelled, FinishedAt: gtime.Now(), }).Update() if err != nil { return err } rowsAffected, _ := result.RowsAffected() if rowsAffected == 0 { return gerror.Newf("task %d cannot be cancelled", taskID) } return nil } // ScheduleTask creates a task scheduled for a future time. func ScheduleTask(ctx context.Context, task *Task, runAt time.Time) (int64, error) { task.ScheduledAt = &runAt return CreateTask(ctx, task) } // RunPendingTasks executes all pending tasks that are due. func RunPendingTasks(ctx context.Context) { var tasks []struct { ID int64 `json:"id"` } err := dao.TskTasks.Ctx(ctx). Where("status", StatusPending). Where("scheduled_at <= ? OR scheduled_at IS NULL", time.Now()). Fields("id"). Scan(&tasks) if err != nil { g.Log().Errorf(ctx, "query pending tasks: %v", err) return } if len(tasks) == 0 { return } workerCount := min(pendingTaskWorkerCount, len(tasks)) taskIDs := make(chan int64) var workers sync.WaitGroup workers.Add(workerCount) for range workerCount { go func() { defer workers.Done() for taskID := range taskIDs { if err := ExecuteTask(ctx, taskID); err != nil { g.Log().Errorf(ctx, "execute task %d: %v", taskID, err) } } }() } for _, task := range tasks { taskIDs <- task.ID } close(taskIDs) workers.Wait() } // RunTaskAsync executes a task in a goroutine. func RunTaskAsync(task *Task) { go func() { defer func() { if r := recover(); r != nil { glogError("panic in async task: %v", r) } }() ctx := gctx.New() id, err := CreateTask(ctx, task) if err != nil { glogError("create async task: %v", err) return } if err := ExecuteTask(ctx, id); err != nil { glogError("execute async task %d: %v", id, err) } }() } // ExecuteTaskAfterCreate 在后台立即执行刚创建的任务(跳过 cron 周期)。 // // 适用于用户主动创建的交互式任务(如数据导出),用户期望提交后立即开始执行, // 而非等待最长 60 秒的 cron 轮询。goroutine 内部使用 gctx.New() 避免被请求上下文 // 取消。ExecuteTask 有 CAS 保护(WHERE status='pending'),即使 cron 也同时竞争 // 也不会重复执行——只有一个会成功 claim,另一个的 RowsAffected == 0 静默退出。 func ExecuteTaskAfterCreate(taskID int64) { go func() { defer func() { if r := recover(); r != nil { glogError("panic in ExecuteTaskAfterCreate task %d: %v", taskID, r) } }() bgCtx := gctx.New() if err := ExecuteTask(bgCtx, taskID); err != nil { glogError("ExecuteTaskAfterCreate task %d: %v", taskID, err) } }() } func glogError(format string, args ...any) { g.Log().Errorf(gctx.New(), format, args...) }