Update parallel_emerge to support --force-remote-binary.

parallel_emerge

 #!/usr/bin/python2.6
 # Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """Program to run emerge in parallel, for significant speedup.
 
 Usage:
  ./parallel_emerge [--board=BOARD] [--workon=PKGS] [--no-workon-deps]
-                   [emerge args] package"
+                   [--force-remote-binary=PKGS] [emerge args] package
 
 Basic operation:
   Runs 'emerge -p --debug' to display dependencies, and stores a
   dependency graph. All non-blocked packages are launched in parallel,
   as 'emerge --nodeps package' with any blocked packages being emerged
   immediately upon deps being met.
 
   For this to work effectively, /usr/lib/portage/pym/portage/locks.py
   must be stubbed out, preventing portage from slowing itself with
   unneccesary locking, as this script ensures that emerge is run in such
@@ skipping 56 matching lines @@
 from _emerge.actions import load_emerge_config
 from _emerge.create_depgraph_params import create_depgraph_params
 from _emerge.depgraph import backtrack_depgraph
 from _emerge.main import emerge_main
 from _emerge.main import parse_opts
 from _emerge.Package import Package
 from _emerge.Scheduler import Scheduler
 from _emerge.stdout_spinner import stdout_spinner
 import portage
 import portage.debug
+import portage.versions
 
 
 def Usage():
   """Print usage."""
   print "Usage:"
   print " ./parallel_emerge [--board=BOARD] [--workon=PKGS] [--no-workon-deps]"
   print "                   [--rebuild] [emerge args] package"
   print
   print "Packages specified as workon packages are always built from source."
   print "Unless --no-workon-deps is specified, packages that depend on these"
