diff --git a/patroni/scripts/aws.py b/patroni/scripts/aws.py index 63f89534..c3be3ca5 100755 --- a/patroni/scripts/aws.py +++ b/patroni/scripts/aws.py @@ -6,13 +6,18 @@ from requests.exceptions import RequestException import sys import boto.ec2 +from patroni.utils import Retry, RetryFailedError + logger = logging.getLogger(__name__) +retry_timeout = 15 + class AWSConnection(object): def __init__(self, cluster_name): self.available = False self.cluster_name = cluster_name if cluster_name is not None else 'unknown' + self._retry = Retry(deadline=retry_timeout, max_delay=5, max_tries=-1, retry_exceptions=(boto.exception,)) try: # get the instance id r = requests.get('http://169.254.169.254/latest/dynamic/instance-identity/document', timeout=0.1) @@ -29,40 +34,36 @@ class AWSConnection(object): return self.available = True + def retry(self, *args, **kwargs): + return self._retry.copy()(*args, **kwargs) + def aws_available(self): return self.available - def _tag_ebs(self, role): + def _tag_ebs(self, conn, role): """ set tags, carrying the cluster name, instance role and instance id for the EBS storage """ - if not self.available: - return False - tags = {'Name': 'spilo_' + self.cluster_name, 'Role': role, 'Instance': self.instance_id} - try: - conn = boto.ec2.connect_to_region(self.region) - volumes = conn.get_all_volumes(filters={'attachment.instance-id': self.instance_id}) - conn.create_tags([v.id for v in volumes], tags) - except Exception as e: - logger.info('could not set tags for EBS storage devices attached: {}'.format(e)) - return False - return True + volumes = conn.get_all_volumes(filters={'attachment.instance-id': self.instance_id}) + conn.create_tags([v.id for v in volumes], tags) - def _tag_ec2(self, role): + def _tag_ec2(self, conn, role): """ tag the current EC2 instance with a cluster role """ - if not self.available: - return False tags = {'Role': role} - try: - conn = boto.ec2.connect_to_region(self.region) - conn.create_tags([self.instance_id], tags) - except Exception as e: - logger.info("could not set tags for EC2 instance %s: %s", self.instance_id, e) - return False - return True + conn.create_tags([self.instance_id], tags) def on_role_change(self, new_role): - ret = self._tag_ec2(new_role) - return self._tag_ebs(new_role) and ret + if not self.available: + return False + try: + conn = self.retry(boto.ec2.connect_to_region, self.region) + self.retry(self._tag_ec2, conn, new_role) + self.retry(self._tag_ebs, conn, new_role) + except RetryFailedError: + logger.warning("Unable to communicate to AWS " + "when setting tags for the EC2 instance {0} " + "and attached EBS volumes".format(self.instance_id)) + return False + return True def main(): diff --git a/tests/test_aws.py b/tests/test_aws.py index be4b918d..457aa1ad 100644 --- a/tests/test_aws.py +++ b/tests/test_aws.py @@ -6,6 +6,7 @@ import unittest from mock import Mock, patch from collections import namedtuple from patroni.scripts.aws import AWSConnection, main as _main +from patroni.utils import RetryFailedError from requests.exceptions import RequestException @@ -16,13 +17,13 @@ class MockEc2Connection(object): def get_all_volumes(self, filters): if self.error: - raise Exception("get_all_volumes") + raise boto.exception("get_all_volumes") oid = namedtuple('Volume', 'id') return [oid(id='a'), oid(id='b')] def create_tags(self, objects, tags): if self.error or len(objects) == 0: - raise Exception("create_tags") + raise boto.exception("create_tags") return True @@ -63,30 +64,20 @@ class TestAWSConnection(unittest.TestCase): self.assertTrue(self.conn.aws_available()) def test_on_role_change(self): - self.assertTrue(self.conn._tag_ebs('master')) - self.assertTrue(self.conn._tag_ec2('master')) self.assertTrue(self.conn.on_role_change('master')) + self.conn.retry = Mock(side_effect=RetryFailedError("retry failed")) + self.assertFalse(self.conn.on_role_change('master')) def test_non_aws(self): self.error = True conn = AWSConnection('test') - self.assertFalse(conn.aws_available()) - self.assertFalse(conn._tag_ebs('master')) - self.assertFalse(conn._tag_ec2('master')) + self.assertFalse(conn.on_role_change("master")) def test_aws_bizare_response(self): self.json_error = True conn = AWSConnection('test') self.assertFalse(conn.aws_available()) - def test_aws_tag_ebs_error(self): - self.error = True - self.assertFalse(self.conn._tag_ebs("master")) - - def test_aws_tag_ec2_error(self): - self.error = True - self.assertFalse(self.conn._tag_ec2("master")) - @patch('sys.exit', Mock()) def test_main(self): self.assertIsNone(_main())