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

     1  package limiter
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  
     7  	"github.com/restic/restic/internal/restic"
     8  )
     9  
    10  // LimitBackend wraps a Backend and applies rate limiting to Load() and Save()
    11  // calls on the backend.
    12  func LimitBackend(be restic.Backend, l Limiter) restic.Backend {
    13  	return rateLimitedBackend{
    14  		Backend: be,
    15  		limiter: l,
    16  	}
    17  }
    18  
    19  type rateLimitedBackend struct {
    20  	restic.Backend
    21  	limiter Limiter
    22  }
    23  
    24  func (r rateLimitedBackend) Save(ctx context.Context, h restic.Handle, rd io.Reader) error {
    25  	return r.Backend.Save(ctx, h, r.limiter.Upstream(rd))
    26  }
    27  
    28  func (r rateLimitedBackend) Load(ctx context.Context, h restic.Handle, length int, offset int64) (io.ReadCloser, error) {
    29  	rc, err := r.Backend.Load(ctx, h, length, offset)
    30  	if err != nil {
    31  		return nil, err
    32  	}
    33  
    34  	return limitedReadCloser{
    35  		original: rc,
    36  		limited:  r.limiter.Downstream(rc),
    37  	}, nil
    38  }
    39  
    40  type limitedReadCloser struct {
    41  	original io.ReadCloser
    42  	limited  io.Reader
    43  }
    44  
    45  func (l limitedReadCloser) Read(b []byte) (n int, err error) {
    46  	return l.limited.Read(b)
    47  }
    48  
    49  func (l limitedReadCloser) Close() error {
    50  	return l.original.Close()
    51  }
    52  
    53  var _ restic.Backend = (*rateLimitedBackend)(nil)