change lru behaviour; #202

shadowsocks/lru_cache.py

@@ -15,12 +15,14 @@ class LRUCache(collections.MutableMapping):
         self.close_callback = close_callback
         self._store = {}
         self._time_to_keys = collections.defaultdict(list)
+        self._keys_to_last_time = {}
         self._last_visits = []
         self.update(dict(*args, **kwargs))  # use the free update to set keys
 
     def __getitem__(self, key):
         # O(logm)
         t = time.time()
+        self._keys_to_last_time[key] = t
         self._time_to_keys[t].append(key)
         heapq.heappush(self._last_visits, t)
         return self._store[key]
@@ -28,6 +30,7 @@ class LRUCache(collections.MutableMapping):
     def __setitem__(self, key, value):
         # O(logm)
         t = time.time()
+        self._keys_to_last_time[key] = t
         self._store[key] = value
         self._time_to_keys[t].append(key)
         heapq.heappush(self._last_visits, t)
@@ -35,6 +38,7 @@ class LRUCache(collections.MutableMapping):
     def __delitem__(self, key):
         # O(1)
         del self._store[key]
+        del self._keys_to_last_time[key]
 
     def __iter__(self):
         return iter(self._store)
@@ -53,13 +57,49 @@ class LRUCache(collections.MutableMapping):
             if self.close_callback is not None:
                 for key in self._time_to_keys[least]:
                     if self._store.__contains__(key):
+                        if now - self._keys_to_last_time[key] > self.timeout:
                             value = self._store[key]
                             self.close_callback(value)
             for key in self._time_to_keys[least]:
                 heapq.heappop(self._last_visits)
                 if self._store.__contains__(key):
+                    if now - self._keys_to_last_time[key] > self.timeout:
                         del self._store[key]
                         c += 1
             del self._time_to_keys[least]
         if c:
             logging.debug('%d keys swept' % c)
+
+
+def test():
+    c = LRUCache(timeout=0.3)
+
+    c['a'] = 1
+    assert c['a'] == 1
+
+    time.sleep(0.5)
+    c.sweep()
+    assert 'a' not in c
+
+    c['a'] = 2
+    c['b'] = 3
+    time.sleep(0.2)
+    c.sweep()
+    assert c['a'] == 2
+    assert c['b'] == 3
+
+    time.sleep(0.2)
+    c.sweep()
+    c['b']
+    time.sleep(0.2)
+    c.sweep()
+    assert 'a' not in c
+    assert c['b'] == 3
+
+    time.sleep(0.5)
+    c.sweep()
+    assert 'a' not in c
+    assert 'b' not in c
+
+if __name__ == '__main__':
+    test()