Retry when tagging EC2 and EBS with Postgres roles. (#418)

Retry if an error happens when setting Role or Name tags for the
EC2 instances or EBS volumes. The maximum retry interval is 15 seconds.

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():

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())