quay/buildman/component/buildcomponent.py

import datetime
import time
import logging
import json
import trollius
import re

from autobahn.wamp.exception import ApplicationError

from buildman.server import BuildJobResult
from buildman.component.basecomponent import BaseComponent
from buildman.jobutil.buildjob import BuildJobLoadException
from buildman.jobutil.buildstatus import StatusHandler
from buildman.jobutil.workererror import WorkerError

from data.database import BUILD_PHASE

HEARTBEAT_DELTA = datetime.timedelta(seconds=30)
HEARTBEAT_TIMEOUT = 10
INITIAL_TIMEOUT = 25

SUPPORTED_WORKER_VERSIONS = ['0.3']

logger = logging.getLogger(__name__)

class ComponentStatus(object):
  """ ComponentStatus represents the possible states of a component. """
  JOINING = 'joining'
  WAITING = 'waiting'
  RUNNING = 'running'
  BUILDING = 'building'
  TIMED_OUT = 'timeout'

class BuildComponent(BaseComponent):
  """ An application session component which conducts one (or more) builds. """
  def __init__(self, config, realm=None, token=None, **kwargs):
    self.expected_token = token
    self.builder_realm = realm

    self.parent_manager = None
    self.registry_hostname = None

    self._component_status = ComponentStatus.JOINING
    self._last_heartbeat = None
    self._current_job = None
    self._build_status = None
    self._image_info = None
    self._worker_version = None

    BaseComponent.__init__(self, config, **kwargs)

  def kind(self):
    return 'builder'

  def onConnect(self):
    self.join(self.builder_realm)

  def onJoin(self, details):
    logger.debug('Registering methods and listeners for component %s', self.builder_realm)
    yield trollius.From(self.register(self._on_ready, u'io.quay.buildworker.ready'))
    yield trollius.From(self.register(self._determine_cache_tag,
                                      u'io.quay.buildworker.determinecachetag'))
    yield trollius.From(self.register(self._ping, u'io.quay.buildworker.ping'))

    yield trollius.From(self.subscribe(self._on_heartbeat, 'io.quay.builder.heartbeat'))
    yield trollius.From(self.subscribe(self._on_log_message, 'io.quay.builder.logmessage'))

    yield trollius.From(self._set_status(ComponentStatus.WAITING))

  def is_ready(self):
    """ Determines whether a build component is ready to begin a build. """
    return self._component_status == ComponentStatus.RUNNING

  @trollius.coroutine
  def start_build(self, build_job):
    """ Starts a build. """
    logger.debug('Starting build for component %s (worker version: %s)',
                 self.builder_realm, self._worker_version)

    self._current_job = build_job
    self._build_status = StatusHandler(self.build_logs, build_job.repo_build.uuid)
    self._image_info = {}

    yield trollius.From(self._set_status(ComponentStatus.BUILDING))

    # Send the notification that the build has started.
    build_job.send_notification('build_start')

    # Parse the build configuration.
    try:
      build_config = build_job.build_config
    except BuildJobLoadException as irbe:
      self._build_failure('Could not load build job information', irbe)

    base_image_information = {}
    buildpack_url = self.user_files.get_file_url(build_job.repo_build.resource_key,
                                                 requires_cors=False)

    # Add the pull robot information, if any.
    if build_job.pull_credentials:
      base_image_information['username'] = build_job.pull_credentials.get('username', '')
      base_image_information['password'] = build_job.pull_credentials.get('password', '')

    # Retrieve the repository's fully qualified name.
    repo = build_job.repo_build.repository
    repository_name = repo.namespace_user.username + '/' + repo.name

    # Parse the build queue item into build arguments.
    #  build_package: URL to the build package to download and untar/unzip.
    #  sub_directory: The location within the build package of the Dockerfile and the build context.
    #  repository: The repository for which this build is occurring.
    #  registry: The registry for which this build is occuring (e.g. 'quay.io', 'staging.quay.io').
    #  pull_token: The token to use when pulling the cache for building.
    #  push_token: The token to use to push the built image.
    #  tag_names: The name(s) of the tag(s) for the newly built image.
    #  base_image: The image name and credentials to use to conduct the base image pull.
    #   repository: The repository to pull (DEPRECATED 0.2)
    #   tag: The tag to pull (DEPRECATED in 0.2)
    #   username: The username for pulling the base image (if any).
    #   password: The password for pulling the base image (if any).
    build_arguments = {
        'build_package': buildpack_url,
        'sub_directory': build_config.get('build_subdir', ''),
        'repository': repository_name,
        'registry': self.registry_hostname,
        'pull_token': build_job.repo_build.access_token.code,
        'push_token': build_job.repo_build.access_token.code,
        'tag_names': build_config.get('docker_tags', ['latest']),
        'base_image': base_image_information,
    }

    # If the trigger has a private key, it's using git, thus we should add
    # git data to the build args.
    if build_job.repo_build.trigger.private_key is not None:
      build_arguments['git'] = {
        'url': build_config['trigger_metadata']['git_url'],
        'sha': build_config['trigger_metadata']['commit_sha'],
        'private_key': build_job.repo_build.trigger.private_key,
      }

    # Invoke the build.
    logger.debug('Invoking build: %s', self.builder_realm)
    logger.debug('With Arguments: %s', build_arguments)

    self.call("io.quay.builder.build", **build_arguments).add_done_callback(self._build_complete)

  @staticmethod
  def _total_completion(statuses, total_images):
    """ Returns the current amount completion relative to the total completion of a build. """
    percentage_with_sizes = float(len(statuses.values())) / total_images
    sent_bytes = sum([status['current'] for status in statuses.values()])
    total_bytes = sum([status['total'] for status in statuses.values()])
    return float(sent_bytes) / total_bytes * percentage_with_sizes

  @staticmethod
  def _process_pushpull_status(status_dict, current_phase, docker_data, images):
    """ Processes the status of a push or pull by updating the provided status_dict and images. """
    if not docker_data:
      return

    num_images = 0
    status_completion_key = ''

    if current_phase == 'pushing':
      status_completion_key = 'push_completion'
      num_images = status_dict['total_commands']
    elif current_phase == 'pulling':
      status_completion_key = 'pull_completion'
    elif current_phase == 'priming-cache':
      status_completion_key = 'cache_completion'
    else:
      return

    if 'progressDetail' in docker_data and 'id' in docker_data:
      image_id = docker_data['id']
      detail = docker_data['progressDetail']

      if 'current' in detail and 'total' in detail:
        images[image_id] = detail
        status_dict[status_completion_key] = \
          BuildComponent._total_completion(images, max(len(images), num_images))

  def _on_log_message(self, phase, json_data):
    """ Tails log messages and updates the build status. """
    # Parse any of the JSON data logged.
    docker_data = {}
    if json_data:
      try:
        docker_data = json.loads(json_data)
      except ValueError:
        pass

    # Extract the current status message (if any).
    fully_unwrapped = ''
    keys_to_extract = ['error', 'status', 'stream']
    for key in keys_to_extract:
      if key in docker_data:
        fully_unwrapped = docker_data[key]
        break

    # Determine if this is a step string.
    current_step = None
    current_status_string = str(fully_unwrapped.encode('utf-8'))

    if current_status_string and phase == BUILD_PHASE.BUILDING:
      step_increment = re.search(r'Step ([0-9]+) :', current_status_string)
      if step_increment:
        current_step = int(step_increment.group(1))

    # Parse and update the phase and the status_dict. The status dictionary contains
    # the pull/push progress, as well as the current step index.
    with self._build_status as status_dict:
      if self._build_status.set_phase(phase):
        logger.debug('Build %s has entered a new phase: %s', self.builder_realm, phase)

      BuildComponent._process_pushpull_status(status_dict, phase, docker_data, self._image_info)

      # If the current message represents the beginning of a new step, then update the
      # current command index.
      if current_step is not None:
        status_dict['current_command'] = current_step

      # If the json data contains an error, then something went wrong with a push or pull.
      if 'error' in docker_data:
        self._build_status.set_error(docker_data['error'])

    if current_step is not None:
      self._build_status.set_command(current_status_string)
    elif phase == BUILD_PHASE.BUILDING:
      self._build_status.append_log(current_status_string)

  @trollius.coroutine
  def _determine_cache_tag(self, command_comments, base_image_name, base_image_tag, base_image_id):
    with self._build_status as status_dict:
      status_dict['total_commands'] = len(command_comments) + 1

    logger.debug('Checking cache on realm %s. Base image: %s:%s (%s)', self.builder_realm,
                 base_image_name, base_image_tag, base_image_id)

    tag_found = self._current_job.determine_cached_tag(base_image_id, command_comments)
    raise trollius.Return(tag_found or '')

  def _build_failure(self, error_message, exception=None):
    """ Handles and logs a failed build. """
    self._build_status.set_error(error_message, {
        'internal_error': str(exception) if exception else None
    })

    build_id = self._current_job.repo_build.uuid
    logger.warning('Build %s failed with message: %s', build_id, error_message)

    # Mark that the build has finished (in an error state)
    trollius.async(self._build_finished(BuildJobResult.ERROR))

  def _build_complete(self, result):
    """ Wraps up a completed build. Handles any errors and calls self._build_finished. """
    try:
      # Retrieve the result. This will raise an ApplicationError on any error that occurred.
      result_value = result.result()
      kwargs = {}

      # Note: If we are hitting an older builder that didn't return ANY map data, then the result
      # value will be a bool instead of a proper CallResult object (because autobahn sucks).
      # Therefore: we have a try-except guard here to ensure we don't hit this pitfall.
      try:
        kwargs = result_value.kwresults
      except:
        pass

      self._build_status.set_phase(BUILD_PHASE.COMPLETE)
      trollius.async(self._build_finished(BuildJobResult.COMPLETE))

      # Send the notification that the build has completed successfully.
      self._current_job.send_notification('build_success', image_id=kwargs.get('image_id'))
    except ApplicationError as aex:
      worker_error = WorkerError(aex.error, aex.kwargs.get('base_error'))

      # Write the error to the log.
      self._build_status.set_error(worker_error.public_message(), worker_error.extra_data(),
                                   internal_error=worker_error.is_internal_error(),
                                   requeued=self._current_job.has_retries_remaining())

      # Send the notification that the build has failed.
      self._current_job.send_notification('build_failure',
                                          error_message=worker_error.public_message())

      # Mark the build as completed.
      if worker_error.is_internal_error():
        trollius.async(self._build_finished(BuildJobResult.INCOMPLETE))
      else:
        trollius.async(self._build_finished(BuildJobResult.ERROR))

  @trollius.coroutine
  def _build_finished(self, job_status):
    """ Alerts the parent that a build has completed and sets the status back to running. """
    yield trollius.From(self.parent_manager.job_completed(self._current_job, job_status, self))
    self._current_job = None

    # Set the component back to a running state.
    yield trollius.From(self._set_status(ComponentStatus.RUNNING))

  @staticmethod
  def _ping():
    """ Ping pong. """
    return 'pong'

  @trollius.coroutine
  def _on_ready(self, token, version):
    self._worker_version = version

    if not version in SUPPORTED_WORKER_VERSIONS:
      logger.warning('Build component (token "%s") is running an out-of-date version: %s', token,
                     version)
      raise trollius.Return(False)

    if self._component_status != ComponentStatus.WAITING:
      logger.warning('Build component (token "%s") is already connected', self.expected_token)
      raise trollius.Return(False)

    if token != self.expected_token:
      logger.warning('Builder token mismatch. Expected: "%s". Found: "%s"', self.expected_token,
                     token)
      raise trollius.Return(False)

    yield trollius.From(self._set_status(ComponentStatus.RUNNING))

    # Start the heartbeat check and updating loop.
    loop = trollius.get_event_loop()
    loop.create_task(self._heartbeat())
    logger.debug('Build worker %s is connected and ready', self.builder_realm)
    raise trollius.Return(True)

  @trollius.coroutine
  def _set_status(self, phase):
    if phase == ComponentStatus.RUNNING:
      yield trollius.From(self.parent_manager.build_component_ready(self))

    self._component_status = phase

  def _on_heartbeat(self):
    """ Updates the last known heartbeat. """
    self._last_heartbeat = datetime.datetime.utcnow()

  @trollius.coroutine
  def _heartbeat(self):
    """ Coroutine that runs every HEARTBEAT_TIMEOUT seconds, both checking the worker's heartbeat
        and updating the heartbeat in the build status dictionary (if applicable). This allows
        the build system to catch crashes from either end.
    """
    yield trollius.From(trollius.sleep(INITIAL_TIMEOUT))

    while True:
      # If the component is no longer running or actively building, nothing more to do.
      if (self._component_status != ComponentStatus.RUNNING and
          self._component_status != ComponentStatus.BUILDING):
        raise trollius.Return()

      # If there is an active build, write the heartbeat to its status.
      build_status = self._build_status
      if build_status is not None:
        with build_status as status_dict:
          status_dict['heartbeat'] = int(time.time())

      # Mark the build item.
      current_job = self._current_job
      if current_job is not None:
        yield trollius.From(self.parent_manager.job_heartbeat(current_job))

      # Check the heartbeat from the worker.
      logger.debug('Checking heartbeat on realm %s', self.builder_realm)
      if (self._last_heartbeat and
          self._last_heartbeat < datetime.datetime.utcnow() - HEARTBEAT_DELTA):
        yield trollius.From(self._timeout())
        raise trollius.Return()

      yield trollius.From(trollius.sleep(HEARTBEAT_TIMEOUT))

  @trollius.coroutine
  def _timeout(self):
    if self._component_status == ComponentStatus.TIMED_OUT:
      raise trollius.Return()

    yield trollius.From(self._set_status(ComponentStatus.TIMED_OUT))
    logger.warning('Build component with realm %s has timed out', self.builder_realm)

    # If we still have a running job, then it has not completed and we need to tell the parent
    # manager.
    if self._current_job is not None:
      self._build_status.set_error('Build worker timed out', internal_error=True,
                                   requeued=self._current_job.has_retries_remaining())

      self.parent_manager.job_completed(self._current_job, BuildJobResult.INCOMPLETE, self)
      self._build_status = None
      self._current_job = None

    # Unregister the current component so that it cannot be invoked again.
    self.parent_manager.build_component_disposed(self, True)