github.com/fawick/restic@v0.1.1-0.20171126184616-c02923fbfc79/internal/limiter/static_limiter.go (about)

     1  package limiter
     2  
     3  import (
     4  	"io"
     5  
     6  	"github.com/juju/ratelimit"
     7  )
     8  
     9  type staticLimiter struct {
    10  	upstream   *ratelimit.Bucket
    11  	downstream *ratelimit.Bucket
    12  }
    13  
    14  // NewStaticLimiter constructs a Limiter with a fixed (static) upload and
    15  // download rate cap
    16  func NewStaticLimiter(uploadKb, downloadKb int) Limiter {
    17  	var (
    18  		upstreamBucket   *ratelimit.Bucket
    19  		downstreamBucket *ratelimit.Bucket
    20  	)
    21  
    22  	if uploadKb > 0 {
    23  		upstreamBucket = ratelimit.NewBucketWithRate(toByteRate(uploadKb), int64(toByteRate(uploadKb)))
    24  	}
    25  
    26  	if downloadKb > 0 {
    27  		downstreamBucket = ratelimit.NewBucketWithRate(toByteRate(downloadKb), int64(toByteRate(downloadKb)))
    28  	}
    29  
    30  	return staticLimiter{
    31  		upstream:   upstreamBucket,
    32  		downstream: downstreamBucket,
    33  	}
    34  }
    35  
    36  func (l staticLimiter) Upstream(r io.Reader) io.Reader {
    37  	return l.limit(r, l.upstream)
    38  }
    39  
    40  func (l staticLimiter) Downstream(r io.Reader) io.Reader {
    41  	return l.limit(r, l.downstream)
    42  }
    43  
    44  func (l staticLimiter) limit(r io.Reader, b *ratelimit.Bucket) io.Reader {
    45  	if b == nil {
    46  		return r
    47  	}
    48  	return ratelimit.Reader(r, b)
    49  }
    50  
    51  func toByteRate(val int) float64 {
    52  	return float64(val) * 1024.
    53  }