Merge pull request #2457 from imachug/segfault

Make ThreadPool a context manager to prevent memory leaks

src/Test/TestNoparallel.py

@@ -149,8 +149,7 @@ class TestNoparallel:
 
     def testMultithreadMix(self, queue_spawn):
         obj1 = ExampleClass()
-        thread_pool = ThreadPool.ThreadPool(10)
-
+        with ThreadPool.ThreadPool(10) as thread_pool:
             s = time.time()
             t1 = queue_spawn(obj1.countBlocking, 5)
             time.sleep(0.01)
@@ -166,4 +165,3 @@ class TestNoparallel:
             time_taken = time.time() - s
             assert obj1.counted == 5
             assert 0.5 < time_taken < 0.7
-        thread_pool.kill()

src/Test/TestThreadPool.py

@@ -9,8 +9,7 @@ from util import ThreadPool
 
 class TestThreadPool:
     def testExecutionOrder(self):
-        pool = ThreadPool.ThreadPool(4)
-
+        with ThreadPool.ThreadPool(4) as pool:
             events = []
 
             @pool.wrap
@@ -33,7 +32,6 @@ class TestThreadPool:
 
             res = blocker()
             assert res == 10000000
-        pool.kill()
 
     def testLockBlockingSameThread(self):
         lock = ThreadPool.Lock()
@@ -60,7 +58,7 @@ class TestThreadPool:
             time.sleep(0.5)
             lock.release()
 
-        pool = ThreadPool.ThreadPool(10)
+        with ThreadPool.ThreadPool(10) as pool:
             threads = [
                 pool.spawn(locker),
                 pool.spawn(locker),
@@ -81,8 +79,7 @@ class TestThreadPool:
 
     def testMainLoopCallerThreadId(self):
         main_thread_id = threading.current_thread().ident
-        pool = ThreadPool.ThreadPool(5)
-
+        with ThreadPool.ThreadPool(5) as pool:
             def getThreadId(*args, **kwargs):
                 return threading.current_thread().ident
 
@@ -94,7 +91,7 @@ class TestThreadPool:
 
     def testMainLoopCallerGeventSpawn(self):
         main_thread_id = threading.current_thread().ident
-        pool = ThreadPool.ThreadPool(5)
+        with ThreadPool.ThreadPool(5) as pool:
             def waiter():
                 time.sleep(1)
                 return threading.current_thread().ident
@@ -116,7 +113,7 @@ class TestThreadPool:
             assert 0.9 < time_taken < 1.2
 
     def testEvent(self):
-        pool = ThreadPool.ThreadPool(5)
+        with ThreadPool.ThreadPool(5) as pool:
             event = ThreadPool.Event()
 
             def setter():
@@ -153,10 +150,9 @@ class TestThreadPool:
             return "ok"
 
         def poolTest():
-            pool = ThreadPool.ThreadPool(5)
+            with ThreadPool.ThreadPool(5) as pool:
                 for i in range(20):
                     pool.spawn(worker)
-            pool.kill()
 
         for i in range(5):
             poolTest()

src/util/ThreadPool.py

@@ -55,6 +55,12 @@ class ThreadPool:
         del self.pool
         self.pool = None
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.kill()
+
 
 lock_pool = gevent.threadpool.ThreadPool(50)
 main_thread_id = threading.current_thread().ident