Keep rate limited reader reads to within the burst capacity

[btrtrc.git] / ratelimitreader.go

diff --git a/ratelimitreader.go b/ratelimitreader.go
index a8e3483b14556b0090cb21009e0abe07bef2526e..d7939a3d3f6c08207ebaf2acecbc6c7cbd40dd96 100644 (file)
--- a/ratelimitreader.go
+++ b/ratelimitreader.go
@@ -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
 }