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>
diff --git a/subcmds/diff.py b/subcmds/diff.py
index c987bf2..8186817 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 @@
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