Refactor update_leader interface (#2690)

pass reference to a last known leader object in order to avoid obtaining it from the `AbstractDCS.cluster` cache.

This change is useful for Consul, Etcd3 and Zookeeper implementations.

patroni/dcs/__init__.py

@@ -965,25 +965,26 @@ class AbstractDCS(abc.ABC):
         return self._last_failsafe
 
     @abc.abstractmethod
-    def _update_leader(self) -> bool:
+    def _update_leader(self, leader: Leader) -> bool:
         """Update leader key (or session) ttl
 
-        :returns: `!True` if leader key (or session) has been updated successfully.
+        :param leader: a reference to a current leader key object
+        :returns: `!True` if leader key (or session) has been updated successfully
 
         You have to use CAS (Compare And Swap) operation in order to update leader key,
         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: Optional[int], slots: Optional[Dict[str, int]] = None,
+    def update_leader(self, leader: Leader, last_lsn: Optional[int],
-                      failsafe: Optional[Dict[str, str]] = None) -> bool:
+                      slots: Optional[Dict[str, int]] = None, failsafe: Optional[Dict[str, str]] = None) -> bool:
         """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."""
 
-        ret = self._update_leader()
+        ret = self._update_leader(leader)
         if ret and last_lsn:
             status: Dict[str, Any] = {self._OPTIME: last_lsn}
             if slots:

patroni/dcs/consul.py

@@ -608,28 +608,24 @@ class Consul(AbstractDCS):
             raise ReturnFalseException
 
     @catch_return_false_exception
-    def _update_leader(self) -> bool:
+    def _update_leader(self, leader: Leader) -> bool:
         retry = self._retry.copy()
 
         self._run_and_handle_exceptions(self._do_refresh_session, True, retry=retry)
 
-        if self._session:
+        if self._session and leader.session != self._session:
-            cluster = self.cluster
+            retry.deadline = retry.stoptime - time.time()
-            leader_session = cluster and isinstance(cluster.leader, Leader) and cluster.leader.session
+            if retry.deadline < 1:
-            if leader_session != self._session:
+                raise ConsulError('update_leader timeout')
-                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 and cluster.leader:
-                    self._run_and_handle_exceptions(self._client.kv.delete, self.leader_path,
-                                                    cas=cluster.leader.version)
 
-                retry.deadline = retry.stoptime - time.time()
+            logger.warning('Recreating the leader key due to session mismatch')
-                if retry.deadline < 0.5:
+            self._run_and_handle_exceptions(self._client.kv.delete, self.leader_path, cas=leader.version)
-                    raise ConsulError('update_leader timeout')
+
-                self._run_and_handle_exceptions(self._client.kv.put, self.leader_path,
+            retry.deadline = retry.stoptime - time.time()
-                                                self._name, acquire=self._session)
+            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)
 

patroni/dcs/etcd.py

@@ -798,7 +798,7 @@ class Etcd(AbstractEtcd):
         return bool(self._client.set(self.failsafe_path, value))
 
     @catch_return_false_exception
-    def _update_leader(self) -> bool:
+    def _update_leader(self, leader: Leader) -> bool:
         return bool(self._run_and_handle_exceptions(self._do_update_leader, retry=None))
 
     @catch_etcd_errors

patroni/dcs/etcd3.py

@@ -877,7 +877,7 @@ class Etcd3(AbstractEtcd):
         return bool(self._client.put(self.failsafe_path, value))
 
     @catch_return_false_exception
-    def _update_leader(self) -> bool:
+    def _update_leader(self, leader: Leader) -> bool:
         retry = self._retry.copy()
 
         def _retry(*args: Any, **kwargs: Any) -> Any:
@@ -886,26 +886,21 @@ class Etcd3(AbstractEtcd):
 
         self._run_and_handle_exceptions(self._do_refresh_lease, True, retry=_retry)
 
-        if self._lease:
+        if self._lease and leader.session != self._lease:
-            cluster = self.cluster
+            retry.deadline = retry.stoptime - time.time()
-            leader_lease = cluster and isinstance(cluster.leader, Leader) and cluster.leader.session
+            if retry.deadline < 1:
-            if leader_lease != self._lease:
+                raise Etcd3Error('update_leader timeout')
-                retry.deadline = retry.stoptime - time.time()
-                if retry.deadline < 1:
-                    raise Etcd3Error('update_leader timeout')
 
-                fields = {'key': base64_encode(self.leader_path),
+            fields = {'key': base64_encode(self.leader_path), 'value': base64_encode(self._name), 'lease': self._lease}
-                          'value': base64_encode(self._name), 'lease': self._lease}
+            # First we try to update lease on existing leader key "hoping" that we still owning it
-                # First we try to update lease on existing leader key "hoping" that we still owning it
+            compare1 = {'key': fields['key'], 'target': 'VALUE', 'value': fields['value']}
-                compare1 = {'key': fields['key'], 'target': 'VALUE', 'value': fields['value']}
+            request_put = {'request_put': fields}
-                request_put = {'request_put': fields}
+            # If the first comparison failed we will try to create the new leader key in a transaction
-                # If the first comparison failed we will try to create the new leader key in a transaction
+            compare2 = {'key': fields['key'], 'target': 'CREATE', 'create_revision': '0'}
-                compare2 = {'key': fields['key'], 'target': 'CREATE', 'create_revision': '0'}
+            request_txn = {'request_txn': {'compare': [compare2], 'success': [request_put]}}
-                request_txn = {'request_txn': {'compare': [compare2], 'success': [request_put]}}
+            ret = self._run_and_handle_exceptions(self._client.txn, compare1, request_put, request_txn, retry=_retry)
-                ret = self._run_and_handle_exceptions(self._client.txn, compare1,
+            return ret.get('succeeded', False)\
-                                                      request_put, request_txn, retry=_retry)
+                or ret.get('responses', [{}])[0].get('response_txn', {}).get('succeeded', False)
-                return ret.get('succeeded', False)\
-                    or ret.get('responses', [{}])[0].get('response_txn', {}).get('succeeded', False)
         return bool(self._lease)
 
     @catch_etcd_errors

patroni/dcs/kubernetes.py

@@ -1129,7 +1129,7 @@ class Kubernetes(AbstractDCS):
         """Unused"""
         raise NotImplementedError  # pragma: no cover
 
-    def _update_leader(self) -> bool:
+    def _update_leader(self, leader: Leader) -> bool:
         """Unused"""
         raise NotImplementedError  # pragma: no cover
 
@@ -1182,8 +1182,8 @@ class Kubernetes(AbstractDCS):
         return bool(_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: Optional[int], slots: Optional[Dict[str, int]] = None,
+    def update_leader(self, leader: Leader, last_lsn: Optional[int],
-                      failsafe: Optional[Dict[str, str]] = None) -> bool:
+                      slots: Optional[Dict[str, int]] = None, failsafe: Optional[Dict[str, str]] = None) -> bool:
         kind = self._kinds.get(self.leader_path)
         kind_annotations = kind and kind.metadata.annotations or {}
 

patroni/dcs/raft.py

@@ -419,7 +419,7 @@ class Raft(AbstractDCS):
     def _write_failsafe(self, value: str) -> bool:
         return self._sync_obj.set(self.failsafe_path, value, timeout=1) is not False
 
-    def _update_leader(self) -> bool:
+    def _update_leader(self, leader: Leader) -> bool:
         ret = self._sync_obj.set(self.leader_path, self._name, ttl=self._ttl,
                                  handle_raft_error=False, prevValue=self._name) is not False
         if not ret and self._sync_obj.get(self.leader_path) is None:

patroni/dcs/zookeeper.py

@@ -443,10 +443,8 @@ class ZooKeeper(AbstractDCS):
     def _write_failsafe(self, value: str) -> bool:
         return self._set_or_create(self.failsafe_path, value) is not False
 
-    def _update_leader(self) -> bool:
+    def _update_leader(self, leader: Leader) -> bool:
-        cluster = self.cluster
+        if self._client.client_id and self._client.client_id[0] != leader.session:
-        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)

patroni/ha.py

@@ -239,8 +239,10 @@ class Ha(object):
                 slots = self.state_handler.slots()
             except Exception:
                 logger.exception('Exception when called state_handler.last_operation()')
+        if TYPE_CHECKING:  # pragma: no cover
+            assert self.cluster.leader is not None
         try:
-            ret = self.dcs.update_leader(last_lsn, slots, self._failsafe_config())
+            ret = self.dcs.update_leader(self.cluster.leader, last_lsn, slots, self._failsafe_config())
         except DCSError:
             raise
         except Exception:

tests/test_consul.py

@@ -176,23 +176,24 @@ class TestConsul(unittest.TestCase):
     @patch.object(consul.Consul.Session, 'renew')
     @patch.object(consul.Consul.KV, 'put', Mock(side_effect=ConsulException))
     def test_update_leader(self, mock_renew):
+        leader = self.c.get_cluster().leader
         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, failsafe={'foo': 'bar'}))
+                self.assertTrue(self.c.update_leader(leader, 12345, failsafe={'foo': 'bar'}))
             with patch.object(consul.Consul.KV, 'put', Mock(side_effect=ConsulException)):
-                self.assertFalse(self.c.update_leader(12345))
+                self.assertFalse(self.c.update_leader(leader, 12345))
-            with patch('time.time', Mock(side_effect=[0, 0, 0, 0, 0, 100, 200, 300])):
+            with patch('time.time', Mock(side_effect=[0, 0, 0, 0, 100, 200, 300])):
-                self.assertRaises(ConsulError, self.c.update_leader, 12345)
+                self.assertRaises(ConsulError, self.c.update_leader, leader, 12345)
         with patch('time.time', Mock(side_effect=[0, 100, 200, 300])):
-            self.assertRaises(ConsulError, self.c.update_leader, 12345)
+            self.assertRaises(ConsulError, self.c.update_leader, leader, 12345)
         with patch.object(consul.Consul.KV, 'delete', Mock(side_effect=ConsulException)):
-            self.assertFalse(self.c.update_leader(12347))
+            self.assertFalse(self.c.update_leader(leader, 12347))
         mock_renew.side_effect = RetryFailedError('')
         self.c._last_session_refresh = 0
-        self.assertRaises(ConsulError, self.c.update_leader, 12346)
+        self.assertRaises(ConsulError, self.c.update_leader, leader, 12346)
         mock_renew.side_effect = ConsulException
-        self.assertFalse(self.c.update_leader(12347))
+        self.assertFalse(self.c.update_leader(leader, 12347))
 
     @patch.object(consul.Consul.KV, 'delete', Mock(return_value=True))
     def test_delete_leader(self):

tests/test_etcd.py

@@ -296,14 +296,15 @@ class TestEtcd(unittest.TestCase):
         self.etcd.write_leader_optime('0')
 
     def test_update_leader(self):
-        self.assertTrue(self.etcd.update_leader(None, failsafe={'foo': 'bar'}))
+        leader = self.etcd.get_cluster().leader
+        self.assertTrue(self.etcd.update_leader(leader, None, failsafe={'foo': 'bar'}))
         with patch.object(etcd.Client, 'write',
                           Mock(side_effect=[etcd.EtcdConnectionFailed, etcd.EtcdClusterIdChanged, Exception])):
-            self.assertRaises(EtcdError, self.etcd.update_leader, None)
+            self.assertRaises(EtcdError, self.etcd.update_leader, leader, None)
-            self.assertFalse(self.etcd.update_leader(None))
+            self.assertFalse(self.etcd.update_leader(leader, None))
-            self.assertRaises(EtcdError, self.etcd.update_leader, None)
+            self.assertRaises(EtcdError, self.etcd.update_leader, leader, None)
         with patch.object(etcd.Client, 'write', Mock(side_effect=etcd.EtcdKeyNotFound)):
-            self.assertFalse(self.etcd.update_leader(None))
+            self.assertFalse(self.etcd.update_leader(leader, None))
 
     def test_initialize(self):
         self.assertFalse(self.etcd.initialize())

tests/test_etcd3.py

@@ -235,19 +235,20 @@ class TestEtcd3(BaseTestEtcd3):
             self.etcd3.touch_member({})
 
     def test__update_leader(self):
+        leader = self.etcd3.get_cluster().leader
         self.etcd3._lease = None
         with patch.object(Etcd3Client, 'txn', Mock(return_value={'succeeded': True})):
-            self.etcd3.update_leader('123', failsafe={'foo': 'bar'})
+            self.etcd3.update_leader(leader, '123', failsafe={'foo': 'bar'})
         self.etcd3._last_lease_refresh = 0
-        self.etcd3.update_leader('124')
+        self.etcd3.update_leader(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.assertRaises(Etcd3Error, self.etcd3.update_leader, leader, '126')
-        self.etcd3._lease = self.etcd3.cluster.leader.session
+        self.etcd3._lease = leader.session
-        self.etcd3.update_leader('124')
+        self.etcd3.update_leader(leader, '124')
         self.etcd3._last_lease_refresh = 0
         with patch.object(PatroniEtcd3Client, 'lease_keepalive', Mock(side_effect=Unknown)):
-            self.assertFalse(self.etcd3.update_leader('125'))
+            self.assertFalse(self.etcd3.update_leader(leader, '125'))
 
     def test_take_leader(self):
         self.assertFalse(self.etcd3.take_leader())

tests/test_kubernetes.py

@@ -340,35 +340,37 @@ class TestKubernetesEndpoints(BaseTestKubernetes):
 
     @patch.object(k8s_client.CoreV1Api, 'patch_namespaced_endpoints', create=True)
     def test_update_leader(self, mock_patch_namespaced_endpoints):
-        self.assertIsNotNone(self.k.update_leader('123', failsafe={'foo': 'bar'}))
+        leader = self.k.get_cluster().leader
+        self.assertIsNotNone(self.k.update_leader(leader, '123', failsafe={'foo': 'bar'}))
         args = mock_patch_namespaced_endpoints.call_args[0]
         self.assertEqual(args[2].subsets[0].addresses[0].target_ref.resource_version, '10')
         self.k._kinds._object_cache['test'].subsets[:] = []
-        self.assertIsNotNone(self.k.update_leader('123'))
+        self.assertIsNotNone(self.k.update_leader(leader, '123'))
         self.k._kinds._object_cache['test'].metadata.annotations['leader'] = 'p-1'
-        self.assertFalse(self.k.update_leader('123'))
+        self.assertFalse(self.k.update_leader(leader, '123'))
 
     @patch.object(k8s_client.CoreV1Api, 'read_namespaced_endpoints', create=True)
     @patch.object(k8s_client.CoreV1Api, 'patch_namespaced_endpoints', create=True)
     def test__update_leader_with_retry(self, mock_patch, mock_read):
+        leader = self.k.get_cluster().leader
         mock_read.return_value = mock_read_namespaced_endpoints()
         mock_patch.side_effect = k8s_client.rest.ApiException(502, '')
-        self.assertFalse(self.k.update_leader('123'))
+        self.assertFalse(self.k.update_leader(leader, '123'))
         mock_patch.side_effect = RetryFailedError('')
-        self.assertRaises(KubernetesError, self.k.update_leader, '123')
+        self.assertRaises(KubernetesError, self.k.update_leader, 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'))
+            self.assertFalse(self.k.update_leader(leader, '123'))
-            self.assertFalse(self.k.update_leader('123'))
+            self.assertFalse(self.k.update_leader(leader, '123'))
-        self.assertFalse(self.k.update_leader('123'))
+        self.assertFalse(self.k.update_leader(leader, '123'))
         mock_patch.side_effect = [k8s_client.rest.ApiException(409, ''), mock_namespaced_kind()]
         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')
+        self.assertRaises(KubernetesError, self.k.update_leader, leader, '123')
         mock_read.side_effect = Exception
-        self.assertFalse(self.k.update_leader('123'))
+        self.assertFalse(self.k.update_leader(leader, '123'))
 
     @patch.object(k8s_client.CoreV1Api, 'patch_namespaced_endpoints',
                   Mock(side_effect=[k8s_client.rest.ApiException(500, ''),

tests/test_raft.py

@@ -146,8 +146,8 @@ class TestRaft(unittest.TestCase):
         self.assertIsInstance(cluster, Cluster)
         self.assertIsInstance(cluster.workers[1], Cluster)
         self.assertTrue(raft._sync_obj.set(raft.status_path, '{"optime":1234567,"slots":{"ls":12345}}'))
-        raft.get_cluster()
+        leader = raft.get_cluster().leader
-        self.assertTrue(raft.update_leader('1', failsafe={'foo': 'bat'}))
+        self.assertTrue(raft.update_leader(leader, '1', failsafe={'foo': 'bat'}))
         self.assertTrue(raft._sync_obj.set(raft.failsafe_path, '{"foo"}'))
         self.assertTrue(raft._sync_obj.set(raft.status_path, '{'))
         raft.get_citus_coordinator()

tests/test_zookeeper.py

@@ -250,14 +250,15 @@ class TestZooKeeper(unittest.TestCase):
             self.zk.take_leader()
 
     def test_update_leader(self):
-        self.assertFalse(self.zk.update_leader(12345))
+        leader = self.zk.get_cluster().leader
+        self.assertFalse(self.zk.update_leader(leader, 12345))
         with patch.object(MockKazooClient, 'delete', Mock(side_effect=RetryFailedError)):
-            self.assertRaises(ZooKeeperError, self.zk.update_leader, 12345)
+            self.assertRaises(ZooKeeperError, self.zk.update_leader, leader, 12345)
         with patch.object(MockKazooClient, 'delete', Mock(side_effect=NoNodeError)):
-            self.assertTrue(self.zk.update_leader(12345, failsafe={'foo': 'bar'}))
+            self.assertTrue(self.zk.update_leader(leader, 12345, failsafe={'foo': 'bar'}))
             with patch.object(MockKazooClient, 'create', Mock(side_effect=[RetryFailedError, Exception])):
-                self.assertRaises(ZooKeeperError, self.zk.update_leader, 12345)
+                self.assertRaises(ZooKeeperError, self.zk.update_leader, leader, 12345)
-                self.assertFalse(self.zk.update_leader(12345))
+                self.assertFalse(self.zk.update_leader(leader, 12345))
 
     @patch.object(Cluster, 'min_version', PropertyMock(return_value=(2, 0)))
     def test_write_leader_optime(self):