From ed92fbbb98a68ca4d332565bf3724e805387ac66 Mon Sep 17 00:00:00 2001
From: "chenyunfei.cs" <chenyunfei.cs@bytedance.com>
Date: Sun, 26 Mar 2023 10:10:09 +0800
Subject: [PATCH] Supporting customize Timer to controller the scheduling
 behavior

---
 cron.go   | 28 +++++++++++++++-------------
 option.go |  8 ++++++++
 timer.go  | 30 ++++++++++++++++++++++++++++++
 3 files changed, 53 insertions(+), 13 deletions(-)
 create mode 100644 timer.go

diff --git a/cron.go b/cron.go
index c7e91766..9575810a 100644
--- a/cron.go
+++ b/cron.go
@@ -24,6 +24,7 @@ type Cron struct {
 	parser    ScheduleParser
 	nextID    EntryID
 	jobWaiter sync.WaitGroup
+	timerFn   func(time.Duration) Timer
 }
 
 // ScheduleParser is an interface for schedule spec parsers that return a Schedule
@@ -97,17 +98,17 @@ func (s byTime) Less(i, j int) bool {
 //
 // Available Settings
 //
-//   Time Zone
-//     Description: The time zone in which schedules are interpreted
-//     Default:     time.Local
+//	Time Zone
+//	  Description: The time zone in which schedules are interpreted
+//	  Default:     time.Local
 //
-//   Parser
-//     Description: Parser converts cron spec strings into cron.Schedules.
-//     Default:     Accepts this spec: https://en.wikipedia.org/wiki/Cron
+//	Parser
+//	  Description: Parser converts cron spec strings into cron.Schedules.
+//	  Default:     Accepts this spec: https://en.wikipedia.org/wiki/Cron
 //
-//   Chain
-//     Description: Wrap submitted jobs to customize behavior.
-//     Default:     A chain that recovers panics and logs them to stderr.
+//	Chain
+//	  Description: Wrap submitted jobs to customize behavior.
+//	  Default:     A chain that recovers panics and logs them to stderr.
 //
 // See "cron.With*" to modify the default behavior.
 func New(opts ...Option) *Cron {
@@ -123,6 +124,7 @@ func New(opts ...Option) *Cron {
 		logger:    DefaultLogger,
 		location:  time.Local,
 		parser:    standardParser,
+		timerFn:   NewStandardTimer,
 	}
 	for _, opt := range opts {
 		opt(c)
@@ -250,18 +252,18 @@ func (c *Cron) run() {
 		// Determine the next entry to run.
 		sort.Sort(byTime(c.entries))
 
-		var timer *time.Timer
+		var timer Timer
 		if len(c.entries) == 0 || c.entries[0].Next.IsZero() {
 			// If there are no entries yet, just sleep - it still handles new entries
 			// and stop requests.
-			timer = time.NewTimer(100000 * time.Hour)
+			timer = c.timerFn(100000 * time.Hour)
 		} else {
-			timer = time.NewTimer(c.entries[0].Next.Sub(now))
+			timer = c.timerFn(c.entries[0].Next.Sub(now))
 		}
 
 		for {
 			select {
-			case now = <-timer.C:
+			case now = <-timer.C():
 				now = now.In(c.location)
 				c.logger.Info("wake", "now", now)
 
diff --git a/option.go b/option.go
index 09e4278e..1fd46e05 100644
--- a/option.go
+++ b/option.go
@@ -43,3 +43,11 @@ func WithLogger(logger Logger) Option {
 		c.logger = logger
 	}
 }
+
+// WithTimerFunc specifies a Timer creator. If not specified, NewStandardTimer will be used by default. Refer to the Timer documentation for details.
+// This option generally used for testing
+func WithTimerFunc(timer func(time.Duration) Timer) Option {
+	return func(c *Cron) {
+		c.timerFn = timer
+	}
+}
diff --git a/timer.go b/timer.go
new file mode 100644
index 00000000..2f9c340b
--- /dev/null
+++ b/timer.go
@@ -0,0 +1,30 @@
+package cron
+
+import "time"
+
+// The Timer works just like time.Timer, and cron schedules Jobs by waiting for the time event emitted by Timer.
+// By default, the standard Timer returned by NewStandardTimer is used.
+// You can also customize a Timer with WithTimerFunc option to control scheduling behavior.
+type Timer interface {
+	C() <-chan time.Time
+	Stop() bool
+}
+
+// NewStandardTimer returns a Timer created using the standard time.Timer.
+func NewStandardTimer(d time.Duration) Timer {
+	return &standardTimer{
+		timer: time.NewTimer(d),
+	}
+}
+
+type standardTimer struct {
+	timer *time.Timer
+}
+
+func (t *standardTimer) C() <-chan time.Time {
+	return t.timer.C
+}
+
+func (t *standardTimer) Stop() bool {
+	return t.timer.Stop()
+}