| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 1 | // Package clockwork contains a simple fake clock for Go. |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 2 | package clockwork |
| 3 | |
| 4 | import ( |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 5 | "context" |
| 6 | "errors" |
| 7 | "slices" |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 8 | "sync" |
| 9 | "time" |
| 10 | ) |
| 11 | |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 12 | // Clock provides an interface that packages can use instead of directly using |
| 13 | // the [time] module, so that chronology-related behavior can be tested. |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 14 | type Clock interface { |
| 15 | After(d time.Duration) <-chan time.Time |
| 16 | Sleep(d time.Duration) |
| 17 | Now() time.Time |
| khenaidoo | 2672188 | 2021-08-11 17:42:52 -0400 | [diff] [blame] | 18 | Since(t time.Time) time.Duration |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 19 | Until(t time.Time) time.Duration |
| khenaidoo | 2672188 | 2021-08-11 17:42:52 -0400 | [diff] [blame] | 20 | NewTicker(d time.Duration) Ticker |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 21 | NewTimer(d time.Duration) Timer |
| 22 | AfterFunc(d time.Duration, f func()) Timer |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 23 | } |
| 24 | |
| 25 | // NewRealClock returns a Clock which simply delegates calls to the actual time |
| 26 | // package; it should be used by packages in production. |
| 27 | func NewRealClock() Clock { |
| 28 | return &realClock{} |
| 29 | } |
| 30 | |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 31 | type realClock struct{} |
| 32 | |
| 33 | func (rc *realClock) After(d time.Duration) <-chan time.Time { |
| 34 | return time.After(d) |
| 35 | } |
| 36 | |
| 37 | func (rc *realClock) Sleep(d time.Duration) { |
| 38 | time.Sleep(d) |
| 39 | } |
| 40 | |
| 41 | func (rc *realClock) Now() time.Time { |
| 42 | return time.Now() |
| 43 | } |
| 44 | |
| khenaidoo | 2672188 | 2021-08-11 17:42:52 -0400 | [diff] [blame] | 45 | func (rc *realClock) Since(t time.Time) time.Duration { |
| 46 | return rc.Now().Sub(t) |
| 47 | } |
| 48 | |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 49 | func (rc *realClock) Until(t time.Time) time.Duration { |
| 50 | return t.Sub(rc.Now()) |
| khenaidoo | 2672188 | 2021-08-11 17:42:52 -0400 | [diff] [blame] | 51 | } |
| 52 | |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 53 | func (rc *realClock) NewTicker(d time.Duration) Ticker { |
| 54 | return realTicker{time.NewTicker(d)} |
| 55 | } |
| 56 | |
| 57 | func (rc *realClock) NewTimer(d time.Duration) Timer { |
| 58 | return realTimer{time.NewTimer(d)} |
| 59 | } |
| 60 | |
| 61 | func (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. |
| 72 | type 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 |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 77 | blockers []*blocker |
| 78 | time time.Time |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 79 | } |
| 80 | |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 81 | // 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. |
| 86 | func NewFakeClock() *FakeClock { |
| 87 | return NewFakeClockAt(time.Now()) |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 88 | } |
| 89 | |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 90 | // NewFakeClockAt returns a FakeClock initialised at the given time.Time. |
| 91 | func NewFakeClockAt(t time.Time) *FakeClock { |
| 92 | return &FakeClock{ |
| 93 | time: t, |
| 94 | } |
| 95 | } |
| 96 | |
| 97 | // blocker is a caller of BlockUntil. |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 98 | type blocker struct { |
| 99 | count int |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 100 | |
| 101 | // ch is closed when the underlying clock has the specified number of blockers. |
| 102 | ch chan struct{} |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 103 | } |
| 104 | |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 105 | // expirer is a timer or ticker that expires at some point in the future. |
| 106 | type 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 |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 117 | // fakeClock, then sends the current time on the returned channel. |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 118 | func (fc *FakeClock) After(d time.Duration) <-chan time.Time { |
| 119 | return fc.NewTimer(d).Chan() |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 120 | } |
| 121 | |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 122 | // Sleep blocks until the given duration has passed on the fakeClock. |
| 123 | func (fc *FakeClock) Sleep(d time.Duration) { |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 124 | <-fc.After(d) |
| 125 | } |
| 126 | |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 127 | // Now returns the current time of the fakeClock |
| 128 | func (fc *FakeClock) Now() time.Time { |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 129 | fc.l.RLock() |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 130 | defer fc.l.RUnlock() |
| 131 | return fc.time |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 132 | } |
| 133 | |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 134 | // Since returns the duration that has passed since the given time on the |
| 135 | // fakeClock. |
| 136 | func (fc *FakeClock) Since(t time.Time) time.Duration { |
| khenaidoo | 2672188 | 2021-08-11 17:42:52 -0400 | [diff] [blame] | 137 | return fc.Now().Sub(t) |
| 138 | } |
| 139 | |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 140 | // Until returns the duration that has to pass from the given time on the fakeClock |
| 141 | // to reach the given time. |
| 142 | func (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. |
| 150 | func (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")) |
| khenaidoo | 2672188 | 2021-08-11 17:42:52 -0400 | [diff] [blame] | 155 | } |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 156 | ft := newFakeTicker(fc, d) |
| 157 | fc.l.Lock() |
| 158 | defer fc.l.Unlock() |
| 159 | fc.setExpirer(ft, d) |
| khenaidoo | 2672188 | 2021-08-11 17:42:52 -0400 | [diff] [blame] | 160 | return ft |
| 161 | } |
| 162 | |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 163 | // NewTimer returns a Timer that will fire only after calls to |
| 164 | // fakeClock.Advance() have moved the clock past the given duration. |
| 165 | func (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. |
| 173 | func (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. |
| 180 | func (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. |
| 192 | func (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. |
| 202 | func (fc *FakeClock) Advance(d time.Duration) { |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 203 | fc.l.Lock() |
| 204 | defer fc.l.Unlock() |
| 205 | end := fc.time.Add(d) |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 206 | // 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) |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 221 | } |
| 222 | } |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 223 | fc.time = end |
| 224 | } |
| 225 | |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 226 | // 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. |
| 232 | func (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. |
| 238 | func (fc *FakeClock) BlockUntilContext(ctx context.Context, n int) error { |
| 239 | b := fc.newBlocker(n) |
| 240 | if b == nil { |
| 241 | return nil |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 242 | } |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 243 | |
| 244 | select { |
| 245 | case <-b.ch: |
| 246 | return nil |
| 247 | case <-ctx.Done(): |
| 248 | return ctx.Err() |
| 249 | } |
| 250 | } |
| 251 | |
| 252 | func (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. |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 260 | b := &blocker{ |
| 261 | count: n, |
| 262 | ch: make(chan struct{}), |
| 263 | } |
| 264 | fc.blockers = append(fc.blockers, b) |
| Abhay Kumar | 40252eb | 2025-10-13 13:25:53 +0000 | [diff] [blame^] | 265 | return b |
| 266 | } |
| 267 | |
| 268 | // stop stops an expirer, returning true if the expirer was stopped. |
| 269 | func (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. |
| 278 | func (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. |
| 294 | func (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 | }) |
| khenaidoo | 59ce9dd | 2019-11-11 13:05:32 -0500 | [diff] [blame] | 319 | } |