From 30119498be91c3f8837a797b1ab212f4520fa0ce Mon Sep 17 00:00:00 2001
From: Albert Larsan <74931857+albertlarsan68@users.noreply.github.com>
Date: Mon, 1 May 2023 13:46:31 +0000
Subject: [PATCH] Make x.py work again in most (all?) cases

Wrap all of x.py in `if __name__ == '__main__':` to avoid problems with `multiprocessing`
Make the pool sizing better
---
 src/bootstrap/bootstrap.py | 20 ++++++++++++++++++-
 x.py                       | 39 ++++++++++++++++++++------------------
 2 files changed, 40 insertions(+), 19 deletions(-)

diff --git a/src/bootstrap/bootstrap.py b/src/bootstrap/bootstrap.py
index dd0b091a2a3..f22cdad7df4 100644
--- a/src/bootstrap/bootstrap.py
+++ b/src/bootstrap/bootstrap.py
@@ -28,6 +28,20 @@ if platform_is_win32():
 else:
     EXE_SUFFIX = ""
 
+def get_cpus():
+    if hasattr(os, "sched_getaffinity"):
+        return len(os.sched_getaffinity(0))
+    if hasattr(os, "cpu_count"):
+        cpus = os.cpu_count()
+        if cpus is not None:
+            return cpus
+    try:
+        return cpu_count()
+    except NotImplementedError:
+        return 1
+
+
+
 def get(base, url, path, checksums, verbose=False):
     with tempfile.NamedTemporaryFile(delete=False) as temp_file:
         temp_path = temp_file.name
@@ -540,11 +554,15 @@ class RustBuild(object):
 
             # Unpack the tarballs in parallle.
             # In Python 2.7, Pool cannot be used as a context manager.
-            p = Pool(min(len(tarballs_download_info), cpu_count()))
+            pool_size = min(len(tarballs_download_info), get_cpus())
+            if self.verbose:
+                print('Choosing a pool size of', pool_size, 'for the unpacking of the tarballs')
+            p = Pool(pool_size)
             try:
                 p.map(unpack_component, tarballs_download_info)
             finally:
                 p.close()
+            p.join()
 
             if self.should_fix_bins_and_dylibs():
                 self.fix_bin_or_dylib("{}/bin/cargo".format(bin_root))
diff --git a/x.py b/x.py
index 5dee953a318..b8cdf67712c 100755
--- a/x.py
+++ b/x.py
@@ -4,26 +4,29 @@
 
 # This file is only a "symlink" to bootstrap.py, all logic should go there.
 
-import os
-import sys
+# Parts of `bootstrap.py` use the `multiprocessing` module, so this entry point
+# must use the normal `if __name__ == '__main__':` convention to avoid problems.
+if __name__ == '__main__':
+    import os
+    import sys
 
-# If this is python2, check if python3 is available and re-execute with that
-# interpreter. Only python3 allows downloading CI LLVM.
-#
-# This matters if someone's system `python` is python2.
-if sys.version_info.major < 3:
-    try:
-        os.execvp("py", ["py", "-3"] + sys.argv)
-    except OSError:
+    # If this is python2, check if python3 is available and re-execute with that
+    # interpreter. Only python3 allows downloading CI LLVM.
+    #
+    # This matters if someone's system `python` is python2.
+    if sys.version_info.major < 3:
         try:
-            os.execvp("python3", ["python3"] + sys.argv)
+            os.execvp("py", ["py", "-3"] + sys.argv)
         except OSError:
-            # Python 3 isn't available, fall back to python 2
-            pass
+            try:
+                os.execvp("python3", ["python3"] + sys.argv)
+            except OSError:
+                # Python 3 isn't available, fall back to python 2
+                pass
 
-rust_dir = os.path.dirname(os.path.abspath(__file__))
-# For the import below, have Python search in src/bootstrap first.
-sys.path.insert(0, os.path.join(rust_dir, "src", "bootstrap"))
+    rust_dir = os.path.dirname(os.path.abspath(__file__))
+    # For the import below, have Python search in src/bootstrap first.
+    sys.path.insert(0, os.path.join(rust_dir, "src", "bootstrap"))
 
-import bootstrap
-bootstrap.main()
+    import bootstrap
+    bootstrap.main()