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:
parent
a389a2979b
commit
8e48fe045e
@ -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
|
||||
|
@ -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
|
||||
|
Loading…
Reference in New Issue
Block a user