tss/schedule/schedule.go

package schedule

import (
	"errors"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/go-logr/logr"
	"github.com/robfig/cron/v3"
	"gopkg.in/yaml.v3"
)

var log = logr.Discard()

func SetLogger(l logr.Logger) {
	log = l
}

type Scheduler interface {
	ShouldDownload(name string, scheduleSpec string) (bool, error)
	Save() error
	Good(name string)
}

type scheduler struct {
	filePath        string
	mutex           sync.Mutex
	lastSuccessTime map[string]time.Time
}

func NewScheduler(filePath string) (Scheduler, error) {
	s := scheduler{filePath: filePath, lastSuccessTime: map[string]time.Time{}}
	data, err := os.ReadFile(filePath)
	if err != nil {
		if !errors.Is(err, os.ErrNotExist) {
			return nil, err
		}
		f, err := os.Create(filePath)
		if err != nil {
			return nil, fmt.Errorf("path:%v does not exist and unable to create: err: %w", filePath, err)
		}
		err = f.Close()
		if err != nil {
			return nil, err
		}
		log.Info("scheduler file does not exist, will be created", "path", filePath)
	} else {
		err = yaml.Unmarshal(data, &s.lastSuccessTime)
		if err != nil {
			return nil, err
		}
	}
	return &s, nil
}

func (s *scheduler) Save() error {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	data, err := yaml.Marshal(&s.lastSuccessTime)
	if err != nil {
		return err
	}
	return os.WriteFile(s.filePath, data, 0644)
}

func (s *scheduler) Good(name string) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	s.lastSuccessTime[name] = time.Now()
}

func (s *scheduler) ShouldDownload(name string, scheduleSpec string) (bool, error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	lastSuccess, ok := s.lastSuccessTime[name]
	if !ok {
		return true, nil
	}
	cron, err := cron.ParseStandard(scheduleSpec)
	if err != nil {
		return false, err
	}
	n := cron.Next(lastSuccess)
	return n.Before(time.Now()), nil
}