extend fake clock

pkg/util/clock.go

@@ -17,6 +17,7 @@ limitations under the License.
 package util
 
 import (
+	"sync"
 	"time"
 )
 
@@ -25,39 +26,115 @@ import (
 type Clock interface {
 	Now() time.Time
 	Since(time.Time) time.Duration
+	After(d time.Duration) <-chan time.Time
 }
 
+var (
+	_ = Clock(RealClock{})
+	_ = Clock(&FakeClock{})
+	_ = Clock(&IntervalClock{})
+)
+
 // RealClock really calls time.Now()
 type RealClock struct{}
 
 // Now returns the current time.
-func (r RealClock) Now() time.Time {
+func (RealClock) Now() time.Time {
 	return time.Now()
 }
 
 // Since returns time since the specified timestamp.
-func (r RealClock) Since(ts time.Time) time.Duration {
+func (RealClock) Since(ts time.Time) time.Duration {
 	return time.Since(ts)
 }
 
+// Same as time.After(d).
+func (RealClock) After(d time.Duration) <-chan time.Time {
+	return time.After(d)
+}
+
 // FakeClock implements Clock, but returns an arbitrary time.
 type FakeClock struct {
-	Time time.Time
+	lock sync.RWMutex
+	time time.Time
+
+	// waiters are waiting for the fake time to pass their specified time
+	waiters []fakeClockWaiter
+}
+
+type fakeClockWaiter struct {
+	targetTime time.Time
+	destChan   chan<- time.Time
+}
+
+func NewFakeClock(t time.Time) *FakeClock {
+	return &FakeClock{
+		time: t,
+	}
 }
 
 // Now returns f's time.
 func (f *FakeClock) Now() time.Time {
-	return f.Time
+	f.lock.RLock()
+	defer f.lock.RUnlock()
+	return f.time
 }
 
 // Since returns time since the time in f.
 func (f *FakeClock) Since(ts time.Time) time.Duration {
-	return f.Time.Sub(ts)
+	f.lock.RLock()
+	defer f.lock.RUnlock()
+	return f.time.Sub(ts)
 }
 
-// Move clock by Duration
+// Fake version of time.After(d).
+func (f *FakeClock) After(d time.Duration) <-chan time.Time {
+	f.lock.Lock()
+	defer f.lock.Unlock()
+	stopTime := f.time.Add(d)
+	ch := make(chan time.Time, 1) // Don't block!
+	f.waiters = append(f.waiters, fakeClockWaiter{
+		targetTime: stopTime,
+		destChan:   ch,
+	})
+	return ch
+}
+
+// Move clock by Duration, notify anyone that's called After
 func (f *FakeClock) Step(d time.Duration) {
-	f.Time = f.Time.Add(d)
+	f.lock.Lock()
+	defer f.lock.Unlock()
+	f.setTimeLocked(f.time.Add(d))
+}
+
+// Sets the time.
+func (f *FakeClock) SetTime(t time.Time) {
+	f.lock.Lock()
+	defer f.lock.Unlock()
+	f.setTimeLocked(t)
+}
+
+// Actually changes the time and checks any waiters. f must be write-locked.
+func (f *FakeClock) setTimeLocked(t time.Time) {
+	f.time = t
+	newWaiters := make([]fakeClockWaiter, 0, len(f.waiters))
+	for i := range f.waiters {
+		w := &f.waiters[i]
+		if !w.targetTime.After(t) {
+			w.destChan <- t
+		} else {
+			newWaiters = append(newWaiters, f.waiters[i])
+		}
+	}
+	f.waiters = newWaiters
+}
+
+// Returns true if After has been called on f but not yet satisfied (so you can
+// write race-free tests).
+func (f *FakeClock) HasWaiters() bool {
+	f.lock.RLock()
+	defer f.lock.RUnlock()
+	return len(f.waiters) > 0
 }
 
 // IntervalClock implements Clock, but each invocation of Now steps the clock forward the specified duration
@@ -76,3 +153,9 @@ func (i *IntervalClock) Now() time.Time {
 func (i *IntervalClock) Since(ts time.Time) time.Duration {
 	return i.Time.Sub(ts)
 }
+
+// Unimplemented, will panic.
+// TODO: make interval clock use FakeClock so this can be implemented.
+func (*IntervalClock) After(d time.Duration) <-chan time.Time {
+	panic("IntervalClock doesn't implement After")
+}