clock.go 7.5 KB

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