package proxy import ( "context" "golang.org/x/time/rate" "io" "time" ) const burstLimit = 1000 * 1000 * 1000 type Reader struct { r io.Reader limiter *rate.Limiter ctx context.Context } type Writer struct { w io.Writer limiter *rate.Limiter ctx context.Context } // NewReader returns a reader that implements io.Reader with rate limiting. func NewReader(r io.Reader) *Reader { return &Reader{ r: r, ctx: context.Background(), } } // SetRateLimit sets rate limit (bytes/sec) to the reader. func (s *Reader) SetRateLimit(bytesPerSec float64) { s.limiter = rate.NewLimiter(rate.Limit(bytesPerSec), burstLimit) s.limiter.AllowN(time.Now(), burstLimit) // spend initial burst } // Read reads bytes into p. func (s *Reader) Read(p []byte) (int, error) { if s.limiter == nil { return s.r.Read(p) } n, err := s.r.Read(p) if err != nil { return n, err } if err := s.limiter.WaitN(s.ctx, n); err != nil { return n, err } return n, nil }