octoleo/restic
2
2
Fork 0
mirror of https://github.com/octoleo/restic.git synced 2024-12-22 02:48:55 +00:00
Code Issues Releases Wiki Activity

Merge pull request #3373 from greatroar/simplify-limiter

Browse Source 
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
internal/limiter/limiter_backend.go
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 {
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
original io.ReadCloser
}
type limitedReadWriteToCloser struct {
limitedReadCloser
writerTo io.WriterTo
limiter Limiter
}
func newDownstreamLimitedReadCloser(rd io.Reader, limiter Limiter, original io.ReadCloser) io.ReadCloser {
lrd := limitedReadCloser{
Reader: limiter.Downstream(rd),
original: original,
}
if _, ok := rd.(io.WriterTo); ok {
return &limitedReadWriteToCloser{
limitedReadCloser: lrd,
writerTo: rd.(io.WriterTo),
limiter: limiter,
func newDownstreamLimitedReader(rd io.Reader, limiter Limiter) io.Reader {
lrd := limiter.Downstream(rd)
if wt, ok := rd.(io.WriterTo); ok {
lrd = &limitedReader{
Reader: lrd,
writerTo: wt,
limiter: limiter,
}
}
return &lrd
return lrd
}
func (l limitedReadCloser) Close() error {
if l.original == nil {
return nil
}
return l.original.Close()
}
func (l limitedReadWriteToCloser) WriteTo(w io.Writer) (int64, error) {
func (l *limitedReader) WriteTo(w io.Writer) (int64, error) {
return l.writerTo.WriteTo(l.limiter.DownstreamWriter(w))
}

17
internal/limiter/static_limiter.go
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) {
type readCloser struct {
io.Reader
io.Closer
}
if req.Body != nil {
req.Body = limitedReadCloser{
Reader: l.Upstream(req.Body),
original: req.Body,
req.Body = &readCloser{
Reader: l.Upstream(req.Body),
Closer: req.Body,
}
}
res, err := rt.RoundTrip(req)
if res != nil && res.Body != nil {
res.Body = limitedReadCloser{
Reader: l.Downstream(res.Body),
original: res.Body,
res.Body = &readCloser{
Reader: l.Downstream(res.Body),
Closer: res.Body,
}
}