clock.go 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231
  1. // Package clock provides an abstraction for system time that enables testing of
  2. // time-sensitive code.
  3. //
  4. // Where you'd use time.Now, instead use clk.Now where clk is an instance of
  5. // Clock.
  6. //
  7. // When running your code in production, pass it a Clock given by
  8. // clock.Default() and when you're running it in your tests, pass it an instance
  9. // of Clock from NewFake().
  10. //
  11. // When you do that, you can use FakeClock's Add and Set methods to control how
  12. // time behaves in your tests. That makes the tests you'll write more reliable
  13. // while also expanding the space of problems you can test.
  14. //
  15. // This code intentionally does not attempt to provide an abstraction over
  16. // time.Ticker and time.Timer because Go does not have the runtime or API hooks
  17. // available to do so reliably. See https://github.com/golang/go/issues/8869
  18. package clock
  19. import (
  20. "sort"
  21. "sync"
  22. "time"
  23. )
  24. // Some in-use reflection-heavy systems, like facebookgo/inject, fail when given
  25. // a value type like sysClock{}. Since it's hidden by an interface, this has
  26. // surprised users. We fixed that by making systemClock a &sysClock.
  27. var systemClock Clock = &sysClock{}
  28. // New returns a Clock that matches the actual system time.
  29. func New() Clock {
  30. // This is a method instead of a public var to prevent folks from
  31. // "making things work" by writing to the var instead of passing
  32. // in a Clock.
  33. return systemClock
  34. }
  35. // Deprecated: Default is just an alias for New but less memorable.
  36. func Default() Clock {
  37. return systemClock
  38. }
  39. // Clock is an abstraction over system time. New instances of it can
  40. // be made with Default and NewFake.
  41. type Clock interface {
  42. // Now returns the Clock's current view of the time. Mutating the
  43. // returned Time will not mutate the clock's time.
  44. Now() time.Time
  45. // Sleep causes the current goroutine to sleep for the given duration.
  46. Sleep(time.Duration)
  47. // After returns a channel that fires after the given duration.
  48. After(time.Duration) <-chan time.Time
  49. // Since is a short hand for Now().Sub(t).
  50. Since(time.Time) time.Duration
  51. // NewTimer makes a Timer based on this clock's time. Using Timers and
  52. // negative durations in the Clock or Timer API is undefined behavior and
  53. // may be changed.
  54. NewTimer(time.Duration) *Timer
  55. }
  56. type sysClock struct{}
  57. func (s *sysClock) Now() time.Time {
  58. return time.Now()
  59. }
  60. func (s *sysClock) Sleep(d time.Duration) {
  61. time.Sleep(d)
  62. }
  63. func (s *sysClock) After(d time.Duration) <-chan time.Time {
  64. return time.After(d)
  65. }
  66. func (s *sysClock) Since(t time.Time) time.Duration {
  67. return time.Since(t)
  68. }
  69. func (s *sysClock) NewTimer(d time.Duration) *Timer {
  70. tt := time.NewTimer(d)
  71. return &Timer{C: tt.C, timer: tt}
  72. }
  73. // NewFake returns a FakeClock to be used in tests that need to
  74. // manipulate time. Its initial value is always the unix epoch in the
  75. // UTC timezone. The FakeClock returned is thread-safe.
  76. func NewFake() FakeClock {
  77. // We're explicit about this time construction to avoid early user
  78. // questions about why the time object doesn't have a Location by
  79. // default.
  80. return &fake{t: time.Unix(0, 0).UTC()}
  81. }
  82. // FakeClock is a Clock with additional controls. The return value of
  83. // Now return can be modified with Add. Use NewFake to get a
  84. // thread-safe FakeClock implementation.
  85. type FakeClock interface {
  86. Clock
  87. // Adjust the time that will be returned by Now.
  88. Add(d time.Duration)
  89. // Set the Clock's time to exactly the time given.
  90. Set(t time.Time)
  91. }
  92. // To prevent mistakes with the API, we hide this behind NewFake. It's
  93. // easy forget to create a pointer to a fake since time.Time (and
  94. // sync.Mutex) are also simple values. The code will appear to work
  95. // but the clock's time will never be adjusted.
  96. type fake struct {
  97. sync.RWMutex
  98. t time.Time
  99. sends sortedSends
  100. }
  101. func (f *fake) Now() time.Time {
  102. f.RLock()
  103. defer f.RUnlock()
  104. return f.t
  105. }
  106. func (f *fake) Sleep(d time.Duration) {
  107. if d < 0 {
  108. // time.Sleep just returns immediately. Do the same.
  109. return
  110. }
  111. f.Add(d)
  112. }
  113. func (f *fake) After(d time.Duration) <-chan time.Time {
  114. return f.NewTimer(d).C
  115. }
  116. func (f *fake) Since(t time.Time) time.Duration {
  117. return f.Now().Sub(t)
  118. }
  119. func (f *fake) NewTimer(d time.Duration) *Timer {
  120. f.Lock()
  121. defer f.Unlock()
  122. ch := make(chan time.Time, 1)
  123. tt := f.t.Add(d)
  124. ft := &fakeTimer{c: ch, clk: f, active: true}
  125. t := &Timer{
  126. C: ch,
  127. fakeTimer: ft,
  128. }
  129. s := f.addSend(tt, ft)
  130. ft.sends = []*send{s}
  131. return t
  132. }
  133. func (f *fake) Add(d time.Duration) {
  134. f.Lock()
  135. defer f.Unlock()
  136. f.t = f.t.Add(d)
  137. f.sendTimes()
  138. }
  139. func (f *fake) Set(t time.Time) {
  140. f.Lock()
  141. defer f.Unlock()
  142. f.t = t
  143. f.sendTimes()
  144. }
  145. // Only to be called while the fake's lock is held
  146. func (f *fake) sendTimes() {
  147. newSends := make(sortedSends, 0)
  148. for _, s := range f.sends {
  149. if !s.active || !s.ft.active {
  150. continue
  151. }
  152. if s.target.Equal(f.t) || s.target.Before(f.t) {
  153. s.ft.active = false
  154. s.active = false
  155. // The select is to drop second sends from resets without a user
  156. // receiving from ft.c.
  157. select {
  158. case s.ft.c <- s.target:
  159. default:
  160. }
  161. }
  162. if s.active {
  163. newSends = append(newSends, s)
  164. }
  165. }
  166. f.sends = newSends
  167. }
  168. // Only to be called while the fake's lock is held
  169. func (f *fake) addSend(target time.Time, ft *fakeTimer) *send {
  170. s := &send{target: target, ft: ft, active: true}
  171. f.sends = append(f.sends, s)
  172. // This will be a small enough slice to be fast. Can be replaced with a more
  173. // complicated container if someone is making many timers.
  174. sort.Sort(f.sends)
  175. return s
  176. }
  177. // send is a struct that represents a scheduled send of a time.Time to its
  178. // fakeTimer's channel. They are actually sent when the relevant fake's time
  179. // goes equal or past their target time, as long as the relevant fakeTimer has
  180. // not been Reset or Stop'ed. When a Timer is Reset, the old sends are
  181. // deactivated and will be removed from the clocks list on the next attempt to
  182. // send.
  183. type send struct {
  184. target time.Time
  185. active bool
  186. ft *fakeTimer
  187. }
  188. type sortedSends []*send
  189. func (s sortedSends) Len() int {
  190. return len(s)
  191. }
  192. func (s sortedSends) Less(i, j int) bool {
  193. return s[i].target.Before(s[j].target)
  194. }
  195. func (s sortedSends) Swap(i, j int) {
  196. s[i], s[j] = s[j], s[i]
  197. }