Improve leader watch functionality (#356)

Previously replicas were always watching for leader key (even if the
postgres was not in the running there). It was not a big issue, but it
was not possible to interrupt such watch in cases if the postgres
started up or stopped successfully. Also it was delaying update_member
call and we had kind of stale information in DCS up to `loop_wait`
seconds. This commit changes such behavior. If the async_executor is
busy by starting/stopping or restarting postgres we will not watch for
leader key but waiting for event from async_executor up to `loop_wait`
seconds. Async executor will fire such event only in case if the
function it was calling returned something what could be evaluated to
boolean True.

Such functionality is really needed to change the way how we are making
decision about necessity of pg_rewind. It will require to have a local
postgres running and for us it is really important to get such
notification as soon as possible.

patroni/__init__.py

@@ -94,7 +94,7 @@ class Patroni(object):
             time.sleep(0.001)
             # Warn user that Patroni is not keeping up
             logger.warning("Loop time exceeded, rescheduling immediately.")
-        elif self.dcs.watch(nap_time):
+        elif self.ha.watch(nap_time):
             self.next_run = time.time()
 
     def run(self):

patroni/api.py

@@ -141,7 +141,7 @@ class RestApiHandler(BaseHTTPRequestHandler):
                 value = json.dumps(data, separators=(',', ':'))
                 if not self.server.patroni.dcs.set_config_value(value, cluster.config.index):
                     return self.send_error(409)
-            self.server.patroni.dcs.event.set()
+            self.server.patroni.ha.wakeup()
             self._write_json_response(200, data)
 
     @check_auth
@@ -323,7 +323,7 @@ class RestApiHandler(BaseHTTPRequestHandler):
                 if _:
                     status_code = _
                 elif self.server.patroni.dcs.manual_failover(leader, candidate, scheduled_at=scheduled_at):
-                    self.server.patroni.dcs.event.set()
+                    self.server.patroni.ha.wakeup()
                     data = 'Failover scheduled'
                     status_code = 202
                 else:
@@ -333,7 +333,7 @@ class RestApiHandler(BaseHTTPRequestHandler):
                 data = self.is_failover_possible(cluster, leader, candidate)
                 if not data:
                     if self.server.patroni.dcs.manual_failover(leader, candidate):
-                        self.server.patroni.dcs.event.set()
+                        self.server.patroni.ha.wakeup()
                         status_code, data = self.poll_failover_result(cluster.leader and cluster.leader.name, candidate)
                     else:
                         data = 'failed to write failover key into DCS'

patroni/async_executor.py

@@ -6,7 +6,8 @@ logger = logging.getLogger(__name__)
 
 class AsyncExecutor(object):
 
-    def __init__(self):
+    def __init__(self, ha_wakeup):
+        self._ha_wakeup = ha_wakeup
         self._thread_lock = RLock()
         self._scheduled_action = None
         self._scheduled_action_lock = RLock()
@@ -32,13 +33,18 @@ class AsyncExecutor(object):
             self._scheduled_action = None
 
     def run(self, func, args=()):
+        wakeup = False
         try:
-            return func(*args) if args else func()
+            # if the func returned something (not None) - wake up main HA loop
+            wakeup = func(*args) if args else func()
+            return wakeup
         except:
             logger.exception('Exception during execution of long running task %s', self.scheduled_action)
         finally:
             with self:
                 self.reset_scheduled_action()
+            if wakeup is not None:
+                self._ha_wakeup()
 
     def run_async(self, func, args=()):
         Thread(target=self.run, args=(func, args)).start()

patroni/ctl.py

@@ -405,20 +405,6 @@ def remove(obj, cluster_name, fmt):
     dcs.delete_cluster()
 
 
-def wait_for_leader(dcs, timeout=30):
-    t_stop = time.time() + timeout
-    timeout /= 2
-
-    while time.time() < t_stop:
-        dcs.watch(timeout)
-        cluster = dcs.get_cluster()
-
-        if cluster.leader:
-            return cluster
-
-    raise PatroniCtlException('Timeout occured')
-
-
 def check_response(response, member_name, action_name, silent_success=False):
     if response.status_code >= 400:
         click.echo('Failed: {0} for member {1}, status code={2}, ({3})'.format(

patroni/dcs/__init__.py

@@ -531,10 +531,11 @@ class AbstractDCS(object):
     def delete_sync_state(self, index=None):
         """"""
 
-    def watch(self, timeout):
+    def watch(self, leader_index, timeout):
         """If the current node is a master it should just sleep.
         Any other node should watch for changes of leader key with a given timeout
 
+        :param leader_index: index of a leader key
         :param timeout: timeout in seconds
         :returns: `!True` if you would like to reschedule the next run of ha cycle"""
 

patroni/dcs/consul.py

@@ -304,24 +304,23 @@ class Consul(AbstractDCS):
     def delete_sync_state(self, index=None):
         return self._client.kv.delete(self.sync_path, cas=index)
 
-    def watch(self, timeout):
+    def watch(self, leader_index, timeout):
         if self.__do_not_watch:
             self.__do_not_watch = False
             return True
 
-        cluster = self.cluster
-        if cluster and cluster.leader and cluster.leader.name != self._name and cluster.leader.index:
+        if leader_index:
             end_time = time.time() + timeout
             while timeout >= 1:
                 try:
-                    idx, _ = self._client.kv.get(self.leader_path, index=cluster.leader.index, wait=str(timeout) + 's')
-                    return str(idx) != str(cluster.leader.index)
+                    idx, _ = self._client.kv.get(self.leader_path, index=leader_index, wait=str(timeout) + 's')
+                    return str(idx) != str(leader_index)
                 except (ConsulException, HTTPException, HTTPError, socket.error, socket.timeout):
                     logging.exception('watch')
 
                 timeout = end_time - time.time()
 
         try:
-            return super(Consul, self).watch(timeout)
+            return super(Consul, self).watch(None, timeout)
         finally:
             self.event.clear()

patroni/dcs/etcd.py

@@ -372,19 +372,17 @@ class Etcd(AbstractDCS):
     def delete_sync_state(self, index=None):
         return self.retry(self._client.delete, self.sync_path, prevIndex=index or 0)
 
-    def watch(self, timeout):
+    def watch(self, leader_index, timeout):
         if self.__do_not_watch:
             self.__do_not_watch = False
             return True
 
-        cluster = self.cluster
-        # watch on leader key changes if it is defined and current node is not lock owner
-        if cluster and cluster.leader and cluster.leader.name != self._name and cluster.leader.index:
+        if leader_index:
             end_time = time.time() + timeout
 
             while timeout >= 1:  # when timeout is too small urllib3 doesn't have enough time to connect
                 try:
-                    self._client.watch(self.leader_path, index=cluster.leader.index, timeout=timeout + 0.5)
+                    self._client.watch(self.leader_path, index=leader_index, timeout=timeout + 0.5)
                     # Synchronous work of all cluster members with etcd is less expensive
                     # than reestablishing http connection every time from every replica.
                     return True
@@ -397,6 +395,6 @@ class Etcd(AbstractDCS):
                 timeout = end_time - time.time()
 
         try:
-            return super(Etcd, self).watch(timeout)
+            return super(Etcd, self).watch(None, timeout)
         finally:
             self.event.clear()

patroni/dcs/zookeeper.py

@@ -326,7 +326,7 @@ class ZooKeeper(AbstractDCS):
     def delete_sync_state(self, index=None):
         return self.set_sync_state_value("{}", index)
 
-    def watch(self, timeout):
-        if super(ZooKeeper, self).watch(timeout):
+    def watch(self, leader_index, timeout):
+        if super(ZooKeeper, self).watch(leader_index, timeout):
             self._fetch_cluster = True
         return self._fetch_cluster

patroni/ha.py

@@ -25,7 +25,7 @@ class Ha(object):
         self.cluster = None
         self.old_cluster = None
         self.recovering = False
-        self._async_executor = AsyncExecutor()
+        self._async_executor = AsyncExecutor(self.wakeup)
 
         # Each member publishes various pieces of information to the DCS using touch_member. This lock protects
         # the state and publishing procedure to have consistent ordering and avoid publishing stale values.
@@ -100,7 +100,7 @@ class Ha(object):
             logger.info('bootstrapped %s', msg)
             cluster = self.dcs.get_cluster()
             node_to_follow = self._get_node_to_follow(cluster)
-            self.state_handler.follow(node_to_follow, cluster.leader, True)
+            return self.state_handler.follow(node_to_follow, cluster.leader, True)
         else:
             logger.error('failed to bootstrap %s', msg)
             self.state_handler.remove_data_directory()
@@ -432,7 +432,7 @@ class Ha(object):
             sleep(2)  # Give a time to somebody to take the leader lock
             cluster = self.dcs.get_cluster()
             node_to_follow = self._get_node_to_follow(cluster)
-            self.state_handler.follow(node_to_follow, cluster.leader, recovery=True, need_rewind=True)
+            return self.state_handler.follow(node_to_follow, cluster.leader, recovery=True, need_rewind=True)
         else:
             self.state_handler.follow(None, None)
 
@@ -666,7 +666,7 @@ class Ha(object):
 
         clone_member = self.cluster.get_clone_member(self.state_handler.name)
         member_role = 'leader' if clone_member == self.cluster.leader else 'replica'
-        self.clone(clone_member, "from {0} '{1}'".format(member_role, clone_member.name))
+        return self.clone(clone_member, "from {0} '{1}'".format(member_role, clone_member.name))
 
     def reinitialize(self):
         with self._async_executor:
@@ -730,9 +730,10 @@ class Ha(object):
             if self._async_executor.busy:
                 return self.handle_long_action_in_progress()
 
-            # we've got here, so any async action has finished. Check if we tried to recover and failed
+            # we've got here, so any async action has finished.
             if self.recovering and not self.state_handler.need_rewind:
                 self.recovering = False
+                # Check if we tried to recover and failed
                 msg = self.post_recover()
                 if msg is not None:
                     return msg
@@ -793,3 +794,20 @@ class Ha(object):
         with self._async_executor:
             info = self._run_cycle()
             return (self.is_paused() and 'PAUSE: ' or '') + info
+
+    def watch(self, timeout):
+        cluster = self.cluster
+        # watch on leader key changes if the postgres is running and leader is known and current node is not lock owner
+        if not self._async_executor.busy and cluster and cluster.leader \
+                and cluster.leader.name != self.state_handler.name:
+            leader_index = cluster.leader.index
+        else:
+            leader_index = None
+
+        return self.dcs.watch(leader_index, timeout)
+
+    def wakeup(self):
+        """Call of this method will trigger the next run of HA loop if there is
+        no "active" leader watch request in progress.
+        This usually happens on the master or if the node is running async action"""
+        self.dcs.event.set()

patroni/postgresql.py

@@ -813,7 +813,7 @@ class Postgresql(object):
             async_executor.schedule('changing primary_conninfo and restarting')
             async_executor.run_async(self._do_follow, (primary_conninfo, leader, recovery))
         else:
-            self._do_follow(primary_conninfo, leader, recovery)
+            return self._do_follow(primary_conninfo, leader, recovery)
 
     def _do_follow(self, primary_conninfo, leader, recovery=False):
         change_role = self.role in ('master', 'demoted')
@@ -862,20 +862,22 @@ class Postgresql(object):
 
             if self.rewind(r) or not self.config.get('remove_data_directory_on_rewind_failure', False):
                 self.write_recovery_conf(primary_conninfo)
-                ret = self.start()
+                self.start()
             else:
                 logger.error('unable to rewind the former master')
                 self.remove_data_directory()
-                ret = True
             self._need_rewind = False
         else:
             self.write_recovery_conf(primary_conninfo)
-            ret = self.start() if recovery else self.restart()
+            if recovery:
+                self.start()
+            else:
+                self.restart()
             self.set_role('replica')
 
         if change_role:
             self.call_nowait(ACTION_ON_ROLE_CHANGE)
-        return ret
+        return True
 
     def save_configuration_files(self):
         """

tests/test_api.py

@@ -67,6 +67,10 @@ class MockHa(object):
     def get_effective_tags():
         return {'nosync': True}
 
+    @staticmethod
+    def wakeup():
+        pass
+
 
 class MockPatroni(object):
 

tests/test_async_executor.py

@@ -8,7 +8,7 @@ from threading import Thread
 class TestAsyncExecutor(unittest.TestCase):
 
     def setUp(self):
-        self.a = AsyncExecutor()
+        self.a = AsyncExecutor(Mock())
 
     @patch.object(Thread, 'start', Mock())
     def test_run_async(self):

tests/test_consul.py

@@ -148,11 +148,11 @@ class TestConsul(unittest.TestCase):
 
     @patch.object(AbstractDCS, 'watch', Mock())
     def test_watch(self):
-        self.c.watch(1)
+        self.c.watch(None, 1)
         self.c._name = ''
-        self.c.watch(1)
+        self.c.watch(6429, 1)
         with patch.object(consul.Consul.KV, 'get', Mock(side_effect=ConsulException)):
-            self.c.watch(1)
+            self.c.watch(6429, 1)
 
     def test_set_retry_timeout(self):
         self.c.set_retry_timeout(10)

tests/test_ctl.py

@@ -7,7 +7,7 @@ import unittest
 from click.testing import CliRunner
 from mock import patch, Mock
 from patroni.ctl import ctl, members, store_config, load_config, output_members, request_patroni, get_dcs, parse_dcs, \
-    wait_for_leader, get_all_members, get_any_member, get_cursor, query_member, configure, PatroniCtlException
+    get_all_members, get_any_member, get_cursor, query_member, configure, PatroniCtlException
 from patroni.dcs.etcd import Client
 from psycopg2 import OperationalError
 from test_etcd import etcd_read, requests_get, socket_getaddrinfo, MockResponse
@@ -305,14 +305,6 @@ class TestCtl(unittest.TestCase):
         result = self.runner.invoke(ctl, ['remove', 'alpha'], input='alpha\nYes I am aware\nleader')
         assert result.exit_code == 0
 
-    @patch('patroni.dcs.AbstractDCS.watch', Mock(return_value=None))
-    @patch('patroni.dcs.AbstractDCS.get_cluster', Mock(return_value=get_cluster_initialized_with_leader()))
-    def test_wait_for_leader(self):
-        self.assertRaises(PatroniCtlException, wait_for_leader, self.e, 0)
-
-        cluster = wait_for_leader(self.e, timeout=2)
-        assert cluster.leader.member.name == 'leader'
-
     @patch('requests.post', Mock(side_effect=requests.exceptions.ConnectionError('foo')))
     def test_request_patroni(self):
         member = get_cluster_initialized_with_leader().leader.member

tests/test_etcd.py

@@ -271,12 +271,12 @@ class TestEtcd(unittest.TestCase):
 
     @patch.object(etcd.Client, 'watch', etcd_watch)
     def test_watch(self):
-        self.etcd.watch(0)
+        self.etcd.watch(None, 0)
         self.etcd.get_cluster()
-        self.etcd.watch(1.5)
-        self.etcd.watch(4.5)
+        self.etcd.watch(20729, 1.5)
+        self.etcd.watch(20729, 4.5)
         with patch.object(AbstractDCS, 'watch', Mock()):
-            self.etcd.watch(9.5)
+            self.etcd.watch(20729, 9.5)
 
     def test_other_exceptions(self):
         self.etcd.retry = Mock(side_effect=AttributeError('foo'))
@@ -284,7 +284,7 @@ class TestEtcd(unittest.TestCase):
 
     def test_set_ttl(self):
         self.etcd.set_ttl(20)
-        self.assertTrue(self.etcd.watch(1))
+        self.assertTrue(self.etcd.watch(None, 1))
 
     def test_sync_state(self):
         self.assertFalse(self.etcd.write_sync_state('leader', None))

tests/test_ha.py

@@ -173,9 +173,6 @@ class TestHa(unittest.TestCase):
         self.ha.cluster = get_cluster_initialized_with_leader()
         self.assertEquals(self.ha.run_cycle(), 'starting as readonly because i had the session lock')
 
-    def test_do_not_recover_in_pause(self):
-        pass
-
     @patch('sys.exit', return_value=1)
     @patch('patroni.ha.Ha.sysid_valid', MagicMock(return_value=True))
     def test_sysid_no_match(self, exit_mock):
@@ -459,6 +456,9 @@ class TestHa(unittest.TestCase):
 
     def test_evaluate_scheduled_restart(self):
         self.p.postmaster_start_time = Mock(return_value=str(postmaster_start_time))
+        # restart already in progres
+        with patch('patroni.async_executor.AsyncExecutor.busy', PropertyMock(return_value=True)):
+            self.assertIsNone(self.ha.evaluate_scheduled_restart())
         # restart while the postmaster has been already restarted, fails
         with patch.object(self.ha,
                           'future_restart_scheduled',
@@ -681,3 +681,10 @@ class TestHa(unittest.TestCase):
         self.ha.has_lock = true
         self.ha.cluster.is_unlocked = false
         self.assertEquals(self.ha.run_cycle(), 'no action.  i am the leader with the lock')
+
+    def test_watch(self):
+        self.ha.cluster = get_cluster_initialized_with_leader()
+        self.ha.watch(0)
+
+    def test_wakup(self):
+        self.ha.wakeup()

tests/test_zookeeper.py

@@ -203,9 +203,9 @@ class TestZooKeeper(unittest.TestCase):
         self.assertTrue(self.zk.delete_cluster())
 
     def test_watch(self):
-        self.zk.watch(0)
-        self.zk.event.isSet = lambda: True
-        self.zk.watch(0)
+        self.zk.watch(None, 0)
+        self.zk.event.isSet = Mock(return_value=True)
+        self.zk.watch(None, 0)
 
     def test__kazoo_connect(self):
         self.zk._client._retry.deadline = 1