90d2c0502b
The GlobalOptions struct now embeds a backend.TransportOptions, so it doesn't need to construct one in open and create. The upload and download limits are similarly now a struct in internal/limiter that is embedded in GlobalOptions.
114 lines
2.5 KiB
Go
114 lines
2.5 KiB
Go
package limiter
|
|
|
|
import (
|
|
"io"
|
|
"net/http"
|
|
|
|
"github.com/juju/ratelimit"
|
|
)
|
|
|
|
type staticLimiter struct {
|
|
upstream *ratelimit.Bucket
|
|
downstream *ratelimit.Bucket
|
|
}
|
|
|
|
// Limits represents static upload and download limits.
|
|
// For both, zero means unlimited.
|
|
type Limits struct {
|
|
UploadKb int
|
|
DownloadKb int
|
|
}
|
|
|
|
// NewStaticLimiter constructs a Limiter with a fixed (static) upload and
|
|
// download rate cap
|
|
func NewStaticLimiter(l Limits) Limiter {
|
|
var (
|
|
upstreamBucket *ratelimit.Bucket
|
|
downstreamBucket *ratelimit.Bucket
|
|
)
|
|
|
|
if l.UploadKb > 0 {
|
|
upstreamBucket = ratelimit.NewBucketWithRate(toByteRate(l.UploadKb), int64(toByteRate(l.UploadKb)))
|
|
}
|
|
|
|
if l.DownloadKb > 0 {
|
|
downstreamBucket = ratelimit.NewBucketWithRate(toByteRate(l.DownloadKb), int64(toByteRate(l.DownloadKb)))
|
|
}
|
|
|
|
return staticLimiter{
|
|
upstream: upstreamBucket,
|
|
downstream: downstreamBucket,
|
|
}
|
|
}
|
|
|
|
func (l staticLimiter) Upstream(r io.Reader) io.Reader {
|
|
return l.limitReader(r, l.upstream)
|
|
}
|
|
|
|
func (l staticLimiter) UpstreamWriter(w io.Writer) io.Writer {
|
|
return l.limitWriter(w, l.upstream)
|
|
}
|
|
|
|
func (l staticLimiter) Downstream(r io.Reader) io.Reader {
|
|
return l.limitReader(r, l.downstream)
|
|
}
|
|
|
|
func (l staticLimiter) DownstreamWriter(w io.Writer) io.Writer {
|
|
return l.limitWriter(w, l.downstream)
|
|
}
|
|
|
|
type roundTripper func(*http.Request) (*http.Response, error)
|
|
|
|
func (rt roundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
|
|
return rt(req)
|
|
}
|
|
|
|
func (l staticLimiter) roundTripper(rt http.RoundTripper, req *http.Request) (*http.Response, error) {
|
|
type readCloser struct {
|
|
io.Reader
|
|
io.Closer
|
|
}
|
|
|
|
if req.Body != nil {
|
|
req.Body = &readCloser{
|
|
Reader: l.Upstream(req.Body),
|
|
Closer: req.Body,
|
|
}
|
|
}
|
|
|
|
res, err := rt.RoundTrip(req)
|
|
|
|
if res != nil && res.Body != nil {
|
|
res.Body = &readCloser{
|
|
Reader: l.Downstream(res.Body),
|
|
Closer: res.Body,
|
|
}
|
|
}
|
|
|
|
return res, err
|
|
}
|
|
|
|
// Transport returns an HTTP transport limited with the limiter l.
|
|
func (l staticLimiter) Transport(rt http.RoundTripper) http.RoundTripper {
|
|
return roundTripper(func(req *http.Request) (*http.Response, error) {
|
|
return l.roundTripper(rt, req)
|
|
})
|
|
}
|
|
|
|
func (l staticLimiter) limitReader(r io.Reader, b *ratelimit.Bucket) io.Reader {
|
|
if b == nil {
|
|
return r
|
|
}
|
|
return ratelimit.Reader(r, b)
|
|
}
|
|
|
|
func (l staticLimiter) limitWriter(w io.Writer, b *ratelimit.Bucket) io.Writer {
|
|
if b == nil {
|
|
return w
|
|
}
|
|
return ratelimit.Writer(w, b)
|
|
}
|
|
|
|
func toByteRate(val int) float64 {
|
|
return float64(val) * 1024.
|
|
}
|