cubefs

1
package clock
2

3
import (
4
	"context"
5
	"sort"
6
	"sync"
7
	"time"
8
)
9

10
// Re-export of time.Duration
11
type Duration = time.Duration
12

13
// Clock represents an interface to the functions in the standard library time
14
// package. Two implementations are available in the clock package. The first
15
// is a real-time clock which simply wraps the time package's functions. The
16
// second is a mock clock which will only change when
17
// programmatically adjusted.
18
type Clock interface {
19
	After(d time.Duration) <-chan time.Time
20
	AfterFunc(d time.Duration, f func()) *Timer
21
	Now() time.Time
22
	Since(t time.Time) time.Duration
23
	Until(t time.Time) time.Duration
24
	Sleep(d time.Duration)
25
	Tick(d time.Duration) <-chan time.Time
26
	Ticker(d time.Duration) *Ticker
27
	Timer(d time.Duration) *Timer
28
	WithDeadline(parent context.Context, d time.Time) (context.Context, context.CancelFunc)
29
	WithTimeout(parent context.Context, t time.Duration) (context.Context, context.CancelFunc)
30
}
31

32
// New returns an instance of a real-time clock.
33
func New() Clock {
34
	return &clock{}
35
}
36

37
// clock implements a real-time clock by simply wrapping the time package functions.
38
type clock struct{}
39

40
func (c *clock) After(d time.Duration) <-chan time.Time { return time.After(d) }
41

42
func (c *clock) AfterFunc(d time.Duration, f func()) *Timer {
43
	return &Timer{timer: time.AfterFunc(d, f)}
44
}
45

46
func (c *clock) Now() time.Time { return time.Now() }
47

48
func (c *clock) Since(t time.Time) time.Duration { return time.Since(t) }
49

50
func (c *clock) Until(t time.Time) time.Duration { return time.Until(t) }
51

52
func (c *clock) Sleep(d time.Duration) { time.Sleep(d) }
53

54
func (c *clock) Tick(d time.Duration) <-chan time.Time { return time.Tick(d) }
55

56
func (c *clock) Ticker(d time.Duration) *Ticker {
57
	t := time.NewTicker(d)
58
	return &Ticker{C: t.C, ticker: t}
59
}
60

61
func (c *clock) Timer(d time.Duration) *Timer {
62
	t := time.NewTimer(d)
63
	return &Timer{C: t.C, timer: t}
64
}
65

66
func (c *clock) WithDeadline(parent context.Context, d time.Time) (context.Context, context.CancelFunc) {
67
	return context.WithDeadline(parent, d)
68
}
69

70
func (c *clock) WithTimeout(parent context.Context, t time.Duration) (context.Context, context.CancelFunc) {
71
	return context.WithTimeout(parent, t)
72
}
73

74
// Mock represents a mock clock that only moves forward programmically.
75
// It can be preferable to a real-time clock when testing time-based functionality.
76
type Mock struct {
77
	mu     sync.Mutex
78
	now    time.Time   // current time
79
	timers clockTimers // tickers & timers
80
}
81

82
// NewMock returns an instance of a mock clock.
83
// The current time of the mock clock on initialization is the Unix epoch.
84
func NewMock() *Mock {
85
	return &Mock{now: time.Unix(0, 0)}
86
}
87

88
// Add moves the current time of the mock clock forward by the specified duration.
89
// This should only be called from a single goroutine at a time.
90
func (m *Mock) Add(d time.Duration) {
91
	// Calculate the final current time.
92
	t := m.now.Add(d)
93

94
	// Continue to execute timers until there are no more before the new time.
95
	for {
96
		if !m.runNextTimer(t) {
97
			break
98
		}
99
	}
100

101
	// Ensure that we end with the new time.
102
	m.mu.Lock()
103
	m.now = t
104
	m.mu.Unlock()
105

106
	// Give a small buffer to make sure that other goroutines get handled.
107
	gosched()
108
}
109

110
// Set sets the current time of the mock clock to a specific one.
111
// This should only be called from a single goroutine at a time.
112
func (m *Mock) Set(t time.Time) {
113
	// Continue to execute timers until there are no more before the new time.
114
	for {
115
		if !m.runNextTimer(t) {
116
			break
117
		}
118
	}
119

120
	// Ensure that we end with the new time.
121
	m.mu.Lock()
122
	m.now = t
123
	m.mu.Unlock()
124

125
	// Give a small buffer to make sure that other goroutines get handled.
126
	gosched()
127
}
128

