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/db"
	"gitlink.org.cn/cloudream/jcs-pub/client/internal/metacache"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/publock"
	"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.Service
}

func New(cfg Config, db *db.DB, spaceMeta *metacache.UserSpaceMeta, stgPool *pool.Pool, evtPub *sysevent.Publisher, pubLock *publock.Service) *TickTock {
	sch, _ := gocron.NewScheduler()
	t := &TickTock{
		cfg:       cfg,
		sch:       sch,
		jobs:      map[string]cronJob{},
		db:        db,
		spaceMeta: spaceMeta,
		stgPool:   stgPool,
		evtPub:    evtPub,
		pubLock:   pubLock,
	}
	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),
	)))

}