utils: Make update-checkout run in parallel.

Author: Doug Coleman

utils/swift_build_support/swift_build_support/shell.py

@@ -20,6 +20,7 @@
 import shutil
 import subprocess
 import sys
+from multiprocessing import Pool, Lock, cpu_count
 from contextlib import contextmanager
 
 from . import diagnostics
@@ -170,3 +171,71 @@ def copytree(src, dest, dry_run=None, echo=True):
     if dry_run:
         return
     shutil.copytree(src, dest)
+
+
+def run(*args, **kwargs):
+    msg = None
+    repo_path = os.getcwd()
+    try:
+        echo_output = kwargs.pop('echo', False)
+        msg = kwargs.pop('msg', False)
+        dry_run = kwargs.pop('dry_run', False)
+        env = kwargs.pop('env', None)
+        allow_non_zero_exit = kwargs.pop('allow_non_zero_exit', False)
+        if dry_run:
+            _echo_command(dry_run, *args, env=env)
+            return(None, 0, args)
+
+        out = subprocess.check_output(*args, stderr=subprocess.STDOUT, **kwargs)
+        lock.acquire()
+
+        if out and echo_output:
+            print(repo_path)
+            _echo_command(dry_run, *args, env=env)
+            print(out)
+        lock.release()
+        return (out, 0, args)
+    except subprocess.CalledProcessError as e:
+        if allow_non_zero_exit:
+            print("non-fatal error running ``%s`` on ``%s``" % (" ".join(*args), repo_path))
+            return (out, 0, args)
+        print("error running ``%s`` on ``%s``" % (" ".join(*args), repo_path))
+        # Workaround for exception bug in 2.7.12
+        # http://bugs.python.org/issue9400
+        eout = Exception(str(e))
+        eout.repo_path = repo_path
+        if msg is not None and msg is not False:
+            print(msg)
+        eout.msg = msg
+        raise eout
+
+
+def run_parallel(fn, pool_args, n_processes=0):
+    def init(l):
+        global lock
+        lock = l
+
+    if n_processes == 0:
+        n_processes = cpu_count()
+
+    l = Lock()
+    pool = Pool(processes=n_processes, initializer=init, initargs=(l,))
+    results = pool.map_async(func=fn, iterable=pool_args).get(9999999)
+    pool.close()
+    pool.join()
+    return results
+
+
+def check_parallel_results(results, op):
+    fail_count = 0
+    if results is None:
+        return 0
+    for r in results:
+        if r is not None:
+            if fail_count == 0:
+                print("======%s FAILURES======" % op)
+            print("%s failed: %s" % (r.repo_path, r))
+            fail_count += 1
+            if r.msg:
+                print(r.msg)
+    return fail_count

utils/update-checkout

@@ -16,9 +16,13 @@ import json
 import os
 import re
 import sys
+import traceback
+import subprocess
 
 
 from functools import reduce
+from multiprocessing import freeze_support
+
 
 sys.path.append(os.path.dirname(__file__))
 
@@ -34,56 +38,53 @@ sys.path.append(os.path.join(SCRIPT_DIR, 'swift_build_support'))
 from swift_build_support import shell  # noqa (E402)
 
 
-def update_single_repository(repo_path, branch, reset_to_remote, should_clean,
-                             cross_repo):
+def update_single_repository(args):
+    repo_path, branch, reset_to_remote, should_clean, cross_repo = args
     if not os.path.isdir(repo_path):
         return
 
