github.com/lrita/ratelimit@v0.0.0-20190723030019-81504bd89bc5/reader.go (about)

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