Merge pull request #447 from rafaelcaricio/timout-to-traffic-change

Only waits for traffic change if timeout provided

senza/aws.py

@@ -4,16 +4,20 @@
 
 import base64
 import collections
+import datetime
 import functools
 import re
+import time
+from contextlib import contextmanager
 from pprint import pformat
+from typing import Optional
 
 import arrow
 import boto3
 import yaml
-from botocore.exceptions import ClientError, BotoCoreError
+from botocore.exceptions import BotoCoreError, ClientError
 from click import FileError
-from clickclick import Action, info
+from clickclick import Action, error, info
 
 from .exceptions import SecurityGroupNotFound
 from .manaus.boto_proxy import BotoClientProxy
@@ -436,3 +440,45 @@ def update_stack_from_template(region: str, template: dict, dry_run: bool):
                 act.ok('NO UPDATE')
             else:
                 act.fatal_error('ClientError: {}'.format(pformat(response)))
+
+
+@contextmanager
+def all_stacks_in_final_state(related_stacks_refs: list, region: str, timeout: Optional[int], interval: int):
+    ''' Wait and check if all related stacks are in a final state before performing code block
+    changes. If there is no timeout, we don't wait anything and just execute the traffic change.
+
+    :param related_stacks_refs: Related stacks to wait
+    :param region: region where stacks are present
+    :param timeout: optional value of how long we should wait for the stack should be `None`
+    :param interval: interval between checks using AWS CF API
+    '''
+    if timeout is None or timeout < 1:
+        yield
+    else:
+        wait_timeout = datetime.datetime.utcnow() + datetime.timedelta(seconds=timeout)
+
+        all_in_final_state = False
+        while not all_in_final_state and wait_timeout > datetime.datetime.utcnow():
+            # assume all stacks are ready
+            all_in_final_state = True
+            related_stacks = list(get_stacks(related_stacks_refs, region))
+
+            if not related_stacks:
+                error("Stack not found!")
+                exit(1)
+
+            for related_stack in related_stacks:
+                current_stack_status = related_stack.StackStatus
+                if current_stack_status.endswith('_IN_PROGRESS'):
+                    # some operation in progress, let's wait some time to try again
+                    all_in_final_state = False
+                    info(
+                        "Waiting for stack {} ({}) to perform requested operation..".format(
+                            related_stack.StackName, current_stack_status))
+                    time.sleep(interval)
+
+        if datetime.datetime.utcnow() > wait_timeout:
+            info("Timeout reached, requested operation not executed.")
+            exit(1)
+        else:
+            yield

senza/cli.py

@@ -29,7 +29,7 @@
                         watchrefresh_option)
 from .aws import (StackReference, get_required_capabilities, get_stacks,
                   get_tag, matches_any, parse_time, resolve_topic_arn,
-                  update_stack_from_template)
+                  update_stack_from_template, all_stacks_in_final_state)
 from .components import evaluate_template, get_component
 from .components.stups_auto_configuration import find_taupage_image
 from .definitions import AccountArguments
@@ -81,7 +81,6 @@
     'ERROR': {'fg': 'red'},
 }
 
-
 TITLES = {
     'creation_time': 'Created',
     'LogicalResourceId': 'Resource ID',
@@ -847,7 +846,7 @@ def events(stack_ref, region, w, watch, output):
 
         with OutputFormat(output):
             print_table(('stack_name version resource_type LogicalResourceId ' +
-                        'ResourceStatus ResourceStatusReason event_time').split(),
+                         'ResourceStatus ResourceStatusReason event_time').split(),
                         rows, styles=STYLES, titles=TITLES, max_column_widths=MAX_COLUMN_WIDTHS)
 
 
