From 649742625102e8537db6d9ea5b355e322a514c34 Mon Sep 17 00:00:00 2001
From: Mark Huang <mark.huang@cerebras.net>
Date: Wed, 30 Sep 2020 12:41:41 -0700
Subject: [PATCH] Fix ForkedFunc to handle hard exit(0)

1. waitfinish() assumes that if the forked child process exited with exitstatus
== 0, that _child() ran to completion and wrote out the retval. This is not true
if the function run in the child process hard exited with status 0,
e.g. os._exit(0) (or even sys.exit(0)).

2. Add new unit test test_exit0() for this case.
---
 py/_process/forkedfunc.py          | 10 ++++++----
 testing/process/test_forkedfunc.py |  9 +++++++++
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/py/_process/forkedfunc.py b/py/_process/forkedfunc.py
index 1c285306..819c3d50 100644
--- a/py/_process/forkedfunc.py
+++ b/py/_process/forkedfunc.py
@@ -89,12 +89,14 @@ def waitfinish(self, waiter=os.waitpid):
             exitstatus = 0
         signal = systemstatus & 0x7f
         if not exitstatus and not signal:
-            retval = self.RETVAL.open('rb')
+            retvalf = self.RETVAL.open('rb')
             try:
-                retval_data = retval.read()
+                retval_data = retvalf.read()
+                retval = marshal.loads(retval_data)
+            except:
+                retval = None
             finally:
-                retval.close()
-            retval = marshal.loads(retval_data)
+                retvalf.close()
         else:
             retval = None
         stdout = self.STDOUT.read()
diff --git a/testing/process/test_forkedfunc.py b/testing/process/test_forkedfunc.py
index ae0d9ab7..7f222108 100644
--- a/testing/process/test_forkedfunc.py
+++ b/testing/process/test_forkedfunc.py
@@ -34,6 +34,15 @@ def func():
     assert not result.out
     assert not result.err
 
+def test_exit0():
+    def exit0():
+        os._exit(0)
+    result = py.process.ForkedFunc(exit0).waitfinish()
+    assert result.exitstatus == 0
+    assert result.signal == 0
+    assert not result.out
+    assert not result.err
+
 def test_execption_in_func():
     def fun():
         raise ValueError(42)