Merge pull request #56 from loads/bug/45

bug: Ensure all containers drained before exiting

ardere/exceptions.py

@@ -8,3 +8,7 @@ class ShutdownPlanException(Exception):
 
 class ValidationException(Exception):
     """Exception to indicate validation error parsing input"""
+
+
+class UndrainedInstancesException(Exception):
+    """There are still ACTIVE or DRAINING instances in the cluster"""

ardere/step_functions.py

@@ -27,6 +27,7 @@
     ServicesStartingException,
     ShutdownPlanException,
     ValidationException,
+    UndrainedInstancesException,
 )
 
 logger = logging.getLogger()
@@ -418,3 +419,30 @@ def cleanup_cluster(self):
         except botocore.exceptions.ClientError:
             pass
         return self.event
+
+    def check_drained(self):
+        """Ensure that all services are shut down before allowing restart
+
+        Step 8
+
+        """
+        client = self.boto.client('ecs')
+        actives = len(
+            client.list_container_instances(
+                cluster=self.event["ecs_name"],
+                maxResults=1,
+                status="ACTIVE",
+            ).get('containerInstanceArns', []))
+        if actives:
+            raise UndrainedInstancesException(
+                "Still {} active.".format(actives))
+        draining = len(
+            client.list_container_instances(
+                cluster=self.event["ecs_name"],
+                maxResults=1,
+                status="DRAINING",
+            ).get('containerInstanceArns', []))
+        if draining:
+            raise UndrainedInstancesException(
+                "Still {} draining.".format(draining))
+        return self.event

handler.py

@@ -39,3 +39,7 @@ def check_for_cluster_done(event, context):
 def cleanup_cluster(event, context):
     runner = AsynchronousPlanRunner(event, context)
     return runner.cleanup_cluster()
+
+
+def check_drain(event, context):
+    return AsynchronousPlanRunner(event, context).check_drained()

serverless.yml

@@ -112,6 +112,8 @@ functions:
     handler: handler.check_for_cluster_done
   cleanup_cluster:
     handler: handler.cleanup_cluster
+  check_drain:
+    handler: handler.check_drain
 
 stepFunctions:
   stateMachines:
@@ -134,6 +136,7 @@ stepFunctions:
               IntervalSeconds: 10
               MaxAttempts: 60
               BackoffRate: 1
+          Catch:
             -
               ErrorEquals:
                 - States.ALL
@@ -201,6 +204,17 @@ stepFunctions:
         "Clean-up Cluster":
           Type: Task
           Resource: cleanup_cluster
+          Next: "Checking Drain"
+        "Checking Drain":
+          Type: Task
+          Resource: check_drain
+          Retry:
+            -
+              ErrorEquals:
+                - UndrainedInstancesException
+              IntervalSeconds: 10
+              MaxAttempts: 10
+              BackoffRate: 1
           End: true
 
 resources:

tests/test_step_functions.py

@@ -273,6 +273,45 @@ def test_cleanup_cluster_error(self):
         self.runner.cleanup_cluster()
         mock_s3.Object.assert_called()
 
+    def test_drain_check_active(self):
+        from ardere.exceptions import UndrainedInstancesException
+
+        mock_client = mock.Mock()
+        mock_client.list_container_instances.return_value = {
+            'containerInstanceArns': [
+                'Some-Arn-01234567890',
+            ],
+            "nextToken": "token-8675309"
+        }
+        self.mock_boto.client.return_value = mock_client
+        assert_raises(UndrainedInstancesException,
+                      self.runner.check_drained)
+
+    def test_drain_check_draining(self):
+        from ardere.exceptions import UndrainedInstancesException
+
+        mock_client = mock.Mock()
+        mock_client.list_container_instances.side_effect = [
+            {},
+            {
+                'containerInstanceArns': [
+                    'Some-Arn-01234567890',
+                ],
+                "nextToken": "token-8675309"
+            }
+        ]
+        self.mock_boto.client.return_value = mock_client
+        assert_raises(UndrainedInstancesException,
+                      self.runner.check_drained)
+
+    def test_drain_check(self):
+        mock_client = mock.Mock()
+        mock_client.list_container_instances.side_effect = [
+            {},
+            {}
+        ]
+        self.mock_boto.client.return_value = mock_client
+        self.runner.check_drained()
 
 class TestValidation(unittest.TestCase):
     def _make_FUT(self):