Merge pull request #774 from degica/wait_for_stopping_tasks

Count STOPPING state ECS tasks
degica · May 10, 2023 · b4ad957 · b4ad957
2 parents 6b9e688 + 433314d
commit b4ad957
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 6 deletions.
diff --git a/drain_instance.py b/drain_instance.py
@@ -3,7 +3,9 @@
 import os
 import time
 import random
+
 from botocore.config import Config
+from botocore.exceptions import ClientError
 
 session = boto3.session.Session()
 config = Config(
@@ -27,6 +29,25 @@ def ciFor(ec2Id):
 
     return None, None
 
+def taskExists(clusterName, ciId):
+    running_tasks = ecs.list_tasks(cluster=clusterName, containerInstance=ciId, desiredStatus='RUNNING')['taskArns']
+    if (len(running_tasks) > 0):
+        return True
+
+    # Assume there are not more than 100 tasks in a container
+    stopping_tasks = ecs.list_tasks(cluster=clusterName, containerInstance=ciId, desiredStatus='STOPPED')['taskArns']
+    for task_arn in stopping_tasks:
+        response = ecs.describe_tasks(
+            cluster=clusterName,
+            tasks=[task_arn]
+        )
+        status = response['tasks'][0]['lastStatus']
+        if status != 'STOPPED':
+            return True
+
+    print('No tasks, will proceed terminating the instance')
+    return False
+
 def lambda_handler(event, context):
     msg = json.loads(event['Records'][0]['Sns']['Message'])
     ec2Id = msg['EC2InstanceId']
@@ -48,13 +69,15 @@ def lambda_handler(event, context):
         if status != 'DRAINING':
             ecs.update_container_instances_state(cluster=clusterName,containerInstances=[ciId],status='DRAINING')
 
-        tasks = ecs.list_tasks(cluster=clusterName, containerInstance=ciId)['taskArns']
-        if len(tasks) > 0:
+        if taskExists(clusterName, ciId):
             time.sleep(5)
             session.client('sns', config=config).publish(TopicArn=topicArn, Message=json.dumps(msg), Subject='Invoking lambda again')
         else:
             session.client('autoscaling', config=config).complete_lifecycle_action(LifecycleHookName=lifecycleHookName, AutoScalingGroupName=asgName, LifecycleActionResult='CONTINUE', InstanceId=ec2Id)
-    except ecs.exceptions.ThrottlingException:
-        sec = random.uniform(3, 5)
-        time.sleep(sec)
-        session.client('sns').publish(TopicArn=topicArn, Message=json.dumps(msg), Subject='Invoking lambda again')
+    except ClientError as exception_obj:
+        if exception_obj.response['Error']['Code'] == 'ThrottlingException':
+            sec = random.uniform(3, 5)
+            time.sleep(sec)
+            session.client('sns').publish(TopicArn=topicArn, Message=json.dumps(msg), Subject='Invoking lambda again')
+        else:
+            raise
diff --git a/lib/barcelona/network/autoscaling_builder.rb b/lib/barcelona/network/autoscaling_builder.rb
@@ -96,6 +96,7 @@ def build_resources
                       "ecs:DescribeContainerInstances",
                       "ecs:UpdateContainerInstancesState",
                       "ecs:ListTasks",
+                      "ecs:DescribeTasks",
                       "logs:CreateLogGroup",
                       "logs:CreateLogStream",
                       "logs:PutLogEvents",

diff --git a/spec/lib/barcelona/network/network_stack_spec.rb b/spec/lib/barcelona/network/network_stack_spec.rb
@@ -195,6 +195,7 @@
                       "ecs:DescribeContainerInstances",
                       "ecs:UpdateContainerInstancesState",
                       "ecs:ListTasks",
+                      "ecs:DescribeTasks",
                       "logs:CreateLogGroup",
                       "logs:CreateLogStream",
                       "logs:PutLogEvents",