traefik/vendor/go.uber.org/ratelimit/internal/clock/clock.go
2019-04-26 11:26:04 +02:00

133 lines
3.6 KiB
Go

// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
package clock
// Forked from github.com/andres-erbsen/clock to isolate a missing nap.
import (
"container/heap"
"sync"
"time"
)
// 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 {
sync.Mutex
now time.Time // current time
timers Timers // 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 duration.
// This should only be called from a single goroutine at a time.
func (m *Mock) Add(d time.Duration) {
m.Lock()
// Calculate the final time.
end := m.now.Add(d)
for len(m.timers) > 0 && m.now.Before(end) {
t := heap.Pop(&m.timers).(*Timer)
m.now = t.next
m.Unlock()
t.Tick()
m.Lock()
}
m.Unlock()
// Give a small buffer to make sure the other goroutines get handled.
nap()
}
// Timer produces a timer that will emit a time some duration after now.
func (m *Mock) Timer(d time.Duration) *Timer {
ch := make(chan time.Time)
t := &Timer{
C: ch,
c: ch,
mock: m,
next: m.now.Add(d),
}
m.addTimer(t)
return t
}
func (m *Mock) addTimer(t *Timer) {
m.Lock()
defer m.Unlock()
heap.Push(&m.timers, t)
}
// After produces a channel that will emit the time after a duration passes.
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.
// A Timer is returned that can be stopped.
func (m *Mock) AfterFunc(d time.Duration, f func()) *Timer {
t := m.Timer(d)
go func() {
<-t.c
f()
}()
nap()
return t
}
// Now returns the current wall time on the mock clock.
func (m *Mock) Now() time.Time {
m.Lock()
defer m.Unlock()
return 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)
}
// Timer represents a single event.
type Timer struct {
C <-chan time.Time
c chan time.Time
next time.Time // next tick time
mock *Mock // mock clock
}
func (t *Timer) Next() time.Time { return t.next }
func (t *Timer) Tick() {
select {
case t.c <- t.next:
default:
}
nap()
}
// Sleep momentarily so that other goroutines can process.
func nap() { time.Sleep(1 * time.Millisecond) }