129
// runNextTimer executes the next timer in chronological order and moves the
130
// current time to the timer's next tick time. The next time is not executed if
131
// its next time is after the max time. Returns true if a timer was executed.
132
func (m *Mock) runNextTimer(max time.Time) bool {
133
	m.mu.Lock()
134

135
	// Sort timers by time.
136
	sort.Sort(m.timers)
137

138
	// If we have no more timers then exit.
139
	if len(m.timers) == 0 {
140
		m.mu.Unlock()
141
		return false
142
	}
143

144
	// Retrieve next timer. Exit if next tick is after new time.
145
	t := m.timers[0]
146
	if t.Next().After(max) {
147
		m.mu.Unlock()
148
		return false
149
	}
150

151
	// Move "now" forward and unlock clock.
152
	m.now = t.Next()
153
	m.mu.Unlock()
154

155
	// Execute timer.
156
	t.Tick(m.now)
157
	return true
158
}
159

160
// After waits for the duration to elapse and then sends the current time on the returned channel.
161
func (m *Mock) After(d time.Duration) <-chan time.Time {
162
	return m.Timer(d).C
163
}
164

165
// AfterFunc waits for the duration to elapse and then executes a function.
166
// A Timer is returned that can be stopped.
167
func (m *Mock) AfterFunc(d time.Duration, f func()) *Timer {
168
	m.mu.Lock()
169
	defer m.mu.Unlock()
170
	ch := make(chan time.Time, 1)
171
	t := &Timer{
172
		c:       ch,
173
		fn:      f,
174
		mock:    m,
175
		next:    m.now.Add(d),
176
		stopped: false,
177
	}
178
	m.timers = append(m.timers, (*internalTimer)(t))
179
	return t
180
}
181

182
// Now returns the current wall time on the mock clock.
183
func (m *Mock) Now() time.Time {
184
	m.mu.Lock()
185
	defer m.mu.Unlock()
186
	return m.now
187
}
188

189
// Since returns time since `t` using the mock clock's wall time.
190
func (m *Mock) Since(t time.Time) time.Duration {
191
	return m.Now().Sub(t)
192
}
193

194
// Until returns time until `t` using the mock clock's wall time.
195
func (m *Mock) Until(t time.Time) time.Duration {
196
	return t.Sub(m.Now())
197
}
198

199
// Sleep pauses the goroutine for the given duration on the mock clock.
200
// The clock must be moved forward in a separate goroutine.
201
func (m *Mock) Sleep(d time.Duration) {
202
	<-m.After(d)
203
}
204

205
// Tick is a convenience function for Ticker().
206
// It will return a ticker channel that cannot be stopped.
207
func (m *Mock) Tick(d time.Duration) <-chan time.Time {
208
	return m.Ticker(d).C
209
}
210

211
// Ticker creates a new instance of Ticker.
212
func (m *Mock) Ticker(d time.Duration) *Ticker {
213
	m.mu.Lock()
214
	defer m.mu.Unlock()
215
	ch := make(chan time.Time, 1)
216
	t := &Ticker{
217
		C:    ch,
218
		c:    ch,
219
		mock: m,
220
		d:    d,
221
		next: m.now.Add(d),
222
	}
223
	m.timers = append(m.timers, (*internalTicker)(t))
224
	return t
225
}
226

227
// Timer creates a new instance of Timer.
228
func (m *Mock) Timer(d time.Duration) *Timer {
229
	m.mu.Lock()
230
	defer m.mu.Unlock()
231
	ch := make(chan time.Time, 1)
232
	t := &Timer{
233
		C:       ch,
234
		c:       ch,
235
		mock:    m,
236
		next:    m.now.Add(d),
237
		stopped: false,
238
	}
239
	m.timers = append(m.timers, (*internalTimer)(t))
240
	return t
241
}
242

243
func (m *Mock) removeClockTimer(t clockTimer) {
244
	for i, timer := range m.timers {
245
		if timer == t {
246
			copy(m.timers[i:], m.timers[i+1:])
247
			m.timers[len(m.timers)-1] = nil
248
			m.timers = m.timers[:len(m.timers)-1]
249
			break
250
		}
251
	}
252
	sort.Sort(m.timers)
253
}
254

