diff --git a/helpers/etcd.py b/helpers/etcd.py index f767175b..ac83fce8 100644 --- a/helpers/etcd.py +++ b/helpers/etcd.py @@ -5,6 +5,8 @@ import os import random import requests import socket +import time +import urllib3 from dns.exception import DNSException from dns import resolver @@ -59,6 +61,11 @@ class Client(etcd.Client): logger.exception('Can not resolve SRV for %s', host) return [] + # try to workarond bug in python-etcd: https://github.com/jplana/python-etcd/issues/81 + def _result_from_response(self, response): + response.data.decode('utf-8') + return super(Client, self)._result_from_response(response) + def _get_machines_cache_from_srv(self, discovery_srv): """Fetch list of etcd-cluster member by resolving _etcd-server._tcp. SRV record. This record should contain list of host and peer ports which could be used to run @@ -136,6 +143,7 @@ class Etcd(AbstractDCS): self.ttl = config['ttl'] self.member_ttl = config.get('member_ttl', 3600) self.client = self.get_etcd_client(config) + self.cluster = None def get_etcd_client(self, config): client = None @@ -174,12 +182,14 @@ class Etcd(AbstractDCS): member = ([m for m in members if m.name == leader.value] or [member])[0] leader = Leader(leader.modifiedIndex, leader.expiration, leader.ttl, member) - return Cluster(initialize, leader, last_leader_operation, members) + self.cluster = Cluster(initialize, leader, last_leader_operation, members) + return self.cluster except etcd.EtcdKeyNotFound: return Cluster(False, None, None, []) except: logger.exception('get_cluster') + self.cluster = None raise EtcdError('Etcd is not responding properly') @catch_etcd_errors @@ -213,3 +223,25 @@ class Etcd(AbstractDCS): @catch_etcd_errors def delete_leader(self): return self.client.delete(self.client_path('/leader'), prevValue=self._name) + + def sleep(self, timeout): + # watch on leader key changes if it is defined and current node is not lock owner + if self.cluster and self.cluster.leader and self.cluster.leader.member.name != self._name: + end_time = time.time() + timeout + index = self.cluster.leader.index + + while index and timeout >= 1: # when timeout is too small urllib3 doesn't have enough time to connect + try: + res = self.client.watch(self.client_path('/leader'), index=index + 1, timeout=timeout) + if res.action not in ['set', 'compareAndSwap'] or res.value != self.cluster.leader.member.name: + return + index = res.modifiedIndex + except urllib3.exceptions.TimeoutError: + self.client.http.clear() + return + except etcd.EtcdException: + index = None + + timeout = end_time - time.time() + + timeout > 0 and super(Etcd, self).sleep(timeout) diff --git a/tests/test_etcd.py b/tests/test_etcd.py index 52de036e..8de4ca46 100644 --- a/tests/test_etcd.py +++ b/tests/test_etcd.py @@ -3,6 +3,7 @@ import dns.resolver import etcd import json import requests +import urllib3 import socket import time import unittest @@ -61,6 +62,20 @@ def requests_get(url, **kwargs): return response +def etcd_watch(key, index=None, timeout=None, recursive=None): + print ('watch', key, index, timeout) + if timeout == 1: + raise urllib3.exceptions.TimeoutError + elif timeout == 5: + return etcd.EtcdResult('delete', {}) + elif timeout == 10: + raise etcd.EtcdException + elif index == 20729: + return etcd.EtcdResult('set', {'value': 'postgresql1', 'modifiedIndex': index + 1}) + elif index == 20731: + return etcd.EtcdResult('set', {'value': 'postgresql2', 'modifiedIndex': index + 1}) + + def etcd_write(key, value, **kwargs): if key == '/service/test/leader': if kwargs.get('prevValue', None) == 'foo' or not kwargs.get('prevExist', True): @@ -237,3 +252,12 @@ class TestEtcd(unittest.TestCase): def test_delete_leader(self): self.etcd.client.delete = etcd_delete self.assertFalse(self.etcd.delete_leader()) + + def test_sleep(self): + self.etcd.client.watch = etcd_watch + self.etcd.sleep(100) + self.etcd.get_cluster() + self.etcd.sleep(1) + self.etcd.sleep(5) + self.etcd.sleep(10) + self.etcd.sleep(100) diff --git a/tests/test_patroni.py b/tests/test_patroni.py index 575c7cf3..c0ea2907 100644 --- a/tests/test_patroni.py +++ b/tests/test_patroni.py @@ -14,7 +14,7 @@ from helpers.zookeeper import ZooKeeper from mock import Mock, patch from patroni import Patroni, main from six.moves import BaseHTTPServer -from test_etcd import Client, etcd_read, etcd_write +from test_etcd import Client, etcd_read, etcd_write, etcd_watch from test_ha import true, false from test_postgresql import Postgresql, subprocess_call, psycopg2_connect from test_zookeeper import MockKazooClient @@ -104,6 +104,7 @@ class TestPatroni(unittest.TestCase): self.p.touch_member = self.touch_member self.p.ha.state_handler.sync_replication_slots = time_sleep self.p.ha.dcs.client.read = etcd_read + self.p.ha.dcs.sleep = time_sleep self.assertRaises(SleepException, self.p.run) self.p.ha.state_handler.is_leader = lambda: False self.p.api.start = nop