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) {
- // 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(fmt.Sprintf("burst exceeded?: %d", 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
}