ratelimitreader.go

   1 package torrent
   2
   3 import (
   4         "context"
   5         "fmt"
   6         "io"
   7         "time"
   8
   9         "golang.org/x/time/rate"
  10 )
  11
  12 type rateLimitedReader struct {
  13         l *rate.Limiter
  14         r io.Reader
  15
  16         // This is the time of the last Read's reservation.
  17         lastRead time.Time
  18 }
  19
  20 func (me *rateLimitedReader) Read(b []byte) (n int, err error) {
  21         const oldStyle = false // Retained for future reference.
  22         if oldStyle {
  23                 // Wait until we can read at all.
  24                 if err := me.l.WaitN(context.Background(), 1); err != nil {
  25                         panic(err)
  26                 }
  27                 // Limit the read to within the burst.
  28                 if me.l.Limit() != rate.Inf && len(b) > me.l.Burst() {
  29                         b = b[:me.l.Burst()]
  30                 }
  31                 n, err = me.r.Read(b)
  32                 // Pay the piper.
  33                 now := time.Now()
  34                 me.lastRead = now
  35                 if !me.l.ReserveN(now, n-1).OK() {
  36                         panic(fmt.Sprintf("burst exceeded?: %d", n-1))
  37                 }
  38         } else {
  39                 // Limit the read to within the burst.
  40                 if me.l.Limit() != rate.Inf && len(b) > me.l.Burst() {
  41                         b = b[:me.l.Burst()]
  42                 }
  43                 n, err = me.r.Read(b)
  44                 now := time.Now()
  45                 r := me.l.ReserveN(now, n)
  46                 if !r.OK() {
  47                         panic(n)
  48                 }
  49                 me.lastRead = now
  50                 time.Sleep(r.Delay())
  51         }
  52         return
  53 }