mirror of https://github.com/restic/restic.git
Merge pull request #3373 from greatroar/simplify-limiter
Simplify internal/limiter
This commit is contained in:
commit
6d8ceefd67
|
@ -42,44 +42,29 @@ func (l limitedRewindReader) Read(b []byte) (int, error) {
|
||||||
|
|
||||||
func (r rateLimitedBackend) Load(ctx context.Context, h restic.Handle, length int, offset int64, consumer func(rd io.Reader) error) error {
|
func (r rateLimitedBackend) Load(ctx context.Context, h restic.Handle, length int, offset int64, consumer func(rd io.Reader) error) error {
|
||||||
return r.Backend.Load(ctx, h, length, offset, func(rd io.Reader) error {
|
return r.Backend.Load(ctx, h, length, offset, func(rd io.Reader) error {
|
||||||
return consumer(newDownstreamLimitedReadCloser(rd, r.limiter, nil))
|
return consumer(newDownstreamLimitedReader(rd, r.limiter))
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
type limitedReadCloser struct {
|
type limitedReader struct {
|
||||||
io.Reader
|
io.Reader
|
||||||
original io.ReadCloser
|
|
||||||
}
|
|
||||||
|
|
||||||
type limitedReadWriteToCloser struct {
|
|
||||||
limitedReadCloser
|
|
||||||
writerTo io.WriterTo
|
writerTo io.WriterTo
|
||||||
limiter Limiter
|
limiter Limiter
|
||||||
}
|
}
|
||||||
|
|
||||||
func newDownstreamLimitedReadCloser(rd io.Reader, limiter Limiter, original io.ReadCloser) io.ReadCloser {
|
func newDownstreamLimitedReader(rd io.Reader, limiter Limiter) io.Reader {
|
||||||
lrd := limitedReadCloser{
|
lrd := limiter.Downstream(rd)
|
||||||
Reader: limiter.Downstream(rd),
|
if wt, ok := rd.(io.WriterTo); ok {
|
||||||
original: original,
|
lrd = &limitedReader{
|
||||||
}
|
Reader: lrd,
|
||||||
if _, ok := rd.(io.WriterTo); ok {
|
writerTo: wt,
|
||||||
return &limitedReadWriteToCloser{
|
|
||||||
limitedReadCloser: lrd,
|
|
||||||
writerTo: rd.(io.WriterTo),
|
|
||||||
limiter: limiter,
|
limiter: limiter,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return &lrd
|
return lrd
|
||||||
}
|
}
|
||||||
|
|
||||||
func (l limitedReadCloser) Close() error {
|
func (l *limitedReader) WriteTo(w io.Writer) (int64, error) {
|
||||||
if l.original == nil {
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
return l.original.Close()
|
|
||||||
}
|
|
||||||
|
|
||||||
func (l limitedReadWriteToCloser) WriteTo(w io.Writer) (int64, error) {
|
|
||||||
return l.writerTo.WriteTo(l.limiter.DownstreamWriter(w))
|
return l.writerTo.WriteTo(l.limiter.DownstreamWriter(w))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
@ -57,19 +57,24 @@ func (rt roundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||||
}
|
}
|
||||||
|
|
||||||
func (l staticLimiter) roundTripper(rt http.RoundTripper, req *http.Request) (*http.Response, error) {
|
func (l staticLimiter) roundTripper(rt http.RoundTripper, req *http.Request) (*http.Response, error) {
|
||||||
|
type readCloser struct {
|
||||||
|
io.Reader
|
||||||
|
io.Closer
|
||||||
|
}
|
||||||
|
|
||||||
if req.Body != nil {
|
if req.Body != nil {
|
||||||
req.Body = limitedReadCloser{
|
req.Body = &readCloser{
|
||||||
Reader: l.Upstream(req.Body),
|
Reader: l.Upstream(req.Body),
|
||||||
original: req.Body,
|
Closer: req.Body,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
res, err := rt.RoundTrip(req)
|
res, err := rt.RoundTrip(req)
|
||||||
|
|
||||||
if res != nil && res.Body != nil {
|
if res != nil && res.Body != nil {
|
||||||
res.Body = limitedReadCloser{
|
res.Body = &readCloser{
|
||||||
Reader: l.Downstream(res.Body),
|
Reader: l.Downstream(res.Body),
|
||||||
original: res.Body,
|
Closer: res.Body,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue