abandon/start: add --jobs support

Use multiprocessing to run in parallel.  When operating on multiple
projects, this can greatly speed things up.  Across 1000 repos, it
goes from ~30sec to ~3sec with the default -j8.

Change-Id: I0dc62d704c022dd02cac0bd67fe79224f4e34095
Reviewed-on: https://gerrit-review.googlesource.com/c/git-repo/+/297484
Tested-by: Mike Frysinger <vapier@google.com>
Reviewed-by: Chris Mcdonald <cjmcdonald@google.com>

1 files changed, 44 insertions, 18 deletions

diff --git a/subcmds/start.py b/subcmds/start.py
index 7684b6d7..25b229f1 100644
--- a/subcmds/start.py
+++ b/subcmds/start.py
@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
+import multiprocessing
 import os
 import sys
 
-from command import Command
+from command import Command, DEFAULT_LOCAL_JOBS, WORKER_BATCH_SIZE
 from git_config import IsImmutable
 from git_command import git
 import gitc_utils
@@ -33,8 +35,10 @@ class Start(Command):
 '%prog' begins a new branch of development, starting from the
 revision specified in the manifest.
 """
+  PARALLEL_JOBS = DEFAULT_LOCAL_JOBS
 
   def _Options(self, p):
+    super()._Options(p)
     p.add_option('--all',
                  dest='all', action='store_true',
                  help='begin branch in all projects')
@@ -51,6 +55,26 @@ revision specified in the manifest.
     if not git.check_ref_format('heads/%s' % nb):
       self.OptionParser.error("'%s' is not a valid name" % nb)
 
+  def _ExecuteOne(self, opt, nb, project):
+    """Start one project."""
+    # If the current revision is immutable, such as a SHA1, a tag or
+    # a change, then we can't push back to it. Substitute with
+    # dest_branch, if defined; or with manifest default revision instead.
+    branch_merge = ''
+    if IsImmutable(project.revisionExpr):
+      if project.dest_branch:
+        branch_merge = project.dest_branch
+      else:
+        branch_merge = self.manifest.default.revisionExpr
+
+    try:
+      ret = project.StartBranch(
+          nb, branch_merge=branch_merge, revision=opt.revision)
+    except Exception as e:
+      print('error: unable to checkout %s: %s' % (project.name, e), file=sys.stderr)
+      ret = False
+    return (ret, project)
+
   def Execute(self, opt, args):
     nb = args[0]
     err = []
@@ -82,11 +106,8 @@ revision specified in the manifest.
       if not os.path.exists(os.getcwd()):
         os.chdir(self.manifest.topdir)
 
-    pm = Progress('Starting %s' % nb, len(all_projects))
-    for project in all_projects:
-      pm.update()
-
-      if self.gitc_manifest:
+      pm = Progress('Syncing %s' % nb, len(all_projects))
+      for project in all_projects:
         gitc_project = self.gitc_manifest.paths[project.relpath]
         # Sync projects that have not been opened.
         if not gitc_project.already_synced:
@@ -99,20 +120,25 @@ revision specified in the manifest.
           sync_buf = SyncBuffer(self.manifest.manifestProject.config)
           project.Sync_LocalHalf(sync_buf)
           project.revisionId = gitc_project.old_revision
+        pm.update()
+      pm.end()
 
-      # If the current revision is immutable, such as a SHA1, a tag or
-      # a change, then we can't push back to it. Substitute with
-      # dest_branch, if defined; or with manifest default revision instead.
-      branch_merge = ''
-      if IsImmutable(project.revisionExpr):
-        if project.dest_branch:
-          branch_merge = project.dest_branch
-        else:
-          branch_merge = self.manifest.default.revisionExpr
+    def _ProcessResults(results):
+      for (result, project) in results:
+        if not result:
+          err.append(project)
+        pm.update()
 
-      if not project.StartBranch(
-              nb, branch_merge=branch_merge, revision=opt.revision):
-        err.append(project)
+    pm = Progress('Starting %s' % nb, len(all_projects))
+    # NB: Multiprocessing is heavy, so don't spin it up for one job.
+    if len(all_projects) == 1 or opt.jobs == 1:
+      _ProcessResults(self._ExecuteOne(opt, nb, x) for x in all_projects)
+    else:
+      with multiprocessing.Pool(opt.jobs) as pool:
+        results = pool.imap_unordered(
+            functools.partial(self._ExecuteOne, opt, nb), all_projects,
+            chunksize=WORKER_BATCH_SIZE)
+        _ProcessResults(results)
     pm.end()
 
     if err: