From 876cfdfb2d1228f8ed43eabf8a02eaff37cdac4d Mon Sep 17 00:00:00 2001
From: Alexander Kukushkin <alexander.kukushkin@zalando.de>
Date: Wed, 29 Jun 2016 15:30:54 +0200
Subject: [PATCH] Fix retry logic in etcd.py

Client class takes care about retrying when connection to the etcd node
fails. It calculates amount of retries and timeout depending on etcd
cluster size.

Etcd class should not retry when EtcdConnectionFailed exception is
raised (this case is already handled in the Client).

Besides that adjust retry timeouts in the Client class.
---
 patroni/dcs/etcd.py | 27 ++++++++++++++++-----------
 tests/test_etcd.py  |  5 +++--
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/patroni/dcs/etcd.py b/patroni/dcs/etcd.py
index 7b2b7477..ed41359e 100644
--- a/patroni/dcs/etcd.py
+++ b/patroni/dcs/etcd.py
@@ -26,7 +26,7 @@ class EtcdError(DCSError):
 class Client(etcd.Client):
 
     def __init__(self, config):
-        super(Client, self).__init__(read_timeout=5)
+        super(Client, self).__init__(read_timeout=config['retry_timeout']/2.0)
         self._config = config
         self._load_machines_cache()
         self._allow_reconnect = True
@@ -50,6 +50,9 @@ class Client(etcd.Client):
             self._update_machines_cache = True
             return [self._base_uri]
 
+    def set_read_timeout(self, timeout):
+        self._read_timeout = timeout/2.0
+
     def _do_http_request(self, request_executor, method, url, fields=None, **kwargs):
         try:
             response = request_executor(method, url, fields=fields, **kwargs)
@@ -70,13 +73,7 @@ class Client(etcd.Client):
         if not path.startswith('/'):
             raise ValueError('Path does not start with /')
 
-        if timeout is None:
-            timeout = self.read_timeout
-
-        if timeout == 0:
-            timeout = None
-
-        kwargs = {'timeout': timeout, 'fields': params, 'redirect': self.allow_redirect,
+        kwargs = {'fields': params, 'redirect': self.allow_redirect,
                   'headers': self._get_headers(), 'preload_content': False}
 
         if method in [self._MGET, self._MDELETE]:
@@ -91,6 +88,14 @@ class Client(etcd.Client):
         if self._update_machines_cache:
             self._load_machines_cache()
 
+        if timeout is None:
+            kwargs['retries'] = 0 if len(self._machines_cache) > 3 else (1 if len(self._machines_cache) > 1 else 2)
+            kwargs['timeout'] = float(self.read_timeout)/(kwargs['retries'] + 1)
+            if kwargs['timeout'] < 1:
+                kwargs['timeout'] = 1
+        else:
+            kwargs.update({'retries': 0, 'timeout': timeout})
+
         response = False
 
         try:
@@ -122,7 +127,7 @@ class Client(etcd.Client):
         for host, port in self.get_srv_record(discovery_srv):
             url = '{0}://{1}:{2}/members'.format(self._protocol, host, port)
             try:
-                response = requests.get(url, timeout=5)
+                response = requests.get(url, timeout=self.read_timeout)
                 if response.ok:
                     for member in response.json():
                         ret.extend(member['clientURLs'])
@@ -195,8 +200,7 @@ class Etcd(AbstractDCS):
         super(Etcd, self).__init__(config)
         self._ttl = int(config.get('ttl') or 30)
         self._retry = Retry(deadline=config['retry_timeout'], max_delay=1, max_tries=-1,
-                            retry_exceptions=(etcd.EtcdConnectionFailed,
-                                              etcd.EtcdLeaderElectionInProgress,
+                            retry_exceptions=(etcd.EtcdLeaderElectionInProgress,
                                               etcd.EtcdWatcherCleared,
                                               etcd.EtcdEventIndexCleared))
         self._client = self.get_etcd_client(config)
@@ -223,6 +227,7 @@ class Etcd(AbstractDCS):
 
     def set_retry_timeout(self, retry_timeout):
         self._retry.deadline = retry_timeout
+        self._client.set_read_timeout(retry_timeout)
 
     @staticmethod
     def member(node):
diff --git a/tests/test_etcd.py b/tests/test_etcd.py
index ebdf6aa2..0fb6ca9c 100644
--- a/tests/test_etcd.py
+++ b/tests/test_etcd.py
@@ -147,7 +147,7 @@ class TestClient(unittest.TestCase):
     def setUp(self):
         with patch.object(etcd.Client, 'machines') as mock_machines:
             mock_machines.__get__ = Mock(return_value=['http://localhost:2379', 'http://localhost:4001'])
-            self.client = Client({'discovery_srv': 'test'})
+            self.client = Client({'discovery_srv': 'test', 'retry_timeout': 3})
             self.client.http.request = http_request
             self.client.http.request_encode_body = http_request
 
@@ -204,7 +204,8 @@ class TestEtcd(unittest.TestCase):
         with patch.object(etcd.Client, 'machines') as mock_machines:
             mock_machines.__get__ = Mock(side_effect=etcd.EtcdException)
             with patch('time.sleep', Mock(side_effect=SleepException())):
-                self.assertRaises(SleepException, self.etcd.get_etcd_client, {'discovery_srv': 'test'})
+                self.assertRaises(SleepException, self.etcd.get_etcd_client,
+                                  {'discovery_srv': 'test', 'retry_timeout': 10})
 
     def test_get_cluster(self):
         self.assertIsInstance(self.etcd.get_cluster(), Cluster)