diff --git a/fs/accounting/accounting.go b/fs/accounting/accounting.go index 5a6be7df4..b42159602 100644 --- a/fs/accounting/accounting.go +++ b/fs/accounting/accounting.go @@ -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 diff --git a/fs/accounting/token_bucket.go b/fs/accounting/token_bucket.go index a8b4bcd2d..07ed2ddac 100644 --- a/fs/accounting/token_bucket.go +++ b/fs/accounting/token_bucket.go @@ -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