Explorar o código

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 %!s(int64=12) %!d(string=hai) anos
pai
achega
db1dd44ef4
Modificáronse 2 ficheiros con 25 adicións e 0 borrados
  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: