diff --git a/clock.go b/clock.go index 2230962..dc621da 100644 --- a/clock.go +++ b/clock.go @@ -74,7 +74,10 @@ func (c *clock) WithTimeout(parent context.Context, t time.Duration) (context.Co // Mock represents a mock clock that only moves forward programmically. // It can be preferable to a real-time clock when testing time-based functionality. type Mock struct { - mu sync.Mutex + // mu protects all other fields in this struct, and the data that they + // point to. + mu sync.Mutex + now time.Time // current time timers clockTimers // tickers & timers } @@ -89,7 +92,9 @@ func NewMock() *Mock { // This should only be called from a single goroutine at a time. func (m *Mock) Add(d time.Duration) { // Calculate the final current time. + m.mu.Lock() t := m.now.Add(d) + m.mu.Unlock() // Continue to execute timers until there are no more before the new time. for { @@ -240,6 +245,8 @@ func (m *Mock) Timer(d time.Duration) *Timer { return t } +// removeClockTimer removes a timer from m.timers. m.mu MUST be held +// when this method is called. func (m *Mock) removeClockTimer(t clockTimer) { for i, timer := range m.timers { if timer == t { @@ -373,7 +380,9 @@ func (t *internalTicker) Tick(now time.Time) { case t.c <- now: default: } + t.mock.mu.Lock() t.next = now.Add(t.d) + t.mock.mu.Unlock() gosched() }