Simplify and probably fix rate limited reader

diff --git a/ratelimitreader.go b/ratelimitreader.go
index 7d9e6d8660ba100e8222fd095293d3acddcd75ea..f1af3125dc1cb5b2e9a12184b59c3b99dbc1c4f4 100644 (file)
--- a/ratelimitreader.go
+++ b/ratelimitreader.go
@@ -1,11 +1,10 @@
 package torrent
 
 import (
-       "context"
-       "fmt"
        "io"
        "time"
 
+       "github.com/anacrolix/missinggo/v2/panicif"
        "golang.org/x/time/rate"
 )
 
@@ -17,37 +16,20 @@ type rateLimitedReader struct {
        lastRead time.Time
 }
 
+func (me *rateLimitedReader) justRead(b []byte) (n int, err error) {
+       n, err = me.r.Read(b)
+       me.lastRead = time.Now()
+       return
+}
+
 func (me *rateLimitedReader) Read(b []byte) (n int, err error) {
-       const oldStyle = false // Retained for future reference.
-       if oldStyle {
-               // 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.
-               now := time.Now()
-               me.lastRead = now
-               if !me.l.ReserveN(now, n-1).OK() {
-                       panic(fmt.Sprintf("burst exceeded?: %d", n-1))
-               }
-       } else {
-               // 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)
-               now := time.Now()
-               r := me.l.ReserveN(now, n)
-               if !r.OK() {
-                       panic(n)
-               }
-               me.lastRead = now
-               time.Sleep(r.Delay())
+       if me.l.Burst() != 0 {
+               b = b[:min(len(b), me.l.Burst())]
        }
+       t := time.Now()
+       n, err = me.justRead(b)
+       r := me.l.ReserveN(t, n)
+       panicif.False(r.OK())
+       time.Sleep(r.DelayFrom(t))
        return
 }