github.com/lrita/ratelimit@v0.0.0-20190723030019-81504bd89bc5/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 r io.Reader 11 bucket *Bucket 12 } 13 14 // Reader returns a reader that is rate limited by 15 // the given token bucket. Each token in the bucket 16 // represents one byte. 17 func Reader(r io.Reader, bucket *Bucket) io.Reader { 18 return &reader{ 19 r: r, 20 bucket: bucket, 21 } 22 } 23 24 func (r *reader) Read(buf []byte) (int, error) { 25 n, err := r.r.Read(buf) 26 if n <= 0 { 27 return n, err 28 } 29 r.bucket.Wait(int64(n)) 30 return n, err 31 } 32 33 type writer struct { 34 w io.Writer 35 bucket *Bucket 36 } 37 38 // Writer returns a reader that is rate limited by 39 // the given token bucket. Each token in the bucket 40 // represents one byte. 41 func Writer(w io.Writer, bucket *Bucket) io.Writer { 42 return &writer{ 43 w: w, 44 bucket: bucket, 45 } 46 } 47 48 func (w *writer) Write(buf []byte) (int, error) { 49 w.bucket.Wait(int64(len(buf))) 50 return w.w.Write(buf) 51 }