@@ skipping 114 matching lines @@
   Typical usage:
     deps = DepGraphGenerator()
     deps.Initialize(sys.argv[1:])
     deps_tree, deps_info = deps.GenDependencyTree()
     deps_graph = deps.GenDependencyGraph(deps_tree, deps_info)
     deps.PrintTree(deps_tree)
     PrintDepsMap(deps_graph)
   """
 
   __slots__ = ["board", "emerge", "mandatory_source", "no_workon_deps",
-               "nomerge", "package_db", "rebuild", "show_output"]
+               "nomerge", "package_db", "rebuild", "show_output",
+               "force_remote_binary", "forced_remote_binary_packages"]
 
   def __init__(self):
     self.board = None
     self.emerge = EmergeData()
     self.mandatory_source = set()
     self.no_workon_deps = False
     self.nomerge = set()
     self.package_db = {}
     self.rebuild = False
     self.show_output = False
+    self.force_remote_binary = set()
+    self.forced_remote_binary_packages = set()
 
   def ParseParallelEmergeArgs(self, argv):
     """Read the parallel emerge arguments from the command-line.
 
     We need to be compatible with emerge arg format.  We scrape arguments that
     are specific to parallel_emerge, and pass through the rest directly to
     emerge.
     Args:
       argv: arguments list
     Returns:
       Arguments that don't belong to parallel_emerge
     """
     emerge_args = []
     for arg in argv:
       # Specifically match arguments that are specific to parallel_emerge, and
       # pass through the rest.
       if arg.startswith("--board="):
         self.board = arg.replace("--board=", "")
       elif arg.startswith("--workon="):
         workon_str = arg.replace("--workon=", "")
         package_list = shlex.split(" ".join(shlex.split(workon_str)))
         self.mandatory_source.update(package_list)
+      elif arg.startswith("--force-remote-binary="):
+        force_remote_binary = arg.replace("--force-remote-binary=", "")
+        force_remote_binary = \
+          shlex.split(" ".join(shlex.split(force_remote_binary)))
+        self.force_remote_binary.update(force_remote_binary)
       elif arg.startswith("--nomerge="):
         nomerge_str = arg.replace("--nomerge=", "")
         package_list = shlex.split(" ".join(shlex.split(nomerge_str)))
         self.nomerge.update(package_list)
       elif arg == "--no-workon-deps":
         self.no_workon_deps = True
       elif arg == "--rebuild":
         self.rebuild = True
       elif arg == "--show-output":
         self.show_output = True
@@ skipping 189 matching lines @@
     #
     # This is used by portage in the _reinstall_for_flags function below.
     forced_flags = set(pkgsettings.useforce).union(pkgsettings.usemask)
 
     depgraph = self.emerge.depgraph
 
     flags = depgraph._reinstall_for_flags(forced_flags, cur_use,
                                           cur_iuse, now_use, now_iuse)
     return not flags
 
-  def GenDependencyTree(self):
+  def GenDependencyTree(self, remote_pkgs):
     """Get dependency tree info from emerge.
 
     TODO(): Update cros_extract_deps to also use this code.
     Returns:
       Dependency tree
     """
     start = time.time()
 
     # Setup emerge options.
     #
     # We treat dependency info a bit differently than emerge itself. Unless
     # you're using --usepkgonly, we disable --getbinpkg and --usepkg here so
     # that emerge will look at the dependencies of the source ebuilds rather
     # than the binary dependencies. This helps ensure that we have the option
     # of merging a package from source, if we want to switch to it with
     # --workon and the dependencies have changed.
     emerge = self.emerge
     emerge_opts = emerge.opts.copy()
-    emerge_opts.pop("--getbinpkg", None)
-    if "--usepkgonly" not in emerge_opts:
-      emerge_opts.pop("--usepkg", None)
-    if self.mandatory_source or self.rebuild:
+    if self.mandatory_source or self.rebuild or self.force_remote_binary:
       # Enable --emptytree so that we get the full tree, which we need for
       # dependency analysis. By default, with this option, emerge optimizes
       # the graph by removing uninstall instructions from the graph. By
       # specifying --tree as well, we tell emerge that it's not safe to remove
       # uninstall instructions because we're planning on analyzing the output.
       emerge_opts["--tree"] = True
       emerge_opts["--emptytree"] = True
 
+      # Tell emerge not to worry about use flags yet. We handle those inside
+      # parallel_emerge itself. Further, when we use the --force-remote-binary
+      # flag, we don't emerge to reject a package just because it has different
+      # use flags.
+      emerge_opts.pop("--newuse", None)
+      emerge_opts.pop("--reinstall", None)
+
     # Create a list of packages to merge
     packages = set(emerge.cmdline_packages[:])
     if self.mandatory_source:
       packages.update(self.mandatory_source)
+    if self.force_remote_binary:
+      forced_pkgs = {}
+      for pkg in remote_pkgs:
+        category, pkgname, _, _ = portage.catpkgsplit(pkg)
+        full_pkgname = "%s/%s" % (category, pkgname)
+        if (pkgname in self.force_remote_binary or
+            full_pkgname in self.force_remote_binary):
+          forced_pkgs.setdefault(full_pkgname, []).append(pkg)
+
+      for pkgs in forced_pkgs.values():
+        forced_package = portage.versions.best(pkgs)
+        packages.add("=%s" % forced_package)
+        self.forced_remote_binary_packages.add(forced_package)
 
     # Tell emerge to be quiet. We print plenty of info ourselves so we don't
     # need any extra output from portage.
     portage.util.noiselimit = -1
 
     # My favorite feature: The silent spinner. It doesn't spin. Ever.
     # I'd disable the colors by default too, but they look kind of cool.
     emerge.spinner = stdout_spinner()
     emerge.spinner.update = emerge.spinner.update_quiet
 
@@ skipping 65 matching lines @@
         # If we're not in emptytree mode, and we're going to replace a package
         # that is already installed, then this operation is possibly optional.
         # ("--selective" mode is handled later, in RemoveInstalledPackages())
         optional = False
         if not emptytree:
           for vardb_pkg in vardb.match_pkgs(pkg.cpv):
             if self.CheckUseFlags(pkgsettings, vardb_pkg, pkg):
               optional = True
               break
 
-        # Add the package to our database.
-        self.package_db[str(pkg.cpv)] = pkg
-
         # Save off info about the package
         deps_info[str(pkg.cpv)] = {"idx": len(deps_info),
                                    "optional": optional}
 
     # Delete the --tree option, because we don't really want to display a
     # tree. We just wanted to get emerge to leave uninstall instructions on
     # the graph. Later, when we display the graph, we'll want standard-looking
     # output, so removing the --tree option is important.
     frozen_config.myopts.pop("--tree", None)
 
     seconds = time.time() - start
     if "--quiet" not in emerge.opts:
       print "Deps calculated in %dm%.1fs" % (seconds / 60, seconds % 60)
 
     return deps_tree, deps_info
 
   def PrintTree(self, deps, depth=""):
     """Print the deps we have seen in the emerge output.
 
     Args:
      deps: Dependency tree structure.
      depth: Allows printing the tree recursively, with indentation.
     """
     for entry in sorted(deps):
       action = deps[entry]["action"]
       print "%s %s (%s)" % (depth, entry, action)
       self.PrintTree(deps[entry]["deps"], depth=depth + "  ")
 
-  def GenDependencyGraph(self, deps_tree, deps_info):
+  def RemotePackageDatabase(self, binhost_url):
+    """Grab the latest binary package database from the prebuilt server.
+
+    We need to know the modification times of the prebuilt packages so that we
+    know when it is OK to use these packages and when we should rebuild them
+    instead.
+
+    Args:
+      binhost_url: Base URL of remote packages (PORTAGE_BINHOST).
+
+    Returns:
+      A dict mapping package identifiers to modification times.
+    """
+
+    if not binhost_url:
+      return {}
+
+    def retry_urlopen(url, tries=3):
+      """Open the specified url, retrying if we run into network errors.
+
+      We do not retry for HTTP errors.
+
+      Args:
+        url: The specified url.
+        tries: The number of times to try.
+
+      Returns:
+        The result of urllib2.urlopen(url).
+      """
+      for i in range(tries):
+        try:
+          return urllib2.urlopen(url)
+        except urllib2.HTTPError as e:
+          raise
+        except urllib2.URLError as e:
+          if i + 1 == tries:
+            raise
+          else:
+            print "Cannot GET %s: %s" % (url, e)
+
+    url = os.path.join(binhost_url, "Packages")
+    try:
+      f = retry_urlopen(url)
+    except urllib2.HTTPError as e:
+      if e.code == 404:
+        return {}
+      else:
+        raise
+    prebuilt_pkgs = {}
+    for line in f:
+      if line.startswith("CPV: "):
+        pkg = line.replace("CPV: ", "").rstrip()
+      elif line.startswith("MTIME: "):
+        prebuilt_pkgs[pkg] = int(line[:-1].replace("MTIME: ", ""))
+    f.close()
+
+    return prebuilt_pkgs
+
+  def GenDependencyGraph(self, deps_tree, deps_info, remote_pkgs):
     """Generate a doubly linked dependency graph.
 
     Args:
       deps_tree: Dependency tree structure.
       deps_info: More details on the dependencies.
     Returns:
       Deps graph in the form of a dict of packages, with each package
       specifying a "needs" list and "provides" list.
     """
     emerge = self.emerge
@@ skipping 28 matching lines @@
     #  - action: What we're planning on doing with this package. Generally,
     #            "merge", "nomerge", or "uninstall"
     #  - mandatory_source:
     #    If true, indicates that this package must be compiled from source.
     #    We set this for "workon" packages, and for packages where the
     #    binaries are known to be out of date.
     #  - mandatory:
     #    If true, indicates that this package must be installed. We don't care
     #    whether it's binary or source, unless the mandatory_source flag is
     #    also set.
+    #  - force_remote_binary:
+    #    If true, indicates that we want to update to the latest remote prebuilt
+    #    of this package. Packages that depend on this package should be built
+    #    from source.
     #
     deps_map = {}
 
     def ReverseTree(packages):
       """Convert tree to digraph.
 
       Take the tree of package -> requirements and reverse it to a digraph of
       buildable packages -> packages they unblock.
       Args:
         packages: Tree(s) of dependencies.
       Returns:
         Unsanitized digraph.
       """
       for pkg in packages:
 
         # Create an entry for the package
         action = packages[pkg]["action"]
         default_pkg = {"needs": {}, "provides": set(), "action": action,
-                       "mandatory_source": False, "mandatory": False}
+                       "mandatory_source": False, "mandatory": False,
+                       "force_remote_binary": False}
         this_pkg = deps_map.setdefault(pkg, default_pkg)
 
         # Create entries for dependencies of this package first.
         ReverseTree(packages[pkg]["deps"])
 
         # Add dependencies to this package.
         for dep, dep_item in packages[pkg]["deps"].iteritems():
           dep_pkg = deps_map[dep]
           dep_type = dep_item["deptype"]
           if dep_type != "runtime_post":
@@ skipping 210 matching lines @@
 
       # Mark this package as non-optional
       deps_info[pkg]["optional"] = False
       this_pkg[merge_type] = True
       for w in this_pkg["provides"].difference(rebuild_blacklist):
         MergeChildren(w, merge_type)
 
       if this_pkg["action"] == "nomerge":
         this_pkg["action"] = "merge"
 
-    def RemotePackageDatabase(binhost_url):
-      """Grab the latest binary package database from the prebuilt server.
-
-      We need to know the modification times of the prebuilt packages so that we
-      know when it is OK to use these packages and when we should rebuild them
-      instead.
-
-      Args:
-        binhost_url: Base URL of remote packages (PORTAGE_BINHOST).
-
-      Returns:
-        A dict mapping package identifiers to modification times.
-      """
-
-      if not binhost_url:
-        return {}
-
-      def retry_urlopen(url, tries=3):
-        """Open the specified url, retrying if we run into network errors.
-
-        We do not retry for HTTP errors.
-
-        Args:
-          url: The specified url.
-          tries: The number of times to try.
-
-        Returns:
-          The result of urllib2.urlopen(url).
-        """
-        for i in range(tries):
-          try:
-            return urllib2.urlopen(url)
-          except urllib2.HTTPError as e:
-            raise
-          except urllib2.URLError as e:
-            if i + 1 == tries:
-              raise
-            else:
-              print "Cannot GET %s: %s" % (url, e)
-
-      url = binhost_url + "/Packages"
-      try:
-        f = retry_urlopen(url)
-      except urllib2.HTTPError as e:
-        if e.code == 404:
-          return {}
-        else:
-          raise
-      prebuilt_pkgs = {}
-      for line in f:
-        if line.startswith("CPV: "):
-          pkg = line.replace("CPV: ", "").rstrip()
-        elif line.startswith("MTIME: "):
-          prebuilt_pkgs[pkg] = int(line[:-1].replace("MTIME: ", ""))
-      f.close()
-
-      return prebuilt_pkgs
-
     def LocalPackageDatabase():
       """Get the modification times of the packages in the local database.
 
       We need to know the modification times of the local packages so that we
       know when they need to be rebuilt.
 
       Returns:
         A dict mapping package identifiers to modification times.
       """
       if self.board:
@@ skipping 32 matching lines @@
         Args:
           pkg: The specified package.
           pkg_db: The package DB to use.
           cache: A dict, where the results are stored.
 
         Returns:
           True iff the prebuilts are ready for pkg and all deps.
         """
         if pkg in cache:
           return cache[pkg]
