package ratelimit
import "io"
type reader struct { r io.Reader
bucket *Bucket
}
// Reader returns a reader that is rate limited by
// the given token bucket. Each token in the bucket
// represents one byte.
func Reader(r io.Reader, bucket *Bucket) io.Reader { return &reader{ r: r,
bucket: bucket,
}
}
func (r *reader) Read(buf []byte) (int, error) { n, err := r.r.Read(buf)
if n <= 0 { return n, err
}
r.bucket.Wait(int64(n))
return n, err
}
type writer struct { w io.Writer
bucket *Bucket
}
// Writer returns a reader that is rate limited by
// the given token bucket. Each token in the bucket
// represents one byte.
func Writer(w io.Writer, bucket *Bucket) io.Writer { return &writer{ w: w,
bucket: bucket,
}
}
func (w *writer) Write(buf []byte) (int, error) { w.bucket.Wait(int64(len(buf)))
return w.w.Write(buf)
}