-    print("--- Updating '" + repo_path + "' ---")
-    with shell.pushd(repo_path, dry_run=False, echo=False):
-        shell.call(["git", "fetch", "--recurse-submodules=yes"], echo=True)
-
-        if should_clean:
-            shell.call(['git', 'clean', '-fdx'],
-                       echo=True)
-            shell.call(['git', 'submodule', 'foreach', '--recursive', 'git',
-                        'clean', '-fdx'], echo=True)
-            shell.call(['git', 'submodule', 'foreach', '--recursive', 'git',
-                        'reset', '--hard', 'HEAD'], echo=True)
-            status = shell.call(['git', 'reset', '--hard', 'HEAD'],
-                                echo=True)
-            if status:
-                print("Please, commit your changes.")
-                print(status)
-                exit(1)
-
-        if branch:
-            status = shell.capture(['git', 'status', '--porcelain', '-uno'],
-                                   echo=False)
-            if status:
-                print("Please, commit your changes.")
-                print(status)
-                exit(1)
-            shell.call(['git', 'checkout', branch], echo=True)
-
-            # If we were asked to reset to the specified branch, do the hard
-            # reset and return.
-            if reset_to_remote and not cross_repo:
-                shell.call(['git', 'reset', '--hard', "origin/%s" % branch],
-                           echo=True)
-                return
-
-        # Prior to Git 2.6, this is the way to do a "git pull
-        # --rebase" that respects rebase.autostash.  See
-        # http://stackoverflow.com/a/30209750/125349
-        if not cross_repo:
-            shell.call(["git", "rebase", "FETCH_HEAD"], echo=True)
-        shell.call(["git", "submodule", "update", "--recursive"],
-                   echo=True)
+    try:
+        print("--- Updating '" + repo_path + "' ---")
+        with shell.pushd(repo_path, dry_run=False, echo=False):
+            shell.run(["git", "fetch", "--recurse-submodules=yes"], echo=True)
+
+            if should_clean:
+                shell.run(['git', 'clean', '-fdx'], echo=True)
+                shell.run(['git', 'submodule', 'foreach', '--recursive', 'git',
+                            'clean', '-fdx'], echo=True)
+                shell.run(['git', 'submodule', 'foreach', '--recursive', 'git',
+                            'reset', '--hard', 'HEAD'], echo=True)
+                shell.run(['git', 'reset', '--hard', 'HEAD'],
+                                    echo=True, msg="Please, commit your changes.")
+
+            if branch:
+                shell.run(['git', 'status', '--porcelain', '-uno'],
+                                       echo=False, msg="Please, commit your changes.")
+                shell.run(['git', 'checkout', branch], echo=True)
+
+                # If we were asked to reset to the specified branch, do the hard
+                # reset and return.
+                if reset_to_remote and not cross_repo:
+                    shell.run(['git', 'reset', '--hard', "origin/%s" % branch],
+                               echo=True)
+                    return
+
+            # Prior to Git 2.6, this is the way to do a "git pull
+            # --rebase" that respects rebase.autostash.  See
+            # http://stackoverflow.com/a/30209750/125349
+            if not cross_repo:
+                shell.run(["git", "rebase", "FETCH_HEAD"], echo=True,
+                    msg="Cannot rebase: You have unstaged changes. Please commit or stash them.")
+            shell.run(["git", "submodule", "update", "--recursive"], echo=True)
+    except:
+        (type, value, tb) = sys.exc_info()
+        print('Error on repo "%s": %s' % (repo_path, traceback.format_exc()))
+        return value
 
 
 def update_all_repositories(args, config, scheme_name, cross_repos_pr):
     repo_branch = scheme_name
+    pool_args = []
     for repo_name in config['repos'].keys():
         cross_repo = False
         if repo_name in args.skip_repository_list:
@@ -105,70 +106,78 @@ def update_all_repositories(args, config, scheme_name, cross_repos_pr):
                     pr_id = cross_repos_pr[remote_repo_id]
                     repo_branch = "ci_pr_{0}".format(pr_id)
                     with shell.pushd(repo_path, dry_run=False, echo=False):
