cros_sdk: Support --board for speeding up initialization.

Adds an optional argument --board to cros_sdk. When used with --create,
this fetches and unpacks a board-specific overlay (toolchains etc) on
top of the vanilla SDK one, which in turn may speed up things later when
setting up the chroot for building this board. Otherwise, this shouldn't
have any detrimental effect on the chroot, nor does it restrict its
generality for use with other boards. This option is ignored if using
--bootstrap.

CQ-DEPEND=CL:274073
BUG=brillo:918
TEST=cros_sdk --create --board link works (contrived overlay tarball).

Change-Id: Ife7e35952a5d258e9613f903e190b656ffffeb2f
Reviewed-on: https://chromium-review.googlesource.com/274074
Reviewed-by: Don Garrett <dgarrett@chromium.org>
Commit-Queue: Gilad Arnold <garnold@chromium.org>
Tested-by: Gilad Arnold <garnold@chromium.org>
diff --git a/scripts/cros_sdk.py b/scripts/cros_sdk.py
index 106dff8..4eb18b7 100644
--- a/scripts/cros_sdk.py
+++ b/scripts/cros_sdk.py
@@ -69,22 +69,46 @@
           for ext in COMPRESSION_PREFERENCE]
 
 
-def FetchRemoteTarballs(storage_dir, urls):
+def GetBoardOverlayUrls(version, board):
+  """Returns the URL(s) for the board-specific toolchain overlay.
+
+  Args:
+    version: The SDK version used, e.g. 2015.05.27.145939. We use the year and
+        month components to point to a subdirectory on the SDK bucket where
+        board overlays are stored (.../2015/05/ in this case).
+    board: The board name.
+
+  Returns:
+    List of alternative download URLs for the board's SDK overlay tarball.
+  """
+  suburl_template = os.path.join(
+      *(version.split('.')[:2] +
+        ['cros-sdk-overlay-%s-%s.tar.%%s' % (board, version)]))
+  return [toolchain.GetSdkURL(suburl=suburl_template % ext)
+          for ext in COMPRESSION_PREFERENCE]
+
+
+def FetchRemoteTarballs(storage_dir, urls, desc, allow_none=False):
   """Fetches a tarball given by url, and place it in |storage_dir|.
 
   Args:
     storage_dir: Path where to save the tarball.
     urls: List of URLs to try to download. Download will stop on first success.
+    desc: A string describing what's the thing we're downloading (for logging).
+    allow_none: Don't fail if none of the URLs worked.
 
   Returns:
-    Full path to the downloaded file
+    Full path to the downloaded file, or None if |allow_none| and no URL worked.
+
+  Raises:
+    ValueError: If |allow_none| is False and none of the URLs worked.
   """
 
   # Note we track content length ourselves since certain versions of curl
   # fail if asked to resume a complete file.
   # pylint: disable=C0301,W0631
   # https://sourceforge.net/tracker/?func=detail&atid=100976&aid=3482927&group_id=976
-  logging.notice('Downloading chroot files.')
+  logging.notice('Downloading %s...', desc)
   for url in urls:
     # http://www.logilab.org/ticket/8766
     # pylint: disable=E1101
@@ -112,7 +136,9 @@
     if successful:
       break
   else:
-    raise Exception('No valid URLs found!')
+    if allow_none:
+      return None
+    raise ValueError('No valid URLs found!')
 
   tarball_dest = os.path.join(storage_dir, tarball_name)
   current_size = 0
@@ -128,26 +154,34 @@
         print_cmd=False)
 
   # Cleanup old tarballs now since we've successfull fetched; only cleanup
-  # the tarballs for our prefix, or unknown ones.
-  ignored_prefix = ('stage3-' if tarball_name.startswith('cros-sdk-')
-                    else 'cros-sdk-')
+  # the tarballs for our prefix, or unknown ones. This gets a bit tricky
+  # because we might have partial overlap between known prefixes.
+  my_prefix = tarball_name.rsplit('-', 1)[0] + '-'
+  all_prefixes = ('stage3-amd64-', 'cros-sdk-', 'cros-sdk-overlay-')
+  ignored_prefixes = [prefix for prefix in all_prefixes if prefix != my_prefix]
   for filename in os.listdir(storage_dir):
