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 }