From 2147005d2c04987ba201428e89ba8c22c1352ec1 Mon Sep 17 00:00:00 2001
From: Charlton Austin <charlton.austin@gmail.com>
Date: Tue, 25 Oct 2016 12:50:58 -0400
Subject: [PATCH] Adding a method of cancelling a build based on etcd message.

---
 buildman/manager/ephemeral.py | 24 ++++++++++++++++++++++++
 data/model/build.py           | 18 ++++++++++++++----
 2 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/buildman/manager/ephemeral.py b/buildman/manager/ephemeral.py
index 460b777ad..b3e8c78e0 100644
--- a/buildman/manager/ephemeral.py
+++ b/buildman/manager/ephemeral.py
@@ -76,6 +76,7 @@ class EphemeralBuilderManager(BaseManager):
     self._etcd_job_prefix = None
     self._etcd_lock_prefix = None
     self._etcd_metric_prefix = None
+    self._etcd_cancel_build_prefix = None
 
     self._ephemeral_api_timeout = DEFAULT_EPHEMERAL_API_TIMEOUT
     self._ephemeral_setup_timeout = DEFAULT_EPHEMERAL_SETUP_TIMEOUT
@@ -356,6 +357,9 @@ class EphemeralBuilderManager(BaseManager):
     self._watch_etcd(self._etcd_realm_prefix, self._handle_realm_change,
                      restarter=self._register_existing_realms)
 
+    self._etcd_cancel_build_prefix = self._manager_config('ETCD_CANCEL_PREFIX', 'cancel/')
+    self._watch_etcd(self._etcd_cancel_build_prefix, self._cancel_build)
+
     self._etcd_lock_prefix = self._manager_config.get('ETCD_LOCK_PREFIX', 'lock/')
     self._etcd_metric_prefix = self._manager_config.get('ETCD_METRIC_PREFIX', 'metric/')
 
@@ -705,3 +709,23 @@ class EphemeralBuilderManager(BaseManager):
     """ Return the number of workers we're managing locally.
     """
     return len(self._component_to_job)
+
+  @coroutine
+  def _cancel_build(self, etcd_result):
+    """ Listens for etcd event and then cancels the build
+    """
+    if etcd_result is None:
+      raise Return(False)
+
+    if etcd_result.action not in (EtcdAction.CREATE, EtcdAction.SET):
+      raise Return(False)
+
+    build_uuid = etcd_result.value
+    build_info = self._build_uuid_to_info.get(build_uuid, None)
+
+    if build_info is None:
+      logger.debug('No build info for "%s" job %s', etcd_result.action, build_uuid)
+      raise Return(False)
+    got_lock = yield From(self._take_etcd_atomic_lock('job-cancelled', build_uuid, build_info.execution_id))
+    if got_lock:
+      yield From(self.kill_builder_executor(build_uuid))
diff --git a/data/model/build.py b/data/model/build.py
index 5a43859ee..2a22f2f77 100644
--- a/data/model/build.py
+++ b/data/model/build.py
@@ -161,11 +161,22 @@ def create_cancel_build_in_queue(build, build_queue):
     if build.phase != BUILD_PHASE.WAITING or not build.queue_id:
       return False
 
-    return build_queue.cancel(build.queue_id)
+    cancelled = build_queue.cancel(build.queue_id)
+    if cancelled:
+      # Delete the build row.
+      build.delete_instance()
+    return cancelled
 
   return cancel_build
 
 
+def create_cancel_build_in_manager(build):
+  """ A function to cancel the build before it starts to push """
+  def cancel_build():
+    return False
+  return cancel_build
+
+
 def cancel_repository_build(build, build_queue):
   """ This tries to cancel the build returns true if request is successful false if it can't be cancelled """
   with db_transaction():
@@ -176,11 +187,10 @@ def cancel_repository_build(build, build_queue):
     except RepositoryBuild.DoesNotExist:
       return False
 
-    cancel_builds = [create_cancel_build_in_queue(build, build_queue), ]
+    cancel_builds = [create_cancel_build_in_queue(build, build_queue),
+                     create_cancel_build_in_manager(build), ]
     for cancelled in cancel_builds:
       if cancelled():
-        # Delete the build row.
-        build.delete_instance()
         return True
 
     return False