Bug 1316140 - Allow use of multiprocessing from config.status on windows. r=mshal draft
authorChris Manchester <cmanchester@mozilla.com>
Tue, 15 Nov 2016 12:09:50 -0800
changeset 439305 a54c06e3e4717db57354c51c0d09f2a3f26611a1
parent 439149 f8ba9c9b401f57b0047ddd6932cb830190865b38
child 537132 f73940914b36f69a048a0505a19faafd5cdedf16
push id35966
push userbmo:cmanchester@mozilla.com
push dateTue, 15 Nov 2016 20:11:44 +0000
reviewersmshal
bugs1316140, 914563
milestone53.0a1
Bug 1316140 - Allow use of multiprocessing from config.status on windows. r=mshal This generalizes the monkey-patching of the main module added in bug 914563 to allow multiprocessing to be used from config.status (which triggers the same bug as when it's used with `mach` as main). MozReview-Commit-ID: AdOdpKzmbsp
build/mach_bootstrap.py
configure.py
mach
python/mozbuild/mozbuild/util.py
--- a/build/mach_bootstrap.py
+++ b/build/mach_bootstrap.py
@@ -200,16 +200,19 @@ def bootstrap(topsrcdir, mozilla_dir=Non
     # creation is much simpler for the "advanced" environment variable use
     # case. For default behavior, we educate users and give them an opportunity
     # to react. We always exit after creating the directory because users don't
     # like surprises.
     sys.path[0:0] = [os.path.join(mozilla_dir, path) for path in SEARCH_PATHS]
     import mach.main
     from mozboot.util import get_state_dir
 
+    from mozbuild.util import patch_main
+    patch_main()
+
     def telemetry_handler(context, data):
         # We have not opted-in to telemetry
         if 'BUILD_SYSTEM_TELEMETRY' not in os.environ:
             return
 
         telemetry_dir = os.path.join(get_state_dir()[0], 'telemetry')
         try:
             os.mkdir(telemetry_dir)
--- a/configure.py
+++ b/configure.py
@@ -74,16 +74,18 @@ def config_status(config):
         for k, v in sanitized_config.iteritems():
             fh.write('%s = encode(%s, encoding)\n' % (k, indented_repr(v)))
         fh.write("__all__ = ['topobjdir', 'topsrcdir', 'defines', "
                  "'non_global_defines', 'substs', 'mozconfig']")
 
         if config.get('MOZ_BUILD_APP') != 'js' or config.get('JS_STANDALONE'):
             fh.write(textwrap.dedent('''
                 if __name__ == '__main__':
+                    from mozbuild.util import patch_main
+                    patch_main()
                     from mozbuild.config_status import config_status
                     args = dict([(name, globals()[name]) for name in __all__])
                     config_status(**args)
             '''))
 
     # Other things than us are going to run this file, so we need to give it
     # executable permissions.
     os.chmod('config.status', 0o755)
--- a/mach
+++ b/mach
@@ -78,73 +78,9 @@ def main(args):
     mach = get_mach()
     if not mach:
         print('Could not run mach: No mach source directory found.')
         sys.exit(1)
     sys.exit(mach.run(args))
 
 
 if __name__ == '__main__':
-    if sys.platform == 'win32':
-        # This is a complete hack to work around the fact that Windows
-        # multiprocessing needs to import the original module (ie: this
-        # file), but only works if it has a .py extension.
-        #
-        # We do this by a sort of two-level function interposing. The first
-        # level interposes forking.get_command_line() with our version defined
-        # in my_get_command_line(). Our version of get_command_line will
-        # replace the command string with the contents of the fork_interpose()
-        # function to be used in the subprocess.
-        #
-        # The subprocess then gets an interposed imp.find_module(), which we
-        # hack up to find 'mach' without the .py extension, since we already
-        # know where it is (it's us!). If we're not looking for 'mach', then
-        # the original find_module will suffice.
-        #
-        # See also: http://bugs.python.org/issue19946
-        # And: https://bugzilla.mozilla.org/show_bug.cgi?id=914563
-        import inspect
-        from multiprocessing import forking
-        global orig_command_line
-
-        def fork_interpose():
-            import imp
-            import os
-            import sys
-            orig_find_module = imp.find_module
-            def my_find_module(name, dirs):
-                if name == 'mach':
-                    path = os.path.join(dirs[0], 'mach')
-                    f = open(path)
-                    return (f, path, ('', 'r', imp.PY_SOURCE))
-                return orig_find_module(name, dirs)
-
-            # Don't allow writing bytecode file for mach module.
-            orig_load_module = imp.load_module
-            def my_load_module(name, file, path, description):
-                # multiprocess.forking invokes imp.load_module manually and
-                # hard-codes the name __parents_main__ as the module name.
-                if name == '__parents_main__':
-                    old_bytecode = sys.dont_write_bytecode
-                    sys.dont_write_bytecode = True
-                    try:
-                        return orig_load_module(name, file, path, description)
-                    finally:
-                        sys.dont_write_bytecode = old_bytecode
-
-                return orig_load_module(name, file, path, description)
-
-            imp.find_module = my_find_module
-            imp.load_module = my_load_module
-            from multiprocessing.forking import main; main()
-
-        def my_get_command_line():
-            fork_code, lineno = inspect.getsourcelines(fork_interpose)
-            # Remove the first line (for 'def fork_interpose():') and the three
-            # levels of indentation (12 spaces).
-            fork_string = ''.join(x[12:] for x in fork_code[1:])
-            cmdline = orig_command_line()
-            cmdline[2] = fork_string
-            return cmdline
-        orig_command_line = forking.get_command_line
-        forking.get_command_line = my_get_command_line
-
     main(sys.argv[1:])
--- a/python/mozbuild/mozbuild/util.py
+++ b/python/mozbuild/mozbuild/util.py
@@ -1257,8 +1257,91 @@ def encode(obj, encoding='utf-8'):
         }
     if isinstance(obj, bytes):
         return obj
     if isinstance(obj, unicode):
         return obj.encode(encoding)
     if isinstance(obj, Iterable):
         return [encode(i, encoding) for i in obj]
     return obj
+
+
+def patch_main():
+    '''This is a hack to work around the fact that Windows multiprocessing needs
+    to import the original main module, and assumes that it corresponds to a file
+    ending in .py.
+
+    We do this by a sort of two-level function interposing. The first
+    level interposes forking.get_command_line() with our version defined
+    in my_get_command_line(). Our version of get_command_line will
+    replace the command string with the contents of the fork_interpose()
+    function to be used in the subprocess.
+
+    The subprocess then gets an interposed imp.find_module(), which we
+    hack up to find the main module name multiprocessing will assume, since we
+    know what this will be based on the main module in the parent. If we're not
+    looking for our main module, then the original find_module will suffice.
+
+    See also: http://bugs.python.org/issue19946
+    And: https://bugzilla.mozilla.org/show_bug.cgi?id=914563
+    '''
+    if sys.platform == 'win32':
+        import inspect
+        import os
+        from multiprocessing import forking
+        global orig_command_line
+
+        # Figure out what multiprocessing will assume our main module
+        # is called (see python/Lib/multiprocessing/forking.py).
+        main_path = getattr(sys.modules['__main__'], '__file__', None)
+        if main_path is None:
+            # If someone deleted or modified __main__, there's nothing left for
+            # us to do.
+            return
+        main_file_name = os.path.basename(main_path)
+        main_module_name, ext = os.path.splitext(main_file_name)
+        if ext == '.py':
+            # If main is a .py file, everything ought to work as expected.
+            return
+
+        def fork_interpose():
+            import imp
+            import os
+            import sys
+            orig_find_module = imp.find_module
+            def my_find_module(name, dirs):
+                if name == main_module_name:
+                    path = os.path.join(dirs[0], main_file_name)
+                    f = open(path)
+                    return (f, path, ('', 'r', imp.PY_SOURCE))
+                return orig_find_module(name, dirs)
+
+            # Don't allow writing bytecode file for the main module.
+            orig_load_module = imp.load_module
+            def my_load_module(name, file, path, description):
+                # multiprocess.forking invokes imp.load_module manually and
+                # hard-codes the name __parents_main__ as the module name.
+                if name == '__parents_main__':
+                    old_bytecode = sys.dont_write_bytecode
+                    sys.dont_write_bytecode = True
+                    try:
+                        return orig_load_module(name, file, path, description)
+                    finally:
+                        sys.dont_write_bytecode = old_bytecode
+
+                return orig_load_module(name, file, path, description)
+
+            imp.find_module = my_find_module
+            imp.load_module = my_load_module
+            from multiprocessing.forking import main; main()
+
+        def my_get_command_line():
+            fork_code, lineno = inspect.getsourcelines(fork_interpose)
+            # Remove the first line (for 'def fork_interpose():') and the three
+            # levels of indentation (12 spaces), add our relevant globals.
+            fork_string = ("main_file_name = '%s'\n" % main_file_name +
+                           "main_module_name = '%s'\n" % main_module_name +
+                           ''.join(x[12:] for x in fork_code[1:]))
+            cmdline = orig_command_line()
+            cmdline[2] = fork_string
+            return cmdline
+        orig_command_line = forking.get_command_line
+        forking.get_command_line = my_get_command_line