Jelajahi Sumber

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 tahun lalu
induk
melakukan
db1dd44ef4
2 mengubah file dengan 25 tambahan dan 0 penghapusan
  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
             x[i] = i
         self.assertListEqual(x.keys(), list(slots[limit:]))
         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):
     def test_least_recently_used(self):
         x = LRUCache(3)
         x = LRUCache(3)
 
 

+ 16 - 0
celery/utils/functional.py

@@ -53,6 +53,22 @@ class LRUCache(UserDict):
     def items(self):
     def items(self):
         return list(self._iterate_items())
         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):
     def __setitem__(self, key, value):
         # remove least recently used key.
         # remove least recently used key.
         with self.mutex:
         with self.mutex: