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