blob: 9567f90060d996625726eacbd4e72a1787875db4 [file] [log] [blame]
Zack Williamse940c7a2019-08-21 14:25:39 -07001/*
2Copyright 2014 The Kubernetes Authors.
3
4Licensed under the Apache License, Version 2.0 (the "License");
5you may not use this file except in compliance with the License.
6You may obtain a copy of the License at
7
8 http://www.apache.org/licenses/LICENSE-2.0
9
10Unless required by applicable law or agreed to in writing, software
11distributed under the License is distributed on an "AS IS" BASIS,
12WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13See the License for the specific language governing permissions and
14limitations under the License.
15*/
16
17package clock
18
19import (
20 "sync"
21 "time"
22)
23
24// Clock allows for injecting fake or real clocks into code that
25// needs to do arbitrary things based on time.
26type Clock interface {
27 Now() time.Time
28 Since(time.Time) time.Duration
29 After(time.Duration) <-chan time.Time
30 NewTimer(time.Duration) Timer
31 Sleep(time.Duration)
32 NewTicker(time.Duration) Ticker
33}
34
35// RealClock really calls time.Now()
36type RealClock struct{}
37
38// Now returns the current time.
39func (RealClock) Now() time.Time {
40 return time.Now()
41}
42
43// Since returns time since the specified timestamp.
44func (RealClock) Since(ts time.Time) time.Duration {
45 return time.Since(ts)
46}
47
48// Same as time.After(d).
49func (RealClock) After(d time.Duration) <-chan time.Time {
50 return time.After(d)
51}
52
53func (RealClock) NewTimer(d time.Duration) Timer {
54 return &realTimer{
55 timer: time.NewTimer(d),
56 }
57}
58
59func (RealClock) NewTicker(d time.Duration) Ticker {
60 return &realTicker{
61 ticker: time.NewTicker(d),
62 }
63}
64
65func (RealClock) Sleep(d time.Duration) {
66 time.Sleep(d)
67}
68
69// FakeClock implements Clock, but returns an arbitrary time.
70type FakeClock struct {
71 lock sync.RWMutex
72 time time.Time
73
74 // waiters are waiting for the fake time to pass their specified time
75 waiters []fakeClockWaiter
76}
77
78type fakeClockWaiter struct {
79 targetTime time.Time
80 stepInterval time.Duration
81 skipIfBlocked bool
82 destChan chan time.Time
83 fired bool
84}
85
86func NewFakeClock(t time.Time) *FakeClock {
87 return &FakeClock{
88 time: t,
89 }
90}
91
92// Now returns f's time.
93func (f *FakeClock) Now() time.Time {
94 f.lock.RLock()
95 defer f.lock.RUnlock()
96 return f.time
97}
98
99// Since returns time since the time in f.
100func (f *FakeClock) Since(ts time.Time) time.Duration {
101 f.lock.RLock()
102 defer f.lock.RUnlock()
103 return f.time.Sub(ts)
104}
105
106// Fake version of time.After(d).
107func (f *FakeClock) After(d time.Duration) <-chan time.Time {
108 f.lock.Lock()
109 defer f.lock.Unlock()
110 stopTime := f.time.Add(d)
111 ch := make(chan time.Time, 1) // Don't block!
112 f.waiters = append(f.waiters, fakeClockWaiter{
113 targetTime: stopTime,
114 destChan: ch,
115 })
116 return ch
117}
118
119// Fake version of time.NewTimer(d).
120func (f *FakeClock) NewTimer(d time.Duration) Timer {
121 f.lock.Lock()
122 defer f.lock.Unlock()
123 stopTime := f.time.Add(d)
124 ch := make(chan time.Time, 1) // Don't block!
125 timer := &fakeTimer{
126 fakeClock: f,
127 waiter: fakeClockWaiter{
128 targetTime: stopTime,
129 destChan: ch,
130 },
131 }
132 f.waiters = append(f.waiters, timer.waiter)
133 return timer
134}
135
136func (f *FakeClock) NewTicker(d time.Duration) Ticker {
137 f.lock.Lock()
138 defer f.lock.Unlock()
139 tickTime := f.time.Add(d)
140 ch := make(chan time.Time, 1) // hold one tick
141 f.waiters = append(f.waiters, fakeClockWaiter{
142 targetTime: tickTime,
143 stepInterval: d,
144 skipIfBlocked: true,
145 destChan: ch,
146 })
147
148 return &fakeTicker{
149 c: ch,
150 }
151}
152
153// Move clock by Duration, notify anyone that's called After, Tick, or NewTimer
154func (f *FakeClock) Step(d time.Duration) {
155 f.lock.Lock()
156 defer f.lock.Unlock()
157 f.setTimeLocked(f.time.Add(d))
158}
159
160// Sets the time.
161func (f *FakeClock) SetTime(t time.Time) {
162 f.lock.Lock()
163 defer f.lock.Unlock()
164 f.setTimeLocked(t)
165}
166
167// Actually changes the time and checks any waiters. f must be write-locked.
168func (f *FakeClock) setTimeLocked(t time.Time) {
169 f.time = t
170 newWaiters := make([]fakeClockWaiter, 0, len(f.waiters))
171 for i := range f.waiters {
172 w := &f.waiters[i]
173 if !w.targetTime.After(t) {
174
175 if w.skipIfBlocked {
176 select {
177 case w.destChan <- t:
178 w.fired = true
179 default:
180 }
181 } else {
182 w.destChan <- t
183 w.fired = true
184 }
185
186 if w.stepInterval > 0 {
187 for !w.targetTime.After(t) {
188 w.targetTime = w.targetTime.Add(w.stepInterval)
189 }
190 newWaiters = append(newWaiters, *w)
191 }
192
193 } else {
194 newWaiters = append(newWaiters, f.waiters[i])
195 }
196 }
197 f.waiters = newWaiters
198}
199
200// Returns true if After has been called on f but not yet satisfied (so you can
201// write race-free tests).
202func (f *FakeClock) HasWaiters() bool {
203 f.lock.RLock()
204 defer f.lock.RUnlock()
205 return len(f.waiters) > 0
206}
207
208func (f *FakeClock) Sleep(d time.Duration) {
209 f.Step(d)
210}
211
212// IntervalClock implements Clock, but each invocation of Now steps the clock forward the specified duration
213type IntervalClock struct {
214 Time time.Time
215 Duration time.Duration
216}
217
218// Now returns i's time.
219func (i *IntervalClock) Now() time.Time {
220 i.Time = i.Time.Add(i.Duration)
221 return i.Time
222}
223
224// Since returns time since the time in i.
225func (i *IntervalClock) Since(ts time.Time) time.Duration {
226 return i.Time.Sub(ts)
227}
228
229// Unimplemented, will panic.
230// TODO: make interval clock use FakeClock so this can be implemented.
231func (*IntervalClock) After(d time.Duration) <-chan time.Time {
232 panic("IntervalClock doesn't implement After")
233}
234
235// Unimplemented, will panic.
236// TODO: make interval clock use FakeClock so this can be implemented.
237func (*IntervalClock) NewTimer(d time.Duration) Timer {
238 panic("IntervalClock doesn't implement NewTimer")
239}
240
241// Unimplemented, will panic.
242// TODO: make interval clock use FakeClock so this can be implemented.
243func (*IntervalClock) NewTicker(d time.Duration) Ticker {
244 panic("IntervalClock doesn't implement NewTicker")
245}
246
247func (*IntervalClock) Sleep(d time.Duration) {
248 panic("IntervalClock doesn't implement Sleep")
249}
250
251// Timer allows for injecting fake or real timers into code that
252// needs to do arbitrary things based on time.
253type Timer interface {
254 C() <-chan time.Time
255 Stop() bool
256 Reset(d time.Duration) bool
257}
258
259// realTimer is backed by an actual time.Timer.
260type realTimer struct {
261 timer *time.Timer
262}
263
264// C returns the underlying timer's channel.
265func (r *realTimer) C() <-chan time.Time {
266 return r.timer.C
267}
268
269// Stop calls Stop() on the underlying timer.
270func (r *realTimer) Stop() bool {
271 return r.timer.Stop()
272}
273
274// Reset calls Reset() on the underlying timer.
275func (r *realTimer) Reset(d time.Duration) bool {
276 return r.timer.Reset(d)
277}
278
279// fakeTimer implements Timer based on a FakeClock.
280type fakeTimer struct {
281 fakeClock *FakeClock
282 waiter fakeClockWaiter
283}
284
285// C returns the channel that notifies when this timer has fired.
286func (f *fakeTimer) C() <-chan time.Time {
287 return f.waiter.destChan
288}
289
290// Stop stops the timer and returns true if the timer has not yet fired, or false otherwise.
291func (f *fakeTimer) Stop() bool {
292 f.fakeClock.lock.Lock()
293 defer f.fakeClock.lock.Unlock()
294
295 newWaiters := make([]fakeClockWaiter, 0, len(f.fakeClock.waiters))
296 for i := range f.fakeClock.waiters {
297 w := &f.fakeClock.waiters[i]
298 if w != &f.waiter {
299 newWaiters = append(newWaiters, *w)
300 }
301 }
302
303 f.fakeClock.waiters = newWaiters
304
305 return !f.waiter.fired
306}
307
308// Reset resets the timer to the fake clock's "now" + d. It returns true if the timer has not yet
309// fired, or false otherwise.
310func (f *fakeTimer) Reset(d time.Duration) bool {
311 f.fakeClock.lock.Lock()
312 defer f.fakeClock.lock.Unlock()
313
314 active := !f.waiter.fired
315
316 f.waiter.fired = false
317 f.waiter.targetTime = f.fakeClock.time.Add(d)
318
319 return active
320}
321
322type Ticker interface {
323 C() <-chan time.Time
324 Stop()
325}
326
327type realTicker struct {
328 ticker *time.Ticker
329}
330
331func (t *realTicker) C() <-chan time.Time {
332 return t.ticker.C
333}
334
335func (t *realTicker) Stop() {
336 t.ticker.Stop()
337}
338
339type fakeTicker struct {
340 c <-chan time.Time
341}
342
343func (t *fakeTicker) C() <-chan time.Time {
344 return t.c
345}
346
347func (t *fakeTicker) Stop() {
348}