Implement --bwlimit to limit data transfer bandwidth

README.md

@@ -125,6 +125,7 @@ the same format as the standard md5sum tool produces.
 General options:
 
 ```
+      --bwlimit=0: Bandwidth limit in kBytes/s, or use suffix k|M|G
       --checkers=8: Number of checkers to run in parallel.
       --config="~/.rclone.conf": Config file.
   -n, --dry-run=false: Do a trial run with no permanent changes

docs/content/docs.md

@@ -107,6 +107,7 @@ the same format as the standard md5sum tool produces.
 General options:
 
 ```
+      --bwlimit=0: Bandwidth limit in kBytes/s, or use suffix k|M|G
       --checkers=8: Number of checkers to run in parallel.
       --transfers=4: Number of file transfers to run in parallel.
       --config="~/.rclone.conf": Config file.

fs/accounting.go

@@ -10,13 +10,24 @@ import (
 	"strings"
 	"sync"
 	"time"
+
+	"github.com/tsenart/tb"
 )
 
 // Globals
 var (
-	Stats = NewStats()
+	Stats       = NewStats()
+	tokenBucket *tb.Bucket
 )
 
+// Start the token bucket if necessary
+func startTokenBucket() {
+	if bwLimit > 0 {
+		tokenBucket = tb.NewBucket(int64(bwLimit), 100*time.Millisecond)
+		Log(nil, "Starting bandwidth limiter at %vBytes/s", &bwLimit)
+	}
+}
+
 // Stringset holds some strings
 type StringSet map[string]bool
 
@@ -178,6 +189,10 @@ func (file *Account) Read(p []byte) (n int, err error) {
 	if err == io.EOF {
 		// FIXME Do something?
 	}
+	// Limit the transfer speed if required
+	if tokenBucket != nil {
+		tokenBucket.Wait(int64(n))
+	}
 	return
 }
 

fs/config.go

@@ -22,6 +22,8 @@ const (
 	configFileName = ".rclone.conf"
 )
 
+type SizeSuffix int64
+
 // Global
 var (
 	// Config file
@@ -40,8 +42,62 @@ var (
 	transfers    = pflag.IntP("transfers", "", 4, "Number of file transfers to run in parallel.")
 	configFile   = pflag.StringP("config", "", ConfigPath, "Config file.")
 	dryRun       = pflag.BoolP("dry-run", "n", false, "Do a trial run with no permanent changes")
+	bwLimit      SizeSuffix
 )
 
+func init() {
+	pflag.VarP(&bwLimit, "bwlimit", "", "Bandwidth limit in kBytes/s, or use suffix K|M|G")
+}
+
+// Turn SizeSuffix into a string
+func (x *SizeSuffix) String() string {
+	switch {
+	case *x == 0:
+		return "0"
+	case *x < 1024*1024:
+		return fmt.Sprintf("%.3fk", float64(*x)/1024)
+	case *x < 1024*1024*1024:
+		return fmt.Sprintf("%.3fM", float64(*x)/1024/1024)
+	default:
+		return fmt.Sprintf("%.3fG", float64(*x)/1024/1024/1024)
+	}
+	panic("shouldn't be reached")
+}
+
+// Set a SizeSuffix
+func (x *SizeSuffix) Set(s string) error {
+	if len(s) == 0 {
+		return fmt.Errorf("Empty string")
+	}
+	suffix := s[len(s)-1]
+	suffixLen := 1
+	var multiplier float64
+	switch suffix {
+	case '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '.':
+		suffixLen = 0
+		multiplier = 1 << 10
+	case 'k', 'K':
+		multiplier = 1 << 10
+	case 'm', 'M':
+		multiplier = 1 << 20
+	case 'g', 'G':
+		multiplier = 1 << 30
+	default:
+		return fmt.Errorf("Bad suffix %q", suffix)
+	}
+	s = s[:len(s)-suffixLen]
+	value, err := strconv.ParseFloat(s, 64)
+	if err != nil {
+		return err
+	}
+	value *= multiplier
+	*x = SizeSuffix(value)
+	return nil
+}
+
+// Check it satisfies the interface
+var _ pflag.Value = (*SizeSuffix)(nil)
+
 // Filesystem config options
 type ConfigInfo struct {
 	Verbose      bool
@@ -97,6 +153,9 @@ func LoadConfig() {
 			log.Fatalf("Failed to read null config file: %v", err)
 		}
 	}
+
+	// Start the token bucket limiter
+	startTokenBucket()
 }
 
 // Save configuration file.

fs/config_test.go

@@ -0,0 +1,55 @@
+package fs
+
+import "testing"
+
+func TestSizeSuffixString(t *testing.T) {
+	for _, test := range []struct {
+		in   float64
+		want string
+	}{
+		{0, "0"},
+		{102, "0.100k"},
+		{1024, "1.000k"},
+		{1024 * 1024, "1.000M"},
+		{1024 * 1024 * 1024, "1.000G"},
+		{10 * 1024 * 1024 * 1024, "10.000G"},
+	} {
+		ss := SizeSuffix(test.in)
+		got := ss.String()
+		if test.want != got {
+			t.Errorf("Want %v got %v", test.want, got)
+		}
+	}
+}
+
+func TestSizeSuffixSet(t *testing.T) {
+	for i, test := range []struct {
+		in   string
+		want int64
+		err  bool
+	}{
+		{"0", 0, false},
+		{"0.1k", 102, false},
+		{"0.1", 102, false},
+		{"1K", 1024, false},
+		{"1", 1024, false},
+		{"2.5", 1024 * 2.5, false},
+		{"1M", 1024 * 1024, false},
+		{"1.g", 1024 * 1024 * 1024, false},
+		{"10G", 10 * 1024 * 1024 * 1024, false},
+		{"", 0, true},
+		{"1p", 0, true},
+		{"1.p", 0, true},
+		{"1p", 0, true},
+	} {
+		ss := SizeSuffix(0)
+		err := ss.Set(test.in)
+		if (err != nil) != test.err {
+			t.Errorf("%d: Expecting error %v but got error %v", i, test.err, err)
+		}
+		got := int64(ss)
+		if test.want != got {
+			t.Errorf("%d: Want %v got %v", i, test.want, got)
+		}
+	}
+}