package clockwork import ( "sync" "time" ) // Clock provides an interface that packages can use instead of directly // using the time module, so that chronology-related behavior can be tested type Clock interface { After(d time.Duration) <-chan time.Time Sleep(d time.Duration) Now() time.Time Since(t time.Time) time.Duration } // FakeClock provides an interface for a clock which can be // manually advanced through time type FakeClock interface { Clock // Advance advances the FakeClock to a new point in time, ensuring any existing // sleepers are notified appropriately before returning Advance(d time.Duration) // BlockUntil will block until the FakeClock has the given number of // sleepers (callers of Sleep or After) BlockUntil(n int) } // NewRealClock returns a Clock which simply delegates calls to the actual time // package; it should be used by packages in production. func NewRealClock() Clock { return &realClock{} } // NewFakeClock returns a FakeClock implementation which can be // manually advanced through time for testing. The initial time of the // FakeClock will be an arbitrary non-zero time. func NewFakeClock() FakeClock { // use a fixture that does not fulfill Time.IsZero() return NewFakeClockAt(time.Date(1984, time.April, 4, 0, 0, 0, 0, time.UTC)) } // NewFakeClockAt returns a FakeClock initialised at the given time.Time. func NewFakeClockAt(t time.Time) FakeClock { return &fakeClock{ time: t, } } type realClock struct{} func (rc *realClock) After(d time.Duration) <-chan time.Time { return time.After(d) } func (rc *realClock) Sleep(d time.Duration) { time.Sleep(d) } func (rc *realClock) Now() time.Time { return time.Now() } func (rc *realClock) Since(t time.Time) time.Duration { return rc.Now().Sub(t) } type fakeClock struct { sleepers []*sleeper blockers []*blocker time time.Time l sync.RWMutex } // sleeper represents a caller of After or Sleep type sleeper struct { until time.Time done chan time.Time } // blocker represents a caller of BlockUntil type blocker struct { count int ch chan struct{} } // After mimics time.After; it waits for the given duration to elapse on the // fakeClock, then sends the current time on the returned channel. func (fc *fakeClock) After(d time.Duration) <-chan time.Time { fc.l.Lock() defer fc.l.Unlock() now := fc.time done := make(chan time.Time, 1) if d.Nanoseconds() == 0 { // special case - trigger immediately done <- now } else { // otherwise, add to the set of sleepers s := &sleeper{ until: now.Add(d), done: done, } fc.sleepers = append(fc.sleepers, s) // and notify any blockers fc.blockers = notifyBlockers(fc.blockers, len(fc.sleepers)) } return done } // notifyBlockers notifies all the blockers waiting until the // given number of sleepers are waiting on the fakeClock. It // returns an updated slice of blockers (i.e. those still waiting) func notifyBlockers(blockers []*blocker, count int) (newBlockers []*blocker) { for _, b := range blockers { if b.count == count { close(b.ch) } else { newBlockers = append(newBlockers, b) } } return } // Sleep blocks until the given duration has passed on the fakeClock func (fc *fakeClock) Sleep(d time.Duration) { <-fc.After(d) } // Time returns the current time of the fakeClock func (fc *fakeClock) Now() time.Time { fc.l.RLock() t := fc.time fc.l.RUnlock() return t } // Since returns the duration that has passed since the given time on the fakeClock func (fc *fakeClock) Since(t time.Time) time.Duration { return fc.Now().Sub(t) } // Advance advances fakeClock to a new point in time, ensuring channels from any // previous invocations of After are notified appropriately before returning func (fc *fakeClock) Advance(d time.Duration) { fc.l.Lock() defer fc.l.Unlock() end := fc.time.Add(d) var newSleepers []*sleeper for _, s := range fc.sleepers { if end.Sub(s.until) >= 0 { s.done <- end } else { newSleepers = append(newSleepers, s) } } fc.sleepers = newSleepers fc.blockers = notifyBlockers(fc.blockers, len(fc.sleepers)) fc.time = end } // BlockUntil will block until the fakeClock has the given number of sleepers // (callers of Sleep or After) func (fc *fakeClock) BlockUntil(n int) { fc.l.Lock() // Fast path: current number of sleepers is what we're looking for if len(fc.sleepers) == n { fc.l.Unlock() return } // Otherwise, set up a new blocker b := &blocker{ count: n, ch: make(chan struct{}), } fc.blockers = append(fc.blockers, b) fc.l.Unlock() <-b.ch }