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

     1  package ratelimit
     2  
     3  import "io"
     4  
     5  type writer struct {
     6  	writer io.Writer
     7  	bucket Bucket
     8  }
     9  
    10  type writeSeeker struct {
    11  	writer io.WriteSeeker
    12  	bucket Bucket
    13  }
    14  
    15  type writeCloser struct {
    16  	writer io.WriteCloser
    17  	bucket Bucket
    18  }
    19  
    20  // NewWriter returns a writer that is rate limited by
    21  // the given token bucket. Each token in the bucket
    22  // represents one byte.
    23  func NewWriter(w io.Writer, bucket Bucket) io.Writer {
    24  	return &writer{
    25  		writer: w,
    26  		bucket: bucket,
    27  	}
    28  }
    29  
    30  func (w *writer) Write(buf []byte) (int, error) {
    31  	w.bucket.Wait(int64(len(buf)))
    32  	return w.writer.Write(buf)
    33  }
    34  
    35  // NewWriter returns a writer that is rate limited by
    36  // the given token bucket. Each token in the bucket
    37  // represents one byte.
    38  func NewWriteSeeker(w io.WriteSeeker, bucket Bucket) io.WriteSeeker {
    39  	return &writeSeeker{
    40  		writer: w,
    41  		bucket: bucket,
    42  	}
    43  }
    44  
    45  func (w *writeSeeker) Write(buf []byte) (int, error) {
    46  	w.bucket.Wait(int64(len(buf)))
    47  	return w.writer.Write(buf)
    48  }
    49  
    50  func (r *writeSeeker) Seek(offset int64, whence int) (int64, error) {
    51  	return r.writer.Seek(offset, whence)
    52  }
    53  
    54  // NewWriter returns a writer that is rate limited by
    55  // the given token bucket. Each token in the bucket
    56  // represents one byte.
    57  func NewWriteCloser(w io.WriteCloser, bucket Bucket) io.WriteCloser {
    58  	return &writeCloser{
    59  		writer: w,
    60  		bucket: bucket,
    61  	}
    62  }
    63  
    64  func (w *writeCloser) Write(buf []byte) (int, error) {
    65  	w.bucket.Wait(int64(len(buf)))
    66  	return w.writer.Write(buf)
    67  }
    68  
    69  func (r *writeCloser) Close() error {
    70  	return r.writer.Close()
    71  }