github.com/anacrolix/torrent@v1.61.0/ratelimitreader.go (about)

     1  package torrent
     2  
     3  import (
     4  	"io"
     5  	"time"
     6  
     7  	"github.com/anacrolix/missinggo/v2/panicif"
     8  	"golang.org/x/time/rate"
     9  )
    10  
    11  func newRateLimitedReader(r io.Reader, l *rate.Limiter) io.Reader {
    12  	if l == nil {
    13  		// Avoids taking Limiter lock to check limit, and allows type assertions to bypass Read.
    14  		return r
    15  	}
    16  	return rateLimitedReader{
    17  		l: l,
    18  		r: r,
    19  	}
    20  }
    21  
    22  type rateLimitedReader struct {
    23  	l *rate.Limiter
    24  	r io.Reader
    25  }
    26  
    27  func (me rateLimitedReader) Read(b []byte) (n int, err error) {
    28  	// Avoid truncating the read if everything is permitted anyway.
    29  	if me.l.Limit() == rate.Inf {
    30  		return me.r.Read(b)
    31  	}
    32  	// If the burst is zero, let the limiter method handle errors.
    33  	if me.l.Burst() != 0 {
    34  		b = b[:min(len(b), me.l.Burst())]
    35  	}
    36  	n, err = me.r.Read(b)
    37  	// golang.org/x/time/rate is completely fucking busted. TODO: Write my own rate limiter.
    38  	r := me.l.ReserveN(time.Now(), n)
    39  	panicif.False(r.OK())
    40  	time.Sleep(r.Delay())
    41  	return
    42  }