JCS-pub/client/internal/ticktock/ticktock.go

package ticktock

import (
	"fmt"

	"github.com/go-co-op/gocron/v2"
	"github.com/samber/lo"
	"gitlink.org.cn/cloudream/common/pkgs/logger"
	"gitlink.org.cn/cloudream/jcs-pub/client/internal/cluster"
	"gitlink.org.cn/cloudream/jcs-pub/client/internal/db"
	"gitlink.org.cn/cloudream/jcs-pub/client/internal/metacache"
	"gitlink.org.cn/cloudream/jcs-pub/client/internal/publock"
	"gitlink.org.cn/cloudream/jcs-pub/client/internal/speedstats"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/storage/pool"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/sysevent"
)

type Job interface {
	Name() string
	Execute(t *TickTock)
}

type cronJob struct {
	cronJob gocron.Job
	job     Job
}

type TickTock struct {
	cfg        Config
	sch        gocron.Scheduler
	jobs       map[string]cronJob
	db         *db.DB
	spaceMeta  *metacache.UserSpaceMeta
	stgPool    *pool.Pool
	evtPub     *sysevent.Publisher
	pubLock    *publock.PubLock
	speedStats *speedstats.SpeedStats
	cluster    *cluster.Cluster
}

func New(cfg Config, db *db.DB, spaceMeta *metacache.UserSpaceMeta, stgPool *pool.Pool, evtPub *sysevent.Publisher, pubLock *publock.PubLock, speedStats *speedstats.SpeedStats, cluster *cluster.Cluster) *TickTock {
	sch, _ := gocron.NewScheduler()
	t := &TickTock{
		cfg:        cfg,
		sch:        sch,
		jobs:       map[string]cronJob{},
		db:         db,
		spaceMeta:  spaceMeta,
		stgPool:    stgPool,
		evtPub:     evtPub,
		pubLock:    pubLock,
		speedStats: speedStats,
		cluster:    cluster,
	}
	t.initJobs()
	return t
}

func (t *TickTock) Start() {
	t.sch.Start()
}

func (t *TickTock) Stop() {
	t.sch.Shutdown()
}

func (t *TickTock) GetJobNames() []string {
	return lo.Keys(t.jobs)
}

func (t *TickTock) RunNow(jobName string) bool {
	j, ok := t.jobs[jobName]
	if !ok {
		logger.Warnf("job %s not found", jobName)
		return false
	}

	j.cronJob.RunNow()
	return true
}

func (t *TickTock) addJob(job Job, duration gocron.JobDefinition) {
	j, err := t.sch.NewJob(duration, gocron.NewTask(job.Execute, t))
	if err != nil {
		panic(fmt.Errorf("add job %s: %w", job.Name(), err))
	}

	t.jobs[job.Name()] = cronJob{
		cronJob: j,
		job:     job,
	}
}

func (t *TickTock) initJobs() {
	t.addJob(&UpdatePackageAccessStatAmount{}, gocron.DailyJob(1, gocron.NewAtTimes(
		gocron.NewAtTime(0, 0, 0),
	)))

	t.addJob(&ChangeRedundancy{}, gocron.DailyJob(1, gocron.NewAtTimes(
		gocron.NewAtTime(0, 5, 0),
	)))

	t.addJob(&CheckShardStore{}, gocron.DailyJob(1, gocron.NewAtTimes(
		gocron.NewAtTime(1, 0, 0),
	)))

	t.addJob(&UserSpaceGC{}, gocron.DailyJob(1, gocron.NewAtTimes(
		gocron.NewAtTime(2, 0, 0),
	)))

	t.addJob(&SpeedStatsStep{}, gocron.DailyJob(1, gocron.NewAtTimes(
		gocron.NewAtTime(0, 0, 0),
		gocron.NewAtTime(6, 0, 0),
		gocron.NewAtTime(12, 0, 0),
		gocron.NewAtTime(18, 0, 0),
	)))

	t.addJob(&StatsLog{}, gocron.DailyJob(1, gocron.NewAtTimes(
		gocron.NewAtTime(0, 0, 0),
	)))
}