From 69b4a9cf216f4ffcb69ea2c39ebd6a01b2d0fec9 Mon Sep 17 00:00:00 2001
From: Mike Frysinger <vapier@google.com>
Date: Tue, 16 Feb 2021 18:18:01 -0500
Subject: [PATCH] diff: add --jobs support

Use multiprocessing to run diff in parallel.

Change-Id: I61e973d9c2cde039d5eebe8d0fe8bb63171ef447
Reviewed-on: https://gerrit-review.googlesource.com/c/git-repo/+/297483
Tested-by: Mike Frysinger <vapier@google.com>
Reviewed-by: Chris Mcdonald <cjmcdonald@google.com>
---
 project.py      |  4 +++-
 subcmds/diff.py | 42 ++++++++++++++++++++++++++++++++++++++----
 2 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/project.py b/project.py
index 52a77f1e..da67c363 100644
--- a/project.py
+++ b/project.py
@@ -832,10 +832,12 @@ class Project(object):
 
     return 'DIRTY'
 
-  def PrintWorkTreeDiff(self, absolute_paths=False):
+  def PrintWorkTreeDiff(self, absolute_paths=False, output_redir=None):
     """Prints the status of the repository to stdout.
     """
     out = DiffColoring(self.config)
+    if output_redir:
+      out.redirect(output_redir)
     cmd = ['diff']
     if out.is_on:
       cmd.append('--color')
diff --git a/subcmds/diff.py b/subcmds/diff.py
index c987bf23..81868176 100644
--- a/subcmds/diff.py
+++ b/subcmds/diff.py
@@ -12,7 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from command import PagedCommand
+import functools
+import io
+import multiprocessing
+
+from command import DEFAULT_LOCAL_JOBS, PagedCommand, WORKER_BATCH_SIZE
 
 
 class Diff(PagedCommand):
@@ -25,15 +29,45 @@ The -u option causes '%prog' to generate diff output with file paths
 relative to the repository root, so the output can be applied
 to the Unix 'patch' command.
 """
+  PARALLEL_JOBS = DEFAULT_LOCAL_JOBS
 
   def _Options(self, p):
+    super()._Options(p)
     p.add_option('-u', '--absolute',
                  dest='absolute', action='store_true',
                  help='Paths are relative to the repository root')
 
+  def _DiffHelper(self, absolute, project):
+    """Obtains the diff for a specific project.
+
+    Args:
+      absolute: Paths are relative to the root.
+      project: Project to get status of.
+
+    Returns:
+      The status of the project.
+    """
+    buf = io.StringIO()
+    ret = project.PrintWorkTreeDiff(absolute, output_redir=buf)
+    return (ret, buf.getvalue())
+
   def Execute(self, opt, args):
     ret = 0
-    for project in self.GetProjects(args):
-      if not project.PrintWorkTreeDiff(opt.absolute):
-        ret = 1
+    all_projects = self.GetProjects(args)
+
+    # NB: Multiprocessing is heavy, so don't spin it up for one job.
+    if len(all_projects) == 1 or opt.jobs == 1:
+      for project in all_projects:
+        if not project.PrintWorkTreeDiff(opt.absolute):
+          ret = 1
+    else:
+      with multiprocessing.Pool(opt.jobs) as pool:
+        states = pool.imap(functools.partial(self._DiffHelper, opt.absolute),
+                           all_projects, WORKER_BATCH_SIZE)
+        for (state, output) in states:
+          if output:
+            print(output, end='')
+          if not state:
+            ret = 1
+
     return ret