quay/buildman/manager/enterprise.py

import logging
import uuid

from buildman.manager.basemanager import BaseManager
from buildman.basecomponent import BaseComponent
from buildman.buildcomponent import BuildComponent

from trollius.coroutines import From

REGISTRATION_REALM = 'registration'
logger = logging.getLogger(__name__)

class DynamicRegistrationComponent(BaseComponent):
  """ Component session that handles dynamic registration of the builder components. """

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

  def onJoin(self, details):
    logger.debug('Registering registration method')
    yield From(self.register(self._worker_register, u'io.quay.buildworker.register'))

  def _worker_register(self):
    realm = self.parent_manager.add_build_component()
    logger.debug('Registering new build component+worker with realm %s', realm)
    return realm


class EnterpriseManager(BaseManager):
  """ Build manager implementation for the Enterprise Registry. """
  build_components = []
  shutting_down = False

  def initialize(self):
    # Add a component which is used by build workers for dynamic registration. Unlike
    # production, build workers in enterprise are long-lived and register dynamically.
    self.register_component(REGISTRATION_REALM, DynamicRegistrationComponent)
  
  def add_build_component(self):
    # Generate a new unique realm ID for the build worker.
    realm = str(uuid.uuid4())
    component = self.register_component(realm, BuildComponent, token="")
    self.build_components.append(component)
    return realm

  def schedule(self, build_job, loop):
    if self.shutting_down:
      return False

    for component in self.build_components:
      if component.is_ready():
        loop.call_soon(component.start_build, build_job)
        return True

    return False

  def shutdown(self):
    self.shutting_down = True

  def job_completed(self, build_job, job_status, build_component):
    self.job_complete_callback(build_job, job_status)

  def build_component_disposed(self, build_component, timed_out):
    self.build_components.remove(build_component)