1package clockwork 2 3import ( 4 "sync" 5 "time" 6) 7 8// Clock provides an interface that packages can use instead of directly 9// using the time module, so that chronology-related behavior can be tested 10type Clock interface { 11 After(d time.Duration) <-chan time.Time 12 Sleep(d time.Duration) 13 Now() time.Time 14 Since(t time.Time) time.Duration 15} 16 17// FakeClock provides an interface for a clock which can be 18// manually advanced through time 19type FakeClock interface { 20 Clock 21 // Advance advances the FakeClock to a new point in time, ensuring any existing 22 // sleepers are notified appropriately before returning 23 Advance(d time.Duration) 24 // BlockUntil will block until the FakeClock has the given number of 25 // sleepers (callers of Sleep or After) 26 BlockUntil(n int) 27} 28 29// NewRealClock returns a Clock which simply delegates calls to the actual time 30// package; it should be used by packages in production. 31func NewRealClock() Clock { 32 return &realClock{} 33} 34 35// NewFakeClock returns a FakeClock implementation which can be 36// manually advanced through time for testing. The initial time of the 37// FakeClock will be an arbitrary non-zero time. 38func NewFakeClock() FakeClock { 39 // use a fixture that does not fulfill Time.IsZero() 40 return NewFakeClockAt(time.Date(1984, time.April, 4, 0, 0, 0, 0, time.UTC)) 41} 42 43// NewFakeClockAt returns a FakeClock initialised at the given time.Time. 44func NewFakeClockAt(t time.Time) FakeClock { 45 return &fakeClock{ 46 time: t, 47 } 48} 49 50type realClock struct{} 51 52func (rc *realClock) After(d time.Duration) <-chan time.Time { 53 return time.After(d) 54} 55 56func (rc *realClock) Sleep(d time.Duration) { 57 time.Sleep(d) 58} 59 60func (rc *realClock) Now() time.Time { 61 return time.Now() 62} 63 64func (rc *realClock) Since(t time.Time) time.Duration { 65 return rc.Now().Sub(t) 66} 67 68type fakeClock struct { 69 sleepers []*sleeper 70 blockers []*blocker 71 time time.Time 72 73 l sync.RWMutex 74} 75 76// sleeper represents a caller of After or Sleep 77type sleeper struct { 78 until time.Time 79 done chan time.Time 80} 81 82// blocker represents a caller of BlockUntil 83type blocker struct { 84 count int 85 ch chan struct{} 86} 87 88// After mimics time.After; it waits for the given duration to elapse on the 89// fakeClock, then sends the current time on the returned channel. 90func (fc *fakeClock) After(d time.Duration) <-chan time.Time { 91 fc.l.Lock() 92 defer fc.l.Unlock() 93 now := fc.time 94 done := make(chan time.Time, 1) 95 if d.Nanoseconds() == 0 { 96 // special case - trigger immediately 97 done <- now 98 } else { 99 // otherwise, add to the set of sleepers 100 s := &sleeper{ 101 until: now.Add(d), 102 done: done, 103 } 104 fc.sleepers = append(fc.sleepers, s) 105 // and notify any blockers 106 fc.blockers = notifyBlockers(fc.blockers, len(fc.sleepers)) 107 } 108 return done 109} 110 111// notifyBlockers notifies all the blockers waiting until the 112// given number of sleepers are waiting on the fakeClock. It 113// returns an updated slice of blockers (i.e. those still waiting) 114func notifyBlockers(blockers []*blocker, count int) (newBlockers []*blocker) { 115 for _, b := range blockers { 116 if b.count == count { 117 close(b.ch) 118 } else { 119 newBlockers = append(newBlockers, b) 120 } 121 } 122 return 123} 124 125// Sleep blocks until the given duration has passed on the fakeClock 126func (fc *fakeClock) Sleep(d time.Duration) { 127 <-fc.After(d) 128} 129 130// Time returns the current time of the fakeClock 131func (fc *fakeClock) Now() time.Time { 132 fc.l.RLock() 133 t := fc.time 134 fc.l.RUnlock() 135 return t 136} 137 138// Since returns the duration that has passed since the given time on the fakeClock 139func (fc *fakeClock) Since(t time.Time) time.Duration { 140 return fc.Now().Sub(t) 141} 142 143// Advance advances fakeClock to a new point in time, ensuring channels from any 144// previous invocations of After are notified appropriately before returning 145func (fc *fakeClock) Advance(d time.Duration) { 146 fc.l.Lock() 147 defer fc.l.Unlock() 148 end := fc.time.Add(d) 149 var newSleepers []*sleeper 150 for _, s := range fc.sleepers { 151 if end.Sub(s.until) >= 0 { 152 s.done <- end 153 } else { 154 newSleepers = append(newSleepers, s) 155 } 156 } 157 fc.sleepers = newSleepers 158 fc.blockers = notifyBlockers(fc.blockers, len(fc.sleepers)) 159 fc.time = end 160} 161 162// BlockUntil will block until the fakeClock has the given number of sleepers 163// (callers of Sleep or After) 164func (fc *fakeClock) BlockUntil(n int) { 165 fc.l.Lock() 166 // Fast path: current number of sleepers is what we're looking for 167 if len(fc.sleepers) == n { 168 fc.l.Unlock() 169 return 170 } 171 // Otherwise, set up a new blocker 172 b := &blocker{ 173 count: n, 174 ch: make(chan struct{}), 175 } 176 fc.blockers = append(fc.blockers, b) 177 fc.l.Unlock() 178 <-b.ch 179} 180