Restore throttle buckets after expiration

diskcache/recipes.py

@@ -265,6 +265,7 @@ def throttle(
 
     Assumes keys will not be evicted. Set the eviction policy to 'none' on the
     cache to guarantee the keys are not evicted.
+    If the bucket expires, the next call starts with a full bucket of tokens.
 
     >>> import diskcache, time
     >>> cache = diskcache.Cache()
@@ -291,8 +292,8 @@ def decorator(func):
         def wrapper(*args, **kwargs):
             while True:
                 with cache.transact(retry=True):
-                    last, tally = cache.get(key)
                     now = time_func()
+                    last, tally = cache.get(key, default=(now, count))
                     tally += (now - last) * rate
                     delay = 0
 

tests/test_recipes.py

@@ -3,6 +3,7 @@
 import shutil
 import threading
 import time
+from unittest import mock
 
 import pytest
 
@@ -28,6 +29,41 @@ def test_averager(cache):
     assert nums.pop() == 9.5
 
 
+@pytest.mark.parametrize('expire', [None, 1])
+@pytest.mark.parametrize('before_first_call', [False, True])
+def test_throttle_after_idle(cache, expire, before_first_call):
+    clock = [time.time()]
+    delays = []
+
+    def sleep(seconds):
+        delays.append(seconds)
+        clock[0] += seconds
+
+    with mock.patch('time.time', lambda: clock[0]):
+
+        @dc.throttle(
+            cache,
+            2,
+            1,
+            name='limiter',
+            expire=expire,
+            time_func=lambda: clock[0],
+            sleep_func=sleep,
+        )
+        def work():
+            return 'ok'
+
+        if not before_first_call:
+            assert work() == 'ok'
+        clock[0] += 10
+
+        assert work() == 'ok'
+        assert work() == 'ok'
+        assert delays == []
+        assert work() == 'ok'
+        assert delays == [0.5]
+
+
 def test_lock(cache):
     state = {'num': 0}
     lock = dc.Lock(cache, 'demo')