status-go/vendor/github.com/benbjohnson/clock/clock.go

423 lines
10 KiB
Go

package clock
import (
"context"
"sort"
"sync"
"time"
)
// Re-export of time.Duration
type Duration = time.Duration
// Clock represents an interface to the functions in the standard library time
// package. Two implementations are available in the clock package. The first
// is a real-time clock which simply wraps the time package's functions. The
// second is a mock clock which will only change when
// programmatically adjusted.
type Clock interface {
After(d time.Duration) <-chan time.Time
AfterFunc(d time.Duration, f func()) *Timer
Now() time.Time
Since(t time.Time) time.Duration
Until(t time.Time) time.Duration
Sleep(d time.Duration)
Tick(d time.Duration) <-chan time.Time
Ticker(d time.Duration) *Ticker
Timer(d time.Duration) *Timer
WithDeadline(parent context.Context, d time.Time) (context.Context, context.CancelFunc)
WithTimeout(parent context.Context, t time.Duration) (context.Context, context.CancelFunc)
}
// New returns an instance of a real-time clock.
func New() Clock {
return &clock{}
}
// clock implements a real-time clock by simply wrapping the time package functions.
type clock struct{}
func (c *clock) After(d time.Duration) <-chan time.Time { return time.After(d) }
func (c *clock) AfterFunc(d time.Duration, f func()) *Timer {
return &Timer{timer: time.AfterFunc(d, f)}
}
func (c *clock) Now() time.Time { return time.Now() }
func (c *clock) Since(t time.Time) time.Duration { return time.Since(t) }
func (c *clock) Until(t time.Time) time.Duration { return time.Until(t) }
func (c *clock) Sleep(d time.Duration) { time.Sleep(d) }
func (c *clock) Tick(d time.Duration) <-chan time.Time { return time.Tick(d) }
func (c *clock) Ticker(d time.Duration) *Ticker {
t := time.NewTicker(d)
return &Ticker{C: t.C, ticker: t}
}
func (c *clock) Timer(d time.Duration) *Timer {
t := time.NewTimer(d)
return &Timer{C: t.C, timer: t}
}
func (c *clock) WithDeadline(parent context.Context, d time.Time) (context.Context, context.CancelFunc) {
return context.WithDeadline(parent, d)
}
func (c *clock) WithTimeout(parent context.Context, t time.Duration) (context.Context, context.CancelFunc) {
return context.WithTimeout(parent, t)
}
// Mock represents a mock clock that only moves forward programmically.
// It can be preferable to a real-time clock when testing time-based functionality.
type Mock struct {
// mu protects all other fields in this struct, and the data that they
// point to.
mu sync.Mutex
now time.Time // current time
timers clockTimers // tickers & timers
}
// NewMock returns an instance of a mock clock.
// The current time of the mock clock on initialization is the Unix epoch.
func NewMock() *Mock {
return &Mock{now: time.Unix(0, 0)}
}
// Add moves the current time of the mock clock forward by the specified duration.
// This should only be called from a single goroutine at a time.
func (m *Mock) Add(d time.Duration) {
// Calculate the final current time.
m.mu.Lock()
t := m.now.Add(d)
m.mu.Unlock()
// Continue to execute timers until there are no more before the new time.
for {
if !m.runNextTimer(t) {
break
}
}
// Ensure that we end with the new time.
m.mu.Lock()
m.now = t
m.mu.Unlock()
// Give a small buffer to make sure that other goroutines get handled.
gosched()
}
// Set sets the current time of the mock clock to a specific one.
// This should only be called from a single goroutine at a time.
func (m *Mock) Set(t time.Time) {
// Continue to execute timers until there are no more before the new time.
for {
if !m.runNextTimer(t) {
break
}
}
// Ensure that we end with the new time.
m.mu.Lock()
m.now = t
m.mu.Unlock()
// Give a small buffer to make sure that other goroutines get handled.
gosched()
}
// WaitForAllTimers sets the clock until all timers are expired
func (m *Mock) WaitForAllTimers() time.Time {
// Continue to execute timers until there are no more
for {
m.mu.Lock()
if len(m.timers) == 0 {
m.mu.Unlock()
return m.Now()
}
sort.Sort(m.timers)
next := m.timers[len(m.timers)-1].Next()
m.mu.Unlock()
m.Set(next)
}
}
// runNextTimer executes the next timer in chronological order and moves the
// current time to the timer's next tick time. The next time is not executed if
// its next time is after the max time. Returns true if a timer was executed.
func (m *Mock) runNextTimer(max time.Time) bool {
m.mu.Lock()
// Sort timers by time.
sort.Sort(m.timers)
// If we have no more timers then exit.
if len(m.timers) == 0 {
m.mu.Unlock()
return false
}
// Retrieve next timer. Exit if next tick is after new time.
t := m.timers[0]
if t.Next().After(max) {
m.mu.Unlock()
return false
}
// Move "now" forward and unlock clock.
m.now = t.Next()
now := m.now
m.mu.Unlock()
// Execute timer.
t.Tick(now)
return true
}
// After waits for the duration to elapse and then sends the current time on the returned channel.
func (m *Mock) After(d time.Duration) <-chan time.Time {
return m.Timer(d).C
}
// AfterFunc waits for the duration to elapse and then executes a function in its own goroutine.
// A Timer is returned that can be stopped.
func (m *Mock) AfterFunc(d time.Duration, f func()) *Timer {
m.mu.Lock()
defer m.mu.Unlock()
ch := make(chan time.Time, 1)
t := &Timer{
c: ch,
fn: f,
mock: m,
next: m.now.Add(d),
stopped: false,
}
m.timers = append(m.timers, (*internalTimer)(t))
return t
}
// Now returns the current wall time on the mock clock.
func (m *Mock) Now() time.Time {
m.mu.Lock()
defer m.mu.Unlock()
return m.now
}
// Since returns time since `t` using the mock clock's wall time.
func (m *Mock) Since(t time.Time) time.Duration {
return m.Now().Sub(t)
}
// Until returns time until `t` using the mock clock's wall time.
func (m *Mock) Until(t time.Time) time.Duration {
return t.Sub(m.Now())
}
// Sleep pauses the goroutine for the given duration on the mock clock.
// The clock must be moved forward in a separate goroutine.
func (m *Mock) Sleep(d time.Duration) {
<-m.After(d)
}
// Tick is a convenience function for Ticker().
// It will return a ticker channel that cannot be stopped.
func (m *Mock) Tick(d time.Duration) <-chan time.Time {
return m.Ticker(d).C
}
// Ticker creates a new instance of Ticker.
func (m *Mock) Ticker(d time.Duration) *Ticker {
m.mu.Lock()
defer m.mu.Unlock()
ch := make(chan time.Time, 1)
t := &Ticker{
C: ch,
c: ch,
mock: m,
d: d,
next: m.now.Add(d),
}
m.timers = append(m.timers, (*internalTicker)(t))
return t
}
// Timer creates a new instance of Timer.
func (m *Mock) Timer(d time.Duration) *Timer {
m.mu.Lock()
ch := make(chan time.Time, 1)
t := &Timer{
C: ch,
c: ch,
mock: m,
next: m.now.Add(d),
stopped: false,
}
m.timers = append(m.timers, (*internalTimer)(t))
now := m.now
m.mu.Unlock()
m.runNextTimer(now)
return t
}
// removeClockTimer removes a timer from m.timers. m.mu MUST be held
// when this method is called.
func (m *Mock) removeClockTimer(t clockTimer) {
for i, timer := range m.timers {
if timer == t {
copy(m.timers[i:], m.timers[i+1:])
m.timers[len(m.timers)-1] = nil
m.timers = m.timers[:len(m.timers)-1]
break
}
}
sort.Sort(m.timers)
}
// clockTimer represents an object with an associated start time.
type clockTimer interface {
Next() time.Time
Tick(time.Time)
}
// clockTimers represents a list of sortable timers.
type clockTimers []clockTimer
func (a clockTimers) Len() int { return len(a) }
func (a clockTimers) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a clockTimers) Less(i, j int) bool { return a[i].Next().Before(a[j].Next()) }
// Timer represents a single event.
// The current time will be sent on C, unless the timer was created by AfterFunc.
type Timer struct {
C <-chan time.Time
c chan time.Time
timer *time.Timer // realtime impl, if set
next time.Time // next tick time
mock *Mock // mock clock, if set
fn func() // AfterFunc function, if set
stopped bool // True if stopped, false if running
}
// Stop turns off the ticker.
func (t *Timer) Stop() bool {
if t.timer != nil {
return t.timer.Stop()
}
t.mock.mu.Lock()
registered := !t.stopped
t.mock.removeClockTimer((*internalTimer)(t))
t.stopped = true
t.mock.mu.Unlock()
return registered
}
// Reset changes the expiry time of the timer
func (t *Timer) Reset(d time.Duration) bool {
if t.timer != nil {
return t.timer.Reset(d)
}
t.mock.mu.Lock()
t.next = t.mock.now.Add(d)
defer t.mock.mu.Unlock()
registered := !t.stopped
if t.stopped {
t.mock.timers = append(t.mock.timers, (*internalTimer)(t))
}
t.stopped = false
return registered
}
type internalTimer Timer
func (t *internalTimer) Next() time.Time { return t.next }
func (t *internalTimer) Tick(now time.Time) {
// a gosched() after ticking, to allow any consequences of the
// tick to complete
defer gosched()
t.mock.mu.Lock()
if t.fn != nil {
// defer function execution until the lock is released, and
defer func() { go t.fn() }()
} else {
t.c <- now
}
t.mock.removeClockTimer((*internalTimer)(t))
t.stopped = true
t.mock.mu.Unlock()
}
// Ticker holds a channel that receives "ticks" at regular intervals.
type Ticker struct {
C <-chan time.Time
c chan time.Time
ticker *time.Ticker // realtime impl, if set
next time.Time // next tick time
mock *Mock // mock clock, if set
d time.Duration // time between ticks
stopped bool // True if stopped, false if running
}
// Stop turns off the ticker.
func (t *Ticker) Stop() {
if t.ticker != nil {
t.ticker.Stop()
} else {
t.mock.mu.Lock()
t.mock.removeClockTimer((*internalTicker)(t))
t.stopped = true
t.mock.mu.Unlock()
}
}
// Reset resets the ticker to a new duration.
func (t *Ticker) Reset(dur time.Duration) {
if t.ticker != nil {
t.ticker.Reset(dur)
return
}
t.mock.mu.Lock()
defer t.mock.mu.Unlock()
if t.stopped {
t.mock.timers = append(t.mock.timers, (*internalTicker)(t))
t.stopped = false
}
t.d = dur
t.next = t.mock.now.Add(dur)
}
type internalTicker Ticker
func (t *internalTicker) Next() time.Time { return t.next }
func (t *internalTicker) Tick(now time.Time) {
select {
case t.c <- now:
default:
}
t.mock.mu.Lock()
t.next = now.Add(t.d)
t.mock.mu.Unlock()
gosched()
}
// Sleep momentarily so that other goroutines can process.
func gosched() { time.Sleep(1 * time.Millisecond) }
var (
// type checking
_ Clock = &Mock{}
)