ratelimitreader.go

   1 package torrent
   2
   3 import (
   4         "fmt"
   5         "io"
   6         "time"
   7
   8         "golang.org/x/net/context"
   9         "golang.org/x/time/rate"
  10 )
  11
  12 type rateLimitedReader struct {
  13         l *rate.Limiter
  14         r io.Reader
  15 }
  16
  17 func (me rateLimitedReader) Read(b []byte) (n int, err error) {
  18         // Wait until we can read at all.
  19         if err := me.l.WaitN(context.Background(), 1); err != nil {
  20                 panic(err)
  21         }
  22         // Limit the read to within the burst.
  23         if me.l.Limit() != rate.Inf && len(b) > me.l.Burst() {
  24                 b = b[:me.l.Burst()]
  25         }
  26         n, err = me.r.Read(b)
  27         // Pay the piper.
  28         if !me.l.ReserveN(time.Now(), n-1).OK() {
  29                 panic(fmt.Sprintf("burst exceeded?: %d", n-1))
  30         }
  31         return
  32 }