255
// clockTimer represents an object with an associated start time.
256
type clockTimer interface {
257
	Next() time.Time
258
	Tick(time.Time)
259
}
260

261
// clockTimers represents a list of sortable timers.
262
type clockTimers []clockTimer
263

264
func (a clockTimers) Len() int           { return len(a) }
265
func (a clockTimers) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
266
func (a clockTimers) Less(i, j int) bool { return a[i].Next().Before(a[j].Next()) }
267

268
// Timer represents a single event.
269
// The current time will be sent on C, unless the timer was created by AfterFunc.
270
type Timer struct {
271
	C       <-chan time.Time
272
	c       chan time.Time
273
	timer   *time.Timer // realtime impl, if set
274
	next    time.Time   // next tick time
275
	mock    *Mock       // mock clock, if set
276
	fn      func()      // AfterFunc function, if set
277
	stopped bool        // True if stopped, false if running
278
}
279

280
// Stop turns off the ticker.
281
func (t *Timer) Stop() bool {
282
	if t.timer != nil {
283
		return t.timer.Stop()
284
	}
285

286
	t.mock.mu.Lock()
287
	registered := !t.stopped
288
	t.mock.removeClockTimer((*internalTimer)(t))
289
	t.stopped = true
290
	t.mock.mu.Unlock()
291
	return registered
292
}
293

294
// Reset changes the expiry time of the timer
295
func (t *Timer) Reset(d time.Duration) bool {
296
	if t.timer != nil {
297
		return t.timer.Reset(d)
298
	}
299

300
	t.mock.mu.Lock()
301
	t.next = t.mock.now.Add(d)
302
	defer t.mock.mu.Unlock()
303

304
	registered := !t.stopped
305
	if t.stopped {
306
		t.mock.timers = append(t.mock.timers, (*internalTimer)(t))
307
	}
308

309
	t.stopped = false
310
	return registered
311
}
312

313
type internalTimer Timer
314

315
func (t *internalTimer) Next() time.Time { return t.next }
316
func (t *internalTimer) Tick(now time.Time) {
317
	// a gosched() after ticking, to allow any consequences of the
318
	// tick to complete
319
	defer gosched()
320

321
	t.mock.mu.Lock()
322
	if t.fn != nil {
323
		// defer function execution until the lock is released, and
324
		defer t.fn()
325
	} else {
326
		t.c <- now
327
	}
328
	t.mock.removeClockTimer((*internalTimer)(t))
329
	t.stopped = true
330
	t.mock.mu.Unlock()
331
}
332

333
// Ticker holds a channel that receives "ticks" at regular intervals.
334
type Ticker struct {
335
	C      <-chan time.Time
336
	c      chan time.Time
337
	ticker *time.Ticker  // realtime impl, if set
338
	next   time.Time     // next tick time
339
	mock   *Mock         // mock clock, if set
340
	d      time.Duration // time between ticks
341
}
342

343
// Stop turns off the ticker.
344
func (t *Ticker) Stop() {
345
	if t.ticker != nil {
346
		t.ticker.Stop()
347
	} else {
348
		t.mock.mu.Lock()
349
		t.mock.removeClockTimer((*internalTicker)(t))
350
		t.mock.mu.Unlock()
351
	}
352
}
353

354
// Reset resets the ticker to a new duration.
355
func (t *Ticker) Reset(dur time.Duration) {
356
	if t.ticker != nil {
357
		t.ticker.Reset(dur)
358
		return
359
	}
360

361
	t.mock.mu.Lock()
362
	defer t.mock.mu.Unlock()
363

364
	t.d = dur
365
	t.next = t.mock.now.Add(dur)
366
}
367

368
type internalTicker Ticker
369

370
func (t *internalTicker) Next() time.Time { return t.next }
371
func (t *internalTicker) Tick(now time.Time) {
372
	select {
373
	case t.c <- now:
374
	default:
375
	}
376
	t.next = now.Add(t.d)
377
	gosched()
378
}
379

380
// Sleep momentarily so that other goroutines can process.
381
func gosched() { time.Sleep(1 * time.Millisecond) }
382

383
var (
384
	// type checking
385
	_ Clock = &Mock{}
386
)
387