Repository: zalando/patroni
Pull request #2484 (2022-12-13 09:45:46)

Raise DCSError when communication with DCS fails
Previously such an exception was raised only from the `get_cluster()` method, and now we will to do the same from the `update_leader()` and `attempt_to_acquire_leader()` methods.

These methods influence Postgres promotion and demotion and we want to make a difference between different types of failures. Specifically, if calls have failed because DCS isn't accessible or due to a timeout.

This commit is extracted from the #2379

coveralls: ## Pull Request Test Coverage Report for [Build 3684181378](https://coveralls.io/builds/55104698)

* **214** of **214**   **(100.0%)**  changed or added relevant lines in **8** files are covered.
* No unchanged relevant lines lost coverage.
* Overall coverage increased (+**0.001%**) to **99.894%**

---



|  Totals | [![Coverage Status](https://coveralls.io/builds/55104698/badge)](https://coveralls.io/builds/55104698) |
| :-- | --: |
| Change from base [Build 3637287276](https://coveralls.io/builds/54923089): |  0.001% |
| Covered Lines: | 11308 |
| Relevant Lines: | 11320 |

---
##### 💛  - [Coveralls](https://coveralls.io)

hughcapet: 👍 
CyberDem0n: 👍

Patch:
diff --git a/patroni/dcs/__init__.py b/patroni/dcs/__init__.py
index 064f1f75..0d414558 100644
--- a/patroni/dcs/__init__.py
+++ b/patroni/dcs/__init__.py
@@ -628,6 +628,20 @@ class Cluster(namedtuple('Cluster', 'initialize,config,leader,last_lsn,members,f
         return next(iter(sorted(filter(lambda v: v, [m.version for m in self.members])) + [None]))
 
 
+class ReturnFalseException(Exception):
+    pass
+
+
+def catch_return_false_exception(func):
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except ReturnFalseException:
+            return False
+
+    return wrapper
+
+
 @six.add_metaclass(abc.ABCMeta)
 class AbstractDCS(object):
 
@@ -801,18 +815,18 @@ class AbstractDCS(object):
         """Update leader key (or session) ttl
 
         :returns: `!True` if leader key (or session) has been updated successfully.
-            If not, `!False` must be returned and current instance would be demoted.
 
         You have to use CAS (Compare And Swap) operation in order to update leader key,
-        for example for etcd `prevValue` parameter must be used."""
+        for example for etcd `prevValue` parameter must be used.
+        If update fails due to DCS not being accessible or because it is not able to
+        process requests (hopefuly temporary), the ~DCSError exception should be raised."""
 
     def update_leader(self, last_lsn, slots=None):
         """Update leader key (or session) ttl and optime/leader
 
         :param last_lsn: absolute WAL LSN in bytes
         :param slots: dict with permanent slots confirmed_flush_lsn
-        :returns: `!True` if leader key (or session) has been updated successfully.
-            If not, `!False` must be returned and current instance would be demoted."""
+        :returns: `!True` if leader key (or session) has been updated successfully."""
 
         ret = self._update_leader()
         if ret and last_lsn:
@@ -831,7 +845,10 @@ class AbstractDCS(object):
         :returns: `!True` if key has been created successfully.
 
         Key must be created atomically. In case if key already exists it should not be
-        overwritten and `!False` must be returned"""
+        overwritten and `!False` must be returned.
+
+        If key creation fails due to DCS not being accessible or because it is not able to
+        process requests (hopefuly temporary), the ~DCSError exception should be raised"""
 
     @abc.abstractmethod
     def set_failover_value(self, value, index=None):
diff --git a/patroni/dcs/consul.py b/patroni/dcs/consul.py
index f4965d90..1d506d4e 100644
--- a/patroni/dcs/consul.py
+++ b/patroni/dcs/consul.py
@@ -14,7 +14,8 @@ from urllib3.exceptions import HTTPError
 from six.moves.urllib.parse import urlencode, urlparse, quote
 from six.moves.http_client import HTTPException
 
-from . import AbstractDCS, Cluster, ClusterConfig, Failover, Leader, Member, SyncState, TimelineHistory
+from . import AbstractDCS, Cluster, ClusterConfig, Failover, Leader, Member,\
+        SyncState, TimelineHistory, ReturnFalseException, catch_return_false_exception
 from ..exceptions import DCSError
 from ..utils import deep_compare, parse_bool, Retry, RetryFailedError, split_host_port, uri, USER_AGENT
 
@@ -286,9 +287,9 @@ class Consul(AbstractDCS):
         except Exception:
             logger.exception('adjust_ttl')
 
-    def _do_refresh_session(self):
+    def _do_refresh_session(self, force=False):
         """:returns: `!True` if it had to create new session"""
-        if self._session and self._last_session_refresh + self._loop_wait > time.time():
+        if not force and self._session and self._last_session_refresh + self._loop_wait > time.time():
             return False
 
         if self._session:
@@ -373,11 +374,6 @@ class Consul(AbstractDCS):
 
             # get leader
             leader = nodes.get(self._LEADER)
-            if not self._ctl and leader and leader['Value'] == self._name \
-                    and self._session != leader.get('Session', 'x'):
-                logger.info('I am leader but not owner of the session. Removing leader node')
-                self._client.kv.delete(self.leader_path, cas=leader['ModifyIndex'])
-                leader = None
 
             if leader:
                 member = Member(-1, leader['Value'], None, {})
@@ -506,22 +502,34 @@ class Consul(AbstractDCS):
         ):
             return self._update_service(new_data)
 
-    @catch_consul_errors
-    def _do_attempt_to_acquire_leader(self, permanent):
+    def _do_attempt_to_acquire_leader(self, permanent, retry):
         try:
             kwargs = {} if permanent else {'acquire': self._session}
-            return self.retry(self._client.kv.put, self.leader_path, self._name, **kwargs)
+            return retry(self._client.kv.put, self.leader_path, self._name, **kwargs)
         except InvalidSession:
-            self._session = None
             logger.error('Our session disappeared from Consul. Will try to get a new one and retry attempt')
-            self.refresh_session()
-            return self.retry(self._client.kv.put, self.leader_path, self._name, acquire=self._session)
+            self._session = None
+            retry.deadline = retry.stoptime - time.time()
+
+            retry(self._do_refresh_session)
 
+            retry.deadline = retry.stoptime - time.time()
+            if retry.deadline < 1:
+                raise ConsulError('_do_attempt_to_acquire_leader timeout')
+
+            return retry(self._client.kv.put, self.leader_path, self._name, acquire=self._session)
+
+    @catch_return_false_exception
     def attempt_to_acquire_leader(self, permanent=False):
-        if not self._session and not permanent:
-            self.refresh_session()
+        retry = self._retry.copy()
+        if not permanent:
+            self._run_and_handle_exceptions(self._do_refresh_session, retry=retry)
+
+            retry.deadline = retry.stoptime - time.time()
+            if retry.deadline < 1:
+                raise ConsulError('attempt_to_acquire_leader timeout')
 
-        ret = self._do_attempt_to_acquire_leader(permanent)
+        ret = self._run_and_handle_exceptions(self._do_attempt_to_acquire_leader, permanent, retry, retry=None)
         if not ret:
             logger.info('Could not take out TTL lock')
 
@@ -546,11 +554,39 @@ class Consul(AbstractDCS):
     def _write_status(self, value):
         return self._client.kv.put(self.status_path, value)
 
-    @catch_consul_errors
+    @staticmethod
+    def _run_and_handle_exceptions(method, *args, **kwargs):
+        retry = kwargs.pop('retry', None)
+        try:
+            return retry(method, *args, **kwargs) if retry else method(*args, **kwargs)
+        except (RetryFailedError, InvalidSession, HTTPException, HTTPError, socket.error, socket.timeout) as e:
+            raise ConsulError(e)
+        except ConsulException:
+            raise ReturnFalseException
+
+    @catch_return_false_exception
     def _update_leader(self):
+        retry = self._retry.copy()
+
+        self._run_and_handle_exceptions(self._do_refresh_session, True, retry=retry)
+
         if self._session:
-            self.retry(self._client.session.renew, self._session)
-            self._last_session_refresh = time.time()
+            cluster = self.cluster
+            leader_session = cluster and isinstance(cluster.leader, Leader) and cluster.leader.session
+            if leader_session != self._session:
+                retry.deadline = retry.stoptime - time.time()
+                if retry.deadline < 1:
+                    raise ConsulError('update_leader timeout')
+                logger.warning('Recreating the leader key due to session mismatch')
+                if cluster.leader:
+                    self._run_and_handle_exceptions(self._client.kv.delete, self.leader_path, cas=cluster.leader.index)
+
+                retry.deadline = retry.stoptime - time.time()
+                if retry.deadline < 0.5:
+                    raise ConsulError('update_leader timeout')
+                self._run_and_handle_exceptions(self._client.kv.put, self.leader_path,
+                                                self._name, acquire=self._session)
+
         return bool(self._session)
 
     @catch_consul_errors
diff --git a/patroni/dcs/etcd.py b/patroni/dcs/etcd.py
index 5be6839d..2921b8e1 100644
--- a/patroni/dcs/etcd.py
+++ b/patroni/dcs/etcd.py
@@ -19,7 +19,8 @@ from six.moves.http_client import HTTPException
 from six.moves.urllib_parse import urlparse
 from threading import Thread
 
-from . import AbstractDCS, Cluster, ClusterConfig, Failover, Leader, Member, SyncState, TimelineHistory
+from . import AbstractDCS, Cluster, ClusterConfig, Failover, Leader, Member,\
+        SyncState, TimelineHistory, ReturnFalseException, catch_return_false_exception
 from ..exceptions import DCSError
 from ..request import get as requests_get
 from ..utils import Retry, RetryFailedError, split_host_port, uri, USER_AGENT
@@ -460,6 +461,18 @@ class AbstractEtcd(AbstractDCS):
         if isinstance(raise_ex, Exception):
             raise raise_ex
 
+    def _run_and_handle_exceptions(self, method, *args, **kwargs):
+        retry = kwargs.pop('retry', self.retry)
+        try:
+            return retry(method, *args, **kwargs) if retry else method(*args, **kwargs)
+        except (RetryFailedError, etcd.EtcdConnectionFailed) as e:
+            raise self._client.ERROR_CLS(e)
+        except etcd.EtcdException as e:
+            self._handle_exception(e)
+            raise ReturnFalseException
+        except Exception as e:
+            self._handle_exception(e, raise_ex=self._client.ERROR_CLS('unexpected error'))
+
     @staticmethod
     def set_socket_options(sock, socket_options):
         if socket_options:
@@ -665,7 +678,7 @@ class Etcd(AbstractEtcd):
     def take_leader(self):
         return self.retry(self._client.write, self.leader_path, self._name, ttl=self._ttl)
 
-    def attempt_to_acquire_leader(self, permanent=False):
+    def _do_attempt_to_acquire_leader(self, permanent=False):
         try:
             return bool(self.retry(self._client.write,
                                    self.leader_path,
@@ -674,9 +687,11 @@ class Etcd(AbstractEtcd):
                                    prevExist=False))
         except etcd.EtcdAlreadyExist:
             logger.info('Could not take out TTL lock')
-        except (RetryFailedError, etcd.EtcdException):
-            pass
-        return False
+            return False
+
+    @catch_return_false_exception
+    def attempt_to_acquire_leader(self, permanent=False):
+        return self._run_and_handle_exceptions(self._do_attempt_to_acquire_leader, permanent=permanent, retry=None)
 
     @catch_etcd_errors
     def set_failover_value(self, value, index=None):
@@ -694,9 +709,16 @@ class Etcd(AbstractEtcd):
     def _write_status(self, value):
         return self._client.set(self.status_path, value)
 
-    @catch_etcd_errors
+    def _do_update_leader(self):
+        try:
+            return self.retry(self._client.write, self.leader_path, self._name,
+                              prevValue=self._name, ttl=self._ttl) is not None
+        except etcd.EtcdKeyNotFound:
+            return self._do_attempt_to_acquire_leader()
+
+    @catch_return_false_exception
     def _update_leader(self):
-        return self.retry(self._client.write, self.leader_path, self._name, prevValue=self._name, ttl=self._ttl)
+        return self._run_and_handle_exceptions(self._do_update_leader, retry=None)
 
     @catch_etcd_errors
     def initialize(self, create_new=True, sysid=""):
diff --git a/patroni/dcs/etcd3.py b/patroni/dcs/etcd3.py
index a507c421..3403f323 100644
--- a/patroni/dcs/etcd3.py
+++ b/patroni/dcs/etcd3.py
@@ -13,7 +13,8 @@ import urllib3
 from threading import Condition, Lock, Thread
 from urllib3.exceptions import ReadTimeoutError, ProtocolError
 
-from . import ClusterConfig, Cluster, Failover, Leader, Member, SyncState, TimelineHistory
+from . import ClusterConfig, Cluster, Failover, Leader, Member,\
+        SyncState, TimelineHistory, ReturnFalseException, catch_return_false_exception
 from .etcd import AbstractEtcdClientWithFailover, AbstractEtcd, catch_etcd_errors
 from ..exceptions import DCSError, PatroniException
 from ..utils import deep_compare, enable_keepalive, iter_response_objects, RetryFailedError, USER_AGENT
@@ -609,8 +610,8 @@ class Etcd3(AbstractEtcd):
         if self.__do_not_watch:
             self._lease = None
 
-    def _do_refresh_lease(self, retry=None):
-        if self._lease and self._last_lease_refresh + self._loop_wait > time.time():
+    def _do_refresh_lease(self, force=False, retry=None):
+        if not force and self._lease and self._last_lease_refresh + self._loop_wait > time.time():
             return False
 
         if self._lease and not self._client.lease_keepalive(self._lease, retry):
@@ -744,21 +745,42 @@ class Etcd3(AbstractEtcd):
     def take_leader(self):
         return self.retry(self._client.put, self.leader_path, self._name, self._lease)
 
-    @catch_etcd_errors
-    def _do_attempt_to_acquire_leader(self, permanent):
+    def _do_attempt_to_acquire_leader(self, permanent, retry):
+        def _retry(*args, **kwargs):
+            kwargs['retry'] = retry
+            return retry(*args, **kwargs)
+
         try:
-            return self.retry(self._client.put, self.leader_path, self._name, None if permanent else self._lease, 0)
+            return _retry(self._client.put, self.leader_path, self._name, None if permanent else self._lease, 0)
         except LeaseNotFound:
-            self._lease = None
             logger.error('Our lease disappeared from Etcd. Will try to get a new one and retry attempt')
-            self.refresh_lease()
-            return self.retry(self._client.put, self.leader_path, self._name, None if permanent else self._lease, 0)
+            self._lease = None
+            retry.deadline = retry.stoptime - time.time()
 
+            _retry(self._do_refresh_lease)
+
+            retry.deadline = retry.stoptime - time.time()
+            if retry.deadline < 1:
+                raise Etcd3Error('_do_attempt_to_acquire_leader timeout')
+
+            return _retry(self._client.put, self.leader_path, self._name, None if permanent else self._lease, 0)
+
+    @catch_return_false_exception
     def attempt_to_acquire_leader(self, permanent=False):
-        if not self._lease and not permanent:
-            self.refresh_lease()
+        retry = self._retry.copy()
+
+        def _retry(*args, **kwargs):
+            kwargs['retry'] = retry
+            return retry(*args, **kwargs)
 
-        ret = self._do_attempt_to_acquire_leader(permanent)
+        if not permanent:
+            self._run_and_handle_exceptions(self._do_refresh_lease, retry=_retry)
+
+            retry.deadline = retry.stoptime - time.time()
+            if retry.deadline < 1:
+                raise Etcd3Error('attempt_to_acquire_leader timeout')
+
+        ret = self._run_and_handle_exceptions(self._do_attempt_to_acquire_leader, permanent, retry, retry=None)
         if not ret:
             logger.info('Could not take out TTL lock')
         return ret
@@ -779,18 +801,29 @@ class Etcd3(AbstractEtcd):
     def _write_status(self, value):
         return self._client.put(self.status_path, value)
 
-    @catch_etcd_errors
+    @catch_return_false_exception
     def _update_leader(self):
-        if not self._lease:
-            self.refresh_lease()
-        elif self.retry(self._client.lease_keepalive, self._lease):
-            self._last_lease_refresh = time.time()
+        retry = self._retry.copy()
+
+        def _retry(*args, **kwargs):
+            kwargs['retry'] = retry
+            return retry(*args, **kwargs)
+
+        self._run_and_handle_exceptions(self._do_refresh_lease, True, retry=_retry)
 
         if self._lease:
             cluster = self.cluster
             leader_lease = cluster and isinstance(cluster.leader, Leader) and cluster.leader.session
             if leader_lease != self._lease:
-                self.take_leader()
+                retry.deadline = retry.stoptime - time.time()
+                if retry.deadline < 1:
+                    raise Etcd3Error('update_leader timeout')
+
+                try:
+                    self._run_and_handle_exceptions(self._client.put, self.leader_path,
+                                                    self._name, self._lease, retry=_retry)
+                except ReturnFalseException:
+                    pass
         return bool(self._lease)
 
     @catch_etcd_errors
diff --git a/patroni/dcs/kubernetes.py b/patroni/dcs/kubernetes.py
index 21625b87..a0fe3dfe 100644
--- a/patroni/dcs/kubernetes.py
+++ b/patroni/dcs/kubernetes.py
@@ -524,16 +524,10 @@ class CoreV1ApiProxy(object):
 
 
 def catch_kubernetes_errors(func):
-    def wrapper(*args, **kwargs):
+    def wrapper(self, *args, **kwargs):
         try:
-            return func(*args, **kwargs)
-        except k8s_client.rest.ApiException as e:
-            if e.status == 403:
-                logger.exception('Permission denied')
-            elif e.status != 409:  # Object exists or conflict in resource_version
-                logger.exception('Unexpected error from Kubernetes API')
-            return False
-        except (RetryFailedError, K8sException):
+            return self._run_and_handle_exceptions(func, self, *args, **kwargs)
+        except KubernetesError:
             return False
     return wrapper
 
@@ -743,6 +737,19 @@ class Kubernetes(AbstractDCS):
         kwargs['_retry'] = retry
         return retry(*args, **kwargs)
 
+    @staticmethod
+    def _run_and_handle_exceptions(method, *args, **kwargs):
+        try:
+            return method(*args, **kwargs)
+        except k8s_client.rest.ApiException as e:
+            if e.status == 403:
+                logger.exception('Permission denied')
+            elif e.status != 409:  # Object exists or conflict in resource_version
+                logger.exception('Unexpected error from Kubernetes API')
+            return False
+        except (RetryFailedError, K8sException) as e:
+            raise KubernetesError(e)
+
     def client_path(self, path):
         return super(Kubernetes, self).client_path(path)[1:].replace('/', '-')
 
@@ -1010,16 +1017,19 @@ class Kubernetes(AbstractDCS):
             else:
                 logger.exception('Permission denied' if e.status == 403 else 'Unexpected error from Kubernetes API')
                 return False
-        except (RetryFailedError, K8sException):
-            return False
+        except (RetryFailedError, K8sException) as e:
+            raise KubernetesError(e)
 
+        # if we are here, that means update failed with 409
         retry.deadline = retry.stoptime - time.time()
         if retry.deadline < 1:
-            return False
+            return False  # No time for retry. Tell ha.py that we have to demote due to failed update.
 
         # Try to get the latest version directly from K8s API instead of relying on async cache
         try:
             kind = _retry(self._api.read_namespaced_kind, self.leader_path, self._namespace)
+        except (RetryFailedError, K8sException) as e:
+            raise KubernetesError(e)
         except Exception as e:
             logger.error('Failed to get the leader object "%s": %r', self.leader_path, e)
             return False
@@ -1037,7 +1047,8 @@ class Kubernetes(AbstractDCS):
         if kind and (kind_annotations.get(self._LEADER) != self._name or kind_resource_version == resource_version):
             return False
 
-        return self.patch_or_create(self.leader_path, annotations, kind_resource_version, ips=ips, retry=_retry)
+        return self._run_and_handle_exceptions(self._patch_or_create, self.leader_path, annotations,
+                                               kind_resource_version, ips=ips, retry=_retry)
 
     def update_leader(self, last_lsn, slots=None):
         kind = self._kinds.get(self.leader_path)
@@ -1074,7 +1085,19 @@ class Kubernetes(AbstractDCS):
                 annotations['acquireTime'] = self._leader_observed_record.get('acquireTime') or now
             annotations['transitions'] = str(transitions)
         ips = [] if self._api.use_endpoints else None
-        ret = self.patch_or_create(self.leader_path, annotations, self._leader_resource_version, ips=ips)
+
+        try:
+            ret = self._patch_or_create(self.leader_path, annotations,
+                                        self._leader_resource_version, retry=self.retry, ips=ips)
+        except k8s_client.rest.ApiException as e:
+            if e.status == 409 and self._leader_resource_version:  # Conflict in resource_version
+                # Terminate watchers, it could be a sign that K8s API is in a failed state
+                self._kinds.kill_stream()
+                self._pods.kill_stream()
+            ret = False
+        except (RetryFailedError, K8sException) as e:
+            raise KubernetesError(e)
+
         if not ret:
             logger.info('Could not take out TTL lock')
         return ret
diff --git a/patroni/dcs/raft.py b/patroni/dcs/raft.py
index 9563f31c..1318c188 100644
--- a/patroni/dcs/raft.py
+++ b/patroni/dcs/raft.py
@@ -11,11 +11,16 @@ from pysyncobj.transport import TCPTransport, CONNECTION_STATE
 from pysyncobj.utility import TcpUtility
 
 from . import AbstractDCS, ClusterConfig, Cluster, Failover, Leader, Member, SyncState, TimelineHistory
+from ..exceptions import DCSError
 from ..utils import validate_directory
 
 logger = logging.getLogger(__name__)
 
 
+class RaftError(DCSError):
+    pass
+
+
 class _TCPTransport(TCPTransport):
 
     def __init__(self, syncObj, selfNode, otherNodes):
@@ -156,7 +161,7 @@ class KVStoreTTL(DynMemberSyncObj):
             elif deadline:
                 timeout = deadline - time.time()
                 if timeout <= 0:
-                    break
+                    raise RaftError('timeout')
             time.sleep(1)
         return False
 
@@ -175,7 +180,7 @@ class KVStoreTTL(DynMemberSyncObj):
             self.__on_set(key, value)
         return True
 
-    def set(self, key, value, ttl=None, **kwargs):
+    def set(self, key, value, ttl=None, handle_raft_error=True, **kwargs):
         old_value = self.__data.get(key, {})
         if not self.__check_requirements(old_value, **kwargs):
             return False
@@ -184,7 +189,12 @@ class KVStoreTTL(DynMemberSyncObj):
         value['created'] = old_value.get('created', value['updated'])
         if ttl:
             value['expire'] = value['updated'] + ttl
-        return self.retry(self._set, key, value, **kwargs)
+        try:
+            return self.retry(self._set, key, value, **kwargs)
+        except RaftError:
+            if not handle_raft_error:
+                raise
+            return False
 
     def __pop(self, key):
         self.__data.pop(key)
@@ -206,7 +216,10 @@ class KVStoreTTL(DynMemberSyncObj):
     def delete(self, key, recursive=False, **kwargs):
         if not recursive and not self.__check_requirements(self.__data.get(key, {}), **kwargs):
             return False
-        return self.retry(self._delete, key, recursive=recursive, **kwargs)
+        try:
+            return self.retry(self._delete, key, recursive=recursive, **kwargs)
+        except RaftError:
+            return False
 
     @staticmethod
     def __values_match(old, new):
@@ -372,14 +385,15 @@ class Raft(AbstractDCS):
         return self._sync_obj.set(self.status_path, value, timeout=1)
 
     def _update_leader(self):
-        ret = self._sync_obj.set(self.leader_path, self._name, ttl=self._ttl, prevValue=self._name)
+        ret = self._sync_obj.set(self.leader_path, self._name, ttl=self._ttl,
+                                 handle_raft_error=False, prevValue=self._name)
         if not ret and self._sync_obj.get(self.leader_path) is None:
             ret = self.attempt_to_acquire_leader()
         return ret
 
     def attempt_to_acquire_leader(self, permanent=False):
-        return self._sync_obj.set(self.leader_path, self._name, prevExist=False,
-                                  ttl=None if permanent else self._ttl)
+        return self._sync_obj.set(self.leader_path, self._name, ttl=None if permanent else self._ttl,
+                                  handle_raft_error=False, prevExist=False)
 
     def set_failover_value(self, value, index=None):
         return self._sync_obj.set(self.failover_path, value, prevIndex=index)
diff --git a/patroni/dcs/zookeeper.py b/patroni/dcs/zookeeper.py
index daba67ae..9101244f 100644
--- a/patroni/dcs/zookeeper.py
+++ b/patroni/dcs/zookeeper.py
@@ -5,9 +5,10 @@ import six
 import time
 
 from kazoo.client import KazooClient, KazooState, KazooRetry
-from kazoo.exceptions import NoNodeError, NodeExistsError, SessionExpiredError
+from kazoo.exceptions import ConnectionClosedError, NoNodeError, NodeExistsError, SessionExpiredError
 from kazoo.handlers.threading import SequentialThreadingHandler
 from kazoo.protocol.states import KeeperState
+from kazoo.retry import RetryFailedError
 from kazoo.security import make_acl
 
 from . import AbstractDCS, ClusterConfig, Cluster, Failover, Leader, Member, SyncState, TimelineHistory
@@ -263,18 +264,10 @@ class ZooKeeper(AbstractDCS):
         # get leader
         leader = self.get_node(self.leader_path) if self._LEADER in nodes else None
         if leader:
-            client_id = self._client.client_id
-            if not self._ctl and leader[0] == self._name and client_id is not None \
-                    and client_id[0] != leader[1].ephemeralOwner:
-                logger.info('I am leader but not owner of the session. Removing leader node')
-                self._client.delete(self.leader_path)
-                leader = None
-
-            if leader:
-                member = Member(-1, leader[0], None, {})
-                member = ([m for m in members if m.name == leader[0]] or [member])[0]
-                leader = Leader(leader[1].version, leader[1].ephemeralOwner, member)
-                self._fetch_cluster = member.index == -1
+            member = Member(-1, leader[0], None, {})
+            member = ([m for m in members if m.name == leader[0]] or [member])[0]
+            leader = Leader(leader[1].version, leader[1].ephemeralOwner, member)
+            self._fetch_cluster = member.index == -1
 
         # get last known leader lsn and slots
         last_lsn, slots = self.get_status(leader)
@@ -325,10 +318,17 @@ class ZooKeeper(AbstractDCS):
         return False
 
     def attempt_to_acquire_leader(self, permanent=False):
-        ret = self._create(self.leader_path, self._name.encode('utf-8'), retry=True, ephemeral=not permanent)
-        if not ret:
-            logger.info('Could not take out TTL lock')
-        return ret
+        try:
+            self._client.retry(self._client.create, self.leader_path, self._name.encode('utf-8'),
+                               makepath=True, ephemeral=not permanent)
+            return True
+        except (ConnectionClosedError, RetryFailedError) as e:
+            raise ZooKeeperError(e)
+        except Exception as e:
+            if not isinstance(e, NodeExistsError):
+                logger.error('Failed to create %s: %r', self.leader_path, e)
+        logger.info('Could not take out TTL lock')
+        return False
 
     def _set_or_create(self, key, value, index=None, retry=False, do_not_create_empty=False):
         value = value.encode('utf-8')
@@ -409,6 +409,28 @@ class ZooKeeper(AbstractDCS):
         return self._set_or_create(self.status_path, value)
 
     def _update_leader(self):
+        cluster = self.cluster
+        session = cluster and isinstance(cluster.leader, Leader) and cluster.leader.session
+        if self._client.client_id and self._client.client_id[0] != session:
+            logger.warning('Recreating the leader ZNode due to ownership mismatch')
+            try:
+                self._client.retry(self._client.delete, self.leader_path)
+            except NoNodeError:
+                pass
+            except (ConnectionClosedError, RetryFailedError) as e:
+                raise ZooKeeperError(e)
+            except Exception as e:
+                logger.error('Failed to remove %s: %r', self.leader_path, e)
+                return False
+
+            try:
+                self._client.retry(self._client.create, self.leader_path,
+                                   self._name.encode('utf-8'), makepath=True, ephemeral=True)
+            except (ConnectionClosedError, RetryFailedError) as e:
+                raise ZooKeeperError(e)
+            except Exception as e:
+                logger.error('Failed to create %s: %r', self.leader_path, e)
+                return False
         return True
 
     def _delete_leader(self):
diff --git a/patroni/ha.py b/patroni/ha.py
index ba4fd2cf..03f31d87 100644
--- a/patroni/ha.py
+++ b/patroni/ha.py
@@ -146,7 +146,13 @@ class Ha(object):
         self._leader_timeline = None if cluster.is_unlocked() else cluster.leader.timeline
 
     def acquire_lock(self):
-        ret = self.dcs.attempt_to_acquire_leader()
+        try:
+            ret = self.dcs.attempt_to_acquire_leader()
+        except DCSError:
+            raise
+        except Exception:
+            logger.exception('Unexpected exception raised from attempt_to_acquire_leader, please report it as a BUG')
+            ret = False
         self.set_is_leader(ret)
         return ret
 
@@ -160,6 +166,8 @@ class Ha(object):
                 logger.exception('Exception when called state_handler.last_operation()')
         try:
             ret = self.dcs.update_leader(last_lsn, slots)
+        except DCSError:
+            raise
         except Exception:
             logger.exception('Unexpected exception raised from update_leader, please report it as a BUG')
             ret = False
@@ -1522,8 +1530,12 @@ class Ha(object):
             dcs_failed = True
             logger.error('Error communicating with DCS')
             if not self.is_paused() and self.state_handler.is_running() and self.state_handler.is_leader():
+                msg = 'demoting self because DCS is not accessible and I was a leader'
+                if not self._async_executor.try_run_async(msg, self.demote, ('offline',)):
+                    return msg
+                logger.warning('AsyncExecutor is busy, demoting from the main thread')
                 self.demote('offline')
-                return 'demoted self because DCS is not accessible and i was a leader'
+                return 'demoted self because DCS is not accessible and I was a leader'
             return 'DCS is not accessible'
         except (psycopg.Error, PostgresConnectionException):
             return 'Error communicating with PostgreSQL. Will try again later'


Test patch:
diff --git a/tests/test_consul.py b/tests/test_consul.py
index 291eda69..03c90273 100644
--- a/tests/test_consul.py
+++ b/tests/test_consul.py
@@ -4,7 +4,7 @@ import unittest
 from consul import ConsulException, NotFound
 from mock import Mock, PropertyMock, patch
 from patroni.dcs.consul import AbstractDCS, Cluster, Consul, ConsulInternalError, \
-                                ConsulError, ConsulClient, HTTPClient, InvalidSessionTTL, InvalidSession
+        ConsulError, ConsulClient, HTTPClient, InvalidSessionTTL, InvalidSession, RetryFailedError
 from . import SleepException
 
 
@@ -122,9 +122,6 @@ class TestConsul(unittest.TestCase):
         self.assertIsInstance(self.c.get_cluster(), Cluster)
         self.c._base_path = '/service/legacy'
         self.assertIsInstance(self.c.get_cluster(), Cluster)
-        self.c._base_path = '/service/good'
-        self.c._session = 'fd4f44fe-2cac-bba5-a60b-304b51ff39b8'
-        self.assertIsInstance(self.c.get_cluster(), Cluster)
 
     @patch.object(consul.Consul.KV, 'delete', Mock(side_effect=[ConsulException, True, True, True]))
     @patch.object(consul.Consul.KV, 'put', Mock(side_effect=[True, ConsulException, InvalidSession]))
@@ -140,11 +137,15 @@ class TestConsul(unittest.TestCase):
         self.c.refresh_session = Mock(side_effect=ConsulError('foo'))
         self.assertFalse(self.c.touch_member({'balbla': 'blabla'}))
 
-    @patch.object(consul.Consul.KV, 'put', Mock(side_effect=InvalidSession))
+    @patch.object(consul.Consul.KV, 'put', Mock(side_effect=[InvalidSession, False, InvalidSession]))
     def test_take_leader(self):
         self.c.set_ttl(20)
-        self.c.refresh_session = Mock()
-        self.c.take_leader()
+        self.c._do_refresh_session = Mock()
+        self.assertFalse(self.c.take_leader())
+        with patch('time.time', Mock(side_effect=[0, 100])):
+            self.assertRaises(ConsulError, self.c.take_leader)
+        with patch('time.time', Mock(side_effect=[0, 0, 0, 0, 0, 0, 100])):
+            self.assertRaises(ConsulError, self.c.take_leader)
 
     @patch.object(consul.Consul.KV, 'put', Mock(return_value=True))
     def test_set_failover_value(self):
@@ -160,10 +161,26 @@ class TestConsul(unittest.TestCase):
         self.c.get_cluster()
         self.c.write_leader_optime('1')
 
-    @patch.object(consul.Consul.Session, 'renew', Mock())
+    @patch.object(consul.Consul.Session, 'renew')
     @patch.object(consul.Consul.KV, 'put', Mock(side_effect=ConsulException))
-    def test_update_leader(self):
-        self.c.update_leader(12345)
+    def test_update_leader(self, mock_renew):
+        self.c._session = 'fd4f44fe-2cac-bba5-a60b-304b51ff39b8'
+        with patch.object(consul.Consul.KV, 'delete', Mock(return_value=True)):
+            with patch.object(consul.Consul.KV, 'put', Mock(return_value=True)):
+                self.assertTrue(self.c.update_leader(12345))
+            with patch.object(consul.Consul.KV, 'put', Mock(side_effect=ConsulException)):
+                self.assertFalse(self.c.update_leader(12345))
+            with patch('time.time', Mock(side_effect=[0, 0, 0, 0, 0, 100, 200, 300])):
+                self.assertRaises(ConsulError, self.c.update_leader, 12345)
+        with patch('time.time', Mock(side_effect=[0, 100, 200, 300])):
+            self.assertRaises(ConsulError, self.c.update_leader, 12345)
+        with patch.object(consul.Consul.KV, 'delete', Mock(side_effect=ConsulException)):
+            self.assertFalse(self.c.update_leader(12347))
+        mock_renew.side_effect = RetryFailedError('')
+        self.c._last_session_refresh = 0
+        self.assertRaises(ConsulError, self.c.update_leader, 12346)
+        mock_renew.side_effect = ConsulException
+        self.assertFalse(self.c.update_leader(12347))
 
     @patch.object(consul.Consul.KV, 'delete', Mock(return_value=True))
     def test_delete_leader(self):
diff --git a/tests/test_etcd.py b/tests/test_etcd.py
index d6ff5f8a..d22d9760 100644
--- a/tests/test_etcd.py
+++ b/tests/test_etcd.py
@@ -276,6 +276,9 @@ class TestEtcd(unittest.TestCase):
         self.assertFalse(self.etcd.attempt_to_acquire_leader())
         self.etcd._base_path = '/service/failed'
         self.assertFalse(self.etcd.attempt_to_acquire_leader())
+        with patch.object(EtcdClient, 'write', Mock(side_effect=[etcd.EtcdConnectionFailed, Exception])):
+            self.assertRaises(EtcdError, self.etcd.attempt_to_acquire_leader)
+            self.assertRaises(EtcdError, self.etcd.attempt_to_acquire_leader)
 
     @patch.object(Cluster, 'min_version', PropertyMock(return_value=(2, 0)))
     def test_write_leader_optime(self):
@@ -284,6 +287,13 @@ class TestEtcd(unittest.TestCase):
 
     def test_update_leader(self):
         self.assertTrue(self.etcd.update_leader(None))
+        with patch.object(etcd.Client, 'write',
+                          Mock(side_effect=[etcd.EtcdConnectionFailed, etcd.EtcdClusterIdChanged, Exception])):
+            self.assertRaises(EtcdError, self.etcd.update_leader, None)
+            self.assertFalse(self.etcd.update_leader(None))
+            self.assertRaises(EtcdError, self.etcd.update_leader, None)
+        with patch.object(etcd.Client, 'write', Mock(side_effect=etcd.EtcdKeyNotFound)):
+            self.assertFalse(self.etcd.update_leader(None))
 
     def test_initialize(self):
         self.assertFalse(self.etcd.initialize())
diff --git a/tests/test_etcd3.py b/tests/test_etcd3.py
index b0ae1d05..e09c59bc 100644
--- a/tests/test_etcd3.py
+++ b/tests/test_etcd3.py
@@ -216,11 +216,26 @@ class TestEtcd3(BaseTestEtcd3):
     def test__update_leader(self):
         self.etcd3._lease = None
         self.etcd3.update_leader('123')
+        self.etcd3._last_lease_refresh = 0
         self.etcd3.update_leader('124')
+        with patch.object(PatroniEtcd3Client, 'lease_keepalive', Mock(return_value=True)),\
+                patch('time.time', Mock(side_effect=[0, 100, 200, 300])):
+            self.assertRaises(Etcd3Error, self.etcd3.update_leader, '126')
+        self.etcd3._last_lease_refresh = 0
+        with patch.object(PatroniEtcd3Client, 'lease_keepalive', Mock(side_effect=Unknown)):
+            self.assertFalse(self.etcd3.update_leader('125'))
+
+    def test_take_leader(self):
+        self.assertFalse(self.etcd3.take_leader())
 
     def test_attempt_to_acquire_leader(self):
-        self.etcd3._lease = None
         self.assertFalse(self.etcd3.attempt_to_acquire_leader())
+        with patch('time.time', Mock(side_effect=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 100, 200])):
+            self.assertRaises(Etcd3Error, self.etcd3.attempt_to_acquire_leader)
+        with patch('time.time', Mock(side_effect=[0, 100, 200, 300, 400])):
+            self.assertRaises(Etcd3Error, self.etcd3.attempt_to_acquire_leader)
+        with patch.object(PatroniEtcd3Client, 'put', Mock(return_value=False)):
+            self.assertFalse(self.etcd3.attempt_to_acquire_leader())
 
     def test_set_ttl(self):
         self.etcd3.set_ttl(20)
diff --git a/tests/test_exhibitor.py b/tests/test_exhibitor.py
index fdd3e993..a908e1fc 100644
--- a/tests/test_exhibitor.py
+++ b/tests/test_exhibitor.py
@@ -30,5 +30,6 @@ class TestExhibitor(unittest.TestCase):
                             'name': 'foo', 'ttl': 30, 'retry_timeout': 10})
 
     @patch.object(ExhibitorEnsembleProvider, 'poll', Mock(return_value=True))
+    @patch.object(MockKazooClient, 'get_children', Mock(side_effect=Exception))
     def test_get_cluster(self):
         self.assertRaises(ZooKeeperError, self.e.get_cluster)
diff --git a/tests/test_ha.py b/tests/test_ha.py
index d7d28bc1..7a7fb5f0 100644
--- a/tests/test_ha.py
+++ b/tests/test_ha.py
@@ -206,7 +206,8 @@ class TestHa(PostgresInit):
 
     def test_update_lock(self):
         self.p.last_operation = Mock(side_effect=PostgresConnectionException(''))
-        self.ha.dcs.update_leader = Mock(side_effect=Exception)
+        self.ha.dcs.update_leader = Mock(side_effect=[DCSError(''), Exception])
+        self.assertRaises(DCSError, self.ha.update_lock)
         self.assertFalse(self.ha.update_lock(True))
 
     @patch.object(Postgresql, 'received_timeline', Mock(return_value=None))
@@ -458,7 +459,9 @@ class TestHa(PostgresInit):
 
     def test_no_etcd_connection_master_demote(self):
         self.ha.load_cluster_from_dcs = Mock(side_effect=DCSError('Etcd is not responding properly'))
-        self.assertEqual(self.ha.run_cycle(), 'demoted self because DCS is not accessible and i was a leader')
+        self.assertEqual(self.ha.run_cycle(), 'demoting self because DCS is not accessible and I was a leader')
+        self.ha._async_executor.schedule('dummy')
+        self.assertEqual(self.ha.run_cycle(), 'demoted self because DCS is not accessible and I was a leader')
 
     @patch('time.sleep', Mock())
     def test_bootstrap_from_another_member(self):
@@ -1285,3 +1288,8 @@ class TestHa(PostgresInit):
         self.ha.fetch_node_status = Mock(return_value=_MemberStatus(self.ha.cluster.members[0],
                                                                     True, True, 0, 2, None, {}, False))
         self.assertFalse(self.ha.is_failover_possible(self.ha.cluster.members))
+
+    def test_acquire_lock(self):
+        self.ha.dcs.attempt_to_acquire_leader = Mock(side_effect=[DCSError('foo'), Exception])
+        self.assertRaises(DCSError, self.ha.acquire_lock)
+        self.assertFalse(self.ha.acquire_lock())
diff --git a/tests/test_kubernetes.py b/tests/test_kubernetes.py
index 79bde32b..9ff376a9 100644
--- a/tests/test_kubernetes.py
+++ b/tests/test_kubernetes.py
@@ -238,6 +238,13 @@ class TestKubernetesConfigMaps(BaseTestKubernetes):
         with patch.object(Kubernetes, '_wait_caches', Mock(side_effect=Exception)):
             self.assertRaises(KubernetesError, self.k.get_cluster)
 
+    def test_attempt_to_acquire_leader(self):
+        with patch.object(k8s_client.CoreV1Api, 'patch_namespaced_config_map', create=True) as mock_patch:
+            mock_patch.side_effect = K8sException
+            self.assertRaises(KubernetesError, self.k.attempt_to_acquire_leader)
+            mock_patch.side_effect = k8s_client.rest.ApiException(409, '')
+            self.assertFalse(self.k.attempt_to_acquire_leader())
+
     def test_take_leader(self):
         self.k.take_leader()
         self.k._leader_observed_record['leader'] = 'test'
@@ -308,7 +315,7 @@ class TestKubernetesEndpoints(BaseTestKubernetes):
         mock_patch.side_effect = k8s_client.rest.ApiException(502, '')
         self.assertFalse(self.k.update_leader('123'))
         mock_patch.side_effect = RetryFailedError('')
-        self.assertFalse(self.k.update_leader('123'))
+        self.assertRaises(KubernetesError, self.k.update_leader, '123')
         mock_patch.side_effect = k8s_client.rest.ApiException(409, '')
         with patch('time.time', Mock(side_effect=[0, 100, 200, 0, 0, 0, 0, 100, 200])):
             self.assertFalse(self.k.update_leader('123'))
@@ -318,6 +325,8 @@ class TestKubernetesEndpoints(BaseTestKubernetes):
         mock_read.return_value.metadata.resource_version = '2'
         self.assertIsNotNone(self.k._update_leader_with_retry({}, '1', []))
         mock_patch.side_effect = k8s_client.rest.ApiException(409, '')
+        mock_read.side_effect = RetryFailedError('')
+        self.assertRaises(KubernetesError, self.k.update_leader, '123')
         mock_read.side_effect = Exception
         self.assertFalse(self.k.update_leader('123'))
 
diff --git a/tests/test_raft.py b/tests/test_raft.py
index 08085299..9797ec0c 100644
--- a/tests/test_raft.py
+++ b/tests/test_raft.py
@@ -4,7 +4,7 @@ import tempfile
 import time
 
 from mock import Mock, PropertyMock, patch
-from patroni.dcs.raft import DynMemberSyncObj, KVStoreTTL, Raft, SyncObjUtility, TCPTransport, _TCPTransport
+from patroni.dcs.raft import DynMemberSyncObj, KVStoreTTL, Raft, RaftError, SyncObjUtility, TCPTransport, _TCPTransport
 from pysyncobj import SyncObjConf, FAIL_REASON
 
 
@@ -79,6 +79,9 @@ class TestKVStoreTTL(unittest.TestCase):
         self.assertFalse(self.so.set('foo', 'bar', prevExist=False, ttl=30))
         self.assertFalse(self.so.retry(self.so._set, 'foo', {'value': 'buz', 'created': 1, 'updated': 1}, prevValue=''))
         self.assertTrue(self.so.retry(self.so._set, 'foo', {'value': 'buz', 'created': 1, 'updated': 1}))
+        with patch.object(KVStoreTTL, 'retry', Mock(side_effect=RaftError(''))):
+            self.assertFalse(self.so.set('foo', 'bar'))
+            self.assertRaises(RaftError, self.so.set, 'foo', 'bar', handle_raft_error=False)
 
     def test_delete(self):
         self.so.autoTickPeriod = 0.2
@@ -87,6 +90,8 @@ class TestKVStoreTTL(unittest.TestCase):
         self.assertFalse(self.so.delete('foo', prevValue='buz'))
         self.assertTrue(self.so.delete('foo', recursive=True))
         self.assertFalse(self.so.retry(self.so._delete, 'foo', prevValue=''))
+        with patch.object(KVStoreTTL, 'retry', Mock(side_effect=RaftError(''))):
+            self.assertFalse(self.so.delete('foo'))
 
     def test_expire(self):
         self.so.set('foo', 'bar', ttl=0.001)
@@ -102,7 +107,7 @@ class TestKVStoreTTL(unittest.TestCase):
             callback(True, return_values.pop(0))
 
         with patch('time.time', Mock(side_effect=[1, 100])):
-            self.assertFalse(self.so.retry(test))
+            self.assertRaises(RaftError, self.so.retry, test)
 
         self.assertTrue(self.so.retry(test))
         self.assertFalse(self.so.retry(test))
diff --git a/tests/test_zookeeper.py b/tests/test_zookeeper.py
index 158be4bb..effcf6de 100644
--- a/tests/test_zookeeper.py
+++ b/tests/test_zookeeper.py
@@ -6,6 +6,7 @@ from kazoo.client import KazooClient, KazooState
 from kazoo.exceptions import NoNodeError, NodeExistsError
 from kazoo.handlers.threading import SequentialThreadingHandler
 from kazoo.protocol.states import KeeperState, ZnodeStat
+from kazoo.retry import RetryFailedError
 from mock import Mock, PropertyMock, patch
 from patroni.dcs.zookeeper import Cluster, Leader, PatroniKazooClient,\
         PatroniSequentialThreadingHandler, ZooKeeper, ZooKeeperError
@@ -173,7 +174,6 @@ class TestZooKeeper(unittest.TestCase):
         self.zk._inner_load_cluster()
 
     def test_get_cluster(self):
-        self.assertRaises(ZooKeeperError, self.zk.get_cluster)
         cluster = self.zk.get_cluster(True)
         self.assertIsInstance(cluster.leader, Leader)
         self.zk.status_watcher(None)
@@ -222,13 +222,25 @@ class TestZooKeeper(unittest.TestCase):
         self.zk.touch_member({'conn_url': 'postgres://repuser:rep-pass@localhost:5434/postgres',
                               'api_url': 'http://127.0.0.1:8009/patroni'})
 
+    @patch.object(MockKazooClient, 'create', Mock(side_effect=[RetryFailedError, Exception]))
+    def test_attempt_to_acquire_leader(self):
+        self.assertRaises(ZooKeeperError, self.zk.attempt_to_acquire_leader)
+        self.assertFalse(self.zk.attempt_to_acquire_leader())
+
     def test_take_leader(self):
         self.zk.take_leader()
         with patch.object(MockKazooClient, 'create', Mock(side_effect=Exception)):
             self.zk.take_leader()
 
     def test_update_leader(self):
-        self.assertTrue(self.zk.update_leader(12345))
+        self.assertFalse(self.zk.update_leader(12345))
+        with patch.object(MockKazooClient, 'delete', Mock(side_effect=RetryFailedError)):
+            self.assertRaises(ZooKeeperError, self.zk.update_leader, 12345)
+        with patch.object(MockKazooClient, 'delete', Mock(side_effect=NoNodeError)):
+            self.assertTrue(self.zk.update_leader(12345))
+            with patch.object(MockKazooClient, 'create', Mock(side_effect=[RetryFailedError, Exception])):
+                self.assertRaises(ZooKeeperError, self.zk.update_leader, 12345)
+                self.assertFalse(self.zk.update_leader(12345))
 
     @patch.object(Cluster, 'min_version', PropertyMock(return_value=(2, 0)))
     def test_write_leader_optime(self):