Use barrier to ensure process is running

portal/process.py

@@ -1,6 +1,5 @@
 import atexit
 import errno
-import os
 import traceback
 
 import cloudpickle
@@ -37,8 +36,10 @@ def __init__(self, fn, *args, name=None, start=False):
     name = name or getattr(fn, '__name__', 'process')
     fn = cloudpickle.dumps(fn)
     options = contextlib.context.options()
+    self.ready = contextlib.context.mp.Barrier(2)
     self.process = contextlib.context.mp.Process(
-        target=self._wrapper, name=name, args=(options, name, fn, args))
+        target=self._wrapper, name=name,
+        args=(options, self.ready, name, fn, args))
     self.started = False
     self.killed = False
     self.thepid = None
@@ -73,6 +74,7 @@ def start(self):
     assert not self.started
     self.started = True
     self.process.start()
+    self.ready.wait()
     self.thepid = self.process.pid
     assert self.thepid is not None
     return self
@@ -83,8 +85,6 @@ def join(self, timeout=None):
     return self
 
   def kill(self, timeout=1):
-    # Cannot early exit if process is not running, because it may just be
-    # starting up.
     assert self.started
     try:
       children = list(psutil.Process(self.pid).children(recursive=True))
@@ -108,9 +108,10 @@ def __repr__(self):
     return 'Process(' + ', '.join(attrs) + ')'
 
   @staticmethod
-  def _wrapper(options, name, fn, args):
+  def _wrapper(options, ready, name, fn, args):
     exitcode = 0
     try:
+      ready.wait()
       contextlib.setup(**options)
       fn = cloudpickle.loads(fn)
       exitcode = fn(*args)

tests/test_process.py

@@ -16,31 +16,39 @@ def test_exitcode(self):
     assert worker.exitcode == 42
 
   def test_error(self):
+
     def fn():
       raise KeyError('foo')
+
     worker = portal.Process(fn, start=True)
     worker.join()
     assert not worker.running
     assert worker.exitcode == 1
 
   def test_error_with_children(self):
+
     def hang():
       while True:
         time.sleep(0.1)
+
     def fn():
       portal.Process(hang, start=True)
       portal.Thread(hang, start=True)
       time.sleep(0.1)
       raise KeyError('foo')
+
     worker = portal.Process(fn, start=True)
     worker.join()
     assert not worker.running
     assert worker.exitcode == 1
 
-  def test_kill(self):
+  @pytest.mark.parametrize('repeat', range(5))
+  def test_kill_basic(self, repeat):
+
     def fn():
       while True:
         time.sleep(0.1)
+
     worker = portal.Process(fn, start=True)
     worker.kill()
     assert not worker.running
@@ -52,14 +60,18 @@ def test_kill_with_subproc(self, repeat):
     queue = portal.context.mp.Queue()
 
     def outer(ready, queue):
-      queue.put(os.getpid())
       portal.Process(inner, ready, queue, start=True)
+      queue.put(os.getpid())
+      queue.close()
+      queue.join_thread()
       ready.wait()
       while True:
         time.sleep(0.1)
 
     def inner(ready, queue):
       queue.put(os.getpid())
+      queue.close()
+      queue.join_thread()
       ready.wait()
       while True:
         time.sleep(0.1)
@@ -90,17 +102,22 @@ def inner(ready):
     assert worker.exitcode < 0
 
   def test_initfn(self):
+
     def init():
       portal.foo = 42
+
     portal.initfn(init)
     ready = portal.context.mp.Event()
     assert portal.foo == 42
+
     def outer(ready):
       assert portal.foo == 42
       portal.Process(inner, ready, start=True).join()
+
     def inner(ready):
       assert portal.foo == 42
       ready.set()
+
     portal.Process(outer, ready, start=True).join()
     ready.wait()
     assert ready.is_set()