subcmds: reduce multiprocessing serialization overhead

Follow the same approach as 39ffd9977e to reduce serialization overhead.

Below benchmarks are tested with 2.7k projects on my workstation
(warm cache). git tracing is disabled for benchmark.

(seconds)              | v2.48 | v2.48 | this CL | this CL
	               |       |  -j32 |         |    -j32
-----------------------------------------------------------
with clean tree state:
branches (none)        |   5.6 |   5.9 |    1.0  |    0.9
status (clean)         |  21.3 |   9.4 |   19.4  |    4.7
diff (none)            |   7.6 |   7.2 |    5.7  |    2.2
prune (none)           |   5.7 |   6.1 |    1.3  |    1.2
abandon (none)         |  19.4 |  18.6 |    0.9  |    0.8
upload (none)          |  19.7 |  18.7 |    0.9  |    0.8
forall -c true         |   7.5 |   7.6 |    0.6  |    0.6
forall -c "git log -1" |  11.3 |  11.1 |    0.6  |    0.6

with branches:
start BRANCH --all     |  21.9 |  20.3 |   13.6  |    2.6
checkout BRANCH        |  29.1 |  27.8 |    1.1  |    1.0
branches (2)           |  28.0 |  28.6 |    1.5  |    1.3
abandon BRANCH         |  29.2 |  27.5 |    9.7  |    2.2

Bug: b/371638995
Change-Id: I53989a3d1e43063587b3f52f852b1c2c56b49412
Reviewed-on: https://gerrit-review.googlesource.com/c/git-repo/+/440221
Reviewed-by: Josip Sokcevic <sokcevic@google.com>
Tested-by: Kuang-che Wu <kcwu@google.com>
Commit-Queue: Kuang-che Wu <kcwu@google.com>

1 files changed, 18 insertions, 12 deletions

diff --git a/subcmds/abandon.py b/subcmds/abandon.py
index e280d69e..3208be6b 100644
--- a/subcmds/abandon.py
+++ b/subcmds/abandon.py
@@ -70,8 +70,10 @@ It is equivalent to "git branch -D <branchname>".
         else:
             args.insert(0, "'All local branches'")
 
-    def _ExecuteOne(self, all_branches, nb, project):
+    @classmethod
+    def _ExecuteOne(cls, all_branches, nb, project_idx):
         """Abandon one project."""
+        project = cls.get_parallel_context()["projects"][project_idx]
         if all_branches:
             branches = project.GetBranches()
         else:
@@ -89,7 +91,7 @@ It is equivalent to "git branch -D <branchname>".
             if status is not None:
                 ret[name] = status
 
-        return (ret, project, errors)
+        return (ret, project_idx, errors)
 
     def Execute(self, opt, args):
         nb = args[0].split()
@@ -102,7 +104,8 @@ It is equivalent to "git branch -D <branchname>".
         _RelPath = lambda p: p.RelPath(local=opt.this_manifest_only)
 
         def _ProcessResults(_pool, pm, states):
-            for results, project, errors in states:
+            for results, project_idx, errors in states:
+                project = all_projects[project_idx]
                 for branch, status in results.items():
                     if status:
                         success[branch].append(project)
@@ -111,15 +114,18 @@ It is equivalent to "git branch -D <branchname>".
                 aggregate_errors.extend(errors)
                 pm.update(msg="")
 
-        self.ExecuteInParallel(
-            opt.jobs,
-            functools.partial(self._ExecuteOne, opt.all, nb),
-            all_projects,
-            callback=_ProcessResults,
-            output=Progress(
-                f"Abandon {nb}", len(all_projects), quiet=opt.quiet
-            ),
-        )
+        with self.ParallelContext():
+            self.get_parallel_context()["projects"] = all_projects
+            self.ExecuteInParallel(
+                opt.jobs,
+                functools.partial(self._ExecuteOne, opt.all, nb),
+                range(len(all_projects)),
+                callback=_ProcessResults,
+                output=Progress(
+                    f"Abandon {nb}", len(all_projects), quiet=opt.quiet
+                ),
+                chunksize=1,
+            )
 
         width = max(
             itertools.chain(