downloader.py

# Copyright (c) 2013 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.

import os
import threading

import build_artifact
import common_util
import gsutil_util
import log_util


class Downloader(log_util.Loggable):
  """Downloader of images to the devsever.

  Given a URL to a build on the archive server:
    - Caches that build and the given artifacts onto the devserver.
    - May also initiate caching of related artifacts in the background.

  Private class members:
    archive_url: a URL where to download build artifacts from.
    static_dir: local filesystem directory to store all artifacts.
    build_dir: the local filesystem directory to store artifacts for the given
      build defined by the archive_url.
  """

  # This filename must be kept in sync with clean_staged_images.py
  _TIMESTAMP_FILENAME = 'staged.timestamp'

  def __init__(self, static_dir, archive_url):
    super(Downloader, self).__init__()
    self._archive_url = archive_url
    self._static_dir = static_dir
    self._build_dir = Downloader.GetBuildDir(static_dir, archive_url)

  @staticmethod
  def ParseUrl(archive_url):
    """Parses archive_url into rel_path and build.

    Parses archive_url into rel_path and build e.g.
    gs://chromeos-image-archive/{rel_path}/{build}.

    Args:
      archive_url: a URL at which build artifacts are archived.

    Returns:
      A tuple of (build relative path, short build name)
    """
    # The archive_url is of the form gs://server/[some_path/target]/...]/build
    # This function discards 'gs://server/' and extracts the [some_path/target]
    # as rel_path and the build as build.
    sub_url = archive_url.partition('://')[2]
    split_sub_url = sub_url.split('/')
    rel_path = '/'.join(split_sub_url[1:-1])
    build = split_sub_url[-1]
    return rel_path, build

  @staticmethod
  def GetBuildDir(static_dir, archive_url):
    """Returns the path to where the artifacts will be staged.

    Args:
      static_dir: The base static dir that will be used.
      archive_url: The gs path to the archive url.
    """
    # Parse archive_url into rel_path (contains the build target) and
    # build e.g. gs://chromeos-image-archive/{rel_path}/{build}.
    rel_path, build = Downloader.ParseUrl(archive_url)
    return os.path.join(static_dir, rel_path, build)

  @staticmethod
  def _TouchTimestampForStaged(directory_path):
    file_name = os.path.join(directory_path, Downloader._TIMESTAMP_FILENAME)
    # Easiest python version of |touch file_name|
    with file(file_name, 'a'):
      os.utime(file_name, None)

  @staticmethod
  def _TryRemoveStageDir(directory_path):
    """If download failed with GSUtilError, try to remove the stage dir.

    If the download attempt failed with GSUtilError and staged.timestamp is the
    only file in that directory. The build could be non-existing, and the
    directory should be removed.

    @param directory_path: directory used to stage the image.

    """
    file_name = os.path.join(directory_path, Downloader._TIMESTAMP_FILENAME)
    if os.path.exists(file_name) and len(os.listdir(directory_path)) == 1:
      os.remove(file_name)
      os.rmdir(directory_path)

  def Download(self, artifacts, async=False):
    """Downloads and caches the |artifacts|.

    Downloads and caches the |artifacts|. Returns once these
    are present on the devserver. A call to this will attempt to cache
    non-specified artifacts in the background following the principle of
    spatial locality.

    @params artifacts: A list of artifact names that correspond to artifacts to
                       stage.
    @params async: True to return without waiting for download to complete.

    """
    common_util.MkDirP(self._build_dir)

    # We are doing some work on this build -- let's touch it to indicate that
    # we shouldn't be cleaning it up anytime soon.
    Downloader._TouchTimestampForStaged(self._build_dir)

    # Create factory to create build_artifacts from artifact names.
    build = self.ParseUrl(self._archive_url)[1]
    factory = build_artifact.ArtifactFactory(self._build_dir, self._archive_url,
                                             artifacts, build)
    background_artifacts = factory.OptionalArtifacts()
    if background_artifacts:
      self._DownloadArtifactsInBackground(background_artifacts)

    required_artifacts = factory.RequiredArtifacts()
    str_repr = [str(a) for a in required_artifacts]
    self._Log('Downloading artifacts %s.', ' '.join(str_repr))
    try:
      if async:
        self._DownloadArtifactsInBackground(required_artifacts)
      else:
        self._DownloadArtifactsSerially(required_artifacts, no_wait=True)
    except gsutil_util.GSUtilError:
      Downloader._TryRemoveStageDir(self._build_dir)
      raise

  def IsStaged(self, artifacts):
    """Check if all artifacts have been downloaded.

    @param artifacts: A list of artifacts to be checked.
    @returns: True if all artifacts are staged.

    """
    # Create factory to create build_artifacts from artifact names.
    build = self.ParseUrl(self._archive_url)[1]
    factory = build_artifact.ArtifactFactory(self._build_dir, self._archive_url,
                                             artifacts, build)
    required_artifacts = factory.RequiredArtifacts()
    return all([artifact.ArtifactStaged() for artifact in required_artifacts])

  def _DownloadArtifactsSerially(self, artifacts, no_wait):
    """Simple function to download all the given artifacts serially.

    @param artifacts: A list of build_artifact.BuildArtifact instances to
                      download.
    @param no_wait: If True, don't block waiting for artifact to exist if we
                    fail to immediately find it.

    """
    for artifact in artifacts:
      artifact.Process(no_wait)

  def _DownloadArtifactsInBackground(self, artifacts):
    """Downloads |artifacts| in the background.

    Downloads |artifacts| in the background. As these are backgrounded
    artifacts, they are done best effort and may not exist.

    Args:
      artifacts: List of build_artifact.BuildArtifact instances to download.
    """
    self._Log('Invoking background download of artifacts for %r', artifacts)
    thread = threading.Thread(target=self._DownloadArtifactsSerially,
                              args=(artifacts, False))
    thread.start()