X-Git-Url: http://www.git.stargrave.org/?a=blobdiff_plain;f=ratelimitreader.go;h=7d9e6d8660ba100e8222fd095293d3acddcd75ea;hb=HEAD;hp=a8e3483b14556b0090cb21009e0abe07bef2526e;hpb=78a8db10da8015467660e3ff009d797fea3ea25f;p=btrtrc.git

diff --git a/ratelimitreader.go b/ratelimitreader.go
index a8e3483b..7d9e6d86 100644
--- a/ratelimitreader.go
+++ b/ratelimitreader.go
@@ -1,25 +1,53 @@
 package torrent
 
 import (
+	"context"
+	"fmt"
 	"io"
 	"time"
 
-	"golang.org/x/net/context"
 	"golang.org/x/time/rate"
 )
 
 type rateLimitedReader struct {
 	l *rate.Limiter
 	r io.Reader
+
+	// This is the time of the last Read's reservation.
+	lastRead time.Time
 }
 
-func (me rateLimitedReader) Read(b []byte) (n int, err error) {
-	if err := me.l.WaitN(context.Background(), 1); err != nil {
-		panic(err)
-	}
-	n, err = me.r.Read(b)
-	if !me.l.ReserveN(time.Now(), n-1).OK() {
-		panic(n - 1)
+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())
 	}
 	return
 }