import logging import etcd import uuid import calendar import os.path from datetime import datetime, timedelta from trollius import From, coroutine, Return, async from concurrent.futures import ThreadPoolExecutor from buildman.manager.basemanager import BaseManager from buildman.manager.executor import PopenExecutor, EC2Executor from buildman.component.buildcomponent import BuildComponent from buildman.asyncutil import AsyncWrapper logger = logging.getLogger(__name__) ETCD_BUILDER_PREFIX = 'building/' ETCD_EXPIRE_RESULT = 'expire' 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._component_to_job = {} self._component_to_builder = {} self._executor = None self._worker_watch_task = None super(EphemeralBuilderManager, self).__init__(*args, **kwargs) def _watch_builders(self): """ Watch the builders key for expirations. """ if not self._shutting_down: workers_future = self._etcd_client.watch(ETCD_BUILDER_PREFIX, recursive=True) workers_future.add_done_callback(self._handle_key_expiration) logger.debug('Scheduling watch task.') self._worker_watch_task = async(workers_future) def _handle_key_expiration(self, changed_key_future): """ Handle when a builder expires """ if self._worker_watch_task is None or self._worker_watch_task.done(): self._watch_builders() if changed_key_future.cancelled(): # Due to lack of interest, tomorrow has been cancelled return etcd_result = changed_key_future.result() if etcd_result.action == ETCD_EXPIRE_RESULT: # Handle the expiration logger.debug('Builder expired, clean up the old build node') async(self._clean_up_old_builder(etcd_result.key, etcd_result._prev_node.value)) 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.public_ip_address) etcd_host = self._manager_config.get('ETCD_HOST', '127.0.0.1') etcd_port = self._manager_config.get('ETCD_PORT', 2379) 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), executor=self._async_thread_executor) self._watch_builders() def setup_time(self): setup_time = self._manager_config.get('MACHINE_SETUP_TIME', 300) logger.debug('Returning setup_time: %s', setup_time) return setup_time def shutdown(self): logger.debug('Shutting down worker.') self._shutting_down = True if self._worker_watch_task is not None: logger.debug('Canceling watch task.') self._worker_watch_task.cancel() self._worker_watch_task = None 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, loop): logger.debug('Calling schedule with job: %s', build_job.job_details['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(ETCD_BUILDER_PREFIX, recursive=True)) workers_alive = sum(1 for child in building.children if not child.dir) except KeyError: workers_alive = 0 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) payload = { 'expiration': calendar.timegm(expiration.timetuple()), } try: yield From(self._etcd_client.write(job_key, payload, prevExist=False, ttl=ttl)) component = self.register_component(realm, BuildComponent, token=token) self._component_to_job[component] = build_job except KeyError: # 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) logger.debug('Starting builder with executor: %s', self._executor) builder_id = yield From(self._executor.start_builder(realm, token)) self._component_to_builder[component] = builder_id # Store the builder in etcd associated with the job id payload['builder_id'] = builder_id yield From(self._etcd_client.write(job_key, payload, prevExist=True)) raise Return(True) def build_component_ready(self, build_component, loop): try: job = self._component_to_job.pop(build_component) logger.debug('Sending build %s to newly ready component on realm %s', job.job_details['build_uuid'], build_component.builder_realm) loop.call_soon(build_component.start_build, job) except KeyError: logger.warning('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.') # TODO make it so that I don't have to unregister the component if it timed out 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 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_response = yield From(self._etcd_client.read(job_key)) ttl = self.heartbeat_period_sec * 2 new_expiration = datetime.utcnow() + timedelta(seconds=ttl) payload = { 'expiration': calendar.timegm(new_expiration.timetuple()), 'builder_id': build_job_response.value['builder_id'], } yield From(self._etcd_client.write(job_key, payload, ttl=ttl)) self.job_heartbeat_callback(build_job) @coroutine def _clean_up_old_builder(self, job_key, job_payload): """ Terminate an old builders once the expiration date has passed. """ logger.debug('Cleaning up the old builder for job: %s', job_key) if 'builder_id' in job_payload: logger.info('Terminating expired build node.') yield From(self._executor.stop_builder(job_payload['builder_id'])) yield From(self._etcd_client.delete(job_key)) @staticmethod def _etcd_job_key(build_job): """ Create a key which is used to track a job in etcd. """ return os.path.join(ETCD_BUILDER_PREFIX, build_job.job_details['build_uuid']) def num_workers(self): """ Return the number of workers we're managing locally. """ return len(self._component_to_builder)