-        if pkg not in pkg_db:
+        if pkg not in pkg_db and pkg not in self.forced_remote_binary_packages:
           cache[pkg] = False
         else:
           cache[pkg] = True
           for dep in deps_map[pkg]["needs"]:
             if not PrebuiltsReady(dep, pkg_db, cache):
               cache[pkg] = False
               break
         return cache[pkg]
 
       def LastModifiedWithDeps(pkg, pkg_db, cache):
@@ skipping 41 matching lines @@
           # least as new as our local packages, we can install them.
           # Otherwise, we need to build from source.
           remote_mtime = LastModifiedWithDeps(pkg, remote_pkgs,
                                               remote_mtime_cache)
           remote_ready = PrebuiltsReady(pkg, remote_pkgs, remote_ready_cache)
           if remote_ready and (local_mtime <= remote_mtime or pkg in cycles):
             MergeChildren(pkg, "mandatory")
           else:
             MergeChildren(pkg, "mandatory_source")
 
-    def UsePrebuiltPackages():
+    def UsePrebuiltPackages(remote_pkgs):
       """Update packages that can use prebuilts to do so."""
       start = time.time()
 
       # The bintree is the database of binary packages. By default, it's
       # empty.
       bintree = emerge.trees[root]["bintree"]
       bindb = bintree.dbapi
       root_config = emerge.root_config
       pkgsettings = emerge.depgraph._frozen_config.pkgsettings[root]
       prebuilt_pkgs = {}
 
       # Populate the DB with packages
       bintree.populate("--getbinpkg" in emerge.opts,
                        "--getbinpkgonly" in emerge.opts)
 
       # Build list of prebuilt packages
       for pkg, info in deps_map.iteritems():
