Keep rate limited reader reads to within the burst capacity

This commit is contained in:
Matt Joiner 2017-08-31 14:32:22 +10:00
parent 3eb422afd0
commit cf022be396
1 changed files with 8 additions and 1 deletions

View File

@ -1,6 +1,7 @@
package torrent
import (
"fmt"
"io"
"time"
@ -14,12 +15,18 @@ type rateLimitedReader struct {
}
func (me rateLimitedReader) Read(b []byte) (n int, err error) {
// Wait until we can read at all.
if err := me.l.WaitN(context.Background(), 1); err != nil {
panic(err)
}
// Limit the read to within the burst.
if me.l.Limit() != rate.Inf && len(b) > me.l.Burst() {
b = b[:me.l.Burst()]
}
n, err = me.r.Read(b)
// Pay the piper.
if !me.l.ReserveN(time.Now(), n-1).OK() {
panic(n - 1)
panic(fmt.Sprintf("burst exceeded?: %d", n-1))
}
return
}