quay/buildman/manager/ephemeral.py

import logging
import etcd
import uuid
import calendar
import os.path
import json

from datetime import datetime, timedelta
from trollius import From, coroutine, Return, async
from concurrent.futures import ThreadPoolExecutor
from urllib3.exceptions import ReadTimeoutError, ProtocolError

from buildman.manager.basemanager import BaseManager
from buildman.manager.executor import PopenExecutor, EC2Executor
from buildman.component.buildcomponent import BuildComponent
from buildman.jobutil.buildjob import BuildJob
from buildman.asyncutil import AsyncWrapper
from util.morecollections import AttrDict


logger = logging.getLogger(__name__)


ETCD_DISABLE_TIMEOUT = 0


class EtcdAction(object):
  GET = 'get'
  SET = 'set'
  EXPIRE = 'expire'
  UPDATE = 'update'
  DELETE = 'delete'
  CREATE = 'create'
  COMPARE_AND_SWAP = 'compareAndSwap'
  COMPARE_AND_DELETE = 'compareAndDelete'


class EphemeralBuilderManager(BaseManager):
  """ Build manager implementation for the Enterprise Registry. """
  _executors = {
      'popen': PopenExecutor,
      'ec2': EC2Executor,
  }

  _etcd_client_klass = etcd.Client

  def __init__(self, *args, **kwargs):
    self._shutting_down = False

    self._manager_config = None
    self._async_thread_executor = None
    self._etcd_client = None

    self._etcd_realm_prefix = None
    self._etcd_builder_prefix = None

    self._component_to_job = {}
    self._job_uuid_to_component = {}
    self._component_to_builder = {}

    self._executor = None

    # Map of etcd keys being watched to the tasks watching them
    self._watch_tasks = {}

    super(EphemeralBuilderManager, self).__init__(*args, **kwargs)

  def _watch_etcd(self, etcd_key, change_callback, recursive=True):
    watch_task_key = (etcd_key, recursive)
    def callback_wrapper(changed_key_future):
      if watch_task_key not in self._watch_tasks or self._watch_tasks[watch_task_key].done():
        self._watch_etcd(etcd_key, change_callback)

      if changed_key_future.cancelled():
        # Due to lack of interest, tomorrow has been cancelled
        return

      try:
        etcd_result = changed_key_future.result()
      except (ReadTimeoutError, ProtocolError, etcd.EtcdException):
        return

      change_callback(etcd_result)

    if not self._shutting_down:
      watch_future = self._etcd_client.watch(etcd_key, recursive=recursive,
                                             timeout=ETCD_DISABLE_TIMEOUT)
      watch_future.add_done_callback(callback_wrapper)
      logger.debug('Scheduling watch of key: %s%s', etcd_key, '/*' if recursive else '')
      self._watch_tasks[watch_task_key] = async(watch_future)

  def _handle_builder_expiration(self, etcd_result):
    if etcd_result is None:
      return

    if etcd_result.action == EtcdAction.EXPIRE:
      # Handle the expiration
      logger.debug('Builder expired, clean up the old build node')
      job_metadata = json.loads(etcd_result._prev_node.value)

      if 'builder_id' in job_metadata:
        logger.info('Terminating expired build node.')
        async(self._executor.stop_builder(job_metadata['builder_id']))

  def _handle_realm_change(self, etcd_result):
    if etcd_result is None:
      return

    if etcd_result.action == EtcdAction.CREATE:
      # We must listen on the realm created by ourselves or another worker
      realm_spec = json.loads(etcd_result.value)
      self._register_realm(realm_spec)

    elif etcd_result.action == EtcdAction.DELETE or etcd_result.action == EtcdAction.EXPIRE:
      # We must stop listening for new connections on the specified realm, if we did not get the
      # connection
      realm_spec = json.loads(etcd_result._prev_node.value)
      build_job = BuildJob(AttrDict(realm_spec['job_queue_item']))
      component = self._job_uuid_to_component.pop(build_job.job_details['build_uuid'], None)
      if component is not None:
        # We were not the manager which the worker connected to, remove the bookkeeping for it
        logger.debug('Unregistering unused component on realm: %s', realm_spec['realm'])
        del self._component_to_job[component]
        del self._component_to_builder[component]
        self.unregister_component(component)

    else:
      logger.warning('Unexpected action (%s) on realm key: %s', etcd_result.action, etcd_result.key)

  def _register_realm(self, realm_spec):
    logger.debug('Registering realm with manager: %s', realm_spec['realm'])
    component = self.register_component(realm_spec['realm'], BuildComponent,
                                        token=realm_spec['token'])
    build_job = BuildJob(AttrDict(realm_spec['job_queue_item']))
    self._component_to_job[component] = build_job
    self._component_to_builder[component] = realm_spec['builder_id']
    self._job_uuid_to_component[build_job.job_details['build_uuid']] = component

  @coroutine
  def _register_existing_realms(self):
    try:
      all_realms = yield From(self._etcd_client.read(self._etcd_realm_prefix, recursive=True))
      for realm in all_realms.children:
        if not realm.dir:
          self._register_realm(json.loads(realm.value))
    except (KeyError, etcd.EtcdKeyError):
      # no realms have been registered yet
      pass

  def initialize(self, manager_config):
    logger.debug('Calling initialize')
    self._manager_config = manager_config

    executor_klass = self._executors.get(manager_config.get('EXECUTOR', ''), PopenExecutor)
    self._executor = executor_klass(manager_config.get('EXECUTOR_CONFIG', {}),
                                    self.manager_hostname)

    etcd_host = self._manager_config.get('ETCD_HOST', '127.0.0.1')
    etcd_port = self._manager_config.get('ETCD_PORT', 2379)
    etcd_auth = self._manager_config.get('ETCD_CERT_AND_KEY', None)
    etcd_ca_cert = self._manager_config.get('ETCD_CA_CERT', None)
    etcd_protocol = 'http' if etcd_auth is None else 'https'
    logger.debug('Connecting to etcd on %s:%s', etcd_host, etcd_port)

    worker_threads = self._manager_config.get('ETCD_WORKER_THREADS', 5)
    self._async_thread_executor = ThreadPoolExecutor(worker_threads)
    self._etcd_client = AsyncWrapper(self._etcd_client_klass(host=etcd_host, port=etcd_port,
                                                             cert=etcd_auth, ca_cert=etcd_ca_cert,
                                                             protocol=etcd_protocol,
                                                             read_timeout=5),
                                     executor=self._async_thread_executor)

    self._etcd_builder_prefix = self._manager_config.get('ETCD_BUILDER_PREFIX', 'building/')
    self._watch_etcd(self._etcd_builder_prefix, self._handle_builder_expiration)

    self._etcd_realm_prefix = self._manager_config.get('ETCD_REALM_PREFIX', 'realm/')
    self._watch_etcd(self._etcd_realm_prefix, self._handle_realm_change)

    # Load components for all realms currently known to the cluster
    async(self._register_existing_realms())

  def setup_time(self):
    setup_time = self._manager_config.get('MACHINE_SETUP_TIME', 300)
    return setup_time

  def shutdown(self):
    logger.debug('Shutting down worker.')
    self._shutting_down = True

    for (etcd_key, _), task in self._watch_tasks.items():
      if not task.done():
        logger.debug('Canceling watch task for %s', etcd_key)
        task.cancel()

    if self._async_thread_executor is not None:
      logger.debug('Shutting down thread pool executor.')
      self._async_thread_executor.shutdown()

  @coroutine
  def schedule(self, build_job):
    build_uuid = build_job.job_details['build_uuid']
    logger.debug('Calling schedule with job: %s', build_uuid)

    # Check if there are worker slots avialable by checking the number of jobs in etcd
    allowed_worker_count = self._manager_config.get('ALLOWED_WORKER_COUNT', 1)
    try:
      building = yield From(self._etcd_client.read(self._etcd_builder_prefix, recursive=True))
      workers_alive = sum(1 for child in building.children if not child.dir)
    except (KeyError, etcd.EtcdKeyError):
      workers_alive = 0
    except etcd.EtcdException:
      logger.exception('Exception when reading job count from etcd')
      raise Return(False)

    logger.debug('Total jobs: %s', workers_alive)

    if workers_alive >= allowed_worker_count:
      logger.info('Too many workers alive, unable to start new worker. %s >= %s', workers_alive,
                  allowed_worker_count)
      raise Return(False)

    job_key = self._etcd_job_key(build_job)

    # First try to take a lock for this job, meaning we will be responsible for its lifeline
    realm = str(uuid.uuid4())
    token = str(uuid.uuid4())
    ttl = self.setup_time()
    expiration = datetime.utcnow() + timedelta(seconds=ttl)

    machine_max_expiration = self._manager_config.get('MACHINE_MAX_TIME', 7200)
    max_expiration = datetime.utcnow() + timedelta(seconds=machine_max_expiration)

    payload = {
        'expiration': calendar.timegm(expiration.timetuple()),
        'max_expiration': calendar.timegm(max_expiration.timetuple()),
    }

    try:
      yield From(self._etcd_client.write(job_key, json.dumps(payload), prevExist=False, ttl=ttl))
    except (KeyError, etcd.EtcdKeyError):
      # The job was already taken by someone else, we are probably a retry
      logger.error('Job already exists in etcd, are timeouts misconfigured or is the queue broken?')
      raise Return(False)
    except etcd.EtcdException:
      logger.exception('Exception when writing job %s to etcd', build_uuid)
      raise Return(False)

    logger.debug('Starting builder with executor: %s', self._executor)

    try:
      builder_id = yield From(self._executor.start_builder(realm, token, build_uuid))
    except:
      logger.exception('Exception when starting builder for job: %s', build_uuid)
      raise Return(False)

    # Store the builder in etcd associated with the job id
    try:
      payload['builder_id'] = builder_id
      yield From(self._etcd_client.write(job_key, json.dumps(payload), prevExist=True, ttl=ttl))
    except etcd.EtcdException:
      logger.exception('Exception when writing job %s to etcd', build_uuid)
      raise Return(False)

    # Store the realm spec which will allow any manager to accept this builder when it connects
    realm_spec = json.dumps({
        'realm': realm,
        'token': token,
        'builder_id': builder_id,
        'job_queue_item': build_job.job_item,
    })

    try:
      yield From(self._etcd_client.write(self._etcd_realm_key(realm), realm_spec, prevExist=False,
                                         ttl=ttl))
    except (KeyError, etcd.EtcdKeyError):
      logger.error('Realm already exists in etcd. UUID collision or something is very very wrong.')
      raise Return(False)
    except etcd.EtcdException:
      logger.exception('Exception when writing realm %s to etcd', realm)
      raise Return(False)

    raise Return(True)

  @coroutine
  def build_component_ready(self, build_component):
    try:
      # Clean up the bookkeeping for allowing any manager to take the job
      job = self._component_to_job.pop(build_component)
      del self._job_uuid_to_component[job.job_details['build_uuid']]
      yield From(self._etcd_client.delete(self._etcd_realm_key(build_component.builder_realm)))

      logger.debug('Sending build %s to newly ready component on realm %s',
                   job.job_details['build_uuid'], build_component.builder_realm)
      yield From(build_component.start_build(job))
    except (KeyError, etcd.EtcdKeyError):
      logger.debug('Builder is asking for more work, but work already completed')

  def build_component_disposed(self, build_component, timed_out):
    logger.debug('Calling build_component_disposed.')
    self.unregister_component(build_component)

  @coroutine
  def job_completed(self, build_job, job_status, build_component):
    logger.debug('Calling job_completed with status: %s', job_status)

    # Kill the ephmeral builder
    yield From(self._executor.stop_builder(self._component_to_builder.pop(build_component)))

    # Release the lock in etcd
    job_key = self._etcd_job_key(build_job)
    yield From(self._etcd_client.delete(job_key))

    self.job_complete_callback(build_job, job_status)

  @coroutine
  def job_heartbeat(self, build_job):
    # Extend the deadline in etcd
    job_key = self._etcd_job_key(build_job)
    build_job_metadata_response = yield From(self._etcd_client.read(job_key))
    build_job_metadata = json.loads(build_job_metadata_response.value)

    max_expiration = datetime.utcfromtimestamp(build_job_metadata['max_expiration'])
    max_expiration_remaining = max_expiration - datetime.utcnow()
    max_expiration_sec = max(0, int(max_expiration_remaining.total_seconds()))

    ttl = min(self.heartbeat_period_sec * 2, max_expiration_sec)
    new_expiration = datetime.utcnow() + timedelta(seconds=ttl)

    payload = {
        'expiration': calendar.timegm(new_expiration.timetuple()),
        'builder_id': build_job_metadata['builder_id'],
        'max_expiration': build_job_metadata['max_expiration'],
    }

    yield From(self._etcd_client.write(job_key, json.dumps(payload), ttl=ttl))

    self.job_heartbeat_callback(build_job)

  def _etcd_job_key(self, build_job):
    """ Create a key which is used to track a job in etcd.
    """
    return os.path.join(self._etcd_builder_prefix, build_job.job_details['build_uuid'])

  def _etcd_realm_key(self, realm):
    """ Create a key which is used to track an incoming connection on a realm.
    """
    return os.path.join(self._etcd_realm_prefix, realm)

  def num_workers(self):
    """ Return the number of workers we're managing locally.
    """
    return len(self._component_to_builder)