-    if filename == tarball_name or filename.startswith(ignored_prefix):
+    if (filename == tarball_name or
+        any([(filename.startswith(p) and
+              not (len(my_prefix) > len(p) and filename.startswith(my_prefix)))
+             for p in ignored_prefixes])):
       continue
-
-    logging.info('Cleaning up old tarball: %s' % (filename,))
+    logging.info('Cleaning up old tarball: %s', filename)
     osutils.SafeUnlink(os.path.join(storage_dir, filename))
 
   return tarball_dest
 
 
-def CreateChroot(chroot_path, sdk_tarball, cache_dir, nousepkg=False,
-                 workspace=None):
+def CreateChroot(chroot_path, sdk_tarball, board_overlay_tarball, cache_dir,
+                 nousepkg=False, workspace=None):
   """Creates a new chroot from a given SDK"""
 
   cmd = MAKE_CHROOT + ['--stage3_path', sdk_tarball,
                        '--chroot', chroot_path,
                        '--cache_dir', cache_dir]
+
+  if board_overlay_tarball:
+    cmd.extend(['--board_overlay_path', board_overlay_tarball])
+
   if nousepkg:
     cmd.append('--nousepkg')
 
@@ -463,6 +497,13 @@
                           % (sdk_latest_version, bootstrap_latest_version)))
   parser.add_option('--workspace', default=None,
                     help='Workspace directory to mount into the chroot.')
+  parser.add_option('--board', default=None,
+                    help='The board we intend to be building in the chroot. '
+                    'Used for downloading an SDK board overlay (if one is '
+                    'found), which may speed up chroot initialization when '
+                    'building for the first time. Otherwise this has no effect '
+                    'and will not restrict the chroot in any way. Ignored if '
+                    'using --bootstrap.')
 
   # Commands.
   group = parser.add_option_group('Commands')
@@ -577,7 +618,7 @@
   if options.buildbot_log_version:
     cros_build_lib.PrintBuildbotStepText(sdk_version)
 
-  # Based on selections, fetch the tarball.
+  # Based on selections, determine the tarball to fetch.
   if options.sdk_url:
     urls = [options.sdk_url]
   elif options.bootstrap:
@@ -585,11 +626,17 @@
   else:
     urls = GetArchStageTarballs(sdk_version)
 
+  # Get URLs for the board-specific overlay, if one is to be used.
+  overlay_urls = None
+  if options.board and not options.bootstrap:
+    overlay_urls = GetBoardOverlayUrls(sdk_version, options.board)
+
   lock_path = os.path.dirname(options.chroot)
   lock_path = os.path.join(
       lock_path, '.%s_lock' % os.path.basename(options.chroot).lstrip('.'))
   with cgroups.SimpleContainChildren('cros_sdk', pid=first_pid):
     with locking.FileLock(lock_path, 'chroot lock') as lock:
+      board_overlay_tarball = None
 
       if options.proxy_sim:
         _ProxySimSetup(options)
@@ -629,11 +676,17 @@
 
       if options.download:
         lock.write_lock()
-        sdk_tarball = FetchRemoteTarballs(sdk_cache, urls)
+        desc = 'stage3 tarball' if options.bootstrap else 'SDK tarball'
+        sdk_tarball = FetchRemoteTarballs(sdk_cache, urls, desc)
+        if overlay_urls:
+          board_overlay_tarball = FetchRemoteTarballs(sdk_cache, overlay_urls,
+                                                      'board overlay tarball',
+                                                      allow_none=True)
 
       if options.create:
         lock.write_lock()
-        CreateChroot(options.chroot, sdk_tarball, options.cache_dir,
+        CreateChroot(options.chroot, sdk_tarball, board_overlay_tarball,
+                     options.cache_dir,
                      nousepkg=(options.bootstrap or options.nousepkg),
                      workspace=options.workspace)