package server import ( log "github.com/sirupsen/logrus" "sync" "time" ) // Scheduler is an interface which implementations can schedule and cancel jobs type Scheduler interface { Cancel(IDs []string) Schedule(in time.Duration, ID string, job func() (nextRunIn time.Duration, reschedule bool)) } // MockScheduler is a mock implementation of Scheduler type MockScheduler struct { CancelFunc func(IDs []string) ScheduleFunc func(in time.Duration, ID string, job func() (nextRunIn time.Duration, reschedule bool)) } // Cancel mocks the Cancel function of the Scheduler interface func (mock *MockScheduler) Cancel(IDs []string) { if mock.CancelFunc != nil { mock.CancelFunc(IDs) return } log.Errorf("MockScheduler doesn't have Cancel function defined ") } // Schedule mocks the Schedule function of the Scheduler interface func (mock *MockScheduler) Schedule(in time.Duration, ID string, job func() (nextRunIn time.Duration, reschedule bool)) { if mock.ScheduleFunc != nil { mock.ScheduleFunc(in, ID, job) return } log.Errorf("MockScheduler doesn't have Schedule function defined") } // DefaultScheduler is a generic structure that allows to schedule jobs (functions) to run in the future and cancel them. type DefaultScheduler struct { // jobs map holds cancellation channels indexed by the job ID jobs map[string]chan struct{} mu *sync.Mutex } // NewDefaultScheduler creates an instance of a DefaultScheduler func NewDefaultScheduler() *DefaultScheduler { return &DefaultScheduler{ jobs: make(map[string]chan struct{}), mu: &sync.Mutex{}, } } func (wm *DefaultScheduler) cancel(ID string) bool { cancel, ok := wm.jobs[ID] if ok { delete(wm.jobs, ID) select { case cancel <- struct{}{}: log.Debugf("cancelled scheduled job %s", ID) default: log.Warnf("couldn't cancel job %s because there was no routine listening on the cancel event", ID) return false } } return ok } // Cancel cancels the scheduled job by ID if present. // If job wasn't found the function returns false. func (wm *DefaultScheduler) Cancel(IDs []string) { wm.mu.Lock() defer wm.mu.Unlock() for _, id := range IDs { wm.cancel(id) } } // Schedule a job to run in some time in the future. If job returns true then it will be scheduled one more time. // If job with the provided ID already exists, a new one won't be scheduled. func (wm *DefaultScheduler) Schedule(in time.Duration, ID string, job func() (nextRunIn time.Duration, reschedule bool)) { wm.mu.Lock() defer wm.mu.Unlock() cancel := make(chan struct{}) if _, ok := wm.jobs[ID]; ok { log.Debugf("couldn't schedule a job %s because it already exists. There are %d total jobs scheduled.", ID, len(wm.jobs)) return } wm.jobs[ID] = cancel log.Debugf("scheduled a job %s to run in %s. There are %d total jobs scheduled.", ID, in.String(), len(wm.jobs)) go func() { select { case <-time.After(in): log.Debugf("time to do a scheduled job %s", ID) runIn, reschedule := job() wm.mu.Lock() defer wm.mu.Unlock() delete(wm.jobs, ID) if reschedule { go wm.Schedule(runIn, ID, job) } case <-cancel: log.Debugf("stopped scheduled job %s ", ID) wm.mu.Lock() defer wm.mu.Unlock() delete(wm.jobs, ID) return } }() }