likedream/internal/framework/schedule/scanner.go

package schedule

import (
	"bytes"
	"context"
	"encoding/gob"
	"errors"
	"fmt"
	cst2 "likedream-goframe/internal/common/cst"
	"likedream-goframe/internal/common/db"
	"likedream-goframe/internal/ent"
	"likedream-goframe/internal/ent/cptask"
	"likedream-goframe/internal/framework/queue"
	"likedream-goframe/internal/utils"
	"slices"
	"strconv"
	"strings"
	"time"

	"entgo.io/ent/dialect/sql"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/os/gctx"
	"github.com/hibiken/asynq"
	"github.com/samber/oops"
)

var globalScanner *TaskScanner

func initScanner() {
	globalScanner = NewTaskScanner()
	go globalScanner.StartScan()
}

type TaskScanner struct {
	ctx        context.Context
	stopSign   chan struct{}
	cancelFunc context.CancelFunc
}

func NewTaskScanner() *TaskScanner {
	ctx, cancelFunc := context.WithCancel(gctx.GetInitCtx())
	return &TaskScanner{
		ctx:        ctx,
		stopSign:   make(chan struct{}),
		cancelFunc: cancelFunc,
	}
}

func waitWithContext(ctx context.Context, d time.Duration) bool {
	t := time.NewTimer(d)
	defer t.Stop()

	select {
	case <-ctx.Done():
		return false
	case <-t.C:
		return true
	}
}

func (ts *TaskScanner) StartScan() {
	for {
		if !waitWithContext(ts.ctx, 30*time.Second) {
			close(ts.stopSign)
			return
		}

		ts.scanCpdailyTask()
	}
}

