From af8e5f0d0fe94c85c7313764c72e195b3e6ba721 Mon Sep 17 00:00:00 2001 From: Alexander Kukushkin Date: Thu, 25 May 2023 14:21:05 +0200 Subject: [PATCH] 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 | 11 ++++++----- patroni/dcs/consul.py | 30 +++++++++++++----------------- patroni/dcs/etcd.py | 2 +- patroni/dcs/etcd3.py | 35 +++++++++++++++-------------------- patroni/dcs/kubernetes.py | 6 +++--- patroni/dcs/raft.py | 2 +- patroni/dcs/zookeeper.py | 6 ++---- patroni/ha.py | 4 +++- tests/test_consul.py | 17 +++++++++-------- tests/test_etcd.py | 11 ++++++----- tests/test_etcd3.py | 13 +++++++------ tests/test_kubernetes.py | 22 ++++++++++++---------- tests/test_raft.py | 4 ++-- tests/test_zookeeper.py | 11 ++++++----- 14 files changed, 86 insertions(+), 88 deletions(-) diff --git a/patroni/dcs/__init__.py b/patroni/dcs/__init__.py index 1ff4a60a..a61c6b59 100644 --- a/patroni/dcs/__init__.py +++ b/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, - failsafe: Optional[Dict[str, str]] = None) -> bool: + def update_leader(self, leader: Leader, last_lsn: Optional[int], + 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: diff --git a/patroni/dcs/consul.py b/patroni/dcs/consul.py index 2383a964..87ae4f6a 100644 --- a/patroni/dcs/consul.py +++ b/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: - 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 and cluster.leader: - self._run_and_handle_exceptions(self._client.kv.delete, self.leader_path, - cas=cluster.leader.version) + if self._session and leader.session != self._session: + retry.deadline = retry.stoptime - time.time() + if retry.deadline < 1: + raise ConsulError('update_leader timeout') - 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) + logger.warning('Recreating the leader key due to session mismatch') + self._run_and_handle_exceptions(self._client.kv.delete, self.leader_path, cas=leader.version) + + 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) diff --git a/patroni/dcs/etcd.py b/patroni/dcs/etcd.py index 190b79fc..3e7f0681 100644 --- a/patroni/dcs/etcd.py +++ b/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 diff --git a/patroni/dcs/etcd3.py b/patroni/dcs/etcd3.py index 6c4312e5..4d0ff75d 100644 --- a/patroni/dcs/etcd3.py +++ b/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: - cluster = self.cluster - leader_lease = cluster and isinstance(cluster.leader, Leader) and cluster.leader.session - if leader_lease != self._lease: - retry.deadline = retry.stoptime - time.time() - if retry.deadline < 1: - raise Etcd3Error('update_leader timeout') + if self._lease and leader.session != self._lease: + retry.deadline = retry.stoptime - time.time() + if retry.deadline < 1: + raise Etcd3Error('update_leader timeout') - fields = {'key': base64_encode(self.leader_path), - 'value': base64_encode(self._name), 'lease': self._lease} - # 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']} - request_put = {'request_put': fields} - # 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'} - 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) - return ret.get('succeeded', False)\ - or ret.get('responses', [{}])[0].get('response_txn', {}).get('succeeded', False) + fields = {'key': base64_encode(self.leader_path), 'value': base64_encode(self._name), 'lease': self._lease} + # 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']} + request_put = {'request_put': fields} + # 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'} + 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) + return ret.get('succeeded', False)\ + or ret.get('responses', [{}])[0].get('response_txn', {}).get('succeeded', False) return bool(self._lease) @catch_etcd_errors diff --git a/patroni/dcs/kubernetes.py b/patroni/dcs/kubernetes.py index 3630ae1e..ac65652d 100644 --- a/patroni/dcs/kubernetes.py +++ b/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, - failsafe: Optional[Dict[str, str]] = None) -> bool: + def update_leader(self, leader: Leader, last_lsn: Optional[int], + 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 {} diff --git a/patroni/dcs/raft.py b/patroni/dcs/raft.py index c068f7f0..c4286f1e 100644 --- a/patroni/dcs/raft.py +++ b/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: diff --git a/patroni/dcs/zookeeper.py b/patroni/dcs/zookeeper.py index 5ec9cd04..ed2b71ed 100644 --- a/patroni/dcs/zookeeper.py +++ b/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: - 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: + def _update_leader(self, leader: Leader) -> bool: + if self._client.client_id and self._client.client_id[0] != leader.session: logger.warning('Recreating the leader ZNode due to ownership mismatch') try: self._client.retry(self._client.delete, self.leader_path) diff --git a/patroni/ha.py b/patroni/ha.py index 1190ad1c..3108f787 100644 --- a/patroni/ha.py +++ b/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: diff --git a/tests/test_consul.py b/tests/test_consul.py index 165d2382..ebbbb90e 100644 --- a/tests/test_consul.py +++ b/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)) - with patch('time.time', Mock(side_effect=[0, 0, 0, 0, 0, 100, 200, 300])): - self.assertRaises(ConsulError, self.c.update_leader, 12345) + self.assertFalse(self.c.update_leader(leader, 12345)) + with patch('time.time', Mock(side_effect=[0, 0, 0, 0, 100, 200, 300])): + 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): diff --git a/tests/test_etcd.py b/tests/test_etcd.py index 6699127a..b5add201 100644 --- a/tests/test_etcd.py +++ b/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.assertFalse(self.etcd.update_leader(None)) - self.assertRaises(EtcdError, self.etcd.update_leader, None) + self.assertRaises(EtcdError, self.etcd.update_leader, leader, None) + self.assertFalse(self.etcd.update_leader(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()) diff --git a/tests/test_etcd3.py b/tests/test_etcd3.py index 4b3a78ab..a737f199 100644 --- a/tests/test_etcd3.py +++ b/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.etcd3._lease = self.etcd3.cluster.leader.session - self.etcd3.update_leader('124') + self.assertRaises(Etcd3Error, self.etcd3.update_leader, leader, '126') + self.etcd3._lease = leader.session + 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()) diff --git a/tests/test_kubernetes.py b/tests/test_kubernetes.py index fbefb448..cfed1559 100644 --- a/tests/test_kubernetes.py +++ b/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('123')) - self.assertFalse(self.k.update_leader('123')) + self.assertFalse(self.k.update_leader(leader, '123')) + self.assertFalse(self.k.update_leader(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, ''), diff --git a/tests/test_raft.py b/tests/test_raft.py index 2029a61e..ab68c2e6 100644 --- a/tests/test_raft.py +++ b/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() - self.assertTrue(raft.update_leader('1', failsafe={'foo': 'bat'})) + leader = raft.get_cluster().leader + 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() diff --git a/tests/test_zookeeper.py b/tests/test_zookeeper.py index c2f3f583..baf9ad10 100644 --- a/tests/test_zookeeper.py +++ b/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.assertFalse(self.zk.update_leader(12345)) + self.assertRaises(ZooKeeperError, self.zk.update_leader, 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):