4 files changed, 213 insertions, 143 deletions

diff --git a/subcmds/forall.py b/subcmds/forall.py
index e2a420a9..7771ec16 100644
--- a/subcmds/forall.py
+++ b/subcmds/forall.py
@@ -14,7 +14,9 @@
 # limitations under the License.
 
 from __future__ import print_function
+import errno
 import fcntl
+import multiprocessing
 import re
 import os
 import select
@@ -31,6 +33,7 @@ _CAN_COLOR = [
   'log',
 ]
 
+
 class ForallColoring(Coloring):
   def __init__(self, config):
     Coloring.__init__(self, config, 'forall')
@@ -87,6 +90,12 @@ revision to a locally executed git command, use REPO_LREV.
 REPO_RREV is the name of the revision from the manifest, exactly
 as written in the manifest.
 
+REPO_COUNT is the total number of projects being iterated.
+
+REPO_I is the current (1-based) iteration count. Can be used in
+conjunction with REPO_COUNT to add a simple progress indicator to your
+command.
+
 REPO__* are any extra environment variables, specified by the
 "annotation" element under any project element.  This can be useful
 for differentiating trees based on user-specific criteria, or simply
@@ -126,9 +135,31 @@ without iterating through the remaining projects.
     g.add_option('-v', '--verbose',
                  dest='verbose', action='store_true',
                  help='Show command error messages')
+    g.add_option('-j', '--jobs',
+                 dest='jobs', action='store', type='int', default=1,
+                 help='number of commands to execute simultaneously')
 
   def WantPager(self, opt):
-    return opt.project_header
+    return opt.project_header and opt.jobs == 1
+
+  def _SerializeProject(self, project):
+    """ Serialize a project._GitGetByExec instance.
+
+    project._GitGetByExec is not pickle-able. Instead of trying to pass it
+    around between processes, make a dict ourselves containing only the
+    attributes that we need.
+
+    """
+    return {
+      'name': project.name,
+      'relpath': project.relpath,
+      'remote_name': project.remote.name,
+      'lrev': project.GetRevisionId(),
+      'rrev': project.revisionExpr,
+      'annotations': dict((a.name, a.value) for a in project.annotations),
+      'gitdir': project.gitdir,
+      'worktree': project.worktree,
+    }
 
   def Execute(self, opt, args):
     if not opt.command:
@@ -167,123 +198,165 @@ without iterating through the remaining projects.
       # pylint: enable=W0631
 
     mirror = self.manifest.IsMirror
-    out = ForallColoring(self.manifest.manifestProject.config)
-    out.redirect(sys.stdout)
-
     rc = 0
-    first = True
 
     if not opt.regex:
       projects = self.GetProjects(args)
     else:
       projects = self.FindProjects(args)
 
