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 }