clock.go 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218
  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. Sleep(d time.Duration)
  25. Tick(d time.Duration) <-chan time.Time
  26. }
  27. var (
  28. _ = Clock(RealClock{})
  29. _ = Clock(&FakeClock{})
  30. _ = Clock(&IntervalClock{})
  31. )
  32. // RealClock really calls time.Now()
  33. type RealClock struct{}
  34. // Now returns the current time.
  35. func (RealClock) Now() time.Time {
  36. return time.Now()
  37. }
  38. // Since returns time since the specified timestamp.
  39. func (RealClock) Since(ts time.Time) time.Duration {
  40. return time.Since(ts)
  41. }
  42. // Same as time.After(d).
  43. func (RealClock) After(d time.Duration) <-chan time.Time {
  44. return time.After(d)
  45. }
  46. func (RealClock) Tick(d time.Duration) <-chan time.Time {
  47. return time.Tick(d)
  48. }
  49. func (RealClock) Sleep(d time.Duration) {
  50. time.Sleep(d)
  51. }
  52. // FakeClock implements Clock, but returns an arbitrary time.
  53. type FakeClock struct {
  54. lock sync.RWMutex
  55. time time.Time
  56. // waiters are waiting for the fake time to pass their specified time
  57. waiters []fakeClockWaiter
  58. }
  59. type fakeClockWaiter struct {
  60. targetTime time.Time
  61. stepInterval time.Duration
  62. skipIfBlocked bool
  63. destChan chan<- time.Time
  64. }
  65. func NewFakeClock(t time.Time) *FakeClock {
  66. return &FakeClock{
  67. time: t,
  68. }
  69. }
  70. // Now returns f's time.
  71. func (f *FakeClock) Now() time.Time {
  72. f.lock.RLock()
  73. defer f.lock.RUnlock()
  74. return f.time
  75. }
  76. // Since returns time since the time in f.
  77. func (f *FakeClock) Since(ts time.Time) time.Duration {
  78. f.lock.RLock()
  79. defer f.lock.RUnlock()
  80. return f.time.Sub(ts)
  81. }
  82. // Fake version of time.After(d).
  83. func (f *FakeClock) After(d time.Duration) <-chan time.Time {
  84. f.lock.Lock()
  85. defer f.lock.Unlock()
  86. stopTime := f.time.Add(d)
  87. ch := make(chan time.Time, 1) // Don't block!
  88. f.waiters = append(f.waiters, fakeClockWaiter{
  89. targetTime: stopTime,
  90. destChan: ch,
  91. })
  92. return ch
  93. }
  94. func (f *FakeClock) Tick(d time.Duration) <-chan time.Time {
  95. f.lock.Lock()
  96. defer f.lock.Unlock()
  97. tickTime := f.time.Add(d)
  98. ch := make(chan time.Time, 1) // hold one tick
  99. f.waiters = append(f.waiters, fakeClockWaiter{
  100. targetTime: tickTime,
  101. stepInterval: d,
  102. skipIfBlocked: true,
  103. destChan: ch,
  104. })
  105. return ch
  106. }
  107. // Move clock by Duration, notify anyone that's called After or Tick
  108. func (f *FakeClock) Step(d time.Duration) {
  109. f.lock.Lock()
  110. defer f.lock.Unlock()
  111. f.setTimeLocked(f.time.Add(d))
  112. }
  113. // Sets the time.
  114. func (f *FakeClock) SetTime(t time.Time) {
  115. f.lock.Lock()
  116. defer f.lock.Unlock()
  117. f.setTimeLocked(t)
  118. }
  119. // Actually changes the time and checks any waiters. f must be write-locked.
  120. func (f *FakeClock) setTimeLocked(t time.Time) {
  121. f.time = t
  122. newWaiters := make([]fakeClockWaiter, 0, len(f.waiters))
  123. for i := range f.waiters {
  124. w := &f.waiters[i]
  125. if !w.targetTime.After(t) {
  126. if w.skipIfBlocked {
  127. select {
  128. case w.destChan <- t:
  129. default:
  130. }
  131. } else {
  132. w.destChan <- t
  133. }
  134. if w.stepInterval > 0 {
  135. for !w.targetTime.After(t) {
  136. w.targetTime = w.targetTime.Add(w.stepInterval)
  137. }
  138. newWaiters = append(newWaiters, *w)
  139. }
  140. } else {
  141. newWaiters = append(newWaiters, f.waiters[i])
  142. }
  143. }
  144. f.waiters = newWaiters
  145. }
  146. // Returns true if After has been called on f but not yet satisfied (so you can
  147. // write race-free tests).
  148. func (f *FakeClock) HasWaiters() bool {
  149. f.lock.RLock()
  150. defer f.lock.RUnlock()
  151. return len(f.waiters) > 0
  152. }
  153. func (f *FakeClock) Sleep(d time.Duration) {
  154. f.Step(d)
  155. }
  156. // IntervalClock implements Clock, but each invocation of Now steps the clock forward the specified duration
  157. type IntervalClock struct {
  158. Time time.Time
  159. Duration time.Duration
  160. }
  161. // Now returns i's time.
  162. func (i *IntervalClock) Now() time.Time {
  163. i.Time = i.Time.Add(i.Duration)
  164. return i.Time
  165. }
  166. // Since returns time since the time in i.
  167. func (i *IntervalClock) Since(ts time.Time) time.Duration {
  168. return i.Time.Sub(ts)
  169. }
  170. // Unimplemented, will panic.
  171. // TODO: make interval clock use FakeClock so this can be implemented.
  172. func (*IntervalClock) After(d time.Duration) <-chan time.Time {
  173. panic("IntervalClock doesn't implement After")
  174. }
  175. // Unimplemented, will panic.
  176. // TODO: make interval clock use FakeClock so this can be implemented.
  177. func (*IntervalClock) Tick(d time.Duration) <-chan time.Time {
  178. panic("IntervalClock doesn't implement Tick")
  179. }
  180. func (*IntervalClock) Sleep(d time.Duration) {
  181. panic("IntervalClock doesn't implement Sleep")
  182. }