-    for project in projects:
-      env = os.environ.copy()
-      def setenv(name, val):
-        if val is None:
-          val = ''
-        env[name] = val.encode()
-
-      setenv('REPO_PROJECT', project.name)
-      setenv('REPO_PATH', project.relpath)
-      setenv('REPO_REMOTE', project.remote.name)
-      setenv('REPO_LREV', project.GetRevisionId())
-      setenv('REPO_RREV', project.revisionExpr)
-      for a in project.annotations:
-        setenv("REPO__%s" % (a.name), a.value)
-
-      if mirror:
-        setenv('GIT_DIR', project.gitdir)
-        cwd = project.gitdir
-      else:
-        cwd = project.worktree
-
-      if not os.path.exists(cwd):
-        if (opt.project_header and opt.verbose) \
-        or not opt.project_header:
-          print('skipping %s/' % project.relpath, file=sys.stderr)
-        continue
-
-      if opt.project_header:
-        stdin = subprocess.PIPE
-        stdout = subprocess.PIPE
-        stderr = subprocess.PIPE
-      else:
-        stdin = None
-        stdout = None
-        stderr = None
-
-      p = subprocess.Popen(cmd,
-                           cwd = cwd,
-                           shell = shell,
-                           env = env,
-                           stdin = stdin,
-                           stdout = stdout,
-                           stderr = stderr)
-
-      if opt.project_header:
-        class sfd(object):
-          def __init__(self, fd, dest):
-            self.fd = fd
-            self.dest = dest
-          def fileno(self):
-            return self.fd.fileno()
-
-        empty = True
-        errbuf = ''
-
-        p.stdin.close()
-        s_in = [sfd(p.stdout, sys.stdout),
-                sfd(p.stderr, sys.stderr)]
-
-        for s in s_in:
-          flags = fcntl.fcntl(s.fd, fcntl.F_GETFL)
-          fcntl.fcntl(s.fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)
-
-        while s_in:
-          in_ready, _out_ready, _err_ready = select.select(s_in, [], [])
-          for s in in_ready:
-            buf = s.fd.read(4096)
-            if not buf:
-              s.fd.close()
-              s_in.remove(s)
-              continue
-
-            if not opt.verbose:
-              if s.fd != p.stdout:
-                errbuf += buf
-                continue
-
-            if empty:
-              if first:
-                first = False
-              else:
-                out.nl()
-
-              if mirror:
-                project_header_path = project.name
-              else:
-                project_header_path = project.relpath
-              out.project('project %s/', project_header_path)
-              out.nl()
-              out.flush()
-              if errbuf:
-                sys.stderr.write(errbuf)
-                sys.stderr.flush()
-                errbuf = ''
-              empty = False
-
-            s.dest.write(buf)
-            s.dest.flush()
-
-      r = p.wait()
-      if r != 0:
-        if r != rc:
-          rc = r
-        if opt.abort_on_errors:
-          print("error: %s: Aborting due to previous error" % project.relpath,
-                file=sys.stderr)
-          sys.exit(r)
+    os.environ['REPO_COUNT'] = str(len(projects))
+
+    pool = multiprocessing.Pool(opt.jobs)
+    try:
+      config = self.manifest.manifestProject.config
+      results_it = pool.imap(
+         DoWorkWrapper,
+         [[mirror, opt, cmd, shell, cnt, config, self._SerializeProject(p)]
+          for cnt, p in enumerate(projects)]
+      )
+      pool.close()
+      for r in results_it:
+        rc = rc or r
+        if r != 0 and opt.abort_on_errors:
+          raise Exception('Aborting due to previous error')
+    except (KeyboardInterrupt, WorkerKeyboardInterrupt):
+      # Catch KeyboardInterrupt raised inside and outside of workers
+      print('Interrupted - terminating the pool')
+      pool.terminate()
+      rc = rc or errno.EINTR
+    except Exception as e:
+      # Catch any other exceptions raised
+      print('Got an error, terminating the pool: %r' % e,
+            file=sys.stderr)
+      pool.terminate()
+      rc = rc or getattr(e, 'errno', 1)
+    finally:
+      pool.join()
     if rc != 0:
       sys.exit(rc)
+
+
+class WorkerKeyboardInterrupt(Exception):
+  """ Keyboard interrupt exception for worker processes. """
+  pass
+
+
+def DoWorkWrapper(args):
+  """ A wrapper around the DoWork() method.
+
+  Catch the KeyboardInterrupt exceptions here and re-raise them as a different,
+  ``Exception``-based exception to stop it flooding the console with stacktraces
+  and making the parent hang indefinitely.
+
+  """
+  project = args.pop()
+  try:
+    return DoWork(project, *args)
+  except KeyboardInterrupt:
+    print('%s: Worker interrupted' % project['name'])
+    raise WorkerKeyboardInterrupt()
+
+
+def DoWork(project, mirror, opt, cmd, shell, cnt, config):
+  env = os.environ.copy()
+  def setenv(name, val):
+    if val is None:
+      val = ''
+    env[name] = val.encode()
+
+  setenv('REPO_PROJECT', project['name'])
+  setenv('REPO_PATH', project['relpath'])
+  setenv('REPO_REMOTE', project['remote_name'])
+  setenv('REPO_LREV', project['lrev'])
+  setenv('REPO_RREV', project['rrev'])
+  setenv('REPO_I', str(cnt + 1))
+  for name in project['annotations']:
+    setenv("REPO__%s" % (name), project['annotations'][name])
+
+  if mirror:
+    setenv('GIT_DIR', project['gitdir'])
+    cwd = project['gitdir']
+  else:
+    cwd = project['worktree']
+
+  if not os.path.exists(cwd):
+    if (opt.project_header and opt.verbose) \
+    or not opt.project_header:
+      print('skipping %s/' % project['relpath'], file=sys.stderr)
+    return
+
+  if opt.project_header:
+    stdin = subprocess.PIPE
+    stdout = subprocess.PIPE
+    stderr = subprocess.PIPE
+  else:
+    stdin = None
+    stdout = None
+    stderr = None
+
+  p = subprocess.Popen(cmd,
+                       cwd=cwd,
+                       shell=shell,
+                       env=env,
+                       stdin=stdin,
+                       stdout=stdout,
+                       stderr=stderr)
+
+  if opt.project_header:
+    out = ForallColoring(config)
+    out.redirect(sys.stdout)
+    class sfd(object):
+      def __init__(self, fd, dest):
+        self.fd = fd
+        self.dest = dest
+      def fileno(self):
+        return self.fd.fileno()
+
+    empty = True
+    errbuf = ''
+
+    p.stdin.close()
+    s_in = [sfd(p.stdout, sys.stdout),
+            sfd(p.stderr, sys.stderr)]
+
+    for s in s_in:
+      flags = fcntl.fcntl(s.fd, fcntl.F_GETFL)
+      fcntl.fcntl(s.fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)
+
+    while s_in:
+      in_ready, _out_ready, _err_ready = select.select(s_in, [], [])
+      for s in in_ready:
+        buf = s.fd.read(4096)
+        if not buf:
+          s.fd.close()
+          s_in.remove(s)
+          continue
+
+        if not opt.verbose:
+          if s.fd != p.stdout:
+            errbuf += buf
+            continue
+
+        if empty and out:
+          if not cnt == 0:
+            out.nl()
+
+          if mirror:
+            project_header_path = project['name']
+          else:
+            project_header_path = project['relpath']
+          out.project('project %s/', project_header_path)
+          out.nl()
+          out.flush()
+          if errbuf:
+            sys.stderr.write(errbuf)
+            sys.stderr.flush()
+            errbuf = ''
+          empty = False
+
+        s.dest.write(buf)
+        s.dest.flush()
+
+  r = p.wait()
+  return r
diff --git a/subcmds/status.py b/subcmds/status.py
index 41c4429a..b42675e0 100644
--- a/subcmds/status.py
+++ b/subcmds/status.py
@@ -113,7 +113,7 @@ the following meanings:
     try:
       state = project.PrintWorkTreeStatus(output)
       if state == 'CLEAN':
