Selaa lähdekoodia

Fix object leak in LRUCache

LRUCache.update() method did not pop the least used items off the dict,
leading to cache constantly growing, and ultimately to OOM kill.
Romuald Brunet 12 vuotta sitten
vanhempi
commit
db1dd44ef4
2 muutettua tiedostoa jossa 25 lisäystä ja 0 poistoa
  1. 9 0
      celery/tests/utilities/test_datastructures.py
  2. 16 0
      celery/utils/functional.py

+ 9 - 0
celery/tests/utilities/test_datastructures.py

@@ -180,6 +180,15 @@ class test_LRUCache(Case):
             x[i] = i
         self.assertListEqual(x.keys(), list(slots[limit:]))
 
+    def test_update_expires(self):
+        limit = 100
+        x = LRUCache(limit=limit)
+        slots = list(xrange(limit * 2))
+        for i in slots:
+            x.update({i: i})
+
+        self.assertListEqual(list(x.keys()), list(slots[limit:]))
+
     def test_least_recently_used(self):
         x = LRUCache(3)
 

+ 16 - 0
celery/utils/functional.py

@@ -53,6 +53,22 @@ class LRUCache(UserDict):
     def items(self):
         return list(self._iterate_items())
 
+    def update(self, *args, **kwargs):
+        with self.mutex:
+            self.data.update(*args, **kwargs)
+
+            if not self.limit:
+                return
+
+            # pop additional items if dict growed too much
+            i = iter(self.data)  # start from bottom (LRU)
+            overflow = len(self.data) - self.limit
+
+            # negative overflow will lead to an empty list
+            to_pop = [next(i) for _ in xrange(overflow)]
+            for item in to_pop:
+                self.data.pop(item)
+
     def __setitem__(self, key, value):
         # remove least recently used key.
         with self.mutex: