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