-        clean_counter.next()
+        next(clean_counter)
     finally:
       sem.release()
 
@@ -141,7 +141,7 @@ the following meanings:
       for project in all_projects:
         state = project.PrintWorkTreeStatus()
         if state == 'CLEAN':
-          counter.next()
+          next(counter)
     else:
       sem = _threading.Semaphore(opt.jobs)
       threads_and_output = []
@@ -164,7 +164,7 @@ the following meanings:
         t.join()
         output.dump(sys.stdout)
         output.close()
-    if len(all_projects) == counter.next():
+    if len(all_projects) == next(counter):
       print('nothing to commit (working directory clean)')
 
     if opt.orphans:
diff --git a/subcmds/sync.py b/subcmds/sync.py
index a0a68960..6f77310f 100644
--- a/subcmds/sync.py
+++ b/subcmds/sync.py
@@ -14,10 +14,10 @@
 # limitations under the License.
 
 from __future__ import print_function
+import json
 import netrc
 from optparse import SUPPRESS_HELP
 import os
-import pickle
 import re
 import shutil
 import socket
@@ -760,7 +760,7 @@ class _FetchTimes(object):
   _ALPHA = 0.5
 
   def __init__(self, manifest):
-    self._path = os.path.join(manifest.repodir, '.repopickle_fetchtimes')
+    self._path = os.path.join(manifest.repodir, '.repo_fetchtimes.json')
     self._times = None
     self._seen = set()
 
@@ -779,22 +779,17 @@ class _FetchTimes(object):
   def _Load(self):
     if self._times is None:
       try:
-        f = open(self._path, 'rb')
-      except IOError:
-        self._times = {}
-        return self._times
-      try:
+        f = open(self._path)
         try:
-          self._times = pickle.load(f)
-        except IOError:
-          try:
-            os.remove(self._path)
-          except OSError:
-            pass
-          self._times = {}
-      finally:
-        f.close()
-    return self._times
+          self._times = json.load(f)
+        finally:
+          f.close()
+      except (IOError, ValueError):
+        try:
+          os.remove(self._path)
+        except OSError:
+          pass
+        self._times = {}
 
   def Save(self):
     if self._times is None:
@@ -808,13 +803,13 @@ class _FetchTimes(object):
       del self._times[name]
 
     try:
-      f = open(self._path, 'wb')
+      f = open(self._path, 'w')
       try:
-        pickle.dump(self._times, f)
-      except (IOError, OSError, pickle.PickleError):
-        try:
-          os.remove(self._path)
-        except OSError:
-          pass
-    finally:
-      f.close()
+        json.dump(self._times, f, indent=2)
+      finally:
+        f.close()
+    except (IOError, TypeError):
+      try:
+        os.remove(self._path)
+      except OSError:
+        pass
diff --git a/subcmds/upload.py b/subcmds/upload.py
index e2fa261e..0ee36df1 100644
--- a/subcmds/upload.py
+++ b/subcmds/upload.py
@@ -25,10 +25,12 @@ from git_command import GitCommand
 from project import RepoHook
 
 from pyversion import is_python3
+# pylint:disable=W0622
 if not is_python3():
-  # pylint:disable=W0622
   input = raw_input
-  # pylint:enable=W0622
+else:
+  unicode = str
+# pylint:enable=W0622
 
 UNUSUAL_COMMIT_THRESHOLD = 5