blob: 85a993467dcfe4a87c96e27f619e5575d1466012 [file] [log] [blame]
Abhay Kumar40252eb2025-10-13 13:25:53 +00001// Package clockwork contains a simple fake clock for Go.
khenaidoo59ce9dd2019-11-11 13:05:32 -05002package clockwork
3
4import (
Abhay Kumar40252eb2025-10-13 13:25:53 +00005 "context"
6 "errors"
7 "slices"
khenaidoo59ce9dd2019-11-11 13:05:32 -05008 "sync"
9 "time"
10)
11
Abhay Kumar40252eb2025-10-13 13:25:53 +000012// Clock provides an interface that packages can use instead of directly using
13// the [time] module, so that chronology-related behavior can be tested.
khenaidoo59ce9dd2019-11-11 13:05:32 -050014type Clock interface {
15 After(d time.Duration) <-chan time.Time
16 Sleep(d time.Duration)
17 Now() time.Time
khenaidoo26721882021-08-11 17:42:52 -040018 Since(t time.Time) time.Duration
Abhay Kumar40252eb2025-10-13 13:25:53 +000019 Until(t time.Time) time.Duration
khenaidoo26721882021-08-11 17:42:52 -040020 NewTicker(d time.Duration) Ticker
Abhay Kumar40252eb2025-10-13 13:25:53 +000021 NewTimer(d time.Duration) Timer
22 AfterFunc(d time.Duration, f func()) Timer
khenaidoo59ce9dd2019-11-11 13:05:32 -050023}
24
25// NewRealClock returns a Clock which simply delegates calls to the actual time
26// package; it should be used by packages in production.
27func NewRealClock() Clock {
28 return &realClock{}
29}
30
khenaidoo59ce9dd2019-11-11 13:05:32 -050031type realClock struct{}
32
33func (rc *realClock) After(d time.Duration) <-chan time.Time {
34 return time.After(d)
35}
36
37func (rc *realClock) Sleep(d time.Duration) {
38 time.Sleep(d)
39}
40
41func (rc *realClock) Now() time.Time {
42 return time.Now()
43}
44
khenaidoo26721882021-08-11 17:42:52 -040045func (rc *realClock) Since(t time.Time) time.Duration {
46 return rc.Now().Sub(t)
47}
48
Abhay Kumar40252eb2025-10-13 13:25:53 +000049func (rc *realClock) Until(t time.Time) time.Duration {
50 return t.Sub(rc.Now())
khenaidoo26721882021-08-11 17:42:52 -040051}
52
Abhay Kumar40252eb2025-10-13 13:25:53 +000053func (rc *realClock) NewTicker(d time.Duration) Ticker {
54 return realTicker{time.NewTicker(d)}
55}
56
57func (rc *realClock) NewTimer(d time.Duration) Timer {
58 return realTimer{time.NewTimer(d)}
59}
60
61func (rc *realClock) AfterFunc(d time.Duration, f func()) Timer {
62 return realTimer{time.AfterFunc(d, f)}
63}
64
65// FakeClock provides an interface for a clock which can be manually advanced
66// through time.
67//
68// FakeClock maintains a list of "waiters," which consists of all callers
69// waiting on the underlying clock (i.e. Tickers and Timers including callers of
70// Sleep or After). Users can call BlockUntil to block until the clock has an
71// expected number of waiters.
72type FakeClock struct {
73 // l protects all attributes of the clock, including all attributes of all
74 // waiters and blockers.
75 l sync.RWMutex
76 waiters []expirer
khenaidoo59ce9dd2019-11-11 13:05:32 -050077 blockers []*blocker
78 time time.Time
khenaidoo59ce9dd2019-11-11 13:05:32 -050079}
80
Abhay Kumar40252eb2025-10-13 13:25:53 +000081// NewFakeClock returns a FakeClock implementation which can be
82// manually advanced through time for testing. The initial time of the
83// FakeClock will be the current system time.
84//
85// Tests that require a deterministic time must use NewFakeClockAt.
86func NewFakeClock() *FakeClock {
87 return NewFakeClockAt(time.Now())
khenaidoo59ce9dd2019-11-11 13:05:32 -050088}
89
Abhay Kumar40252eb2025-10-13 13:25:53 +000090// NewFakeClockAt returns a FakeClock initialised at the given time.Time.
91func NewFakeClockAt(t time.Time) *FakeClock {
92 return &FakeClock{
93 time: t,
94 }
95}
96
97// blocker is a caller of BlockUntil.
khenaidoo59ce9dd2019-11-11 13:05:32 -050098type blocker struct {
99 count int
Abhay Kumar40252eb2025-10-13 13:25:53 +0000100
101 // ch is closed when the underlying clock has the specified number of blockers.
102 ch chan struct{}
khenaidoo59ce9dd2019-11-11 13:05:32 -0500103}
104
Abhay Kumar40252eb2025-10-13 13:25:53 +0000105// expirer is a timer or ticker that expires at some point in the future.
106type expirer interface {
107 // expire the expirer at the given time, returning the desired duration until
108 // the next expiration, if any.
109 expire(now time.Time) (next *time.Duration)
110
111 // Get and set the expiration time.
112 expiration() time.Time
113 setExpiration(time.Time)
114}
115
116// After mimics [time.After]; it waits for the given duration to elapse on the
khenaidoo59ce9dd2019-11-11 13:05:32 -0500117// fakeClock, then sends the current time on the returned channel.
Abhay Kumar40252eb2025-10-13 13:25:53 +0000118func (fc *FakeClock) After(d time.Duration) <-chan time.Time {
119 return fc.NewTimer(d).Chan()
khenaidoo59ce9dd2019-11-11 13:05:32 -0500120}
121
Abhay Kumar40252eb2025-10-13 13:25:53 +0000122// Sleep blocks until the given duration has passed on the fakeClock.
123func (fc *FakeClock) Sleep(d time.Duration) {
khenaidoo59ce9dd2019-11-11 13:05:32 -0500124 <-fc.After(d)
125}
126
Abhay Kumar40252eb2025-10-13 13:25:53 +0000127// Now returns the current time of the fakeClock
128func (fc *FakeClock) Now() time.Time {
khenaidoo59ce9dd2019-11-11 13:05:32 -0500129 fc.l.RLock()
Abhay Kumar40252eb2025-10-13 13:25:53 +0000130 defer fc.l.RUnlock()
131 return fc.time
khenaidoo59ce9dd2019-11-11 13:05:32 -0500132}
133
Abhay Kumar40252eb2025-10-13 13:25:53 +0000134// Since returns the duration that has passed since the given time on the
135// fakeClock.
136func (fc *FakeClock) Since(t time.Time) time.Duration {
khenaidoo26721882021-08-11 17:42:52 -0400137 return fc.Now().Sub(t)
138}
139
Abhay Kumar40252eb2025-10-13 13:25:53 +0000140// Until returns the duration that has to pass from the given time on the fakeClock
141// to reach the given time.
142func (fc *FakeClock) Until(t time.Time) time.Duration {
143 return t.Sub(fc.Now())
144}
145
146// NewTicker returns a Ticker that will expire only after calls to
147// FakeClock.Advance() have moved the clock past the given duration.
148//
149// The duration d must be greater than zero; if not, NewTicker will panic.
150func (fc *FakeClock) NewTicker(d time.Duration) Ticker {
151 // Maintain parity with
152 // https://cs.opensource.google/go/go/+/refs/tags/go1.20.3:src/time/tick.go;l=23-25
153 if d <= 0 {
154 panic(errors.New("non-positive interval for NewTicker"))
khenaidoo26721882021-08-11 17:42:52 -0400155 }
Abhay Kumar40252eb2025-10-13 13:25:53 +0000156 ft := newFakeTicker(fc, d)
157 fc.l.Lock()
158 defer fc.l.Unlock()
159 fc.setExpirer(ft, d)
khenaidoo26721882021-08-11 17:42:52 -0400160 return ft
161}
162
Abhay Kumar40252eb2025-10-13 13:25:53 +0000163// NewTimer returns a Timer that will fire only after calls to
164// fakeClock.Advance() have moved the clock past the given duration.
165func (fc *FakeClock) NewTimer(d time.Duration) Timer {
166 t, _ := fc.newTimer(d, nil)
167 return t
168}
169
170// AfterFunc mimics [time.AfterFunc]; it returns a Timer that will invoke the
171// given function only after calls to fakeClock.Advance() have moved the clock
172// past the given duration.
173func (fc *FakeClock) AfterFunc(d time.Duration, f func()) Timer {
174 t, _ := fc.newTimer(d, f)
175 return t
176}
177
178// newTimer returns a new timer using an optional afterFunc and the time that
179// timer expires.
180func (fc *FakeClock) newTimer(d time.Duration, afterfunc func()) (*fakeTimer, time.Time) {
181 ft := newFakeTimer(fc, afterfunc)
182 fc.l.Lock()
183 defer fc.l.Unlock()
184 fc.setExpirer(ft, d)
185 return ft, ft.expiration()
186}
187
188// newTimerAtTime is like newTimer, but uses a time instead of a duration.
189//
190// It is used to ensure FakeClock's lock is held constant through calling
191// fc.After(t.Sub(fc.Now())). It should not be exposed externally.
192func (fc *FakeClock) newTimerAtTime(t time.Time, afterfunc func()) *fakeTimer {
193 ft := newFakeTimer(fc, afterfunc)
194 fc.l.Lock()
195 defer fc.l.Unlock()
196 fc.setExpirer(ft, t.Sub(fc.time))
197 return ft
198}
199
200// Advance advances fakeClock to a new point in time, ensuring waiters and
201// blockers are notified appropriately before returning.
202func (fc *FakeClock) Advance(d time.Duration) {
khenaidoo59ce9dd2019-11-11 13:05:32 -0500203 fc.l.Lock()
204 defer fc.l.Unlock()
205 end := fc.time.Add(d)
Abhay Kumar40252eb2025-10-13 13:25:53 +0000206 // Expire the earliest waiter until the earliest waiter's expiration is after
207 // end.
208 //
209 // We don't iterate because the callback of the waiter might register a new
210 // waiter, so the list of waiters might change as we execute this.
211 for len(fc.waiters) > 0 && !end.Before(fc.waiters[0].expiration()) {
212 w := fc.waiters[0]
213 fc.waiters = fc.waiters[1:]
214
215 // Use the waiter's expiration as the current time for this expiration.
216 now := w.expiration()
217 fc.time = now
218 if d := w.expire(now); d != nil {
219 // Set the new expiration if needed.
220 fc.setExpirer(w, *d)
khenaidoo59ce9dd2019-11-11 13:05:32 -0500221 }
222 }
khenaidoo59ce9dd2019-11-11 13:05:32 -0500223 fc.time = end
224}
225
Abhay Kumar40252eb2025-10-13 13:25:53 +0000226// BlockUntil blocks until the FakeClock has the given number of waiters.
227//
228// Prefer BlockUntilContext in new code, which offers context cancellation to
229// prevent deadlock.
230//
231// Deprecated: New code should prefer BlockUntilContext.
232func (fc *FakeClock) BlockUntil(n int) {
233 fc.BlockUntilContext(context.TODO(), n)
234}
235
236// BlockUntilContext blocks until the fakeClock has the given number of waiters
237// or the context is cancelled.
238func (fc *FakeClock) BlockUntilContext(ctx context.Context, n int) error {
239 b := fc.newBlocker(n)
240 if b == nil {
241 return nil
khenaidoo59ce9dd2019-11-11 13:05:32 -0500242 }
Abhay Kumar40252eb2025-10-13 13:25:53 +0000243
244 select {
245 case <-b.ch:
246 return nil
247 case <-ctx.Done():
248 return ctx.Err()
249 }
250}
251
252func (fc *FakeClock) newBlocker(n int) *blocker {
253 fc.l.Lock()
254 defer fc.l.Unlock()
255 // Fast path: we already have >= n waiters.
256 if len(fc.waiters) >= n {
257 return nil
258 }
259 // Set up a new blocker to wait for more waiters.
khenaidoo59ce9dd2019-11-11 13:05:32 -0500260 b := &blocker{
261 count: n,
262 ch: make(chan struct{}),
263 }
264 fc.blockers = append(fc.blockers, b)
Abhay Kumar40252eb2025-10-13 13:25:53 +0000265 return b
266}
267
268// stop stops an expirer, returning true if the expirer was stopped.
269func (fc *FakeClock) stop(e expirer) bool {
270 fc.l.Lock()
271 defer fc.l.Unlock()
272 return fc.stopExpirer(e)
273}
274
275// stopExpirer stops an expirer, returning true if the expirer was stopped.
276//
277// The caller must hold fc.l.
278func (fc *FakeClock) stopExpirer(e expirer) bool {
279 idx := slices.Index(fc.waiters, e)
280 if idx == -1 {
281 return false
282 }
283 // Remove element, maintaining order, setting inaccessible elements to nil so
284 // they can be garbage collected.
285 copy(fc.waiters[idx:], fc.waiters[idx+1:])
286 fc.waiters[len(fc.waiters)-1] = nil
287 fc.waiters = fc.waiters[:len(fc.waiters)-1]
288 return true
289}
290
291// setExpirer sets an expirer to expire at a future point in time.
292//
293// The caller must hold fc.l.
294func (fc *FakeClock) setExpirer(e expirer, d time.Duration) {
295 if d.Nanoseconds() <= 0 {
296 // Special case for timers with duration <= 0: trigger immediately, never
297 // reset.
298 //
299 // Tickers never get here, they panic if d is < 0.
300 e.expire(fc.time)
301 return
302 }
303 // Add the expirer to the set of waiters and notify any blockers.
304 e.setExpiration(fc.time.Add(d))
305 fc.waiters = append(fc.waiters, e)
306 slices.SortFunc(fc.waiters, func(a, b expirer) int {
307 return a.expiration().Compare(b.expiration())
308 })
309
310 // Notify blockers of our new waiter.
311 count := len(fc.waiters)
312 fc.blockers = slices.DeleteFunc(fc.blockers, func(b *blocker) bool {
313 if b.count <= count {
314 close(b.ch)
315 return true
316 }
317 return false
318 })
khenaidoo59ce9dd2019-11-11 13:05:32 -0500319}