Issue #9573: os.fork now works when triggered as a side effect of import (the wisdom of actually relying on this remains questionable!)

ncoghlan · ncoghlan · commit b2ddf7979d22 · 2010-12-02T04:11:46.000Z
diff --git a/Lib/test/test_fork1.py b/Lib/test/test_fork1.py
@@ -8,13 +8,14 @@
 import time
 
 from test.fork_wait import ForkWait
-from test.support import run_unittest, reap_children, get_attribute, import_module
+from test.support import (run_unittest, reap_children, get_attribute,
+                          import_module, verbose)
+
 threading = import_module('threading')
 
 # Skip test if fork does not exist.
 get_attribute(os, 'fork')
 
-
 class ForkTest(ForkWait):
     def wait_impl(self, cpid):
         for i in range(10):
@@ -28,7 +29,8 @@ def wait_impl(self, cpid):
         self.assertEqual(spid, cpid)
         self.assertEqual(status, 0, "cause = %d, exit = %d" % (status&0xff, status>>8))
 
-    def test_import_lock_fork(self):
+    def test_threaded_import_lock_fork(self):
+        """Check fork() in main thread works while a subthread is doing an import"""
         import_started = threading.Event()
         fake_module_name = "fake test module"
         partial_module = "partial"
@@ -45,11 +47,16 @@ def importer():
         import_started.wait()
         pid = os.fork()
         try:
+            # PyOS_BeforeFork should have waited for the import to complete
+            # before forking, so the child can recreate the import lock
+            # correctly, but also won't see a partially initialised module
             if not pid:
                 m = __import__(fake_module_name)
                 if m == complete_module:
                     os._exit(0)
                 else:
+                    if verbose > 1:
+                        print("Child encountered partial module")
                     os._exit(1)
             else:
                 t.join()
@@ -63,6 +70,39 @@ def importer():
             except OSError:
                 pass
 
+
+    def test_nested_import_lock_fork(self):
+        """Check fork() in main thread works while the main thread is doing an import"""
+        # Issue 9573: this used to trigger RuntimeError in the child process
+        def fork_with_import_lock(level):
+            release = 0
+            in_child = False
+            try:
+                try:
+                    for i in range(level):
+                        imp.acquire_lock()
+                        release += 1
+                    pid = os.fork()
+                    in_child = not pid
+                finally:
+                    for i in range(release):
+                        imp.release_lock()
+            except RuntimeError:
+                if in_child:
+                    if verbose > 1:
+                        print("RuntimeError in child")
+                    os._exit(1)
+                raise
+            if in_child:
+                os._exit(0)
+            self.wait_impl(pid)
+
+        # Check this works with various levels of nested
+        # import in the main thread
+        for level in range(5):
+            fork_with_import_lock(level)
+
+
 def test_main():
     run_unittest(ForkTest)
     reap_children()
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -46,6 +46,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #9573: os.fork() now works correctly when triggered as a side effect
+  of a module import
+
 - Issue #10464: netrc now correctly handles lines with embedded '#' characters.
 
 - Added itertools.accumulate().
diff --git a/Python/import.c b/Python/import.c
@@ -325,8 +325,17 @@ _PyImport_ReInitLock(void)
 {
     if (import_lock != NULL)
         import_lock = PyThread_allocate_lock();
-    import_lock_thread = -1;
-    import_lock_level = 0;
+    if (import_lock_level > 1) {
+        /* Forked as a side effect of import */
+        long me = PyThread_get_thread_ident();
+        PyThread_acquire_lock(import_lock, 0);
+	/* XXX: can the previous line fail? */
+        import_lock_thread = me;
+        import_lock_level--;
+    } else {
+        import_lock_thread = -1;
+        import_lock_level = 0;
+    }
 }
 
 #endif
