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>

1 files changed, 38 insertions, 4 deletions

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