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