-        if info and not info["mandatory_source"] and info["action"] == "merge":
+        if info and info["action"] == "merge":
+          if (not info["force_remote_binary"] and info["mandatory_source"] or
+              "--usepkgonly" not in emerge.opts and pkg not in remote_pkgs):
+            continue
+
           db_keys = list(bindb._aux_cache_keys)
           try:
             db_vals = bindb.aux_get(pkg, db_keys + ["MTIME"])
           except KeyError:
             # No binary package
             continue
+
           mtime = int(db_vals.pop() or 0)
           metadata = zip(db_keys, db_vals)
           db_pkg = Package(built=True, cpv=pkg, installed=False,
                            metadata=metadata, onlydeps=False, mtime=mtime,
                            operation="merge", root_config=root_config,
                            type_name="binary")
           prebuilt_pkgs[pkg] = db_pkg
 
       # Calculate what packages need to be rebuilt due to changes in use flags.
       for pkg, db_pkg in prebuilt_pkgs.iteritems():
-        db_pkg_src = self.package_db[pkg]
-        if not self.CheckUseFlags(pkgsettings, db_pkg, db_pkg_src):
+        db_pkg_src = self.package_db.get(pkg)
+        if db_pkg_src and not self.CheckUseFlags(pkgsettings, db_pkg,
+                                                 db_pkg_src):
           MergeChildren(pkg, "mandatory_source")
 
       # Convert eligible packages to binaries.
       for pkg, info in deps_map.iteritems():
