1 package clock
2 3 import (
4 "context"
5 "sort"
6 "sync"
7 "time"
8 )
9 10 // Re-export of time.Duration
11 type Duration = time.Duration
12 13 // Clock represents an interface to the functions in the standard library time
14 // package. Two implementations are available in the clock package. The first
15 // is a real-time clock which simply wraps the time package's functions. The
16 // second is a mock clock which will only change when
17 // programmatically adjusted.
18 type Clock interface {
19 After(d time.Duration) <-chan time.Time
20 AfterFunc(d time.Duration, f func()) *Timer
21 Now() time.Time
22 Since(t time.Time) time.Duration
23 Until(t time.Time) time.Duration
24 Sleep(d time.Duration)
25 Tick(d time.Duration) <-chan time.Time
26 Ticker(d time.Duration) *Ticker
27 Timer(d time.Duration) *Timer
28 WithDeadline(parent context.Context, d time.Time) (context.Context, context.CancelFunc)
29 WithTimeout(parent context.Context, t time.Duration) (context.Context, context.CancelFunc)
30 }
31 32 // New returns an instance of a real-time clock.
33 func New() Clock {
34 return &clock{}
35 }
36 37 // clock implements a real-time clock by simply wrapping the time package functions.
38 type clock struct{}
39 40 func (c *clock) After(d time.Duration) <-chan time.Time { return time.After(d) }
41 42 func (c *clock) AfterFunc(d time.Duration, f func()) *Timer {
43 return &Timer{timer: time.AfterFunc(d, f)}
44 }
45 46 func (c *clock) Now() time.Time { return time.Now() }
47 48 func (c *clock) Since(t time.Time) time.Duration { return time.Since(t) }
49 50 func (c *clock) Until(t time.Time) time.Duration { return time.Until(t) }
51 52 func (c *clock) Sleep(d time.Duration) { time.Sleep(d) }
53 54 func (c *clock) Tick(d time.Duration) <-chan time.Time { return time.Tick(d) }
55 56 func (c *clock) Ticker(d time.Duration) *Ticker {
57 t := time.NewTicker(d)
58 return &Ticker{C: t.C, ticker: t}
59 }
60 61 func (c *clock) Timer(d time.Duration) *Timer {
62 t := time.NewTimer(d)
63 return &Timer{C: t.C, timer: t}
64 }
65 66 func (c *clock) WithDeadline(parent context.Context, d time.Time) (context.Context, context.CancelFunc) {
67 return context.WithDeadline(parent, d)
68 }
69 70 func (c *clock) WithTimeout(parent context.Context, t time.Duration) (context.Context, context.CancelFunc) {
71 return context.WithTimeout(parent, t)
72 }
73 74 // Mock represents a mock clock that only moves forward programmically.
75 // It can be preferable to a real-time clock when testing time-based functionality.
76 type Mock struct {
77 // mu protects all other fields in this struct, and the data that they
78 // point to.
79 mu sync.Mutex
80 81 now time.Time // current time
82 timers clockTimers // tickers & timers
83 }
84 85 // NewMock returns an instance of a mock clock.
86 // The current time of the mock clock on initialization is the Unix epoch.
87 func NewMock() *Mock {
88 return &Mock{now: time.Unix(0, 0)}
89 }
90 91 // Add moves the current time of the mock clock forward by the specified duration.
92 // This should only be called from a single goroutine at a time.
93 func (m *Mock) Add(d time.Duration) {
94 // Calculate the final current time.
95 m.mu.Lock()
96 t := m.now.Add(d)
97 m.mu.Unlock()
98 99 // Continue to execute timers until there are no more before the new time.
100 for {
101 if !m.runNextTimer(t) {
102 break
103 }
104 }
105 106 // Ensure that we end with the new time.
107 m.mu.Lock()
108 m.now = t
109 m.mu.Unlock()
110 111 // Give a small buffer to make sure that other goroutines get handled.
112 gosched()
113 }
114 115 // Set sets the current time of the mock clock to a specific one.
116 // This should only be called from a single goroutine at a time.
117 func (m *Mock) Set(t time.Time) {
118 // Continue to execute timers until there are no more before the new time.
119 for {
120 if !m.runNextTimer(t) {
121 break
122 }
123 }
124 125 // Ensure that we end with the new time.
126 m.mu.Lock()
127 m.now = t
128 m.mu.Unlock()
129 130 // Give a small buffer to make sure that other goroutines get handled.
131 gosched()
132 }
133 134 // WaitForAllTimers sets the clock until all timers are expired
135 func (m *Mock) WaitForAllTimers() time.Time {
136 // Continue to execute timers until there are no more
137 for {
138 m.mu.Lock()
139 if len(m.timers) == 0 {
140 m.mu.Unlock()
141 return m.Now()
142 }
143 144 sort.Sort(m.timers)
145 next := m.timers[len(m.timers)-1].Next()
146 m.mu.Unlock()
147 m.Set(next)
148 }
149 }
150 151 // runNextTimer executes the next timer in chronological order and moves the
152 // current time to the timer's next tick time. The next time is not executed if
153 // its next time is after the max time. Returns true if a timer was executed.
154 func (m *Mock) runNextTimer(max time.Time) bool {
155 m.mu.Lock()
156 157 // Sort timers by time.
158 sort.Sort(m.timers)
159 160 // If we have no more timers then exit.
161 if len(m.timers) == 0 {
162 m.mu.Unlock()
163 return false
164 }
165 166 // Retrieve next timer. Exit if next tick is after new time.
167 t := m.timers[0]
168 if t.Next().After(max) {
169 m.mu.Unlock()
170 return false
171 }
172 173 // Move "now" forward and unlock clock.
174 m.now = t.Next()
175 now := m.now
176 m.mu.Unlock()
177 178 // Execute timer.
179 t.Tick(now)
180 return true
181 }
182 183 // After waits for the duration to elapse and then sends the current time on the returned channel.
184 func (m *Mock) After(d time.Duration) <-chan time.Time {
185 return m.Timer(d).C
186 }
187 188 // AfterFunc waits for the duration to elapse and then executes a function in its own goroutine.
189 // A Timer is returned that can be stopped.
190 func (m *Mock) AfterFunc(d time.Duration, f func()) *Timer {
191 m.mu.Lock()
192 defer m.mu.Unlock()
193 ch := make(chan time.Time, 1)
194 t := &Timer{
195 c: ch,
196 fn: f,
197 mock: m,
198 next: m.now.Add(d),
199 stopped: false,
200 }
201 m.timers = append(m.timers, (*internalTimer)(t))
202 return t
203 }
204 205 // Now returns the current wall time on the mock clock.
206 func (m *Mock) Now() time.Time {
207 m.mu.Lock()
208 defer m.mu.Unlock()
209 return m.now
210 }
211 212 // Since returns time since `t` using the mock clock's wall time.
213 func (m *Mock) Since(t time.Time) time.Duration {
214 return m.Now().Sub(t)
215 }
216 217 // Until returns time until `t` using the mock clock's wall time.
218 func (m *Mock) Until(t time.Time) time.Duration {
219 return t.Sub(m.Now())
220 }
221 222 // Sleep pauses the goroutine for the given duration on the mock clock.
223 // The clock must be moved forward in a separate goroutine.
224 func (m *Mock) Sleep(d time.Duration) {
225 <-m.After(d)
226 }
227 228 // Tick is a convenience function for Ticker().
229 // It will return a ticker channel that cannot be stopped.
230 func (m *Mock) Tick(d time.Duration) <-chan time.Time {
231 return m.Ticker(d).C
232 }
233 234 // Ticker creates a new instance of Ticker.
235 func (m *Mock) Ticker(d time.Duration) *Ticker {
236 m.mu.Lock()
237 defer m.mu.Unlock()
238 ch := make(chan time.Time, 1)
239 t := &Ticker{
240 C: ch,
241 c: ch,
242 mock: m,
243 d: d,
244 next: m.now.Add(d),
245 }
246 m.timers = append(m.timers, (*internalTicker)(t))
247 return t
248 }
249 250 // Timer creates a new instance of Timer.
251 func (m *Mock) Timer(d time.Duration) *Timer {
252 m.mu.Lock()
253 ch := make(chan time.Time, 1)
254 t := &Timer{
255 C: ch,
256 c: ch,
257 mock: m,
258 next: m.now.Add(d),
259 stopped: false,
260 }
261 m.timers = append(m.timers, (*internalTimer)(t))
262 now := m.now
263 m.mu.Unlock()
264 m.runNextTimer(now)
265 return t
266 }
267 268 // removeClockTimer removes a timer from m.timers. m.mu MUST be held
269 // when this method is called.
270 func (m *Mock) removeClockTimer(t clockTimer) {
271 for i, timer := range m.timers {
272 if timer == t {
273 copy(m.timers[i:], m.timers[i+1:])
274 m.timers[len(m.timers)-1] = nil
275 m.timers = m.timers[:len(m.timers)-1]
276 break
277 }
278 }
279 sort.Sort(m.timers)
280 }
281 282 // clockTimer represents an object with an associated start time.
283 type clockTimer interface {
284 Next() time.Time
285 Tick(time.Time)
286 }
287 288 // clockTimers represents a list of sortable timers.
289 type clockTimers []clockTimer
290 291 func (a clockTimers) Len() int { return len(a) }
292 func (a clockTimers) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
293 func (a clockTimers) Less(i, j int) bool { return a[i].Next().Before(a[j].Next()) }
294 295 // Timer represents a single event.
296 // The current time will be sent on C, unless the timer was created by AfterFunc.
297 type Timer struct {
298 C <-chan time.Time
299 c chan time.Time
300 timer *time.Timer // realtime impl, if set
301 next time.Time // next tick time
302 mock *Mock // mock clock, if set
303 fn func() // AfterFunc function, if set
304 stopped bool // True if stopped, false if running
305 }
306 307 // Stop turns off the ticker.
308 func (t *Timer) Stop() bool {
309 if t.timer != nil {
310 return t.timer.Stop()
311 }
312 313 t.mock.mu.Lock()
314 registered := !t.stopped
315 t.mock.removeClockTimer((*internalTimer)(t))
316 t.stopped = true
317 t.mock.mu.Unlock()
318 return registered
319 }
320 321 // Reset changes the expiry time of the timer
322 func (t *Timer) Reset(d time.Duration) bool {
323 if t.timer != nil {
324 return t.timer.Reset(d)
325 }
326 327 t.mock.mu.Lock()
328 t.next = t.mock.now.Add(d)
329 defer t.mock.mu.Unlock()
330 331 registered := !t.stopped
332 if t.stopped {
333 t.mock.timers = append(t.mock.timers, (*internalTimer)(t))
334 }
335 336 t.stopped = false
337 return registered
338 }
339 340 type internalTimer Timer
341 342 func (t *internalTimer) Next() time.Time { return t.next }
343 func (t *internalTimer) Tick(now time.Time) {
344 // a gosched() after ticking, to allow any consequences of the
345 // tick to complete
346 defer gosched()
347 348 t.mock.mu.Lock()
349 if t.fn != nil {
350 // defer function execution until the lock is released, and
351 defer func() { go t.fn() }()
352 } else {
353 t.c <- now
354 }
355 t.mock.removeClockTimer((*internalTimer)(t))
356 t.stopped = true
357 t.mock.mu.Unlock()
358 }
359 360 // Ticker holds a channel that receives "ticks" at regular intervals.
361 type Ticker struct {
362 C <-chan time.Time
363 c chan time.Time
364 ticker *time.Ticker // realtime impl, if set
365 next time.Time // next tick time
366 mock *Mock // mock clock, if set
367 d time.Duration // time between ticks
368 stopped bool // True if stopped, false if running
369 }
370 371 // Stop turns off the ticker.
372 func (t *Ticker) Stop() {
373 if t.ticker != nil {
374 t.ticker.Stop()
375 } else {
376 t.mock.mu.Lock()
377 t.mock.removeClockTimer((*internalTicker)(t))
378 t.stopped = true
379 t.mock.mu.Unlock()
380 }
381 }
382 383 // Reset resets the ticker to a new duration.
384 func (t *Ticker) Reset(dur time.Duration) {
385 if t.ticker != nil {
386 t.ticker.Reset(dur)
387 return
388 }
389 390 t.mock.mu.Lock()
391 defer t.mock.mu.Unlock()
392 393 if t.stopped {
394 t.mock.timers = append(t.mock.timers, (*internalTicker)(t))
395 t.stopped = false
396 }
397 398 t.d = dur
399 t.next = t.mock.now.Add(dur)
400 }
401 402 type internalTicker Ticker
403 404 func (t *internalTicker) Next() time.Time { return t.next }
405 func (t *internalTicker) Tick(now time.Time) {
406 select {
407 case t.c <- now:
408 default:
409 }
410 t.mock.mu.Lock()
411 t.next = now.Add(t.d)
412 t.mock.mu.Unlock()
413 gosched()
414 }
415 416 // Sleep momentarily so that other goroutines can process.
417 func gosched() { time.Sleep(1 * time.Millisecond) }
418 419 var (
420 // type checking
421 _ Clock = &Mock{}
422 )
423