func (ts *TaskScanner) StopScan() {
	ts.cancelFunc()
	timeoutTimer := time.NewTimer(10 * time.Second)
	defer timeoutTimer.Stop()

	select {
	case <-timeoutTimer.C:
		g.Log().Warningf(ts.ctx, "task scanner stop timeout")
	case <-ts.stopSign:
	}
}
func (ts *TaskScanner) scanCpdailyTask() {

	currentTime := time.Now().In(cst2.TimeZoneShangHai)

	// 只查询30分钟以内的任务
	minExecuteTime := currentTime.Add(-30 * time.Minute).Unix()

	allTasks := db.Global.CpTask.Query().
		Where(
			cptask.NextExecuteTimeLTE(currentTime.Unix()),
			cptask.NextExecuteTimeNotNil(),
			cptask.NextExecuteTimeGT(0),
			cptask.StepEQ(cptask.StepComplete),
			cptask.TaskStatusEQ(cptask.TaskStatusNotStart),
			func(selector *sql.Selector) {
				selector.Where(sql.ExprP("success_count < frequency"))
			}).
		Select(cptask.FieldID, cptask.FieldExecuteHour, cptask.FieldExecuteMinute, cptask.FieldNextExecuteTime).
		AllX(ts.ctx)

	var failedIDS []int64

	for _, task := range allTasks {
		// 30分钟之前的任务不再执行，但是要重置下次执行时间
		if *task.NextExecuteTime < minExecuteTime {
			originExecuteTime := time.Unix(*task.NextExecuteTime, 0).In(cst2.TimeZoneShangHai).Format("2006-01-02 15:04:05")
			g.Log().Debugf(ts.ctx, "任务ID %d 已过期，原来预计执行时间：%s，跳过执行", task.ID, originExecuteTime)
			continue
		}

		payload := &CpdailyTaskPayload{
			CpTaskID: task.ID,
		}

		payloadWriter := bytes.NewBuffer(nil)

		e := gob.NewEncoder(payloadWriter).Encode(payload)
		utils.OkOrPanic(e, oops.WithContext(ts.ctx).With("payload", payload).In("CpdailyTaskScanner.scanCpdailyTask").Wrap)

		year, month, day := currentTime.Date()
		execTime := time.Date(year, month, day, *task.ExecuteHour, *task.ExecuteMinute, 0, 0, cst2.TimeZoneShangHai)

		e = utils.Catch(func() {
			queue.AddTask(ts.ctx, asynq.NewTask(cst2.TaskCpdaily, payloadWriter.Bytes()),
				asynq.TaskID(fmt.Sprintf(cst2.TaskCpdailyId, task.ID)),
				asynq.MaxRetry(4),
				asynq.Unique(10*time.Second),
				asynq.ProcessAt(execTime),
				asynq.Deadline(execTime.Add(12*time.Minute)),
				asynq.Queue(cst2.TaskQueueCpdaily),
			)
			// 更新未开始的任务状态为队列中
			db.Global.CpTask.Update().Where(cptask.ID(task.ID), cptask.TaskStatusEQ(cptask.TaskStatusNotStart)).SetTaskStatus(cptask.TaskStatusQueue).Exec(ts.ctx)
		})
		if e != nil {
			if errors.Is(e, asynq.ErrDuplicateTask) {
				// 任务已在队列中，不更新下次执行时间，等任务执行完成后再处理
				g.Log().Debugf(ts.ctx, "任务ID %d 已在队列中，跳过本次扫描", task.ID)
				failedIDS = append(failedIDS, task.ID)
				continue
			} else if errors.Is(e, context.Canceled) {
				return
			} else {
				utils.LogError(ts.ctx, "扫描今日校园任务时，入列失败", oops.With("task_id", task.ID).Wrap(e))
				failedIDS = append(failedIDS, task.ID)
			}
		}
	}

	if len(failedIDS) > 0 {
		allTasks = slices.DeleteFunc(allTasks, func(task *ent.CpTask) bool {
			return slices.Contains(failedIDS, task.ID)
		})
	}

	for _, task := range allTasks {
		nextExecuteAt := utils.GetCpdailyNextExecuteTime(task)

		e := db.Global.CpTask.UpdateOneID(task.ID).Where(cptask.NextExecuteTimeNotNil()).SetNextExecuteTime(nextExecuteAt).Exec(ts.ctx)
		utils.LogError(ts.ctx, "更新任务下次执行时间失败", oops.With("task_id", task.ID, "next_execute_at", nextExecuteAt).Wrap(e))
	}

	// 重置30分钟之前的任务到下次执行时间
	// 只处理已完成配置且状态为未开始的任务，避免重复处理正在执行或已失败的任务
	expiredTasks := db.Global.CpTask.Query().
		Where(cptask.NextExecuteTimeLT(minExecuteTime), cptask.NextExecuteTimeNotNil(), cptask.NextExecuteTimeGT(0), cptask.StepEQ(cptask.StepComplete), cptask.TaskStatusEQ(cptask.TaskStatusNotStart)).
		Select(cptask.FieldID, cptask.FieldExecuteHour, cptask.FieldExecuteMinute).
		AllX(ts.ctx)

	expiredIDs := slices.Collect(func(yield func(string) bool) {
		for _, task := range expiredTasks {
			if !yield(strconv.FormatInt(task.ID, 10)) {
				return
			}
		}
	})

	if len(expiredIDs) > 0 {
		g.Log().Warningf(ts.ctx, "以下任务ID已过期，重置了下次执行时间: %v", strings.Join(expiredIDs, ", "))
	}

	for _, task := range expiredTasks {
		nextExecuteAt := utils.GetCpdailyNextExecuteTime(task)

		e := db.Global.CpTask.UpdateOneID(task.ID).Where(cptask.NextExecuteTimeNotNil()).SetNextExecuteTime(nextExecuteAt).Exec(ts.ctx)
		utils.LogError(ts.ctx, "重置过期任务下次执行时间失败", oops.With("task_id", task.ID, "next_execute_at", nextExecuteAt).Wrap(e))
	}
}