-        if (info and not info["mandatory_source"] and
-            info["action"] == "merge" and pkg in prebuilt_pkgs):
-          self.package_db[pkg] = prebuilt_pkgs[pkg]
+        if info and info["action"] == "merge" and pkg in prebuilt_pkgs:
+          if not info["mandatory_source"] or info["force_remote_binary"]:
+            self.package_db[pkg] = prebuilt_pkgs[pkg]
 
       seconds = time.time() - start
       if "--quiet" not in emerge.opts:
         print "Prebuilt DB populated in %dm%.1fs" % (seconds / 60, seconds % 60)
 
       return prebuilt_pkgs
 
     def AddRemainingPackages():
       """Fill in packages that don't have entries in the package db.
 
       Every package we are installing needs an entry in the package db.
       This function should only be called after we have removed the
       packages that are not being merged from our deps_map.
       """
       for pkg in deps_map:
         if pkg not in self.package_db:
           if deps_map[pkg]["action"] != "merge":
             # We should only fill in packages that are being merged. If
             # there's any other packages here, something funny is going on.
             print "Missing entry for %s in package db" % pkg
             sys.exit(1)
 
           db_pkg = emerge.depgraph._pkg(pkg, "ebuild", emerge.root_config)
           self.package_db[pkg] = db_pkg
 
     ReverseTree(deps_tree)
     BuildFinalPackageSet()
     AddSecretDeps()
 
