11 "github.com/bradfitz/iter"
12 "github.com/stretchr/testify/assert"
13 "github.com/stretchr/testify/require"
14 "golang.org/x/time/rate"
17 func writeN(ws []io.Writer, n int) error {
19 for _, w := range ws[1:] {
21 wn, err := w.Write(b[:n1])
30 wn, err := ws[0].Write(b[:n])
39 func TestRateLimitReaders(t *testing.T) {
48 control := rate.NewLimiter(bytesPerSecond, burst)
49 shared := rate.NewLimiter(bytesPerSecond, burst)
54 wg := sync.WaitGroup{}
57 // When the read was allowed.
60 reads := make(chan read)
61 done := make(chan struct{})
62 for range iter.N(numReaders) {
69 r := rateLimitedReader{
73 b := make([]byte, readSize)
77 case reads <- read{n, r.lastRead}:
91 for _, c := range cs {
102 for range iter.N(writeRounds) {
103 err := writeN(ws, bytesPerRound)
105 log.Printf("error writing: %s", err)
108 written += bytesPerRound
115 started := time.Now()
116 for r := range reads {
117 totalBytesRead += r.N
118 require.False(t, r.At.IsZero())
119 // Copy what the reader should have done with its reservation.
120 res := control.ReserveN(r.At, r.N)
121 // If we don't have to wait with the control, the reader has gone too
124 log.Printf("%d bytes not allowed at %s", r.N, time.Since(started))
128 assert.EqualValues(t, writeRounds*bytesPerRound, totalBytesRead)