From a287e3ced70f2666d675c52e671be5881e2c1c21 Mon Sep 17 00:00:00 2001 From: Nick Craig-Wood Date: Thu, 19 Feb 2015 19:26:00 +0000 Subject: [PATCH] Implement --bwlimit to limit data transfer bandwidth --- README.md | 1 + docs/content/docs.md | 1 + fs/accounting.go | 17 ++++++++++++- fs/config.go | 59 ++++++++++++++++++++++++++++++++++++++++++++ fs/config_test.go | 55 +++++++++++++++++++++++++++++++++++++++++ 5 files changed, 132 insertions(+), 1 deletion(-) create mode 100644 fs/config_test.go diff --git a/README.md b/README.md index b098bde52..0e20c2867 100644 --- a/README.md +++ b/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 diff --git a/docs/content/docs.md b/docs/content/docs.md index c948676e8..0e0600935 100644 --- a/docs/content/docs.md +++ b/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. diff --git a/fs/accounting.go b/fs/accounting.go index d8631dc46..93f38c613 100644 --- a/fs/accounting.go +++ b/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 } diff --git a/fs/config.go b/fs/config.go index cb0272b18..a68c89605 100644 --- a/fs/config.go +++ b/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. diff --git a/fs/config_test.go b/fs/config_test.go new file mode 100644 index 000000000..f0989ece9 --- /dev/null +++ b/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) + } + } +}