From 5c6e7adb0594bd6facd52634efba005bd814969f Mon Sep 17 00:00:00 2001
From: Brenden Matthews <brenden@diddyinc.com>
Date: Sun, 22 May 2016 08:22:41 -0700
Subject: [PATCH] Add a pre-kill hook to ZDD deploy. (#200)

---
 bluegreen_deploy.py            | 22 ++++++++++++++++++++++
 tests/test_bluegreen_deploy.py | 17 +++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/bluegreen_deploy.py b/bluegreen_deploy.py
index 02efafe2..0dea57b6 100755
--- a/bluegreen_deploy.py
+++ b/bluegreen_deploy.py
@@ -15,6 +15,7 @@
 import six.moves.urllib as urllib
 import socket
 import sys
+import subprocess
 
 
 logger = logging.getLogger('bluegreen_deploy')
@@ -296,6 +297,15 @@ def deployment_in_progress(app):
     return len(app['deployments']) > 0
 
 
+def execute_pre_kill_hook(args, old_app, tasks_to_kill):
+    if args.pre_kill_hook is not None:
+        logger.info("Calling pre-kill hook '{}'".format(args.pre_kill_hook))
+
+        subprocess.check_call([args.pre_kill_hook,
+                              json.dumps(old_app),
+                              json.dumps(tasks_to_kill)])
+
+
 def swap_bluegreen_apps(args, new_app, old_app):
     old_app = fetch_marathon_app(args, old_app['id'])
     new_app = fetch_marathon_app(args, new_app['id'])
@@ -310,6 +320,8 @@ def swap_bluegreen_apps(args, new_app, old_app):
         return safe_delete_app(args, old_app)
 
     if len(tasks_to_kill) > 0:
+        execute_pre_kill_hook(args, old_app, tasks_to_kill)
+
         logger.info("There are {} draining listeners, "
                     "about to kill the following tasks:\n  - {}"
                     .format(len(tasks_to_kill),
@@ -608,6 +620,16 @@ def get_arg_parser():
                         " for HAProxy to drain connections",
                         type=int, default=300
                         )
+    parser.add_argument("--pre-kill-hook",
+                        help="A path to an executable (such as a script) "
+                        "which will be called before killing any tasks marked "
+                        "for draining at each step. The script will be called "
+                        "with 2 arguments: the old app definition (in JSON), "
+                        "and the list of tasks which will be killed. An exit "
+                        "code of 0 indicates the deploy may continue. "
+                        "If the hook returns a non-zero exit code, the deploy "
+                        "will stop, and an operator must intervene."
+                        )
     parser = set_logging_args(parser)
     parser = set_marathon_auth_args(parser)
     return parser
diff --git a/tests/test_bluegreen_deploy.py b/tests/test_bluegreen_deploy.py
index f7dae363..ffa4cc24 100644
--- a/tests/test_bluegreen_deploy.py
+++ b/tests/test_bluegreen_deploy.py
@@ -14,6 +14,7 @@ class Arguments:
     initial_instances = 1
     marathon_auth_credential_file = None
     auth_credentials = None
+    pre_kill_hook = None
 
 
 class MyResponse:
@@ -122,6 +123,22 @@ def test_parse_haproxy_stats(self):
             assert results[2].pxname == 'http-out'
             assert results[2].svname == 'IPv4-cached'
 
+    @mock.patch('subprocess.check_call')
+    def test_pre_kill_hook(self, mock):
+        # TODO(BM): This test is naive. An end-to-end test would be nice.
+        args = Arguments()
+        args.pre_kill_hook = 'myhook'
+        app = {
+            'id': 'myApp'
+        }
+        tasks_to_kill = ['task1', 'task2']
+
+        bluegreen_deploy.execute_pre_kill_hook(args, app, tasks_to_kill)
+
+        mock.assert_called_with([args.pre_kill_hook,
+                                 '{"id": "myApp"}',
+                                 '["task1", "task2"]'])
+
     @mock.patch('bluegreen_deploy.fetch_combined_haproxy_stats',
                 mock.Mock(side_effect=lambda _: _load_listeners()))
     def test_fetch_app_listeners(self):