Add support for rolling update of stateful Elastigroups (#572)

Author: a1exsh

senza/cli.py

@@ -1933,11 +1933,16 @@ def patch_spotinst_elastigroup(properties, elastigroup_id, region, stack_name):
     metavar="PERCENTAGE",
     type=click.IntRange(0, 100, clamp=True),
     help="Percentage (int value) of the ElastiGroup cluster that is respawned in each step."
-    " Valid only for ElastiGroups. The default value for this of 20.",
+    " Valid only for stateless ElastiGroups. The default value for this of 20.",
+)
+@click.option(
+    "--batch-per-subnet",
+    is_flag=True,
+    help="Recycle ElastiGroup instances in batches per subnet. Valid only for stateful ElastiGroups.",
 )
 @region_option
 @stacktrace_visible_option
-def respawn_instances(stack_ref, inplace, force, batch_size_percentage, region):
+def respawn_instances(stack_ref, inplace, force, batch_size_percentage, batch_per_subnet, region):
     """Replace all EC2 instances in Auto Scaling Group(s)
 
     Performs a rolling update to prevent downtimes."""
@@ -1954,7 +1959,7 @@ def respawn_instances(stack_ref, inplace, force, batch_size_percentage, region):
             )
         elif group["type"] == ELASTIGROUP_RESOURCE_TYPE:
             respawn.respawn_elastigroup(
-                group["resource_id"], group["stack_name"], region, batch_size_percentage
+                group["resource_id"], group["stack_name"], region, batch_size_percentage, batch_per_subnet
             )
 
 

senza/respawn.py

@@ -1,6 +1,7 @@
 """
 Functions to scale and respawn Auto Scale Groups
 """
+import collections
 import time
 
 from clickclick import Action, info
@@ -14,6 +15,7 @@
 ELASTIGROUP_TERMINATED_DEPLOY_STATUS = ["stopped", "failed"]
 
 DEFAULT_BATCH_SIZE = 20
+WAIT_FOR_ELASTIGROUP_SEC = 10
 
 
 def get_auto_scaling_group(asg, asg_name: str):
