feat: multiple At() times (#297)

go-co-op · Feb 22, 2022 · b35e4e8 · b35e4e8
1 parent 8c87d62
commit b35e4e8
Show file tree

Hide file tree

Showing 6 changed files with 278 additions and 72 deletions.
diff --git a/README.md b/README.md
@@ -33,6 +33,14 @@ s.Every(5).Days().Do(func(){ ... })
 
 s.Every(1).Month(1, 2, 3).Do(func(){ ... })
 
+// set time
+s.Every(1).Day().At("10:30").Do(func(){ ... })
+
+// set multiple times
+s.Every(1).Day().At("10:30;08:00").Do(func(){ ... })
+
+s.Every(1).Day().At("10:30").At("08:00").Do(func(){ ... })
+
 // Schedule each last day of the month
 s.Every(1).MonthLastDay().Do(func(){ ... })
 

diff --git a/example_test.go b/example_test.go
@@ -88,8 +88,22 @@ func ExampleJob_ScheduledAtTime() {
 	job, _ := s.Every(1).Day().At("10:30").Do(task)
 	s.StartAsync()
 	fmt.Println(job.ScheduledAtTime())
+
+	// if multiple times are set, the earliest time will be returned
+	job1, _ := s.Every(1).Day().At("10:30;08:00").Do(task)
+	fmt.Println(job1.ScheduledAtTime())
 	// Output:
 	// 10:30
+	// 8:0
+}
+
+func ExampleJob_ScheduledAtTimes() {
+	s := gocron.NewScheduler(time.UTC)
+	job, _ := s.Every(1).Day().At("10:30;08:00").Do(task)
+	s.StartAsync()
+	fmt.Println(job.ScheduledAtTimes())
+	// Output:
+	// [8:0 10:30]
 }
 
 func ExampleJob_ScheduledTime() {
@@ -171,6 +185,9 @@ func ExampleScheduler_At() {
 	s := gocron.NewScheduler(time.UTC)
 	_, _ = s.Every(1).Day().At("10:30").Do(task)
 	_, _ = s.Every(1).Monday().At("10:30:01").Do(task)
+	// multiple
+	_, _ = s.Every(1).Monday().At("10:30;18:00").Do(task)
+	_, _ = s.Every(1).Monday().At("10:30").At("18:00").Do(task)
 }
 
 func ExampleScheduler_ChangeLocation() {

diff --git a/go.mod b/go.mod
@@ -3,10 +3,13 @@ module github.com/go-co-op/gocron
 go 1.17
 
 require (
-	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/robfig/cron/v3 v3.0.1
 	github.com/stretchr/testify v1.7.0
 	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
+)
+
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
 	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect
 )
diff --git a/job.go b/job.go
@@ -3,6 +3,7 @@ package gocron
 import (
 	"context"
 	"fmt"
+	"sort"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -15,21 +16,21 @@ import (
 type Job struct {
 	mu sync.RWMutex
 	jobFunction
-	interval          int            // pause interval * unit between runs
-	duration          time.Duration  // time duration between runs
-	unit              schedulingUnit // time units, e.g. 'minutes', 'hours'...
-	startsImmediately bool           // if the Job should run upon scheduler start
-	atTime            time.Duration  // optional time at which this Job runs when interval is day
-	startAtTime       time.Time      // optional time at which the Job starts
-	error             error          // error related to Job
-	lastRun           time.Time      // datetime of last run
-	nextRun           time.Time      // datetime of next run
-	scheduledWeekdays []time.Weekday // Specific days of the week to start on
-	daysOfTheMonth    []int          // Specific days of the month to run the job
-	tags              []string       // allow the user to tag Jobs with certain labels
-	runCount          int            // number of times the job ran
-	timer             *time.Timer    // handles running tasks at specific time
-	cronSchedule      cron.Schedule  // stores the schedule when a task uses cron
+	interval          int             // pause interval * unit between runs
+	duration          time.Duration   // time duration between runs
+	unit              schedulingUnit  // time units, e.g. 'minutes', 'hours'...
+	startsImmediately bool            // if the Job should run upon scheduler start
+	atTimes           []time.Duration // optional time(s) at which this Job runs when interval is day
+	startAtTime       time.Time       // optional time at which the Job starts
+	error             error           // error related to Job
+	lastRun           time.Time       // datetime of last run
+	nextRun           time.Time       // datetime of next run
+	scheduledWeekdays []time.Weekday  // Specific days of the week to start on
+	daysOfTheMonth    []int           // Specific days of the month to run the job
+	tags              []string        // allow the user to tag Jobs with certain labels
+	runCount          int             // number of times the job ran
+	timer             *time.Timer     // handles running tasks at specific time
+	cronSchedule      cron.Schedule   // stores the schedule when a task uses cron
 }
 
 type jobFunction struct {
@@ -113,12 +114,63 @@ func (j *Job) setTimer(t *time.Timer) {
 	j.timer = t
 }
 
-func (j *Job) getAtTime() time.Duration {
-	return j.atTime
+func (j *Job) getFirstAtTime() time.Duration {
+	var t time.Duration
+	if len(j.atTimes) > 0 {
+		t = j.atTimes[0]
+	}
+
+	return t
+}
+
+func (j *Job) getAtTime(lastRun time.Time) time.Duration {
+	var r time.Duration
+	if len(j.atTimes) == 0 {
+		return r
+	}
+
+	if len(j.atTimes) == 1 {
+		return j.atTimes[0]
+	}
+
+	if lastRun.IsZero() {
+		r = j.atTimes[0]
+	} else {
+		for _, d := range j.atTimes {
+			nt := time.Date(lastRun.Year(), lastRun.Month(), lastRun.Day(), 0, 0, 0, 0, lastRun.Location()).Add(d)
+			if nt.After(lastRun) {
+				r = d
+				break
+			}
+		}
+	}
+
+	return r
 }
 
-func (j *Job) setAtTime(t time.Duration) {
-	j.atTime = t
+func (j *Job) addAtTime(t time.Duration) {
+	if len(j.atTimes) == 0 {
+		j.atTimes = append(j.atTimes, t)
+		return
+	}
+	exist := false
+	index := sort.Search(len(j.atTimes), func(i int) bool {
+		atTime := j.atTimes[i]
+		b := atTime >= t
+		if b {
+			exist = atTime == t
+		}
+		return b
+	})
+
+	// ignore if present
+	if exist {
+		return
+	}
+
+	j.atTimes = append(j.atTimes, time.Duration(0))
+	copy(j.atTimes[index+1:], j.atTimes[index:])
+	j.atTimes[index] = t
 }
 
 func (j *Job) getStartAtTime() time.Time {
@@ -208,9 +260,24 @@ func (j *Job) ScheduledTime() time.Time {
 	return j.nextRun
 }
 
-// ScheduledAtTime returns the specific time of day the Job will run at
+// ScheduledAtTime returns the specific time of day the Job will run at.
+// If multiple times are set, the earliest time will be returned.
 func (j *Job) ScheduledAtTime() string {
-	return fmt.Sprintf("%d:%d", j.atTime/time.Hour, (j.atTime%time.Hour)/time.Minute)
+	if len(j.atTimes) == 0 {
+		return "0:0"
+	}
+
+	return fmt.Sprintf("%d:%d", j.getFirstAtTime()/time.Hour, (j.getFirstAtTime()%time.Hour)/time.Minute)
+}
+
+// ScheduledAtTimes returns the specific times of day the Job will run at
+func (j *Job) ScheduledAtTimes() []string {
+	r := make([]string, len(j.atTimes))
+	for i, t := range j.atTimes {
+		r[i] = fmt.Sprintf("%d:%d", t/time.Hour, (t%time.Hour)/time.Minute)
+	}
+
+	return r
 }
 
 // Weekday returns which day of the week the Job will run on and

diff --git a/scheduler.go b/scheduler.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"reflect"
 	"sort"
+	"strings"
 	"sync"
 	"time"
 
@@ -266,56 +267,69 @@ func (s *Scheduler) calculateMonths(job *Job, lastRun time.Time) nextRun {
 
 		return nextRunResult
 	}
-	next := lastRunRoundedMidnight.Add(job.getAtTime()).AddDate(0, job.interval, 0)
-	return nextRun{duration: until(lastRunRoundedMidnight, next), dateTime: next}
+	next := lastRunRoundedMidnight.Add(job.getFirstAtTime()).AddDate(0, job.interval, 0)
+	return nextRun{duration: until(lastRun, next), dateTime: next}
 }
 
 func calculateNextRunForLastDayOfMonth(s *Scheduler, job *Job, lastRun time.Time) nextRun {
 	// Calculate the last day of the next month, by adding job.interval+1 months (i.e. the
 	// first day of the month after the next month), and subtracting one day, unless the
 	// last run occurred before the end of the month.
 	addMonth := job.interval
+	atTime := job.getAtTime(lastRun)
 	if testDate := lastRun.AddDate(0, 0, 1); testDate.Month() != lastRun.Month() &&
-		!s.roundToMidnight(lastRun).Add(job.getAtTime()).After(lastRun) {
+		!s.roundToMidnight(lastRun).Add(atTime).After(lastRun) {
 		// Our last run was on the last day of this month.
 		addMonth++
+		atTime = job.getFirstAtTime()
 	}
+
 	next := time.Date(lastRun.Year(), lastRun.Month(), 1, 0, 0, 0, 0, s.Location()).
-		Add(job.getAtTime()).
+		Add(atTime).
 		AddDate(0, addMonth, 0).
 		AddDate(0, 0, -1)
 	return nextRun{duration: until(lastRun, next), dateTime: next}
 }
 
 func calculateNextRunForMonth(s *Scheduler, job *Job, lastRun time.Time, dayOfMonth int) nextRun {
-
-	jobDay := time.Date(lastRun.Year(), lastRun.Month(), dayOfMonth, 0, 0, 0, 0, s.Location()).Add(job.getAtTime())
+	atTime := job.getAtTime(lastRun)
+	natTime := atTime
+	jobDay := time.Date(lastRun.Year(), lastRun.Month(), dayOfMonth, 0, 0, 0, 0, s.Location()).Add(atTime)
 	difference := absDuration(lastRun.Sub(jobDay))
 	next := lastRun
 	if jobDay.Before(lastRun) { // shouldn't run this month; schedule for next interval minus day difference
 		next = next.AddDate(0, job.interval, -0)
 		next = next.Add(-difference)
+		natTime = job.getFirstAtTime()
 	} else {
 		if job.interval == 1 && !jobDay.Equal(lastRun) { // every month counts current month
 			next = next.AddDate(0, job.interval-1, 0)
 		} else { // should run next month interval
 			next = next.AddDate(0, job.interval, 0)
+			natTime = job.getFirstAtTime()
 		}
 		next = next.Add(difference)
 	}
+	if atTime != natTime {
+		next = next.Add(-atTime).Add(natTime)
+	}
 	return nextRun{duration: until(lastRun, next), dateTime: next}
 }
 
 func (s *Scheduler) calculateWeekday(job *Job, lastRun time.Time) nextRun {
 	daysToWeekday := s.remainingDaysToWeekday(lastRun, job)
 	totalDaysDifference := s.calculateTotalDaysDifference(lastRun, daysToWeekday, job)
-	next := s.roundToMidnight(lastRun).Add(job.getAtTime()).AddDate(0, 0, totalDaysDifference)
+	acTime := job.getAtTime(lastRun)
+	if totalDaysDifference > 0 {
+		acTime = job.getFirstAtTime()
+	}
+	next := s.roundToMidnight(lastRun).Add(acTime).AddDate(0, 0, totalDaysDifference)
 	return nextRun{duration: until(lastRun, next), dateTime: next}
 }
 
 func (s *Scheduler) calculateWeeks(job *Job, lastRun time.Time) nextRun {
 	totalDaysDifference := int(job.interval) * 7
-	next := s.roundToMidnight(lastRun).Add(job.getAtTime()).AddDate(0, 0, totalDaysDifference)
+	next := s.roundToMidnight(lastRun).Add(job.getFirstAtTime()).AddDate(0, 0, totalDaysDifference)
 	return nextRun{duration: until(lastRun, next), dateTime: next}
 }
 
@@ -332,7 +346,7 @@ func (s *Scheduler) calculateTotalDaysDifference(lastRun time.Time, daysToWeekda
 	}
 
 	if daysToWeekday == 0 { // today, at future time or already passed
-		lastRunAtTime := time.Date(lastRun.Year(), lastRun.Month(), lastRun.Day(), 0, 0, 0, 0, s.Location()).Add(job.getAtTime())
+		lastRunAtTime := time.Date(lastRun.Year(), lastRun.Month(), lastRun.Day(), 0, 0, 0, 0, s.Location()).Add(job.getAtTime(lastRun))
 		if lastRun.Before(lastRunAtTime) {
 			return 0
 		}
@@ -344,7 +358,7 @@ func (s *Scheduler) calculateTotalDaysDifference(lastRun time.Time, daysToWeekda
 func (s *Scheduler) calculateDays(job *Job, lastRun time.Time) nextRun {
 
 	if job.interval == 1 {
-		lastRunDayPlusJobAtTime := time.Date(lastRun.Year(), lastRun.Month(), lastRun.Day(), 0, 0, 0, 0, s.Location()).Add(job.getAtTime())
+		lastRunDayPlusJobAtTime := s.roundToMidnight(lastRun).Add(job.getAtTime(lastRun))
 
 		// handle occasional occurrence of job running to quickly / too early such that last run was within a second of now
 		lastRunUnix, nowUnix := job.LastRun().Unix(), s.now().Unix()
@@ -353,11 +367,11 @@ func (s *Scheduler) calculateDays(job *Job, lastRun time.Time) nextRun {
 		}
 
 		if shouldRunToday(lastRun, lastRunDayPlusJobAtTime) {
-			return nextRun{duration: until(lastRun, s.roundToMidnight(lastRun).Add(job.getAtTime())), dateTime: s.roundToMidnight(lastRun).Add(job.getAtTime())}
+			return nextRun{duration: until(lastRun, lastRunDayPlusJobAtTime), dateTime: lastRunDayPlusJobAtTime}
 		}
 	}
 
-	nextRunAtTime := s.roundToMidnight(lastRun).Add(job.getAtTime()).AddDate(0, 0, job.interval).In(s.Location())
+	nextRunAtTime := s.roundToMidnight(lastRun).Add(job.getFirstAtTime()).AddDate(0, 0, job.interval).In(s.Location())
 	return nextRun{duration: until(lastRun, nextRunAtTime), dateTime: nextRunAtTime}
 }
 
@@ -382,11 +396,11 @@ func in(scheduleWeekdays []time.Weekday, weekday time.Weekday) bool {
 }
 
 func (s *Scheduler) calculateDuration(job *Job) time.Duration {
-	lastRun := job.LastRun()
 	if job.neverRan() && shouldRunAtSpecificTime(job) { // ugly. in order to avoid this we could prohibit setting .At() and allowing only .StartAt() when dealing with Duration types
-		atTime := time.Date(lastRun.Year(), lastRun.Month(), lastRun.Day(), 0, 0, 0, 0, s.Location()).Add(job.getAtTime())
-		if lastRun.Before(atTime) || lastRun.Equal(atTime) {
-			return time.Until(s.roundToMidnight(lastRun).Add(job.getAtTime()))
+		now := s.time.Now(s.location)
+		next := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, s.Location()).Add(job.getFirstAtTime())
+		if now.Before(next) || now.Equal(next) {
+			return next.Sub(now)
 		}
 	}
 
@@ -404,7 +418,7 @@ func (s *Scheduler) calculateDuration(job *Job) time.Duration {
 }
 
 func shouldRunAtSpecificTime(job *Job) bool {
-	return job.getAtTime() != 0
+	return job.getAtTime(job.lastRun) != 0
 }
 
 func (s *Scheduler) remainingDaysToWeekday(lastRun time.Time, job *Job) int {
@@ -424,7 +438,7 @@ func (s *Scheduler) remainingDaysToWeekday(lastRun time.Time, job *Job) int {
 	})
 	// check atTime
 	if equals {
-		if s.roundToMidnight(lastRun).Add(job.getAtTime()).After(lastRun) {
+		if s.roundToMidnight(lastRun).Add(job.getAtTime(lastRun)).After(lastRun) {
 			return 0
 		}
 		index++
@@ -764,7 +778,7 @@ func (s *Scheduler) Do(jobFun interface{}, params ...interface{}) (*Job, error)
 	job := s.getCurrentJob()
 
 	jobUnit := job.getUnit()
-	if job.atTime != 0 && (jobUnit <= hours || jobUnit >= duration) {
+	if job.getAtTime(job.lastRun) != 0 && (jobUnit <= hours || jobUnit >= duration) {
 		job.error = wrapOrError(job.error, ErrAtTimeNotSupported)
 	}
 
@@ -821,15 +835,17 @@ func (s *Scheduler) At(i interface{}) *Scheduler {
 
 	switch t := i.(type) {
 	case string:
-		hour, min, sec, err := parseTime(t)
-		if err != nil {
-			job.error = wrapOrError(job.error, err)
-			return s
+		for _, tt := range strings.Split(t, ";") {
+			hour, min, sec, err := parseTime(tt)
+			if err != nil {
+				job.error = wrapOrError(job.error, err)
+				return s
+			}
+			// save atTime start as duration from midnight
+			job.addAtTime(time.Duration(hour)*time.Hour + time.Duration(min)*time.Minute + time.Duration(sec)*time.Second)
 		}
-		// save atTime start as duration from midnight
-		job.setAtTime(time.Duration(hour)*time.Hour + time.Duration(min)*time.Minute + time.Duration(sec)*time.Second)
 	case time.Time:
-		job.setAtTime(time.Duration(t.Hour())*time.Hour + time.Duration(t.Minute())*time.Minute + time.Duration(t.Second())*time.Second + time.Duration(t.Nanosecond())*time.Nanosecond)
+		job.addAtTime(time.Duration(t.Hour())*time.Hour + time.Duration(t.Minute())*time.Minute + time.Duration(t.Second())*time.Second + time.Duration(t.Nanosecond())*time.Nanosecond)
 	default:
 		job.error = wrapOrError(job.error, ErrUnsupportedTimeFormat)
 	}