Improve tests for the ephemeral build manager.

buildman/manager/ephemeral.py

@@ -2,9 +2,11 @@ import logging
 import etcd
 import uuid
 import calendar
+import os.path
 
 from datetime import datetime, timedelta
-from trollius import From, coroutine, Return
+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
@@ -16,25 +18,11 @@ logger = logging.getLogger(__name__)
 
 
 ETCD_BUILDER_PREFIX = 'building/'
-
-
-def clear_etcd(client):
-  """ Debugging method used to clear out the section of etcd we are using to track jobs in flight.
-  """
-  try:
-    building = client.read(ETCD_BUILDER_PREFIX, recursive=True)
-    for child in building.leaves:
-      if not child.dir:
-        logger.warning('Deleting key: %s', child.key)
-        client.delete(child.key)
-  except KeyError:
-    pass
+ETCD_EXPIRE_RESULT = 'expire'
 
 
 class EphemeralBuilderManager(BaseManager):
   """ Build manager implementation for the Enterprise Registry. """
-  shutting_down = False
-
   _executors = {
       'popen': PopenExecutor,
       'ec2': EC2Executor,
@@ -43,7 +31,10 @@ class EphemeralBuilderManager(BaseManager):
   _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 = {}
@@ -51,8 +42,35 @@ class EphemeralBuilderManager(BaseManager):
 
     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
@@ -65,7 +83,11 @@ class EphemeralBuilderManager(BaseManager):
     etcd_port = self._manager_config.get('ETCD_PORT', 2379)
     logger.debug('Connecting to etcd on %s:%s', etcd_host, etcd_port)
 
-    self._etcd_client = AsyncWrapper(self._etcd_client_klass(host=etcd_host, port=etcd_port))
+    self._async_thread_executor = ThreadPoolExecutor(self._manager_config.get('ETCD_WORKERS', 5))
+    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)
@@ -73,8 +95,17 @@ class EphemeralBuilderManager(BaseManager):
     return setup_time
 
   def shutdown(self):
-    logger.debug('Calling shutdown.')
-    raise NotImplementedError
+    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):
@@ -161,9 +192,8 @@ class EphemeralBuilderManager(BaseManager):
 
     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 '{0}{1}'.format(ETCD_BUILDER_PREFIX, build_job.job_details['build_uuid'])
+    return os.path.join(ETCD_BUILDER_PREFIX, build_job.job_details['build_uuid'])