-                        shell.call(["git", "checkout", v['repos'][repo_name]],
+                        shell.run(["git", "checkout", v['repos'][repo_name]],
                                    echo=True)
-                        shell.capture(["git", "branch", "-D", repo_branch],
+                        shell.run(["git", "branch", "-D", repo_branch],
                                       echo=True, allow_non_zero_exit=True)
-                        shell.call(["git", "fetch", "origin",
+                        shell.run(["git", "fetch", "origin",
                                     "pull/{0}/merge:{1}"
                                     .format(pr_id, repo_branch)], echo=True)
                 break
-        update_single_repository(repo_path,
-                                 repo_branch,
-                                 args.reset_to_remote,
-                                 args.clean,
-                                 cross_repo)
+        pool_args.append([repo_path, repo_branch, args.reset_to_remote, args.clean, cross_repo])
+
+    return shell.run_parallel(update_single_repository, pool_args, args.n_processes)
 
+def obtain_additional_swift_sources(pool_args):
+    args, repo_name, repo_info, repo_branch, remote, with_ssh, scheme_name, skip_history, skip_repository_list = pool_args
 
-def obtain_additional_swift_sources(
-        config, with_ssh, scheme_name, skip_history, skip_repository_list):
     with shell.pushd(SWIFT_SOURCE_ROOT, dry_run=False,
                      echo=False):
-        for repo_name, repo_info in config['repos'].items():
-            if repo_name in skip_repository_list:
-                print("--- Skipping '" + repo_name + "' ---")
-                continue
+        print("--- Cloning '" + repo_name + "' ---")
+
+        if skip_history:
+            shell.run(['git', 'clone', '--recursive', '--depth', '1',
+                        remote, repo_name], echo=True)
+        else:
+            shell.run(['git', 'clone', '--recursive', remote,
+                        repo_name], echo=True)
+        if scheme_name:
+            src_path = os.path.join(SWIFT_SOURCE_ROOT, repo_name, ".git")
+            shell.run(['git', '--git-dir', src_path, '--work-tree',
+                        os.path.join(SWIFT_SOURCE_ROOT, repo_name),
+                        'checkout', repo_branch], echo=False)
+    with shell.pushd(os.path.join(SWIFT_SOURCE_ROOT, repo_name),
+                     dry_run=False, echo=False):
+        shell.run(["git", "submodule", "update", "--recursive"],
+                   echo=False)
 
-            if os.path.isdir(os.path.join(repo_name, ".git")):
-                continue
+def obtain_all_additional_swift_sources(
+    args, config, with_ssh, scheme_name, skip_history, skip_repository_list):
 
-            print("--- Cloning '" + repo_name + "' ---")
+    pool_args = []
+    for repo_name, repo_info in config['repos'].items():
+        if repo_name in skip_repository_list:
+            print("--- Skipping '" + repo_name + "' ---")
+            continue
+
+        if os.path.isdir(os.path.join(repo_name, ".git")):
+            continue
 
-            # If we have a url override, use that url instead of
-            # interpolating.
-            remote_repo_info = repo_info['remote']
-            if 'url' in remote_repo_info:
-                remote = remote_repo_info['url']
+        # If we have a url override, use that url instead of
+        # interpolating.
+        remote_repo_info = repo_info['remote']
+        if 'url' in remote_repo_info:
+            remote = remote_repo_info['url']
+        else:
+            remote_repo_id = remote_repo_info['id']
+            if with_ssh is True or 'https-clone-pattern' not in config:
+                remote = config['ssh-clone-pattern'] % remote_repo_id
             else:
-                remote_repo_id = remote_repo_info['id']
-                if with_ssh is True or 'https-clone-pattern' not in config:
-                    remote = config['ssh-clone-pattern'] % remote_repo_id
-                else:
-                    remote = config['https-clone-pattern'] % remote_repo_id
-
-            if skip_history:
-                shell.call(['git', 'clone', '--recursive', '--depth', '1',
-                            remote, repo_name], echo=True)
+                remote = config['https-clone-pattern'] % remote_repo_id
+
+        repo_branch = None
+        if scheme_name:
+            for v in config['branch-schemes'].values():
+                if scheme_name not in v['aliases']:
+                    continue
+                repo_branch = v['repos'][repo_name]
+                break
             else:
-                shell.call(['git', 'clone', '--recursive', remote,
-                            repo_name], echo=True)
-            if scheme_name:
-                for v in config['branch-schemes'].values():
-                    if scheme_name not in v['aliases']:
-                        continue
-                    repo_branch = v['repos'][repo_name]
-                    break
-                else:
-                    repo_branch = scheme_name
-                src_path = os.path.join(SWIFT_SOURCE_ROOT, repo_name,
-                                        ".git")
-                shell.call(['git', '--git-dir', src_path, '--work-tree',
-                            os.path.join(SWIFT_SOURCE_ROOT, repo_name),
-                            'checkout', repo_branch], echo=False)
-    with shell.pushd(os.path.join(SWIFT_SOURCE_ROOT, repo_name),
-                     dry_run=False, echo=False):
-        shell.call(["git", "submodule", "update", "--recursive"],
-                   echo=False)
+                repo_branch = scheme_name
+
+        pool_args.append([args, repo_name, repo_info, repo_branch, remote, with_ssh, scheme_name, skip_history, skip_repository_list])
+
+    return shell.run_parallel(obtain_additional_swift_sources, pool_args, args.n_processes)
 
 
 def validate_config(config):
@@ -192,7 +201,6 @@ def validate_config(config):
     raise RuntimeError('Configuration file has schemes with duplicate '
                        'aliases?!')
 
-
 def main():
     parser = argparse.ArgumentParser(
         formatter_class=argparse.RawDescriptionHelpFormatter,
@@ -242,6 +250,12 @@ By default, updates your checkouts of Swift, SourceKit, LLDB, and SwiftPM.""")
         free-form GitHub-style comment.""",
         metavar='GITHUB-COMMENT',
         dest='github_comment')
+    parser.add_argument(
+        "-j", "--jobs",
+        type=int,
+        help="Number of threads to run at once",
+        default=0,
+        dest="n_processes")
     args = parser.parse_args()
 
     clone = args.clone
@@ -262,19 +276,25 @@ By default, updates your checkouts of Swift, SourceKit, LLDB, and SwiftPM.""")
         repos_with_pr = [pr.replace('/pull/', '#') for pr in repos_with_pr]
         cross_repos_pr = dict(pr.split('#') for pr in repos_with_pr)
 
+    clone_results = None
     if clone or clone_with_ssh:
         # If branch is None, default to using the default branch alias
         # specified by our configuration file.
         if scheme is None:
             scheme = config['default-branch-scheme']
 
-        obtain_additional_swift_sources(
-            config, clone_with_ssh, scheme, skip_history,
-            args.skip_repository_list)
-
-    update_all_repositories(args, config, scheme, cross_repos_pr)
+        clone_results = obtain_all_additional_swift_sources(args, config,
+                    clone_with_ssh, scheme, skip_history, args.skip_repository_list)
 
-    return 0
+    update_results = update_all_repositories(args, config, scheme, cross_repos_pr)
+    return (clone_results, update_results)
 
 if __name__ == "__main__":
-    sys.exit(main())
+    freeze_support()
+    clone_results, update_results = main()
+    fail_count = 0
+    fail_count += shell.check_parallel_results(clone_results, "CLONE")
+    fail_count += shell.check_parallel_results(update_results, "UPDATE")
+    if fail_count > 0:
+        print("update-checkout failed, fix errors and try again")
+    sys.exit(fail_count)