Get rid of internal/clock - Consolidate AfterFunc and Sleep in tests

go.mod

@@ -3,6 +3,7 @@ module go.uber.org/ratelimit
 go 1.14
 
 require (
+	github.com/andres-erbsen/clock v0.0.0-20160526145045-9e14626cd129
 	github.com/stretchr/testify v1.6.1
 	go.uber.org/atomic v1.7.0
 )

go.sum

@@ -1,3 +1,5 @@
+github.com/andres-erbsen/clock v0.0.0-20160526145045-9e14626cd129 h1:MzBOUgng9orim59UnfUTLRjMpd09C5uEVQ6RPGeCaVI=
+github.com/andres-erbsen/clock v0.0.0-20160526145045-9e14626cd129/go.mod h1:rFgpPQZYZ8vdbc+48xibu8ALc3yeyd64IhHS+PU6Yyg=
 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=

ratelimit.go

@@ -26,7 +26,7 @@ import (
 	"sync/atomic"
 	"unsafe"
 
-	"go.uber.org/ratelimit/internal/clock"
+	"github.com/andres-erbsen/clock"
 )
 
 // Note: This file is inspired by:

ratelimit_test.go

@@ -8,8 +8,8 @@ import (
 
 	"go.uber.org/atomic"
 	"go.uber.org/ratelimit"
-	"go.uber.org/ratelimit/internal/clock"
 
+	"github.com/andres-erbsen/clock"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -22,6 +22,7 @@ type runner interface {
 	// getClock returns the test clock.
 	getClock() ratelimit.Clock
 	// afterFunc executes a func at a given time.
+	// not using clock.AfterFunc because andres-erbsen/clock misses a nap there.
 	afterFunc(d time.Duration, fn func())
 }
 
@@ -52,7 +53,7 @@ func runTest(t *testing.T, fn func(runner)) {
 
 // startTaking tries to Take() on passed in limiters in a loop/goroutine.
 func (r *runnerImpl) startTaking(rls ...ratelimit.Limiter) {
-	go func() {
+	r.goWait(func() {
 		for {
 			for _, rl := range rls {
 				rl.Take()
@@ -64,7 +65,7 @@ func (r *runnerImpl) startTaking(rls ...ratelimit.Limiter) {
 			default:
 			}
 		}
-	}()
+	})
 }
 
 // assertCountAt asserts the limiters have Taken() a number of times at a given time.
@@ -86,7 +87,26 @@ func (r *runnerImpl) afterFunc(d time.Duration, fn func()) {
 	if d > r.maxDuration {
 		r.maxDuration = d
 	}
-	r.clock.AfterFunc(d, fn)
+
+	r.goWait(func() {
+		select {
+		case <-r.doneCh:
+			return
+		case <-r.clock.After(d):
+		}
+		fn()
+	})
+}
+
+// goWait runs a function in a goroutine and makes sure the gouritine was scheduled.
+func (r *runnerImpl) goWait(fn func()) {
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		wg.Done()
+		fn()
+	}()
+	wg.Wait()
 }
 
 func Example() {
@@ -136,7 +156,6 @@ func TestRateLimiter(t *testing.T) {
 		r.assertCountAt(2*time.Second, 200)
 		r.assertCountAt(3*time.Second, 300)
 	})
-
 }
 
 func TestDelayedRateLimiter(t *testing.T) {
@@ -158,5 +177,4 @@ func TestDelayedRateLimiter(t *testing.T) {
 
 		r.assertCountAt(30*time.Second, 1200)
 	})
-
 }