github.com/divyam234/rclone@v1.64.1/fs/accounting/token_bucket.go (about) 1 package accounting 2 3 import ( 4 "context" 5 "errors" 6 "fmt" 7 "sync" 8 "time" 9 10 "github.com/divyam234/rclone/fs" 11 "github.com/divyam234/rclone/fs/rc" 12 "golang.org/x/time/rate" 13 ) 14 15 // TokenBucket holds the global token bucket limiter 16 var TokenBucket tokenBucket 17 18 // TokenBucketSlot is the type to select which token bucket to use 19 type TokenBucketSlot int 20 21 // Slots for the token bucket 22 const ( 23 TokenBucketSlotAccounting TokenBucketSlot = iota 24 TokenBucketSlotTransportRx 25 TokenBucketSlotTransportTx 26 TokenBucketSlots 27 ) 28 29 type buckets [TokenBucketSlots]*rate.Limiter 30 31 // tokenBucket holds info about the rate limiters in use 32 type tokenBucket struct { 33 mu sync.RWMutex // protects the token bucket variables 34 curr buckets 35 prev buckets 36 toggledOff bool 37 currLimit fs.BwTimeSlot 38 } 39 40 // Return true if limit is disabled 41 // 42 // Call with lock held 43 func (bs *buckets) _isOff() bool { 44 return bs[0] == nil 45 } 46 47 // Disable the limits 48 // 49 // Call with lock held 50 func (bs *buckets) _setOff() { 51 for i := range bs { 52 bs[i] = nil 53 } 54 } 55 56 const defaultMaxBurstSize = 4 * 1024 * 1024 // must be bigger than the biggest request 57 58 // make a new empty token bucket with the bandwidth given 59 func newEmptyTokenBucket(bandwidth fs.SizeSuffix) *rate.Limiter { 60 // Relate maxBurstSize to bandwidth limit 61 // 4M gives 2.5 Gb/s on Windows 62 // Use defaultMaxBurstSize up to 2GBit/s (256MiB/s) then scale 63 maxBurstSize := (bandwidth * defaultMaxBurstSize) / (256 * 1024 * 1024) 64 if maxBurstSize < defaultMaxBurstSize { 65 maxBurstSize = defaultMaxBurstSize 66 } 67 // fs.Debugf(nil, "bandwidth=%v maxBurstSize=%v", bandwidth, maxBurstSize) 68 tb := rate.NewLimiter(rate.Limit(bandwidth), int(maxBurstSize)) 69 if tb != nil { 70 // empty the bucket 71 err := tb.WaitN(context.Background(), int(maxBurstSize)) 72 if err != nil { 73 fs.Errorf(nil, "Failed to empty token bucket: %v", err) 74 } 75 } 76 return tb 77 } 78 79 // make a new empty token bucket with the bandwidth(s) given 80 func newTokenBucket(bandwidth fs.BwPair) (tbs buckets) { 81 bandwidthAccounting := fs.SizeSuffix(-1) 82 if bandwidth.Tx > 0 { 83 tbs[TokenBucketSlotTransportTx] = newEmptyTokenBucket(bandwidth.Tx) 84 bandwidthAccounting = bandwidth.Tx 85 } 86 if bandwidth.Rx > 0 { 87 tbs[TokenBucketSlotTransportRx] = newEmptyTokenBucket(bandwidth.Rx) 88 if bandwidth.Rx > bandwidthAccounting { 89 bandwidthAccounting = bandwidth.Rx 90 } 91 } 92 // Limit core bandwidth to max of Rx and Tx if both are limited 93 if bandwidth.Tx > 0 && bandwidth.Rx > 0 { 94 tbs[TokenBucketSlotAccounting] = newEmptyTokenBucket(bandwidthAccounting) 95 } 96 return tbs 97 } 98 99 // StartTokenBucket starts the token bucket if necessary 100 func (tb *tokenBucket) StartTokenBucket(ctx context.Context) { 101 tb.mu.Lock() 102 defer tb.mu.Unlock() 103 ci := fs.GetConfig(ctx) 104 tb.currLimit = ci.BwLimit.LimitAt(time.Now()) 105 if tb.currLimit.Bandwidth.IsSet() { 106 tb.curr = newTokenBucket(tb.currLimit.Bandwidth) 107 fs.Infof(nil, "Starting bandwidth limiter at %v Byte/s", &tb.currLimit.Bandwidth) 108 } 109 110 // Start the SIGUSR2 signal handler to toggle bandwidth. 111 // This function does nothing in windows systems. 112 tb.startSignalHandler() 113 } 114 115 // StartTokenTicker creates a ticker to update the bandwidth limiter every minute. 116 func (tb *tokenBucket) StartTokenTicker(ctx context.Context) { 117 ci := fs.GetConfig(ctx) 118 // If the timetable has a single entry or was not specified, we don't need 119 // a ticker to update the bandwidth. 120 if len(ci.BwLimit) <= 1 { 121 return 122 } 123 124 ticker := time.NewTicker(time.Minute) 125 go func() { 126 for range ticker.C { 127 limitNow := ci.BwLimit.LimitAt(time.Now()) 128 tb.mu.Lock() 129 130 if tb.currLimit.Bandwidth != limitNow.Bandwidth { 131 // If bwlimit is toggled off, the change should only 132 // become active on the next toggle, which causes 133 // an exchange of tb.curr <-> tb.prev 134 var targetBucket *buckets 135 if tb.toggledOff { 136 targetBucket = &tb.prev 137 } else { 138 targetBucket = &tb.curr 139 } 140 141 // Set new bandwidth. If unlimited, set tokenbucket to nil. 142 if limitNow.Bandwidth.IsSet() { 143 *targetBucket = newTokenBucket(limitNow.Bandwidth) 144 if tb.toggledOff { 145 fs.Logf(nil, "Scheduled bandwidth change. "+ 146 "Limit will be set to %v Byte/s when toggled on again.", &limitNow.Bandwidth) 147 } else { 148 fs.Logf(nil, "Scheduled bandwidth change. Limit set to %v Byte/s", &limitNow.Bandwidth) 149 } 150 } else { 151 targetBucket._setOff() 152 fs.Logf(nil, "Scheduled bandwidth change. Bandwidth limits disabled") 153 } 154 155 tb.currLimit = limitNow 156 } 157 158 tb.mu.Unlock() 159 } 160 }() 161 } 162 163 // LimitBandwidth sleeps for the correct amount of time for the passage 164 // of n bytes according to the current bandwidth limit 165 func (tb *tokenBucket) LimitBandwidth(i TokenBucketSlot, n int) { 166 tb.mu.RLock() 167 168 // Limit the transfer speed if required 169 if tb.curr[i] != nil { 170 err := tb.curr[i].WaitN(context.Background(), n) 171 if err != nil { 172 fs.Errorf(nil, "Token bucket error: %v", err) 173 } 174 } 175 176 tb.mu.RUnlock() 177 } 178 179 // SetBwLimit sets the current bandwidth limit 180 func (tb *tokenBucket) SetBwLimit(bandwidth fs.BwPair) { 181 tb.mu.Lock() 182 defer tb.mu.Unlock() 183 if bandwidth.IsSet() { 184 tb.curr = newTokenBucket(bandwidth) 185 fs.Logf(nil, "Bandwidth limit set to %v", bandwidth) 186 } else { 187 tb.curr._setOff() 188 fs.Logf(nil, "Bandwidth limit reset to unlimited") 189 } 190 } 191 192 // read and set the bandwidth limits 193 func (tb *tokenBucket) rcBwlimit(ctx context.Context, in rc.Params) (out rc.Params, err error) { 194 if in["rate"] != nil { 195 bwlimit, err := in.GetString("rate") 196 if err != nil { 197 return out, err 198 } 199 var bws fs.BwTimetable 200 err = bws.Set(bwlimit) 201 if err != nil { 202 return out, fmt.Errorf("bad bwlimit: %w", err) 203 } 204 if len(bws) != 1 { 205 return out, errors.New("need exactly 1 bandwidth setting") 206 } 207 bw := bws[0] 208 tb.SetBwLimit(bw.Bandwidth) 209 } 210 tb.mu.RLock() 211 bytesPerSecond := int64(-1) 212 if tb.curr[TokenBucketSlotAccounting] != nil { 213 bytesPerSecond = int64(tb.curr[TokenBucketSlotAccounting].Limit()) 214 } 215 var bp = fs.BwPair{Tx: -1, Rx: -1} 216 if tb.curr[TokenBucketSlotTransportTx] != nil { 217 bp.Tx = fs.SizeSuffix(tb.curr[TokenBucketSlotTransportTx].Limit()) 218 } 219 if tb.curr[TokenBucketSlotTransportRx] != nil { 220 bp.Rx = fs.SizeSuffix(tb.curr[TokenBucketSlotTransportRx].Limit()) 221 } 222 tb.mu.RUnlock() 223 out = rc.Params{ 224 "rate": bp.String(), 225 "bytesPerSecond": bytesPerSecond, 226 "bytesPerSecondTx": int64(bp.Tx), 227 "bytesPerSecondRx": int64(bp.Rx), 228 } 229 return out, nil 230 } 231 232 // Remote control for the token bucket 233 func init() { 234 rc.Add(rc.Call{ 235 Path: "core/bwlimit", 236 Fn: func(ctx context.Context, in rc.Params) (out rc.Params, err error) { 237 return TokenBucket.rcBwlimit(ctx, in) 238 }, 239 Title: "Set the bandwidth limit.", 240 Help: ` 241 This sets the bandwidth limit to the string passed in. This should be 242 a single bandwidth limit entry or a pair of upload:download bandwidth. 243 244 Eg 245 246 rclone rc core/bwlimit rate=off 247 { 248 "bytesPerSecond": -1, 249 "bytesPerSecondTx": -1, 250 "bytesPerSecondRx": -1, 251 "rate": "off" 252 } 253 rclone rc core/bwlimit rate=1M 254 { 255 "bytesPerSecond": 1048576, 256 "bytesPerSecondTx": 1048576, 257 "bytesPerSecondRx": 1048576, 258 "rate": "1M" 259 } 260 rclone rc core/bwlimit rate=1M:100k 261 { 262 "bytesPerSecond": 1048576, 263 "bytesPerSecondTx": 1048576, 264 "bytesPerSecondRx": 131072, 265 "rate": "1M" 266 } 267 268 269 If the rate parameter is not supplied then the bandwidth is queried 270 271 rclone rc core/bwlimit 272 { 273 "bytesPerSecond": 1048576, 274 "bytesPerSecondTx": 1048576, 275 "bytesPerSecondRx": 1048576, 276 "rate": "1M" 277 } 278 279 The format of the parameter is exactly the same as passed to --bwlimit 280 except only one bandwidth may be specified. 281 282 In either case "rate" is returned as a human-readable string, and 283 "bytesPerSecond" is returned as a number. 284 `, 285 }) 286 }