github.com/angenalZZZ/gofunc@v0.0.0-20210507121333-48ff1be3917b/http/ratelimit/reader.go (about)

     1  package ratelimit
     2  
     3  import "io"
     4  
     5  type reader struct {
     6  	r      io.Reader
     7  	bucket *Bucket
     8  }
     9  
    10  // Reader returns a reader that is rate limited by
    11  // the given token bucket. Each token in the bucket
    12  // represents one byte.
    13  func Reader(r io.Reader, bucket *Bucket) io.Reader {
    14  	return &reader{
    15  		r:      r,
    16  		bucket: bucket,
    17  	}
    18  }
    19  
    20  func (r *reader) Read(buf []byte) (int, error) {
    21  	n, err := r.r.Read(buf)
    22  	if n <= 0 {
    23  		return n, err
    24  	}
    25  	r.bucket.Wait(int64(n))
    26  	return n, err
    27  }
    28  
    29  type writer struct {
    30  	w      io.Writer
    31  	bucket *Bucket
    32  }
    33  
    34  // Writer returns a reader that is rate limited by
    35  // the given token bucket. Each token in the bucket
    36  // represents one byte.
    37  func Writer(w io.Writer, bucket *Bucket) io.Writer {
    38  	return &writer{
    39  		w:      w,
    40  		bucket: bucket,
    41  	}
    42  }
    43  
    44  func (w *writer) Write(buf []byte) (int, error) {
    45  	w.bucket.Wait(int64(len(buf)))
    46  	return w.w.Write(buf)
    47  }