Basic rate limiting implementation.

Added `--limit-upload` and `--limit-download` flags to rate limit
backups and restores.
This commit is contained in:
rmdashrf
2017-10-08 11:28:03 -07:00
parent 8ceb22fe8a
commit 32637a0328
12 changed files with 1179 additions and 4 deletions

View File

@@ -0,0 +1,53 @@
package limiter
import (
"io"
"github.com/juju/ratelimit"
)
type staticLimiter struct {
upstream *ratelimit.Bucket
downstream *ratelimit.Bucket
}
// NewStaticLimiter constructs a Limiter with a fixed (static) upload and
// download rate cap
func NewStaticLimiter(uploadKb, downloadKb int) Limiter {
var (
upstreamBucket *ratelimit.Bucket
downstreamBucket *ratelimit.Bucket
)
if uploadKb > 0 {
upstreamBucket = ratelimit.NewBucketWithRate(toByteRate(uploadKb), int64(toByteRate(uploadKb)))
}
if downloadKb > 0 {
downstreamBucket = ratelimit.NewBucketWithRate(toByteRate(downloadKb), int64(toByteRate(downloadKb)))
}
return staticLimiter{
upstream: upstreamBucket,
downstream: downstreamBucket,
}
}
func (l staticLimiter) Upstream(r io.Reader) io.Reader {
return l.limit(r, l.upstream)
}
func (l staticLimiter) Downstream(r io.Reader) io.Reader {
return l.limit(r, l.downstream)
}
func (l staticLimiter) limit(r io.Reader, b *ratelimit.Bucket) io.Reader {
if b == nil {
return r
}
return ratelimit.Reader(r, b)
}
func toByteRate(val int) float64 {
return float64(val) * 1024.
}