ratelimitreader.go | 9 ++++++++- diff --git a/ratelimitreader.go b/ratelimitreader.go index a8e3483b14556b0090cb21009e0abe07bef2526e..d7939a3d3f6c08207ebaf2acecbc6c7cbd40dd96 100644 --- a/ratelimitreader.go +++ b/ratelimitreader.go @@ -1,6 +1,7 @@ package torrent import ( + "fmt" "io" "time" @@ -14,12 +15,18 @@ r io.Reader } 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 }