From 5c16d031ff51bcc9a34e2bb75c4105948b1dfe47 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Fri, 11 Jul 2025 00:14:52 +1000
Subject: [PATCH] Simplify and probably fix rate limited reader

---
 ratelimitreader.go | 46 ++++++++++++++--------------------------------
 1 file changed, 14 insertions(+), 32 deletions(-)

diff --git a/ratelimitreader.go b/ratelimitreader.go
index 7d9e6d86..f1af3125 100644
--- 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
 }
-- 
2.51.0