@@ -982,7 +981,6 @@ def instances(stack_ref, all, terminated, docker_image, piu, odd_host, region,
             if not stack_refs or matches_any(cf_stack_name, stack_refs):
                 instance_health = get_instance_health(cf_stack_name, region)
                 if instance.state['Name'].upper() != 'TERMINATED' or terminated:
-
                     docker_source = get_instance_docker_image_source(instance) if docker_image else ''
 
                     rows.append({'stack_name': stack_name or '',
@@ -1148,50 +1146,31 @@ def domains(stack_ref, region, output, w, watch):
 @click.argument('stack_name')
 @click.argument('stack_version', required=False)
 @click.argument('percentage', type=FloatRange(0, 100, clamp=True), required=False)
+@click.option('-t', '--timeout', default=None,
+              type=click.IntRange(1, 600, clamp=True),
+              help=('Timeout for waiting for stacks to be ready to perform the '
+                    'traffic change (by default it does not wait)'))
 @click.option('-i', '--interval', default=5,
               type=click.IntRange(1, 600, clamp=True),
               help='Time between checks (default: 5s)')
 @region_option
 @output_option
 @stacktrace_visible_option
-def traffic(stack_name, stack_version, percentage, region, output, interval):
-    """Route traffic to a specific stack (weighted DNS record)"""
+def traffic(stack_name, stack_version, percentage, region, output, timeout, interval):
+    '''Route traffic to a specific stack (weighted DNS record)'''
 
     stack_refs = get_stack_refs([stack_name, stack_version])
-    related_stack_refs = get_stack_refs([stack_name])
     region = get_region(region)
     check_credentials(region)
 
     with OutputFormat(output):
-        for ref in stack_refs:
+        for ref in stack_refs:  # it's expected to always iterate only once
             if percentage is None:
                 print_version_traffic(ref, region)
             else:
-                all_stacks_in_final_state = False
-                while not all_stacks_in_final_state:
-                    # assume all stacks are ready
-                    all_stacks_in_final_state = True
-                    related_stacks = list(get_stacks(related_stack_refs, region))
-
-                    if len(related_stacks) > 0:
-                        for related_stack in related_stacks:
-                            current_stack_status = related_stack.StackStatus
-
-                            if current_stack_status.endswith('_COMPLETE') or current_stack_status.endswith('_FAILED'):
-                                continue
-                            elif current_stack_status.endswith('_IN_PROGRESS'):
-                                # some operation in progress, let's wait some time to try again
-                                all_stacks_in_final_state = False
-                                info(
-                                    "Waiting for stack {} ({}) to perform traffic change..".format(
-                                        related_stack.StackName, current_stack_status))
-                                time.sleep(interval)
-                    else:
-                        error("Stack not found!")
-                        exit(1)
-
-                # change traffic after all related stacks are in a final state
-                change_version_traffic(ref, percentage, region)
+                related_stacks_refs = get_stack_refs([stack_name])
+                with all_stacks_in_final_state(related_stacks_refs, region, timeout=timeout, interval=interval):
+                    change_version_traffic(ref, percentage, region)
 
 
 @cli.command()
@@ -1473,7 +1452,7 @@ def scale(stack_ref, region, desired_capacity):
         group = get_auto_scaling_group(asg, asg_name)
         current_capacity = group['DesiredCapacity']
         with Action('Scaling {} from {} to {} instances..'.format(
-                    asg_name, current_capacity, desired_capacity)) as act:
+                asg_name, current_capacity, desired_capacity)) as act:
             if current_capacity == desired_capacity:
                 act.ok('NO CHANGES')
             else:

senza/traffic.py

@@ -259,6 +259,7 @@ def dns_name(self):
 
 
 def get_stack_versions(stack_name: str, region: str) -> Iterator[StackVersion]:
+    '''Get stack versions by name and region.'''
     cf = boto3.resource('cloudformation', region)
     for stack in get_stacks([StackReference(name=stack_name, version=None)], region):
         if stack.StackStatus in ('ROLLBACK_COMPLETE', 'CREATE_FAILED'):

tests/test_cli.py

@@ -1419,14 +1419,14 @@ def _fake_progress_of_stack_changes(stack_refs, *args, **kwargs) -> List:
                 return [
                     SenzaStackSummary({'StackName': 'myapp',
                                        'StackStatus': stacks_state_progress_queue.popleft()})
-                    for ref in stack_refs]
+                    for _ in stack_refs]
             else:
                 return []
 
-        monkeypatch.setattr('senza.cli.get_stacks', _fake_progress_of_stack_changes)
+        monkeypatch.setattr('senza.aws.get_stacks', _fake_progress_of_stack_changes)
 
         with runner.isolated_filesystem():
-            sub_command = ['traffic', '--region=aa-fakeregion-1', 'myapp', target_stack_version, '100']
+            sub_command = ['traffic', '--region=aa-fakeregion-1', 'myapp', target_stack_version, '100', '-t', '200']
             return runner.invoke(cli, sub_command, catch_exceptions=False)
 
     mocked_change_version_traffic = MagicMock(name='mocked_change_version_traffic')
@@ -1453,7 +1453,7 @@ def _reset_mocks_ctx():
     with _reset_mocks_ctx():
         result = _run_for_stacks_states_changes(['UPDATE_IN_PROGRESS', 'UPDATE_COMPLETE'])
 
-        assert 'Waiting for stack myapp (UPDATE_IN_PROGRESS) to perform traffic change..' in result.output
+        assert 'Waiting for stack myapp (UPDATE_IN_PROGRESS) to perform requested operation..' in result.output
         mocked_time_sleep.assert_called_once_with(5)
         mocked_change_version_traffic.assert_called_once_with(get_stack_refs(['myapp', 'v1'])[0], 100.0,
                                                               'aa-fakeregion-1')
@@ -1462,7 +1462,7 @@ def _reset_mocks_ctx():
     with _reset_mocks_ctx():
         result = _run_for_stacks_states_changes(['CREATE_IN_PROGRESS', 'CREATE_FAILED'])
 
-        assert 'Waiting for stack myapp (CREATE_IN_PROGRESS) to perform traffic change..' in result.output
+        assert 'Waiting for stack myapp (CREATE_IN_PROGRESS) to perform requested operation..' in result.output
         mocked_time_sleep.assert_called_once_with(5)
 
     # test stack ready to change