1
0
mirror of https://github.com/rclone/rclone.git synced 2025-01-08 12:34:53 +02:00

accounting: add per ctx token bucket

expose a token bucket wrapper type, so that client code can set a
limiter in context to limit the transfer bandwidth in context scope.
this is useful when rclone is used as imported library.
This commit is contained in:
Feng Shao 2024-12-31 13:58:47 +08:00
parent a389a2979b
commit 8e48fe045e
2 changed files with 26 additions and 0 deletions

View File

@ -344,6 +344,16 @@ func (acc *Account) limitPerFileBandwidth(n int) {
}
}
// Account for n bytes from the current context bandwidth limit (if any)
func (acc *Account) limitPerContextBandwidth(n int) {
value := acc.ctx.Value(TokenBucketWrapperPerContextKey)
if value != nil {
if tbw, ok := value.(*TokenBucketWrapper); ok {
tbw.LimitBandwidth(TokenBucketSlotAccounting, n)
}
}
}
// Account the read and limit bandwidth
func (acc *Account) accountRead(n int) {
// Update Stats
@ -356,6 +366,7 @@ func (acc *Account) accountRead(n int) {
TokenBucket.LimitBandwidth(TokenBucketSlotAccounting, n)
acc.limitPerFileBandwidth(n)
acc.limitPerContextBandwidth(n)
}
// read bytes from the io.Reader passed in and account them

View File

@ -15,6 +15,11 @@ import (
// TokenBucket holds the global token bucket limiter
var TokenBucket tokenBucket
type TokenBucketWrapperPerContextKeyType struct{}
// Context key for per context token bucket
var TokenBucketWrapperPerContextKey = TokenBucketWrapperPerContextKeyType{}
// TokenBucketSlot is the type to select which token bucket to use
type TokenBucketSlot int
@ -37,6 +42,16 @@ type tokenBucket struct {
currLimit fs.BwTimeSlot
}
type TokenBucketWrapper struct {
tokenBucket
}
func NewTokenBucketWrapper() *TokenBucketWrapper {
return &TokenBucketWrapper{
tokenBucket: tokenBucket{},
}
}
// Return true if limit is disabled
//
// Call with lock held