Merge pull request swiftlang#84911 from charles-zablit/charles-zablit/update-checkout/threadpool

[update-checkout] switch to using a ThreadPool for scheduling git commands

Author: charles-zablit

utils/swift_build_support/swift_build_support/shell.py

@@ -216,7 +216,7 @@ def remove(path, dry_run=None, echo=True):
 
 
 def run(*args, **kwargs):
-    repo_path = os.getcwd()
+    repo_path = kwargs.pop('repo_path', os.getcwd())
     echo_output = kwargs.pop('echo', False)
     dry_run = kwargs.pop('dry_run', False)
     env = kwargs.get('env', None)
@@ -249,9 +249,9 @@ def run(*args, **kwargs):
         lock.release()
 
     if ret != 0:
-        eout = Exception()
+        eout = Exception(f"[{repo_path}] '{args}' failed with '{output.decode('utf-8')}'")
         eout.ret = ret
-        eout.args = args
+        eout.arguments = args
         eout.repo_path = repo_path
         eout.stderr = output
         raise eout

utils/update_checkout/tests/test_clone.py

@@ -41,7 +41,10 @@ def test_clone_with_additional_scheme(self):
                             '--verbose'])
 
         # Test that we're actually checking out the 'extra' scheme based on the output
-        self.assertIn(b"git checkout refs/heads/main", output)
+        self.assertIn(
+            f"git -C {os.path.join(self.source_root, 'repo1')} checkout refs/heads/main",
+            output.decode("utf-8"),
+        )
 
     def test_manager_not_called_on_long_socket(self):
         fake_tmpdir = '/tmp/very/' + '/long' * 20 + '/tmp'

utils/update_checkout/update_checkout/git_command.py

@@ -0,0 +1,37 @@
+from typing import List, Any, Optional, Union
+
+from swift_build_support.swift_build_support import shell
+
+
+class Git:
+    @staticmethod
+    def run(repo_path: Optional[str], args: List[str], **kwargs):
+        cmd = ["git"]
+        if repo_path is not None:
+            cmd += ["-C", repo_path]
+        kwargs["repo_path"] = repo_path
+        # FIXME: The way we are passing args below is broken. shell.run takes 
+        # *args as list arguments and sometimes treats them as one positional 
+        # argument or as a list of arguments.
+        return shell.run(cmd+args, **kwargs)
+
+    @staticmethod
+    def capture(
+        repo_path: str,
+        args: List[str],
+        stderr=None,
+        env=None,
+        dry_run=None,
+        echo=True,
+        optional=False,
+        allow_non_zero_exit=False,
+    ) -> Union[str, Any, None]:
+        return shell.capture(
+            ["git", "-C", repo_path] + args,
+            stderr=stderr,
+            env=env,
+            dry_run=dry_run,
+            echo=echo,
+            optional=optional,
+            allow_non_zero_exit=allow_non_zero_exit,
+        )

utils/update_checkout/update_checkout/parallel_runner.py

@@ -1,9 +1,10 @@
 from multiprocessing.managers import ListProxy, ValueProxy
 import sys
-from multiprocessing import Pool, cpu_count, Manager
+from multiprocessing import cpu_count, Manager
 import time
 from typing import Callable, List, Any, Union
 from threading import Lock, Thread, Event
+from concurrent.futures import ThreadPoolExecutor
 import shutil
 
 from .runner_arguments import RunnerArguments, AdditionalSwiftSourcesArguments
@@ -44,18 +45,17 @@ def __init__(
         self,
         fn: Callable,
         pool_args: List[Union[RunnerArguments, AdditionalSwiftSourcesArguments]],
-        n_processes: int = 0,
+        n_threads: int = 0,
     ):
-        if n_processes == 0:
-            # Limit the number of processes as the performance regresses after
-            # if the number is too high.
-            n_processes = min(cpu_count() * 2, 16)
-        self._n_processes = n_processes
+        if n_threads == 0:
+            # Limit the number of threads as the performance regresses if the
+            # number is too high.
+            n_threads = min(cpu_count() * 2, 16)
+        self._n_threads = n_threads
         self._monitor_polling_period = 0.1
         self._terminal_width = shutil.get_terminal_size().columns
         self._pool_args = pool_args
         self._fn = fn
-        self._pool = Pool(processes=self._n_processes)
         self._output_prefix = pool_args[0].output_prefix
         self._nb_repos = len(pool_args)
         self._stop_event = Event()
@@ -70,24 +70,15 @@ def __init__(
             )
 
     def run(self) -> List[Any]:
-        print(
-            "Running ``%s`` with up to %d processes."
-            % (self._fn.__name__, self._n_processes)
-        )
+        print(f"Running ``{self._fn.__name__}`` with up to {self._n_threads} processes.")
         if self._verbose:
-            results = self._pool.map_async(
-                func=self._fn, iterable=self._pool_args
-            ).get(timeout=1800)
-            self._pool.close()
-            self._pool.join()
+            with ThreadPoolExecutor(max_workers=self._n_threads) as pool:
+                results = list(pool.map(self._fn, self._pool_args, timeout=1800))
         else:
             monitor_thread = Thread(target=self._monitor, daemon=True)
             monitor_thread.start()
-            results = self._pool.map_async(
-                func=self._monitored_fn, iterable=self._pool_args
-            ).get(timeout=1800)
-            self._pool.close()
-            self._pool.join()
+            with ThreadPoolExecutor(max_workers=self._n_threads) as pool:
+                results = list(pool.map(self._monitored_fn, self._pool_args, timeout=1800))
             self._stop_event.set()
             monitor_thread.join()
         return results
@@ -131,14 +122,18 @@ def check_results(results, op) -> int:
         if results is None:
             return 0
         for r in results:
-            if r is not None:
-                if fail_count == 0:
-                    print("======%s FAILURES======" % op)
-                fail_count += 1
-                if isinstance(r, str):
-                    print(r)
-                    continue
-                print("%s failed (ret=%d): %s" % (r.repo_path, r.ret, r))
-                if r.stderr:
-                    print(r.stderr.decode())
+            if r is None:
+                continue
+            if fail_count == 0:
+                print("======%s FAILURES======" % op)
+            fail_count += 1
+            if isinstance(r, str):
+                print(r)
+                continue
+            if not hasattr(r, "repo_path"):
+                # TODO: create a proper Exception class with these attributes
+                continue
+            print("%s failed (ret=%d): %s" % (r.repo_path, r.ret, r))
+            if r.stderr:
+                print(r.stderr.decode())
         return fail_count

utils/update_checkout/update_checkout/runner_arguments.py

@@ -22,7 +22,8 @@ class UpdateArguments(RunnerArguments):
 
 @dataclass
 class AdditionalSwiftSourcesArguments(RunnerArguments):
-    args: RunnerArguments
+    args: Any
+    "Arguments passed during CLI invocation." # TODO: Properly type this.
     repo_info: str
     repo_branch: str
     remote: str