github.com/larrabee/ratelimit@v1.0.6-0.20191102113931-712217ec4fdc/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  	reader io.Reader
    11  	bucket Bucket
    12  }
    13  
    14  type readSeeker struct {
    15  	reader io.ReadSeeker
    16  	bucket Bucket
    17  }
    18  
    19  type readCloser struct {
    20  	reader io.ReadCloser
    21  	bucket Bucket
    22  }
    23  
    24  // NewReader returns a reader that is rate limited by
    25  // the given token bucket. Each token in the bucket
    26  // represents one byte.
    27  func NewReader(r io.Reader, bucket Bucket) io.Reader {
    28  	return &reader{
    29  		reader: r,
    30  		bucket: bucket,
    31  	}
    32  }
    33  
    34  func (r *reader) Read(buf []byte) (int, error) {
    35  	n, err := r.reader.Read(buf)
    36  	if n <= 0 {
    37  		return n, err
    38  	}
    39  	r.bucket.Wait(int64(n))
    40  	return n, err
    41  }
    42  
    43  // NewReader returns a readSeeker that is rate limited by
    44  // the given token bucket. Each token in the bucket
    45  // represents one byte.
    46  func NewReadSeeker(r io.ReadSeeker, bucket Bucket) io.ReadSeeker {
    47  	return &readSeeker{
    48  		reader: r,
    49  		bucket: bucket,
    50  	}
    51  }
    52  
    53  func (r *readSeeker) Read(buf []byte) (int, error) {
    54  	n, err := r.reader.Read(buf)
    55  	if n <= 0 {
    56  		return n, err
    57  	}
    58  	r.bucket.Wait(int64(n))
    59  	return n, err
    60  }
    61  
    62  func (r *readSeeker) Seek(offset int64, whence int) (int64, error) {
    63  	return r.reader.Seek(offset, whence)
    64  }
    65  
    66  // NewReader returns a readCloser that is rate limited by
    67  // the given token bucket. Each token in the bucket
    68  // represents one byte.
    69  func NewReadCloser(r io.ReadCloser, bucket Bucket) io.ReadCloser {
    70  	return &readCloser{
    71  		reader: r,
    72  		bucket: bucket,
    73  	}
    74  }
    75  
    76  func (r *readCloser) Read(buf []byte) (int, error) {
    77  	n, err := r.reader.Read(buf)
    78  	if n <= 0 {
    79  		return n, err
    80  	}
    81  	r.bucket.Wait(int64(n))
    82  	return n, err
    83  }
    84  
    85  func (r *readCloser) Close() error {
    86  	return r.reader.Close()
    87  }