rclone: Respect bandwith limits

internal/backend/rclone/backend.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
+	"io"
 	"math/rand"
 	"net"
 	"net/http"
@@ -18,6 +19,7 @@ import (
 	"github.com/restic/restic/internal/backend/rest"
 	"github.com/restic/restic/internal/debug"
 	"github.com/restic/restic/internal/errors"
+	"github.com/restic/restic/internal/limiter"
 	"golang.org/x/net/context/ctxhttp"
 	"golang.org/x/net/http2"
 )
@@ -81,8 +83,38 @@ func run(command string, args ...string) (*StdioConn, *exec.Cmd, *sync.WaitGroup
 	return c, cmd, &wg, bg, nil
 }
 
+// wrappedConn adds bandwidth limiting capabilities to the StdioConn by
+// wrapping the Read/Write methods.
+type wrappedConn struct {
+	*StdioConn
+	io.Reader
+	io.Writer
+}
+
+func (c wrappedConn) Read(p []byte) (int, error) {
+	return c.Reader.Read(p)
+}
+
+func (c wrappedConn) Write(p []byte) (int, error) {
+	return c.Writer.Write(p)
+}
+
+func wrapConn(c *StdioConn, lim limiter.Limiter) wrappedConn {
+	wc := wrappedConn{
+		StdioConn: c,
+		Reader:    c,
+		Writer:    c,
+	}
+	if lim != nil {
+		wc.Reader = lim.Downstream(c)
+		wc.Writer = lim.UpstreamWriter(c)
+	}
+
+	return wc
+}
+
 // New initializes a Backend and starts the process.
-func New(cfg Config) (*Backend, error) {
+func New(cfg Config, lim limiter.Limiter) (*Backend, error) {
 	var (
 		args []string
 		err  error
@@ -118,11 +150,16 @@ func New(cfg Config) (*Backend, error) {
 	arg0, args := args[0], args[1:]
 
 	debug.Log("running command: %v %v", arg0, args)
-	conn, cmd, wg, bg, err := run(arg0, args...)
+	stdioConn, cmd, wg, bg, err := run(arg0, args...)
 	if err != nil {
 		return nil, err
 	}
 
+	var conn net.Conn = stdioConn
+	if lim != nil {
+		conn = wrapConn(stdioConn, lim)
+	}
+
 	dialCount := 0
 	tr := &http2.Transport{
 		AllowHTTP: true, // this is not really HTTP, just stdin/stdout
@@ -141,7 +178,7 @@ func New(cfg Config) (*Backend, error) {
 		tr:     tr,
 		cmd:    cmd,
 		waitCh: waitCh,
-		conn:   conn,
+		conn:   stdioConn,
 		wg:     wg,
 	}
 
@@ -202,8 +239,8 @@ func New(cfg Config) (*Backend, error) {
 }
 
 // Open starts an rclone process with the given config.
-func Open(cfg Config) (*Backend, error) {
-	be, err := New(cfg)
+func Open(cfg Config, lim limiter.Limiter) (*Backend, error) {
+	be, err := New(cfg, lim)
 	if err != nil {
 		return nil, err
 	}
@@ -229,7 +266,7 @@ func Open(cfg Config) (*Backend, error) {
 
 // Create initializes a new restic repo with clone.
 func Create(cfg Config) (*Backend, error) {
-	be, err := New(cfg)
+	be, err := New(cfg, nil)
 	if err != nil {
 		return nil, err
 	}

internal/backend/rclone/backend_test.go

@@ -39,7 +39,7 @@ func newTestSuite(t testing.TB) *test.Suite {
 		Open: func(config interface{}) (restic.Backend, error) {
 			t.Logf("Open()")
 			cfg := config.(rclone.Config)
-			return rclone.Open(cfg)
+			return rclone.Open(cfg, nil)
 		},
 
 		// CleanupFn removes data created during the tests.

internal/limiter/limiter.go

@@ -12,6 +12,10 @@ type Limiter interface {
 	// uploads.
 	Upstream(r io.Reader) io.Reader
 
+	// UpstreamWriter returns a rate limited writer that is intended to be used
+	// in uploads.
+	UpstreamWriter(w io.Writer) io.Writer
+
 	// Downstream returns a rate limited reader that is intended to be used
 	// for downloads.
 	Downstream(r io.Reader) io.Reader

internal/limiter/static_limiter.go

@@ -35,11 +35,15 @@ func NewStaticLimiter(uploadKb, downloadKb int) Limiter {
 }
 
 func (l staticLimiter) Upstream(r io.Reader) io.Reader {
-	return l.limit(r, l.upstream)
+	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.limit(r, l.downstream)
+	return l.limitReader(r, l.downstream)
 }
 
 type roundTripper func(*http.Request) (*http.Response, error)
@@ -75,13 +79,20 @@ func (l staticLimiter) Transport(rt http.RoundTripper) http.RoundTripper {
 	})
 }
 
-func (l staticLimiter) limit(r io.Reader, b *ratelimit.Bucket) io.Reader {
+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.
 }