Merge pull request #3373 from greatroar/simplify-limiter

Simplify internal/limiter
This commit is contained in:
MichaelEischer 2021-05-14 20:27:31 +02:00 committed by GitHub
commit 6d8ceefd67
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 22 additions and 32 deletions

View File

@ -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))
} }

View File

@ -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,
} }
} }