Browse Source

Reset the whole MPFuture backend in tests

Aleksandr Borzunov 4 years ago
parent
commit
8b6f13e2a7
2 changed files with 11 additions and 4 deletions
  1. 9 0
      hivemind/utils/mpfuture.py
  2. 2 4
      tests/conftest.py

+ 9 - 0
hivemind/utils/mpfuture.py

@@ -127,11 +127,20 @@ class MPFuture(base.Future, Generic[ResultType]):
                     )
                     cls._pipe_waiter_thread.start()
 
+    @classmethod
+    def reset_backend(cls):
+        cls._initialization_lock = mp.Lock()
+        cls._update_lock = mp.Lock()
+        cls._active_pid = None
+
     @classmethod
     def _process_updates_in_background(cls, receiver_pipe: mp.connection.Connection):
         pid = os.getpid()
         while True:
             try:
+                if cls._pipe_waiter_thread is not threading.current_thread():
+                    break  # Backend was reset, a new background thread has started
+
                 uid, msg_type, payload = receiver_pipe.recv()
                 future = None
                 future_ref = cls._active_futures.get(uid)

+ 2 - 4
tests/conftest.py

@@ -1,5 +1,4 @@
 import gc
-import multiprocessing as mp
 from contextlib import suppress
 
 import psutil
@@ -29,6 +28,5 @@ def cleanup_children():
             with suppress(psutil.NoSuchProcess):
                 child.kill()
 
-    # Killing child processes may leave the global MPFuture locks acquired, so we recreate them
-    MPFuture._initialization_lock = mp.Lock()
-    MPFuture._update_lock = mp.Lock()
+    # Killing child processes may leave the global MPFuture state broken, so we reset it
+    MPFuture.reset_backend()