ratelimitreader.go 1.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. package torrent
  2. import (
  3. "context"
  4. "fmt"
  5. "io"
  6. "time"
  7. "golang.org/x/time/rate"
  8. )
  9. type rateLimitedReader struct {
  10. l *rate.Limiter
  11. r io.Reader
  12. // This is the time of the last Read's reservation.
  13. lastRead time.Time
  14. }
  15. func (me *rateLimitedReader) Read(b []byte) (n int, err error) {
  16. const oldStyle = false // Retained for future reference.
  17. if oldStyle {
  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. now := time.Now()
  29. me.lastRead = now
  30. if !me.l.ReserveN(now, n-1).OK() {
  31. panic(fmt.Sprintf("burst exceeded?: %d", n-1))
  32. }
  33. } else {
  34. // Limit the read to within the burst.
  35. if me.l.Limit() != rate.Inf && len(b) > me.l.Burst() {
  36. b = b[:me.l.Burst()]
  37. }
  38. n, err = me.r.Read(b)
  39. now := time.Now()
  40. r := me.l.ReserveN(now, n)
  41. if !r.OK() {
  42. panic(n)
  43. }
  44. me.lastRead = now
  45. time.Sleep(r.Delay())
  46. }
  47. return
  48. }