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>

2 files changed, 83 insertions, 38 deletions

diff --git a/subcmds/abandon.py b/subcmds/abandon.py
index 359c431b..b82a2dbf 100644
--- a/subcmds/abandon.py
+++ b/subcmds/abandon.py
@@ -13,9 +13,12 @@
 # limitations under the License.
 
 from collections import defaultdict
+import functools
+import itertools
+import multiprocessing
 import sys
 
-from command import Command
+from command import Command, DEFAULT_LOCAL_JOBS, WORKER_BATCH_SIZE
 from git_command import git
 from progress import Progress
 
@@ -31,8 +34,10 @@ deleting it (and all its history) from your local repository.
 
 It is equivalent to "git branch -D <branchname>".
 """
+  PARALLEL_JOBS = DEFAULT_LOCAL_JOBS
 
   def _Options(self, p):
+    super()._Options(p)
     p.add_option('-q', '--quiet',
                  action='store_true', default=False,
                  help='be quiet')
@@ -51,35 +56,49 @@ It is equivalent to "git branch -D <branchname>".
     else:
       args.insert(0, "'All local branches'")
 
+  def _ExecuteOne(self, opt, nb, project):
+    """Abandon one project."""
+    if opt.all:
+      branches = project.GetBranches()
+    else:
+      branches = [nb]
+
+    ret = {}
+    for name in branches:
+      status = project.AbandonBranch(name)
+      if status is not None:
+        ret[name] = status
+    return (ret, project)
+
   def Execute(self, opt, args):
     nb = args[0]
     err = defaultdict(list)
     success = defaultdict(list)
     all_projects = self.GetProjects(args[1:])
 
-    pm = Progress('Abandon %s' % nb, len(all_projects))
-    for project in all_projects:
-      pm.update()
-
-      if opt.all:
-        branches = list(project.GetBranches().keys())
-      else:
-        branches = [nb]
-
-      for name in branches:
-        status = project.AbandonBranch(name)
-        if status is not None:
+    def _ProcessResults(states):
+      for (results, project) in states:
+        for branch, status in results.items():
           if status:
-            success[name].append(project)
+            success[branch].append(project)
           else:
-            err[name].append(project)
-    pm.end()
+            err[branch].append(project)
+        pm.update()
 
-    width = 25
-    for name in branches:
-      if width < len(name):
-        width = len(name)
+    pm = Progress('Abandon %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:
+        states = pool.imap_unordered(
+            functools.partial(self._ExecuteOne, opt, nb), all_projects,
+            chunksize=WORKER_BATCH_SIZE)
+        _ProcessResults(states)
+    pm.end()
 
+    width = max(itertools.chain(
+        [25], (len(x) for x in itertools.chain(success, err))))
     if err:
       for br in err.keys():
         err_msg = "error: cannot abandon %s" % br
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: