Zack Williams | e940c7a | 2019-08-21 14:25:39 -0700 | [diff] [blame] | 1 | /* |
| 2 | Copyright 2014 The Kubernetes Authors. |
| 3 | |
| 4 | Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | you may not use this file except in compliance with the License. |
| 6 | You may obtain a copy of the License at |
| 7 | |
| 8 | http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | |
| 10 | Unless required by applicable law or agreed to in writing, software |
| 11 | distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | See the License for the specific language governing permissions and |
| 14 | limitations under the License. |
| 15 | */ |
| 16 | |
| 17 | package clock |
| 18 | |
| 19 | import ( |
| 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. |
| 26 | type 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() |
| 36 | type RealClock struct{} |
| 37 | |
| 38 | // Now returns the current time. |
| 39 | func (RealClock) Now() time.Time { |
| 40 | return time.Now() |
| 41 | } |
| 42 | |
| 43 | // Since returns time since the specified timestamp. |
| 44 | func (RealClock) Since(ts time.Time) time.Duration { |
| 45 | return time.Since(ts) |
| 46 | } |
| 47 | |
| 48 | // Same as time.After(d). |
| 49 | func (RealClock) After(d time.Duration) <-chan time.Time { |
| 50 | return time.After(d) |
| 51 | } |
| 52 | |
| 53 | func (RealClock) NewTimer(d time.Duration) Timer { |
| 54 | return &realTimer{ |
| 55 | timer: time.NewTimer(d), |
| 56 | } |
| 57 | } |
| 58 | |
| 59 | func (RealClock) NewTicker(d time.Duration) Ticker { |
| 60 | return &realTicker{ |
| 61 | ticker: time.NewTicker(d), |
| 62 | } |
| 63 | } |
| 64 | |
| 65 | func (RealClock) Sleep(d time.Duration) { |
| 66 | time.Sleep(d) |
| 67 | } |
| 68 | |
| 69 | // FakeClock implements Clock, but returns an arbitrary time. |
| 70 | type 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 | |
| 78 | type fakeClockWaiter struct { |
| 79 | targetTime time.Time |
| 80 | stepInterval time.Duration |
| 81 | skipIfBlocked bool |
| 82 | destChan chan time.Time |
| 83 | fired bool |
| 84 | } |
| 85 | |
| 86 | func NewFakeClock(t time.Time) *FakeClock { |
| 87 | return &FakeClock{ |
| 88 | time: t, |
| 89 | } |
| 90 | } |
| 91 | |
| 92 | // Now returns f's time. |
| 93 | func (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. |
| 100 | func (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). |
| 107 | func (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). |
| 120 | func (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 | |
| 136 | func (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 |
| 154 | func (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. |
| 161 | func (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. |
| 168 | func (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). |
| 202 | func (f *FakeClock) HasWaiters() bool { |
| 203 | f.lock.RLock() |
| 204 | defer f.lock.RUnlock() |
| 205 | return len(f.waiters) > 0 |
| 206 | } |
| 207 | |
| 208 | func (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 |
| 213 | type IntervalClock struct { |
| 214 | Time time.Time |
| 215 | Duration time.Duration |
| 216 | } |
| 217 | |
| 218 | // Now returns i's time. |
| 219 | func (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. |
| 225 | func (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. |
| 231 | func (*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. |
| 237 | func (*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. |
| 243 | func (*IntervalClock) NewTicker(d time.Duration) Ticker { |
| 244 | panic("IntervalClock doesn't implement NewTicker") |
| 245 | } |
| 246 | |
| 247 | func (*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. |
| 253 | type 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. |
| 260 | type realTimer struct { |
| 261 | timer *time.Timer |
| 262 | } |
| 263 | |
| 264 | // C returns the underlying timer's channel. |
| 265 | func (r *realTimer) C() <-chan time.Time { |
| 266 | return r.timer.C |
| 267 | } |
| 268 | |
| 269 | // Stop calls Stop() on the underlying timer. |
| 270 | func (r *realTimer) Stop() bool { |
| 271 | return r.timer.Stop() |
| 272 | } |
| 273 | |
| 274 | // Reset calls Reset() on the underlying timer. |
| 275 | func (r *realTimer) Reset(d time.Duration) bool { |
| 276 | return r.timer.Reset(d) |
| 277 | } |
| 278 | |
| 279 | // fakeTimer implements Timer based on a FakeClock. |
| 280 | type fakeTimer struct { |
| 281 | fakeClock *FakeClock |
| 282 | waiter fakeClockWaiter |
| 283 | } |
| 284 | |
| 285 | // C returns the channel that notifies when this timer has fired. |
| 286 | func (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. |
| 291 | func (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. |
| 310 | func (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 | |
| 322 | type Ticker interface { |
| 323 | C() <-chan time.Time |
| 324 | Stop() |
| 325 | } |
| 326 | |
| 327 | type realTicker struct { |
| 328 | ticker *time.Ticker |
| 329 | } |
| 330 | |
| 331 | func (t *realTicker) C() <-chan time.Time { |
| 332 | return t.ticker.C |
| 333 | } |
| 334 | |
| 335 | func (t *realTicker) Stop() { |
| 336 | t.ticker.Stop() |
| 337 | } |
| 338 | |
| 339 | type fakeTicker struct { |
| 340 | c <-chan time.Time |
| 341 | } |
| 342 | |
| 343 | func (t *fakeTicker) C() <-chan time.Time { |
| 344 | return t.c |
| 345 | } |
| 346 | |
| 347 | func (t *fakeTicker) Stop() { |
| 348 | } |