status-go/vendor/github.com/anacrolix/torrent/ratelimitreader.go

54 lines
1.1 KiB
Go

package torrent
import (
"context"
"fmt"
"io"
"time"
"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) {
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
}