bucket improvement

3 files changed, 17 insertions, 20 deletions

diff --git a/module/network/Bucket.py b/module/network/Bucket.py
index c7eb62a54..04b706ff2 100644
--- a/module/network/Bucket.py
+++ b/module/network/Bucket.py
@@ -22,32 +22,30 @@ from threading import Lock
 
 class Bucket:
     def __init__(self):
-        self.tokens = 0
+        self.content = 0
         self.rate = 0
         self.lastDrip = time()
         self.lock = Lock()
     
     def setRate(self, rate):
-        self.lock.acquire()
         self.rate = rate
-        self.lock.release()
     
-    def consume(self, amount):
-        """ consume specified amount, return False if not enough tokens in bucket """
-        if not self.rate: return True
+    def add(self, amount):
         self.lock.acquire()
-        if amount < self.getTokens():
-            self.tokens -= amount
-            self.lock.release()
-            return True
+        self.drip()
+        allowable = min(amount, self.rate - self.content)
+        if allowable < 3072:
+            allowable = 0
 
+        self.content += allowable
         self.lock.release()
-        return False
+        return allowable
 
-    def getTokens(self):
-        if self.tokens < self.rate:
+    def drip(self):
+        if self.rate:
             now = time()
-            delta = self.rate * (now - self.lastDrip)
-            self.tokens = min(self.rate, self.tokens + delta)
+            deltaT = now - self.lastDrip
+            self.content = long(max(0, self.content - deltaT * self.rate))
             self.lastDrip = now
-        return self.tokens
+        else:
+            self.content = 0
diff --git a/module/network/HTTPChunk.py b/module/network/HTTPChunk.py
index c7d981880..0c87283f1 100644
--- a/module/network/HTTPChunk.py
+++ b/module/network/HTTPChunk.py
@@ -19,7 +19,6 @@
 
 from HTTPBase import HTTPBase
 from urllib2 import HTTPError
-from threading import Lock
 from helper import *
 from time import sleep
 from traceback import print_exc
@@ -76,8 +75,8 @@ class HTTPChunk(HTTPBase):
             if self.noRangeHeader:
                 count = min(count, self.range[1] - self.arrived)
             if self.bucket:
-                allow = self.bucket.consume(count)
-                if not allow:
+                count = self.bucket.add(count)
+                if not count:
                     sleep(0.01)
                     continue
             
diff --git a/module/network/HTTPDownload.py b/module/network/HTTPDownload.py
index 1be1a4552..d0e2eeb1f 100644
--- a/module/network/HTTPDownload.py
+++ b/module/network/HTTPDownload.py
@@ -272,7 +272,7 @@ if __name__ == "__main__":
     from Bucket import Bucket
     bucket = Bucket()
     bucket.setRate(200*1024)
-    #bucket = None
+    bucket = None
     
     url = "http://speedtest.netcologne.de/test_100mb.bin"