github.com/10XDev/rclone@v1.52.3-0.20200626220027-16af9ab76b2a/fs/accounting/token_bucket.go (about) 1 package accounting 2 3 import ( 4 "context" 5 "sync" 6 "time" 7 8 "github.com/pkg/errors" 9 "github.com/rclone/rclone/fs" 10 "github.com/rclone/rclone/fs/rc" 11 "golang.org/x/time/rate" 12 ) 13 14 // Globals 15 var ( 16 tokenBucketMu sync.Mutex // protects the token bucket variables 17 tokenBucket *rate.Limiter 18 prevTokenBucket = tokenBucket 19 bwLimitToggledOff = false 20 currLimitMu sync.Mutex // protects changes to the timeslot 21 currLimit fs.BwTimeSlot 22 ) 23 24 const maxBurstSize = 4 * 1024 * 1024 // must be bigger than the biggest request 25 26 // make a new empty token bucket with the bandwidth given 27 func newTokenBucket(bandwidth fs.SizeSuffix) *rate.Limiter { 28 newTokenBucket := rate.NewLimiter(rate.Limit(bandwidth), maxBurstSize) 29 // empty the bucket 30 err := newTokenBucket.WaitN(context.Background(), maxBurstSize) 31 if err != nil { 32 fs.Errorf(nil, "Failed to empty token bucket: %v", err) 33 } 34 return newTokenBucket 35 } 36 37 // StartTokenBucket starts the token bucket if necessary 38 func StartTokenBucket() { 39 currLimitMu.Lock() 40 currLimit := fs.Config.BwLimit.LimitAt(time.Now()) 41 currLimitMu.Unlock() 42 43 if currLimit.Bandwidth > 0 { 44 tokenBucket = newTokenBucket(currLimit.Bandwidth) 45 fs.Infof(nil, "Starting bandwidth limiter at %vBytes/s", &currLimit.Bandwidth) 46 47 // Start the SIGUSR2 signal handler to toggle bandwidth. 48 // This function does nothing in windows systems. 49 startSignalHandler() 50 } 51 } 52 53 // StartTokenTicker creates a ticker to update the bandwidth limiter every minute. 54 func StartTokenTicker() { 55 // If the timetable has a single entry or was not specified, we don't need 56 // a ticker to update the bandwidth. 57 if len(fs.Config.BwLimit) <= 1 { 58 return 59 } 60 61 ticker := time.NewTicker(time.Minute) 62 go func() { 63 for range ticker.C { 64 limitNow := fs.Config.BwLimit.LimitAt(time.Now()) 65 currLimitMu.Lock() 66 67 if currLimit.Bandwidth != limitNow.Bandwidth { 68 tokenBucketMu.Lock() 69 70 // If bwlimit is toggled off, the change should only 71 // become active on the next toggle, which causes 72 // an exchange of tokenBucket <-> prevTokenBucket 73 var targetBucket **rate.Limiter 74 if bwLimitToggledOff { 75 targetBucket = &prevTokenBucket 76 } else { 77 targetBucket = &tokenBucket 78 } 79 80 // Set new bandwidth. If unlimited, set tokenbucket to nil. 81 if limitNow.Bandwidth > 0 { 82 *targetBucket = newTokenBucket(limitNow.Bandwidth) 83 if bwLimitToggledOff { 84 fs.Logf(nil, "Scheduled bandwidth change. "+ 85 "Limit will be set to %vBytes/s when toggled on again.", &limitNow.Bandwidth) 86 } else { 87 fs.Logf(nil, "Scheduled bandwidth change. Limit set to %vBytes/s", &limitNow.Bandwidth) 88 } 89 } else { 90 *targetBucket = nil 91 fs.Logf(nil, "Scheduled bandwidth change. Bandwidth limits disabled") 92 } 93 94 currLimit = limitNow 95 tokenBucketMu.Unlock() 96 } 97 currLimitMu.Unlock() 98 } 99 }() 100 } 101 102 // limitBandwith sleeps for the correct amount of time for the passage 103 // of n bytes according to the current bandwidth limit 104 func limitBandwidth(n int) { 105 tokenBucketMu.Lock() 106 107 // Limit the transfer speed if required 108 if tokenBucket != nil { 109 err := tokenBucket.WaitN(context.Background(), n) 110 if err != nil { 111 fs.Errorf(nil, "Token bucket error: %v", err) 112 } 113 } 114 115 tokenBucketMu.Unlock() 116 } 117 118 // SetBwLimit sets the current bandwidth limit 119 func SetBwLimit(bandwidth fs.SizeSuffix) { 120 tokenBucketMu.Lock() 121 defer tokenBucketMu.Unlock() 122 if bandwidth > 0 { 123 tokenBucket = newTokenBucket(bandwidth) 124 fs.Logf(nil, "Bandwidth limit set to %v", bandwidth) 125 } else { 126 tokenBucket = nil 127 fs.Logf(nil, "Bandwidth limit reset to unlimited") 128 } 129 } 130 131 // Remote control for the token bucket 132 func init() { 133 rc.Add(rc.Call{ 134 Path: "core/bwlimit", 135 Fn: func(ctx context.Context, in rc.Params) (out rc.Params, err error) { 136 if in["rate"] != nil { 137 bwlimit, err := in.GetString("rate") 138 if err != nil { 139 return out, err 140 } 141 var bws fs.BwTimetable 142 err = bws.Set(bwlimit) 143 if err != nil { 144 return out, errors.Wrap(err, "bad bwlimit") 145 } 146 if len(bws) != 1 { 147 return out, errors.New("need exactly 1 bandwidth setting") 148 } 149 bw := bws[0] 150 SetBwLimit(bw.Bandwidth) 151 } 152 bytesPerSecond := int64(-1) 153 if tokenBucket != nil { 154 bytesPerSecond = int64(tokenBucket.Limit()) 155 } 156 out = rc.Params{ 157 "rate": fs.SizeSuffix(bytesPerSecond).String(), 158 "bytesPerSecond": bytesPerSecond, 159 } 160 return out, nil 161 }, 162 Title: "Set the bandwidth limit.", 163 Help: ` 164 This sets the bandwidth limit to that passed in. 165 166 Eg 167 168 rclone rc core/bwlimit rate=off 169 { 170 "bytesPerSecond": -1, 171 "rate": "off" 172 } 173 rclone rc core/bwlimit rate=1M 174 { 175 "bytesPerSecond": 1048576, 176 "rate": "1M" 177 } 178 179 180 If the rate parameter is not supplied then the bandwidth is queried 181 182 rclone rc core/bwlimit 183 { 184 "bytesPerSecond": 1048576, 185 "rate": "1M" 186 } 187 188 The format of the parameter is exactly the same as passed to --bwlimit 189 except only one bandwidth may be specified. 190 191 In either case "rate" is returned as a human readable string, and 192 "bytesPerSecond" is returned as a number. 193 `, 194 }) 195 }