diff --git a/Makefile b/Makefile index 7eb6d24..7ff707d 100644 --- a/Makefile +++ b/Makefile @@ -53,4 +53,4 @@ staticcheck: bin/staticcheck .PHONY: test test: - go test -race ./... + go test -v -race ./... diff --git a/go.mod b/go.mod index ce7cfc9..77c1eda 100644 --- a/go.mod +++ b/go.mod @@ -3,7 +3,6 @@ module go.uber.org/ratelimit go 1.18 require ( - github.com/benbjohnson/clock v1.3.0 github.com/stretchr/testify v1.6.1 go.uber.org/atomic v1.7.0 ) diff --git a/go.sum b/go.sum index 471801e..8c610cb 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,3 @@ -github.com/benbjohnson/clock v1.3.0 h1:ip6w0uFQkncKQ979AypyG0ER7mqUSBdKLOgAle/AT8A= -github.com/benbjohnson/clock v1.3.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= diff --git a/ratelimit.go b/ratelimit.go index 7370526..7afce5b 100644 --- a/ratelimit.go +++ b/ratelimit.go @@ -22,8 +22,6 @@ package ratelimit // import "go.uber.org/ratelimit" import ( "time" - - "github.com/benbjohnson/clock" ) // Note: This file is inspired by: @@ -45,6 +43,16 @@ type Clock interface { Sleep(time.Duration) } +type internalClock struct{} + +func (i *internalClock) Now() time.Time { + return time.Now() +} + +func (i *internalClock) Sleep(duration time.Duration) { + time.Sleep(duration) +} + // config configures a limiter. type config struct { clock Clock @@ -60,7 +68,7 @@ func New(rate int, opts ...Option) Limiter { // buildConfig combines defaults with options. func buildConfig(opts []Option) config { c := config{ - clock: clock.New(), + clock: &internalClock{}, slack: 10, per: time.Second, } diff --git a/ratelimit_mock_time_test.go b/ratelimit_mock_time_test.go new file mode 100644 index 0000000..e88f5a2 --- /dev/null +++ b/ratelimit_mock_time_test.go @@ -0,0 +1,114 @@ +package ratelimit + +/** +This fake time implementation is a modification of time mocking +the mechanism used by Ian Lance Taylor in https://github.com/golang/time project +https://github.com/golang/time/commit/579cf78fd858857c0d766e0d63eb2b0ccf29f436 + +Modified parts: + - timers are sorted on every addition, and then we relly of that order, + we could use heap data structure, but sorting is OK for now. + - advance accepts backoffDuration to sleep without lock held after every timer triggering + - advanceUnlocked method yields the processor, after every timer triggering, + allowing other goroutines to run +*/ + +import ( + "runtime" + "sort" + "sync" + "time" +) + +// testTime is a fake time used for testing. +type testTime struct { + mu sync.Mutex + cur time.Time // current fake time + timers []testTimer // fake timers +} + +// makeTestTime hooks the testTimer into the package. +func makeTestTime() *testTime { + return &testTime{ + cur: time.Now(), + } +} + +// testTimer is a fake timer. +type testTimer struct { + when time.Time + ch chan<- time.Time +} + +// now returns the current fake time. +func (tt *testTime) now() time.Time { + tt.mu.Lock() + defer tt.mu.Unlock() + return tt.cur +} + +// newTimer creates a fake timer. It returns the channel, +// a function to stop the timer (which we don't care about), +// and a function to advance to the next timer. +func (tt *testTime) newTimer(dur time.Duration) (<-chan time.Time, func() bool) { + tt.mu.Lock() + defer tt.mu.Unlock() + ch := make(chan time.Time, 1) + timer := testTimer{ + when: tt.cur.Add(dur), + ch: ch, + } + tt.timers = append(tt.timers, timer) + sort.Slice(tt.timers, func(i, j int) bool { + return tt.timers[i].when.Before(tt.timers[j].when) + }) + return ch, func() bool { return true } +} + +// advance advances the fake time. +func (tt *testTime) advance(dur time.Duration, backoffDuration time.Duration) { + tt.mu.Lock() + defer tt.mu.Unlock() + + targetTime := tt.cur.Add(dur) + for { + if len(tt.timers) == 0 || tt.timers[0].when.After(targetTime) { + tt.cur = targetTime + return + } + if tt.advanceUnlocked(tt.timers[0].when.Sub(tt.cur)) && backoffDuration > 0 { + // after every timer triggering, we release our mutex + // and give time for other goroutines to run + tt.mu.Unlock() + time.Sleep(backoffDuration) + tt.mu.Lock() + } + } +} + +// advanceUnlock advances the fake time, assuming it is already locked. +func (tt *testTime) advanceUnlocked(dur time.Duration) bool { + tt.cur = tt.cur.Add(dur) + if len(tt.timers) == 0 || tt.timers[0].when.After(tt.cur) { + return false + } + + i := 0 + for i < len(tt.timers) { + if tt.timers[i].when.After(tt.cur) { + break + } + tt.timers[i].ch <- tt.cur + i++ + // calculate how many goroutines we currently have in runtime + // and yield the processor, after every timer triggering, + // allowing all other goroutines to run + numOfAllRunningGoroutines := runtime.NumGoroutine() + for j := 0; j < numOfAllRunningGoroutines; j++ { + runtime.Gosched() + } + } + + tt.timers = tt.timers[i:] + return true +} diff --git a/ratelimit_test.go b/ratelimit_test.go index 7268f87..ac18c71 100644 --- a/ratelimit_test.go +++ b/ratelimit_test.go @@ -7,10 +7,20 @@ import ( "go.uber.org/atomic" - "github.com/benbjohnson/clock" "github.com/stretchr/testify/assert" ) +const advanceBackoffDuration = 5 * time.Millisecond + +func (tt *testTime) Now() time.Time { + return tt.now() +} + +func (tt *testTime) Sleep(duration time.Duration) { + timer, _ := tt.newTimer(duration) + <-timer +} + type testRunner interface { // createLimiter builds a limiter with given options. createLimiter(int, ...Option) Limiter @@ -23,13 +33,13 @@ type testRunner interface { // not using clock.AfterFunc because andres-erbsen/clock misses a nap there. afterFunc(d time.Duration, fn func()) // some tests want raw access to the clock. - getClock() *clock.Mock + getClock() *testTime } type runnerImpl struct { t *testing.T - clock *clock.Mock + clock *testTime constructor func(int, ...Option) Limiter count atomic.Int32 // maxDuration is the time we need to move into the future for a test. @@ -66,13 +76,9 @@ func runTest(t *testing.T, fn func(testRunner)) { for _, tt := range impls { t.Run(tt.name, func(t *testing.T) { - // Set a non-default time.Time since some limiters (int64 in particular) use - // the default value as "non-initialized" state. - clockMock := clock.NewMock() - clockMock.Set(time.Now()) r := runnerImpl{ t: t, - clock: clockMock, + clock: makeTestTime(), constructor: tt.constructor, doneCh: make(chan struct{}), } @@ -80,7 +86,7 @@ func runTest(t *testing.T, fn func(testRunner)) { defer r.wg.Wait() fn(&r) - r.clock.Add(r.maxDuration) + r.clock.advance(r.maxDuration, advanceBackoffDuration) }) } } @@ -91,7 +97,7 @@ func (r *runnerImpl) createLimiter(rate int, opts ...Option) Limiter { return r.constructor(rate, opts...) } -func (r *runnerImpl) getClock() *clock.Mock { +func (r *runnerImpl) getClock() *testTime { return r.clock } @@ -102,6 +108,7 @@ func (r *runnerImpl) startTaking(rls ...Limiter) { for _, rl := range rls { rl.Take() } + r.clock.advance(time.Nanosecond, 0) r.count.Inc() select { case <-r.doneCh: @@ -126,14 +133,14 @@ func (r *runnerImpl) afterFunc(d time.Duration, fn func()) { if d > r.maxDuration { r.maxDuration = d } - + timer, _ := r.clock.newTimer(d) r.goWait(func() { select { case <-r.doneCh: return - case <-r.clock.After(d): + case <-timer: + fn() } - fn() }) } @@ -237,17 +244,17 @@ func TestInitial(t *testing.T) { have []time.Duration startWg sync.WaitGroup ) - startWg.Add(3) + startWg.Add(3) for i := 0; i < 3; i++ { go func() { startWg.Done() results <- rl.Take() }() } - startWg.Wait() - clk.Add(time.Second) + + r.getClock().advance(time.Second, advanceBackoffDuration) for i := 0; i < 3; i++ { ts := <-results @@ -262,7 +269,7 @@ func TestInitial(t *testing.T) { time.Millisecond * 100, }, have, - "bad timestamps for inital takes", + "bad timestamps for initial takes", ) }) }) @@ -342,6 +349,7 @@ func TestSlack(t *testing.T) { for _, tt := range tests { t.Run(tt.msg, func(t *testing.T) { + t.Parallel() runTest(t, func(r testRunner) { slow := r.createLimiter(10, WithoutSlack) fast := r.createLimiter(100, tt.opt...)