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 package torrent
import ( import (
"fmt"
"io" "io"
"time" "time"
@ -14,12 +15,18 @@ type rateLimitedReader struct {
} }
func (me rateLimitedReader) Read(b []byte) (n int, err error) { 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 { if err := me.l.WaitN(context.Background(), 1); err != nil {
panic(err) 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) n, err = me.r.Read(b)
// Pay the piper.
if !me.l.ReserveN(time.Now(), n-1).OK() { if !me.l.ReserveN(time.Now(), n-1).OK() {
panic(n - 1) panic(fmt.Sprintf("burst exceeded?: %d", n-1))
} }
return return
} }