package scheduler import ( "context" "strconv" "sync" "time" "github.com/pkg/errors" "github.com/robfig/cron/v3" "github.com/rs/zerolog/log" ) type Scheduler struct { crontab *cron.Cron activeJobs map[cron.EntryID]context.CancelFunc mu sync.Mutex } type PermanentError struct { err error } func NewPermanentError(err error) *PermanentError { return &PermanentError{err: err} } func (e *PermanentError) Error() string { return e.err.Error() } func NewScheduler(ctx context.Context) *Scheduler { crontab := cron.New(cron.WithChain(cron.Recover(cron.DefaultLogger))) crontab.Start() s := &Scheduler{ crontab: crontab, activeJobs: make(map[cron.EntryID]context.CancelFunc), } if ctx != nil { go func() { <-ctx.Done() s.Shutdown() }() } return s } // Shutdown stops the scheduler and waits for it to stop if it is running; otherwise does nothing. func (s *Scheduler) Shutdown() error { if s.crontab == nil { return nil } log.Debug().Msg("stopping scheduler") ctx := s.crontab.Stop() <-ctx.Done() s.mu.Lock() for _, job := range s.crontab.Entries() { if cancel, ok := s.activeJobs[job.ID]; ok { cancel() } } s.mu.Unlock() err := ctx.Err() if errors.Is(err, context.Canceled) { return nil } return err } // StopJob stops the job from being run in the future func (s *Scheduler) StopJob(jobID string) error { id, err := strconv.Atoi(jobID) if err != nil { return errors.Wrapf(err, "failed convert jobID %q to int", jobID) } entryID := cron.EntryID(id) s.mu.Lock() if cancel, ok := s.activeJobs[entryID]; ok { cancel() delete(s.activeJobs, entryID) } s.mu.Unlock() return nil } // StartJobEvery schedules a new periodic job with a given duration. // Returns job id that could be used to stop the given job. // When job run returns an error, that job won't be run again. func (s *Scheduler) StartJobEvery(duration time.Duration, job func() error) string { entryID := new(cron.EntryID) cancelFn := func() { log.Debug().Msg("job cancelled, stopping") s.crontab.Remove(*entryID) } jobFn := cron.FuncJob(func() { err := job() if err == nil { return } var permErr *PermanentError if errors.As(err, &permErr) { log.Error().Err(permErr).Msg("job returned a permanent error, it will be stopped") cancelFn() return } log.Error().Err(err).Msg("job returned an error, it will be rescheduled") }) *entryID = s.crontab.Schedule(cron.Every(duration), jobFn) s.mu.Lock() s.activeJobs[*entryID] = cancelFn s.mu.Unlock() return strconv.Itoa(int(*entryID)) }