@@ -177,19 +179,117 @@ def respawn_auto_scaling_group(
 
 
 def respawn_elastigroup(
-    elastigroup_id: str, stack_name: str, region: str, batch_size: int
+    elastigroup_id: str, stack_name: str, region: str, batch_size: int, batch_per_subnet: bool
 ):
     """
     Respawn all instances in the ElastiGroup.
     """
 
+    spotinst_account = elastigroup_api.get_spotinst_account_data(region, stack_name)
+
+    stateful_instances = elastigroup_api.get_stateful_instances(elastigroup_id, spotinst_account)
+    if stateful_instances:
+        if batch_size is not None:
+            raise Exception("Batch size is not supported when respawning stateful ElastiGroups")
+
+        respawn_stateful_elastigroup(elastigroup_id, stack_name, region, batch_per_subnet, stateful_instances, spotinst_account)
+    else:
+        if batch_per_subnet:
+            raise Exception("Batch per subnet is not supported when respawning stateless ElastiGroups")
+
+        respawn_stateless_elastigroup(elastigroup_id, stack_name, batch_size, spotinst_account)
+
+
+def respawn_stateful_elastigroup(
+    elastigroup_id: str, stack_name: str, region: str, batch_per_subnet: bool, stateful_instances: list, spotinst_account
+):
+    """
+    Recycles stateful instances of the ElastiGroup.
+    """
+
+    if not stateful_elastigroup_ready(stateful_instances):
+        raise Exception(
+            "Stateful ElastiGroup {} is not ready: some instances are not in the ACTIVE state".format(elastigroup_id)
+        )
+
+    info(
+        "Recycling {} stateful instances for ElastiGroup {} (ID {})".format(
+            len(stateful_instances), stack_name, elastigroup_id
+        )
+    )
+
+    if batch_per_subnet:
+        instances_by_subnet = stateful_elastigroup_instances_by_subnet(region, stateful_instances)
+        for subnet, subnet_instances in sorted(instances_by_subnet.items(), key=lambda item: item[0]):
+            info("Recycling ALL stateful instances in subnet: {}".format(subnet))
+
+            for instance in sorted(subnet_instances, key=lambda i: i['privateIp']):
+                time.sleep(WAIT_FOR_ELASTIGROUP_SEC)
+                recycle_stateful_elastigroup_instance(elastigroup_id, instance, spotinst_account)
+
+            wait_for_stateful_elastigroup(elastigroup_id, spotinst_account)
+
+    else:
+        for instance in sorted(stateful_instances, key=lambda i: i['privateIp']):
+            recycle_stateful_elastigroup_instance(elastigroup_id, instance, spotinst_account)
+            wait_for_stateful_elastigroup(elastigroup_id, spotinst_account)
+
+
+def stateful_elastigroup_instances_by_subnet(region: str, stateful_instances: list):
+    instances_by_subnet = collections.defaultdict(list)
+    instances_by_ec2_id = {i['instanceId']: i for i in stateful_instances}
+
+    ec2 = BotoClientProxy("ec2", region)
+    ec2_instances = ec2.describe_instances(InstanceIds=list(instances_by_ec2_id.keys()))
+    for r in ec2_instances['Reservations']:
+        for i in r['Instances']:
+            subnet = "{} | {}".format(
+                i['Placement']['AvailabilityZone'], i['SubnetId']
+            )
+            instance = instances_by_ec2_id[i['InstanceId']]
+            instances_by_subnet[subnet].append(instance)
+
+    return instances_by_subnet
+
+
+def stateful_elastigroup_ready(stateful_instances: list):
+    return all(i['state'] == elastigroup_api.STATEFUL_STATE_ACTIVE for i in stateful_instances)
+
+
+def wait_for_stateful_elastigroup(elastigroup_id: str, spotinst_account):
+    """
+    Waits for all stateful instances of the ElastiGroup to be in the ACTIVE state.
+    """
+    with Action("Waiting for all stateful instances to be in the ACTIVE state") as act:
+        while True:
+            time.sleep(WAIT_FOR_ELASTIGROUP_SEC)
+            act.progress()
+            stateful_instances = elastigroup_api.get_stateful_instances(elastigroup_id, spotinst_account)
+            if stateful_elastigroup_ready(stateful_instances):
+                break
+
+
+def recycle_stateful_elastigroup_instance(elastigroup_id: str, instance: dict, spotinst_account):
+    info(
+        "Recycling stateful instance {} ({} | {})".format(
+            instance['id'], instance['instanceId'], instance['privateIp']
+        )
+    )
+    elastigroup_api.recycle_stateful_instance(elastigroup_id, instance['id'], spotinst_account)
+
+
+def respawn_stateless_elastigroup(
+    elastigroup_id: str, stack_name: str, batch_size: int, spotinst_account
+):
+    """
+    Start a deployment of the ElastiGroup and wait for it to complete.
+    """
+
     if batch_size is None or batch_size < 1:
         batch_size = DEFAULT_BATCH_SIZE
 
-    spotinst_account = elastigroup_api.get_spotinst_account_data(region, stack_name)
-
     info(
-        "Redeploying the cluster for ElastiGroup {} (ID {})".format(
+        "Redeploying instances for ElastiGroup {} (ID {})".format(
             stack_name, elastigroup_id
         )
     )

senza/spotinst/components/elastigroup_api.py

@@ -14,6 +14,7 @@
 DEPLOY_STRATEGY_REPLACE = 'REPLACE_SERVER'
 DEFAULT_CONNECT_TIMEOUT = 9
 DEFAULT_READ_TIMEOUT = 30
+STATEFUL_STATE_ACTIVE = 'ACTIVE'
 
 
 class SpotInstAccountData:
@@ -112,6 +113,56 @@ def get_elastigroup(elastigroup_id, spotinst_account_data):
     return groups
 
 
+def get_stateful_instances(elastigroup_id, spotinst_account_data):
+    '''
+    Returns a list containing the description of the stateful instances of an ElastiGroup.
+    Exceptions will be thrown for HTTP errors.
+
+    https://docs.spot.io/spotinst-api/elastigroup/amazon-web-services/stateful-api/list-stateful-instances/
+    '''
+    headers = {
+        "Authorization": "Bearer {}".format(spotinst_account_data.access_token),
+        "Content-Type": "application/json"
+    }
+
+    response = requests.get(
+        '{}/aws/ec2/group/{}/statefulInstance?accountId={}'.format(
+            SPOTINST_API_URL, elastigroup_id, spotinst_account_data.account_id
+        ),
+        headers=headers, timeout=(DEFAULT_CONNECT_TIMEOUT, DEFAULT_READ_TIMEOUT))
+    response.raise_for_status()
+    data = response.json()
+    stateful_instances = data.get("response", {}).get("items", [])
+
+    return stateful_instances
+
+
+def recycle_stateful_instance(elastigroup_id, stateful_instance_id, spotinst_account_data):
+    '''
+    Triggers recycling of a single stateful instance of an ElastiGroup.
+    Returns operation status.
+
+    Exceptions will be thrown for HTTP errors.
+
+    https://docs.spot.io/spotinst-api/elastigroup/amazon-web-services/stateful-api/recycle-stateful-instance/
+    '''
+    headers = {
+        "Authorization": "Bearer {}".format(spotinst_account_data.access_token),
+        "Content-Type": "application/json"
+    }
+
+    response = requests.put(
+        '{}/aws/ec2/group/{}/statefulInstance/{}/recycle?accountId={}'.format(
+            SPOTINST_API_URL, elastigroup_id, stateful_instance_id, spotinst_account_data.account_id
+        ),
+        headers=headers, timeout=(DEFAULT_CONNECT_TIMEOUT, DEFAULT_READ_TIMEOUT))
+    response.raise_for_status()
+    data = response.json()
+    status = data.get("response", {}).get("status", {})
+
+    return status
+
+
 def patch_elastigroup(properties, elastigroup_id, spotinst_account_data):
     '''
     Patch specific properties of the ElastiGroup.

tests/test_elastigroup_api.py

@@ -4,8 +4,9 @@
 from pytest import raises
 
 from senza.components.elastigroup import ELASTIGROUP_RESOURCE_TYPE
-from senza.spotinst.components.elastigroup_api import update_capacity, get_elastigroup, patch_elastigroup, deploy, \
-    deploy_status, SPOTINST_API_URL, SpotInstAccountData, get_spotinst_account_data
+from senza.spotinst.components.elastigroup_api import update_capacity, get_elastigroup, get_stateful_instances, \
+    recycle_stateful_instance, patch_elastigroup, deploy, deploy_status, SPOTINST_API_URL, SpotInstAccountData, \
+    get_spotinst_account_data
 
 
 def test_update_capacity():
@@ -65,6 +66,57 @@ def test_get_elastigroup():
         assert group['name'] == 'my-app-1'
 
 
+def test_get_stateful_instances():
+    instances = {
+        'response': {
+            'items': [{
+                'id': 'ssi-abc123',
+                'instanceId': 'i-321cba',
+            }, {
+                'id': 'ssi-def456',
+                'instanceId': 'i-456def',
+            }]
+        }
+    }
+
+    elastigroup_id = 'sig-xfy'
+    spotinst_account_data = SpotInstAccountData('act-zwk', 'fake-token')
+    with responses.RequestsMock() as rsps:
+        rsps.add(rsps.GET,
+                 '{}/aws/ec2/group/{}/statefulInstance?accountId={}'.format(
+                     SPOTINST_API_URL, elastigroup_id, spotinst_account_data.account_id
+                 ),
+                 status=200,
+                 json=instances)
+
+        instances = get_stateful_instances(elastigroup_id, spotinst_account_data)
+        assert instances[0]['id'] == 'ssi-abc123'
+        assert instances[1]['id'] == 'ssi-def456'
+
+
+def test_recycle_stateful_instance():
+    recycle_response = {
+        'response': {
+            'status': {
+                'code': 200
+            }
+        }
+    }
+    elastigroup_id = 'sig-xfy'
+    stateful_instance_id = 'ssi-abcdef1'
+    spotinst_account_data = SpotInstAccountData('act-zwk', 'fake-token')
+    with responses.RequestsMock() as rsps:
+        rsps.add(rsps.PUT,
+                 '{}/aws/ec2/group/{}/statefulInstance/{}/recycle?accountId={}'.format(
+                     SPOTINST_API_URL, elastigroup_id, stateful_instance_id, spotinst_account_data.account_id
+                 ),
+                 status=200,
+                 json=recycle_response)
+
+        recycle_response = recycle_stateful_instance(elastigroup_id, stateful_instance_id, spotinst_account_data)
+        assert recycle_response['code'] == 200
+
+
 def test_patch_elastigroup():
     patch = {
         'ImageId': 'image-foo',