Merge 8e48fe045e91d1a73fb25c74c6c820b70a2700df into 0b9671313b14ffe839ecbd7dd2ae5ac7f6f05db8

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

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