blob: 1018051f4af775f38b0daa21c3f9b73614ce9387 [file] [log] [blame]
khenaidoo59ce9dd2019-11-11 13:05:32 -05001package clockwork
2
3import (
4 "sync"
5 "time"
6)
7
8// Clock provides an interface that packages can use instead of directly
9// using the time module, so that chronology-related behavior can be tested
10type Clock interface {
11 After(d time.Duration) <-chan time.Time
12 Sleep(d time.Duration)
13 Now() time.Time
khenaidoo26721882021-08-11 17:42:52 -040014 Since(t time.Time) time.Duration
15 NewTicker(d time.Duration) Ticker
khenaidoo59ce9dd2019-11-11 13:05:32 -050016}
17
18// FakeClock provides an interface for a clock which can be
19// manually advanced through time
20type FakeClock interface {
21 Clock
22 // Advance advances the FakeClock to a new point in time, ensuring any existing
23 // sleepers are notified appropriately before returning
24 Advance(d time.Duration)
25 // BlockUntil will block until the FakeClock has the given number of
26 // sleepers (callers of Sleep or After)
27 BlockUntil(n int)
28}
29
30// NewRealClock returns a Clock which simply delegates calls to the actual time
31// package; it should be used by packages in production.
32func NewRealClock() Clock {
33 return &realClock{}
34}
35
36// NewFakeClock returns a FakeClock implementation which can be
37// manually advanced through time for testing. The initial time of the
38// FakeClock will be an arbitrary non-zero time.
39func NewFakeClock() FakeClock {
40 // use a fixture that does not fulfill Time.IsZero()
41 return NewFakeClockAt(time.Date(1984, time.April, 4, 0, 0, 0, 0, time.UTC))
42}
43
44// NewFakeClockAt returns a FakeClock initialised at the given time.Time.
45func NewFakeClockAt(t time.Time) FakeClock {
46 return &fakeClock{
47 time: t,
48 }
49}
50
51type realClock struct{}
52
53func (rc *realClock) After(d time.Duration) <-chan time.Time {
54 return time.After(d)
55}
56
57func (rc *realClock) Sleep(d time.Duration) {
58 time.Sleep(d)
59}
60
61func (rc *realClock) Now() time.Time {
62 return time.Now()
63}
64
khenaidoo26721882021-08-11 17:42:52 -040065func (rc *realClock) Since(t time.Time) time.Duration {
66 return rc.Now().Sub(t)
67}
68
69func (rc *realClock) NewTicker(d time.Duration) Ticker {
70 return &realTicker{time.NewTicker(d)}
71}
72
khenaidoo59ce9dd2019-11-11 13:05:32 -050073type fakeClock struct {
74 sleepers []*sleeper
75 blockers []*blocker
76 time time.Time
77
78 l sync.RWMutex
79}
80
81// sleeper represents a caller of After or Sleep
82type sleeper struct {
83 until time.Time
84 done chan time.Time
85}
86
87// blocker represents a caller of BlockUntil
88type blocker struct {
89 count int
90 ch chan struct{}
91}
92
93// After mimics time.After; it waits for the given duration to elapse on the
94// fakeClock, then sends the current time on the returned channel.
95func (fc *fakeClock) After(d time.Duration) <-chan time.Time {
96 fc.l.Lock()
97 defer fc.l.Unlock()
98 now := fc.time
99 done := make(chan time.Time, 1)
khenaidoo26721882021-08-11 17:42:52 -0400100 if d.Nanoseconds() <= 0 {
khenaidoo59ce9dd2019-11-11 13:05:32 -0500101 // special case - trigger immediately
102 done <- now
103 } else {
104 // otherwise, add to the set of sleepers
105 s := &sleeper{
106 until: now.Add(d),
107 done: done,
108 }
109 fc.sleepers = append(fc.sleepers, s)
110 // and notify any blockers
111 fc.blockers = notifyBlockers(fc.blockers, len(fc.sleepers))
112 }
113 return done
114}
115
116// notifyBlockers notifies all the blockers waiting until the
117// given number of sleepers are waiting on the fakeClock. It
118// returns an updated slice of blockers (i.e. those still waiting)
119func notifyBlockers(blockers []*blocker, count int) (newBlockers []*blocker) {
120 for _, b := range blockers {
121 if b.count == count {
122 close(b.ch)
123 } else {
124 newBlockers = append(newBlockers, b)
125 }
126 }
127 return
128}
129
130// Sleep blocks until the given duration has passed on the fakeClock
131func (fc *fakeClock) Sleep(d time.Duration) {
132 <-fc.After(d)
133}
134
135// Time returns the current time of the fakeClock
136func (fc *fakeClock) Now() time.Time {
137 fc.l.RLock()
138 t := fc.time
139 fc.l.RUnlock()
140 return t
141}
142
khenaidoo26721882021-08-11 17:42:52 -0400143// Since returns the duration that has passed since the given time on the fakeClock
144func (fc *fakeClock) Since(t time.Time) time.Duration {
145 return fc.Now().Sub(t)
146}
147
148func (fc *fakeClock) NewTicker(d time.Duration) Ticker {
149 ft := &fakeTicker{
150 c: make(chan time.Time, 1),
151 stop: make(chan bool, 1),
152 clock: fc,
153 period: d,
154 }
155 ft.runTickThread()
156 return ft
157}
158
khenaidoo59ce9dd2019-11-11 13:05:32 -0500159// Advance advances fakeClock to a new point in time, ensuring channels from any
160// previous invocations of After are notified appropriately before returning
161func (fc *fakeClock) Advance(d time.Duration) {
162 fc.l.Lock()
163 defer fc.l.Unlock()
164 end := fc.time.Add(d)
165 var newSleepers []*sleeper
166 for _, s := range fc.sleepers {
167 if end.Sub(s.until) >= 0 {
168 s.done <- end
169 } else {
170 newSleepers = append(newSleepers, s)
171 }
172 }
173 fc.sleepers = newSleepers
174 fc.blockers = notifyBlockers(fc.blockers, len(fc.sleepers))
175 fc.time = end
176}
177
178// BlockUntil will block until the fakeClock has the given number of sleepers
179// (callers of Sleep or After)
180func (fc *fakeClock) BlockUntil(n int) {
181 fc.l.Lock()
182 // Fast path: current number of sleepers is what we're looking for
183 if len(fc.sleepers) == n {
184 fc.l.Unlock()
185 return
186 }
187 // Otherwise, set up a new blocker
188 b := &blocker{
189 count: n,
190 ch: make(chan struct{}),
191 }
192 fc.blockers = append(fc.blockers, b)
193 fc.l.Unlock()
194 <-b.ch
195}