timing.go 5.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234
  1. // Forked from github.com/StefanKopieczek/gossip by @StefanKopieczek
  2. package timing
  3. import (
  4. "sync"
  5. "time"
  6. )
  7. // Controls whether library calls should be mocked, or whether we should use the standard Go time library.
  8. // If we're in Mock Mode, then time does not pass as normal, but only progresses when Elapse is called.
  9. // False by default, indicating that we just call through to standard Go functions.
  10. var MockMode = false
  11. var currentTimeMock = time.Unix(0, 0)
  12. var mockTimers = make([]*mockTimer, 0)
  13. var mockTimerMu = new(sync.Mutex)
  14. // Interface over Golang's built-in Timers, allowing them to be swapped out for mocked timers.
  15. type Timer interface {
  16. // Returns a channel which sends the current time immediately when the timer expires.
  17. // Equivalent to time.Timer.C; however, we have to use a method here instead of a member since this is an interface.
  18. C() <-chan time.Time
  19. // Resets the timer such that it will expire in duration 'd' after the current time.
  20. // Returns true if the timer had been active, and false if it had expired or been stopped.
  21. Reset(d time.Duration) bool
  22. // Stops the timer, preventing it from firing.
  23. // Returns true if the timer had been active, and false if it had expired or been stopped.
  24. Stop() bool
  25. }
  26. // Implementation of Timer that just wraps time.Timer.
  27. type realTimer struct {
  28. *time.Timer
  29. }
  30. func (t *realTimer) C() <-chan time.Time {
  31. return t.Timer.C
  32. }
  33. func (t *realTimer) Reset(d time.Duration) bool {
  34. t.Stop()
  35. return t.Timer.Reset(d)
  36. }
  37. func (t *realTimer) Stop() bool {
  38. // return t.Timer.Stop()
  39. if !t.Timer.Stop() {
  40. select {
  41. case <-t.Timer.C:
  42. return true
  43. default:
  44. return false
  45. }
  46. }
  47. return true
  48. }
  49. // Implementation of Timer that mocks time.Timer, firing when the total elapsed time (as controlled by Elapse)
  50. // exceeds the duration specified when the timer was constructed.
  51. type mockTimer struct {
  52. EndTime time.Time
  53. Chan chan time.Time
  54. fired bool
  55. toRun func()
  56. }
  57. func (t *mockTimer) C() <-chan time.Time {
  58. return t.Chan
  59. }
  60. func (t *mockTimer) Reset(d time.Duration) bool {
  61. wasActive := removeMockTimer(t)
  62. t.EndTime = currentTimeMock.Add(d)
  63. if d > 0 {
  64. mockTimerMu.Lock()
  65. mockTimers = append(mockTimers, t)
  66. mockTimerMu.Unlock()
  67. } else {
  68. // The new timer has an expiry time of 0.
  69. // Fire it right away, and don't bother tracking it.
  70. t.Chan <- currentTimeMock
  71. }
  72. return wasActive
  73. }
  74. func (t *mockTimer) Stop() bool {
  75. if !removeMockTimer(t) {
  76. select {
  77. case <-t.Chan:
  78. return true
  79. default:
  80. return false
  81. }
  82. }
  83. return true
  84. }
  85. // Creates a new Timer; either a wrapper around a standard Go time.Timer, or a mocked-out Timer,
  86. // depending on whether MockMode is set.
  87. func NewTimer(d time.Duration) Timer {
  88. if MockMode {
  89. t := mockTimer{currentTimeMock.Add(d), make(chan time.Time, 1), false, nil}
  90. if d == 0 {
  91. t.Chan <- currentTimeMock
  92. } else {
  93. mockTimerMu.Lock()
  94. mockTimers = append(mockTimers, &t)
  95. mockTimerMu.Unlock()
  96. }
  97. return &t
  98. } else {
  99. return &realTimer{time.NewTimer(d)}
  100. }
  101. }
  102. // See built-in time.After() function.
  103. func After(d time.Duration) <-chan time.Time {
  104. return NewTimer(d).C()
  105. }
  106. // See built-in time.AfterFunc() function.
  107. func AfterFunc(d time.Duration, f func()) Timer {
  108. if MockMode {
  109. mockTimerMu.Lock()
  110. t := mockTimer{currentTimeMock.Add(d), make(chan time.Time, 1), false, f}
  111. mockTimerMu.Unlock()
  112. if d == 0 {
  113. go f()
  114. t.Chan <- currentTimeMock
  115. } else {
  116. mockTimerMu.Lock()
  117. mockTimers = append(mockTimers, &t)
  118. mockTimerMu.Unlock()
  119. }
  120. return &t
  121. } else {
  122. return &realTimer{time.AfterFunc(d, f)}
  123. }
  124. }
  125. // See built-in time.Sleep() function.
  126. func Sleep(d time.Duration) {
  127. <-After(d)
  128. }
  129. // Increment the current time by the given Duration.
  130. // This function can only be called in Mock Mode, otherwise we will panic.
  131. func Elapse(d time.Duration) {
  132. requireMockMode()
  133. mockTimerMu.Lock()
  134. currentTimeMock = currentTimeMock.Add(d)
  135. mockTimerMu.Unlock()
  136. // Fire any timers whose time has come up.
  137. mockTimerMu.Lock()
  138. for _, t := range mockTimers {
  139. t.fired = false
  140. if !t.EndTime.After(currentTimeMock) {
  141. if t.toRun != nil {
  142. go t.toRun()
  143. }
  144. // Clear the channel if something is already in it.
  145. select {
  146. case <-t.Chan:
  147. default:
  148. }
  149. t.Chan <- currentTimeMock
  150. t.fired = true
  151. }
  152. }
  153. mockTimerMu.Unlock()
  154. // Stop tracking any fired timers.
  155. remainingTimers := make([]*mockTimer, 0)
  156. mockTimerMu.Lock()
  157. for _, t := range mockTimers {
  158. if !t.fired {
  159. remainingTimers = append(remainingTimers, t)
  160. }
  161. }
  162. mockTimers = remainingTimers
  163. mockTimerMu.Unlock()
  164. }
  165. // Returns the current time.
  166. // If Mock Mode is set, this will be the sum of all Durations passed into Elapse calls;
  167. // otherwise it will be the true system time.
  168. func Now() time.Time {
  169. if MockMode {
  170. return currentTimeMock
  171. } else {
  172. return time.Now()
  173. }
  174. }
  175. // Shortcut method to enforce that Mock Mode is enabled.
  176. func requireMockMode() {
  177. if !MockMode {
  178. panic("This method requires MockMode to be enabled")
  179. }
  180. }
  181. // Utility method to remove a mockTimer from the list of outstanding timers.
  182. func removeMockTimer(t *mockTimer) bool {
  183. // First, find the index of the timer in our list.
  184. found := false
  185. var idx int
  186. var elt *mockTimer
  187. mockTimerMu.Lock()
  188. for idx, elt = range mockTimers {
  189. if elt == t {
  190. found = true
  191. break
  192. }
  193. }
  194. mockTimerMu.Unlock()
  195. if found {
  196. mockTimerMu.Lock()
  197. // We found the given timer. Remove it.
  198. mockTimers = append(mockTimers[:idx], mockTimers[idx+1:]...)
  199. mockTimerMu.Unlock()
  200. return true
  201. } else {
  202. // The timer was not present, indicating that it was already expired.
  203. return false
  204. }
  205. }