+    # Mark that we want to use remote binaries only for a particular package.
+    vardb = emerge.depgraph._frozen_config.trees[root]["vartree"].dbapi
+    for pkg in self.force_remote_binary:
+      for db_pkg in final_db.match_pkgs(pkg):
+        match = deps_map.get(str(db_pkg.cpv))
+        if match:
+          match["force_remote_binary"] = True
+
+          rebuild_blacklist.add(str(db_pkg.cpv))
+          if not vardb.match_pkgs(db_pkg.cpv):
+            MergeChildren(str(db_pkg.cpv), "mandatory")
+
     if self.no_workon_deps:
       for pkg in self.mandatory_source.copy():
         for db_pkg in final_db.match_pkgs(pkg):
           deps_map[str(db_pkg.cpv)]["mandatory_source"] = True
     else:
       for pkg in self.mandatory_source.copy():
         for db_pkg in final_db.match_pkgs(pkg):
           MergeChildren(str(db_pkg.cpv), "mandatory_source")
 
       cycles = FindCycles()
       if self.rebuild:
         local_pkgs = LocalPackageDatabase()
-        remote_pkgs = {}
-        if "--getbinpkg" in emerge.opts:
-          binhost = emerge.settings["PORTAGE_BINHOST"]
-          remote_pkgs = RemotePackageDatabase(binhost)
         AutoRebuildDeps(local_pkgs, remote_pkgs, cycles)
 
     # We need to remove installed packages so that we can use the dependency
     # ordering of the install process to show us what cycles to crack. Once
     # we've done that, we also need to recalculate our list of cycles so that
     # we don't include the installed packages in our cycles.
     RemoveInstalledPackages()
     SanitizeTree()
     if deps_map:
       if "--usepkg" in emerge.opts:
-        UsePrebuiltPackages()
+        UsePrebuiltPackages(remote_pkgs)
       AddRemainingPackages()
     return deps_map
 
   def PrintInstallPlan(self, deps_map):
     """Print an emerge-style install plan.
 
     The install plan lists what packages we're installing, in order.
     It's useful for understanding what parallel_emerge is doing.
 
     Args:
@@ skipping 533 matching lines @@
   if "--quiet" not in emerge.opts:
     cmdline_packages = " ".join(emerge.cmdline_packages)
     nomerge_packages = " ".join(deps.nomerge)
     print "Starting fast-emerge."
     print " Building package %s on %s" % (cmdline_packages,
                                           deps.board or "root")
     if nomerge_packages:
       print " Skipping package %s on %s" % (nomerge_packages,
                                             deps.board or "root")
 
-  deps_tree, deps_info = deps.GenDependencyTree()
+  remote_pkgs = {}
+  if "--getbinpkg" in emerge.opts:
+    binhost = emerge.settings["PORTAGE_BINHOST"]
+    remote_pkgs = deps.RemotePackageDatabase(binhost)
+
+  deps_tree, deps_info = deps.GenDependencyTree(remote_pkgs)
 
   # You want me to be verbose? I'll give you two trees! Twice as much value.
   if "--tree" in emerge.opts and "--verbose" in emerge.opts:
     deps.PrintTree(deps_tree)
 
-  deps_graph = deps.GenDependencyGraph(deps_tree, deps_info)
+  deps_graph = deps.GenDependencyGraph(deps_tree, deps_info, remote_pkgs)
 
   # OK, time to print out our progress so far.
   deps.PrintInstallPlan(deps_graph)
   if "--tree" in emerge.opts:
     PrintDepsMap(deps_graph)
 
   # Are we upgrading portage? If so, and there are more packages to merge,
   # schedule a restart of parallel_emerge to merge the rest. This ensures that
   # we pick up all updates to portage settings before merging any more
   # packages.
@@ skipping 35 matching lines @@
   # need to upgrade the rest of the packages. So we'll go ahead and do that.
   if portage_upgrade:
     args = sys.argv[1:] + ["--nomerge=sys-apps/portage"]
     os.execvp(os.path.realpath(sys.argv[0]), args)
 
   print "Done"
   sys.exit(0)
 
 if __name__ == "__main__":
   main()