11 "github.com/stretchr/testify/assert"
12 "github.com/stretchr/testify/require"
13 "golang.org/x/time/rate"
16 func writeN(ws []io.Writer, n int) error {
18 for _, w := range ws[1:] {
20 wn, err := w.Write(b[:n1])
29 wn, err := ws[0].Write(b[:n])
38 func TestRateLimitReaders(t *testing.T) {
47 control := rate.NewLimiter(bytesPerSecond, burst)
48 shared := rate.NewLimiter(bytesPerSecond, burst)
53 wg := sync.WaitGroup{}
56 // When the read was allowed.
59 reads := make(chan read)
60 done := make(chan struct{})
61 for i := 0; i < numReaders; i += 1 {
68 r := rateLimitedReader{
72 b := make([]byte, readSize)
76 case reads <- read{n, r.lastRead}:
90 for _, c := range cs {
101 for i := 0; i < writeRounds; i += 1 {
102 err := writeN(ws, bytesPerRound)
104 log.Printf("error writing: %s", err)
107 written += bytesPerRound
114 started := time.Now()
115 for r := range reads {
116 totalBytesRead += r.N
117 require.False(t, r.At.IsZero())
118 // Copy what the reader should have done with its reservation.
119 res := control.ReserveN(r.At, r.N)
120 // If we don't have to wait with the control, the reader has gone too
123 log.Printf("%d bytes not allowed at %s", r.N, time.Since(started))
127 assert.EqualValues(t, writeRounds*bytesPerRound, totalBytesRead)