forall: improve pool logic

Use a pool contextmanager to take care of the messy details like
properly cleaning it up when aborting.

Change-Id: I264ebb591c2e67c9a975b6dcc0f14b29cc66a874
Reviewed-on: https://gerrit-review.googlesource.com/c/git-repo/+/297243
Reviewed-by: Chris Mcdonald <cjmcdonald@google.com>
Tested-by: Mike Frysinger <vapier@google.com>
This commit is contained in:
Mike Frysinger 2021-02-16 01:56:30 -05:00
parent 7c871163c8
commit 15e807cf3c

View File

@ -21,7 +21,7 @@ import sys
import subprocess import subprocess
from color import Coloring from color import Coloring
from command import DEFAULT_LOCAL_JOBS, Command, MirrorSafeCommand from command import DEFAULT_LOCAL_JOBS, Command, MirrorSafeCommand, WORKER_BATCH_SIZE
import platform_utils import platform_utils
_CAN_COLOR = [ _CAN_COLOR = [
@ -234,31 +234,26 @@ without iterating through the remaining projects.
os.environ['REPO_COUNT'] = str(len(projects)) os.environ['REPO_COUNT'] = str(len(projects))
pool = multiprocessing.Pool(opt.jobs, InitWorker)
try: try:
config = self.manifest.manifestProject.config config = self.manifest.manifestProject.config
with multiprocessing.Pool(opt.jobs, InitWorker) as pool:
results_it = pool.imap( results_it = pool.imap(
DoWorkWrapper, DoWorkWrapper,
self.ProjectArgs(projects, mirror, opt, cmd, shell, config)) self.ProjectArgs(projects, mirror, opt, cmd, shell, config),
pool.close() chunksize=WORKER_BATCH_SIZE)
for r in results_it: for r in results_it:
rc = rc or r rc = rc or r
if r != 0 and opt.abort_on_errors: if r != 0 and opt.abort_on_errors:
raise Exception('Aborting due to previous error') raise Exception('Aborting due to previous error')
except (KeyboardInterrupt, WorkerKeyboardInterrupt): except (KeyboardInterrupt, WorkerKeyboardInterrupt):
# Catch KeyboardInterrupt raised inside and outside of workers # Catch KeyboardInterrupt raised inside and outside of workers
print('Interrupted - terminating the pool')
pool.terminate()
rc = rc or errno.EINTR rc = rc or errno.EINTR
except Exception as e: except Exception as e:
# Catch any other exceptions raised # Catch any other exceptions raised
print('Got an error, terminating the pool: %s: %s' % print('Got an error, terminating the pool: %s: %s' %
(type(e).__name__, e), (type(e).__name__, e),
file=sys.stderr) file=sys.stderr)
pool.terminate()
rc = rc or getattr(e, 'errno', 1) rc = rc or getattr(e, 'errno', 1)
finally:
pool.join()
if rc != 0: if rc != 0:
sys.exit(rc) sys.exit(rc)