diff options
Diffstat (limited to 'module/network/Bucket.py')
| -rw-r--r-- | module/network/Bucket.py | 28 |
1 files changed, 14 insertions, 14 deletions
diff --git a/module/network/Bucket.py b/module/network/Bucket.py index 35e27bcd4..c7eb62a54 100644 --- a/module/network/Bucket.py +++ b/module/network/Bucket.py @@ -22,7 +22,7 @@ from threading import Lock class Bucket: def __init__(self): - self.content = 0 + self.tokens = 0 self.rate = 0 self.lastDrip = time() self.lock = Lock() @@ -32,22 +32,22 @@ class Bucket: self.rate = rate self.lock.release() - def add(self, amount): + def consume(self, amount): + """ consume specified amount, return False if not enough tokens in bucket """ + if not self.rate: return True self.lock.acquire() - self.drip() - allowable = min(amount, self.rate - self.content) - if allowable > 0: - sleep(0.005) #@XXX: high sysload without?! + if amount < self.getTokens(): + self.tokens -= amount + self.lock.release() + return True - self.content += allowable self.lock.release() - return allowable + return False - def drip(self): - if self.rate == 0: - self.content = 0 - else: + def getTokens(self): + if self.tokens < self.rate: now = time() - deltaT = now - self.lastDrip - self.content = long(max(0, self.content - deltaT * self.rate)) + delta = self.rate * (now - self.lastDrip) + self.tokens = min(self.rate, self.tokens + delta) self.lastDrip = now + return self.tokens |