From 038b5aed72655815c1e3a802e13a6b6efc41b782 Mon Sep 17 00:00:00 2001
From: Alexander Kukushkin <cyberdemn@gmail.com>
Date: Tue, 22 Nov 2016 16:22:30 +0100
Subject: [PATCH] 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          |  2 +-
 patroni/api.py               |  6 +++---
 patroni/async_executor.py    | 10 ++++++++--
 patroni/ctl.py               | 14 --------------
 patroni/dcs/__init__.py      |  3 ++-
 patroni/dcs/consul.py        | 11 +++++------
 patroni/dcs/etcd.py          | 10 ++++------
 patroni/dcs/zookeeper.py     |  4 ++--
 patroni/ha.py                | 28 +++++++++++++++++++++++-----
 patroni/postgresql.py        | 12 +++++++-----
 tests/test_api.py            |  4 ++++
 tests/test_async_executor.py |  2 +-
 tests/test_consul.py         |  6 +++---
 tests/test_ctl.py            | 10 +---------
 tests/test_etcd.py           | 10 +++++-----
 tests/test_ha.py             | 13 ++++++++++---
 tests/test_zookeeper.py      |  6 +++---
 17 files changed, 82 insertions(+), 69 deletions(-)

diff --git a/patroni/__init__.py b/patroni/__init__.py
index 0d83a50b..029ae6ff 100644
--- a/patroni/__init__.py
+++ b/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):
diff --git a/patroni/api.py b/patroni/api.py
index 92f931a7..0f1de6e1 100644
--- a/patroni/api.py
+++ b/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'
diff --git a/patroni/async_executor.py b/patroni/async_executor.py
index 04b6a5a8..b5c56ee7 100644
--- a/patroni/async_executor.py
+++ b/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()
diff --git a/patroni/ctl.py b/patroni/ctl.py
index 07ed22b1..0cedb7db 100644
--- a/patroni/ctl.py
+++ b/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(
diff --git a/patroni/dcs/__init__.py b/patroni/dcs/__init__.py
index 8bf3505a..ce4020af 100644
--- a/patroni/dcs/__init__.py
+++ b/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"""
 
diff --git a/patroni/dcs/consul.py b/patroni/dcs/consul.py
index c9c866b8..94407b10 100644
--- a/patroni/dcs/consul.py
+++ b/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()
diff --git a/patroni/dcs/etcd.py b/patroni/dcs/etcd.py
index 97dd3741..c1c7258d 100644
--- a/patroni/dcs/etcd.py
+++ b/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()
diff --git a/patroni/dcs/zookeeper.py b/patroni/dcs/zookeeper.py
index caaed510..c68bd3d8 100644
--- a/patroni/dcs/zookeeper.py
+++ b/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
diff --git a/patroni/ha.py b/patroni/ha.py
index af340264..a25ee50d 100644
--- a/patroni/ha.py
+++ b/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()
diff --git a/patroni/postgresql.py b/patroni/postgresql.py
index e8970177..3f1a0719 100644
--- a/patroni/postgresql.py
+++ b/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):
         """
diff --git a/tests/test_api.py b/tests/test_api.py
index 76d75a30..3a76d506 100644
--- a/tests/test_api.py
+++ b/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):
 
diff --git a/tests/test_async_executor.py b/tests/test_async_executor.py
index 17d7f94a..6f867428 100644
--- a/tests/test_async_executor.py
+++ b/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):
diff --git a/tests/test_consul.py b/tests/test_consul.py
index 135f1843..04700b3d 100644
--- a/tests/test_consul.py
+++ b/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)
diff --git a/tests/test_ctl.py b/tests/test_ctl.py
index 645df8f6..9c7e457f 100644
--- a/tests/test_ctl.py
+++ b/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
diff --git a/tests/test_etcd.py b/tests/test_etcd.py
index bb6567fa..3f38bc86 100644
--- a/tests/test_etcd.py
+++ b/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))
diff --git a/tests/test_ha.py b/tests/test_ha.py
index 035d39ad..c5b01884 100644
--- a/tests/test_ha.py
+++ b/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()
diff --git a/tests/test_zookeeper.py b/tests/test_zookeeper.py
index 3b7e8073..9166138a 100644
--- a/tests/test_zookeeper.py
+++ b/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