github.com/larrabee/ratelimit@v1.0.6-0.20191102113931-712217ec4fdc/readwrite.go (about)

     1  package ratelimit
     2  
     3  import "io"
     4  
     5  type readWriter struct {
     6  	rw     io.ReadWriter
     7  	bucket Bucket
     8  }
     9  
    10  type readWriteCloser struct {
    11  	rw     io.ReadWriteCloser
    12  	bucket Bucket
    13  }
    14  
    15  type readWriteSeeker struct {
    16  	rw     io.ReadWriteSeeker
    17  	bucket Bucket
    18  }
    19  
    20  // NewReadWriter returns a readwriter that is rate limited by
    21  // the given token bucket. Each token in the bucket
    22  // represents one byte.
    23  func NewReadWriter(w io.ReadWriter, bucket Bucket) io.ReadWriter {
    24  	return &readWriter{
    25  		rw:     w,
    26  		bucket: bucket,
    27  	}
    28  }
    29  
    30  func (w *readWriter) Write(buf []byte) (int, error) {
    31  	w.bucket.Wait(int64(len(buf)))
    32  	return w.rw.Write(buf)
    33  }
    34  
    35  func (r *readWriter) Read(buf []byte) (int, error) {
    36  	n, err := r.rw.Read(buf)
    37  	if n <= 0 {
    38  		return n, err
    39  	}
    40  	r.bucket.Wait(int64(n))
    41  	return n, err
    42  }
    43  
    44  // NewReadWriter returns a readwriter that is rate limited by
    45  // the given token bucket. Each token in the bucket
    46  // represents one byte.
    47  func NewReadWriteCloser(w io.ReadWriteCloser, bucket Bucket) io.ReadWriteCloser {
    48  	return &readWriteCloser{
    49  		rw:     w,
    50  		bucket: bucket,
    51  	}
    52  }
    53  
    54  func (w *readWriteCloser) Write(buf []byte) (int, error) {
    55  	w.bucket.Wait(int64(len(buf)))
    56  	return w.rw.Write(buf)
    57  }
    58  
    59  func (r *readWriteCloser) Read(buf []byte) (int, error) {
    60  	n, err := r.rw.Read(buf)
    61  	if n <= 0 {
    62  		return n, err
    63  	}
    64  	r.bucket.Wait(int64(n))
    65  	return n, err
    66  }
    67  
    68  func (r *readWriteCloser) Close() error {
    69  	return r.rw.Close()
    70  }
    71  
    72  // NewReadWriter returns a readwriter that is rate limited by
    73  // the given token bucket. Each token in the bucket
    74  // represents one byte.
    75  func NewReadWriteSeeker(w io.ReadWriteSeeker, bucket Bucket) io.ReadWriteSeeker {
    76  	return &readWriteSeeker{
    77  		rw:     w,
    78  		bucket: bucket,
    79  	}
    80  }
    81  
    82  func (w *readWriteSeeker) Write(buf []byte) (int, error) {
    83  	w.bucket.Wait(int64(len(buf)))
    84  	return w.rw.Write(buf)
    85  }
    86  
    87  func (r *readWriteSeeker) Read(buf []byte) (int, error) {
    88  	n, err := r.rw.Read(buf)
    89  	if n <= 0 {
    90  		return n, err
    91  	}
    92  	r.bucket.Wait(int64(n))
    93  	return n, err
    94  }
    95  
    96  func (r *readWriteSeeker) Seek(offset int64, whence int) (int64, error) {
    97  	return r.rw.Seek(offset, whence)
    98  }