From e3ef9ac306994d10ce7fcc8b995f16da8f784b17 Mon Sep 17 00:00:00 2001
From: Alexander Kukushkin <cyberdemn@gmail.com>
Date: Mon, 14 Dec 2020 15:12:57 +0100
Subject: [PATCH] Fix issues with zookeeper (#1792)

1. The `ttl` was incorrectly returned 1000 times higher then it should
2. The `watch()` method must return True if the parent method returned True. Not doing so resulted in the incorrect calculation of sleep time.
3. Move mock of exhibitor api to the features/environment.py. It simplifies testing with behave.
---
 .github/workflows/install_deps.py |  9 ---------
 features/environment.py           | 21 ++++++++++++++++++++-
 patroni/dcs/zookeeper.py          | 10 ++++++----
 tests/test_zookeeper.py           |  2 +-
 4 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/.github/workflows/install_deps.py b/.github/workflows/install_deps.py
index b795aeb2..1a942fe1 100644
--- a/.github/workflows/install_deps.py
+++ b/.github/workflows/install_deps.py
@@ -154,13 +154,6 @@ users:
     return 0
 
 
-def setup_exhibitor():
-    response = '{"servers":["127.0.0.1"],"port":2181}'
-    response = 'HTTP/1.0 200 OK\\nContent-Length: {0}\\n\\n{1}'.format(len(response), response)
-    s = subprocess.Popen("while true; do echo '{0}'| nc -l 8181 > /dev/null; done".format(response), shell=True)
-    return 0 if s.poll() is None else s.returncode
-
-
 def main():
     what = os.environ.get('DCS', sys.argv[1] if len(sys.argv) > 1 else 'all')
     r = install_requirements(what)
@@ -178,8 +171,6 @@ def main():
         return install_etcd()
     elif what == 'kubernetes':
         return setup_kubernetes()
-    elif what == 'exhibitor':
-        return setup_exhibitor()
     return 0
 
 
diff --git a/features/environment.py b/features/environment.py
index b2ff6c86..7402658a 100644
--- a/features/environment.py
+++ b/features/environment.py
@@ -13,6 +13,8 @@ import threading
 import time
 import yaml
 
+from six.moves.BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
+
 
 @six.add_metaclass(abc.ABCMeta)
 class AbstractController(object):
@@ -559,11 +561,28 @@ class ZooKeeperController(AbstractDcsController):
             return False
 
 
+class MockExhibitor(BaseHTTPRequestHandler):
+
+    def do_GET(self):
+        self.send_response(200)
+        self.end_headers()
+        self.wfile.write(b'{"servers":["127.0.0.1"],"port":2181}')
+
+    def log_message(self, fmt, *args):
+        pass
+
+
 class ExhibitorController(ZooKeeperController):
 
     def __init__(self, context):
         super(ExhibitorController, self).__init__(context, False)
-        os.environ.update({'PATRONI_EXHIBITOR_HOSTS': 'localhost', 'PATRONI_EXHIBITOR_PORT': '8181'})
+        port = 8181
+        exhibitor = HTTPServer(('', port), MockExhibitor)
+        exhibitor.daemon_thread = True
+        exhibitor_thread = threading.Thread(target=exhibitor.serve_forever)
+        exhibitor_thread.daemon = True
+        exhibitor_thread.start()
+        os.environ.update({'PATRONI_EXHIBITOR_HOSTS': 'localhost', 'PATRONI_EXHIBITOR_PORT': str(port)})
 
 
 class RaftController(AbstractDcsController):
diff --git a/patroni/dcs/zookeeper.py b/patroni/dcs/zookeeper.py
index 30b1e222..a7732cd4 100644
--- a/patroni/dcs/zookeeper.py
+++ b/patroni/dcs/zookeeper.py
@@ -123,13 +123,14 @@ class ZooKeeper(AbstractDCS):
         # the same time, set_ttl method will reestablish connection and return
         # `!True`, otherwise we will close existing connection and let kazoo
         # open the new one.
-        if not self.set_ttl(int(config['ttl'] * 1000)) and loop_wait_changed:
+        if not self.set_ttl(config['ttl']) and loop_wait_changed:
             self._client._connection._socket.close()
 
     def set_ttl(self, ttl):
         """It is not possible to change ttl (session_timeout) in zookeeper without
         destroying old session and creating the new one. This method returns `!True`
         if session_timeout has been changed (`restart()` has been called)."""
+        ttl = int(ttl * 1000)
         if self._client._session_timeout != ttl:
             self._client._session_timeout = ttl
             self._client.restart()
@@ -137,7 +138,7 @@ class ZooKeeper(AbstractDCS):
 
     @property
     def ttl(self):
-        return self._client._session_timeout
+        return self._client._session_timeout / 1000.0
 
     def set_retry_timeout(self, retry_timeout):
         retry = self._client.retry if isinstance(self._client.retry, KazooRetry) else self._client._retry
@@ -372,6 +373,7 @@ class ZooKeeper(AbstractDCS):
         return self.set_sync_state_value("{}", index)
 
     def watch(self, leader_index, timeout):
-        if super(ZooKeeper, self).watch(leader_index, timeout) and not self._fetch_optime:
+        ret = super(ZooKeeper, self).watch(leader_index, timeout)
+        if ret and not self._fetch_optime:
             self._fetch_cluster = True
-        return self._fetch_cluster
+        return ret or self._fetch_cluster
diff --git a/tests/test_zookeeper.py b/tests/test_zookeeper.py
index 3df1ad3e..8d00894a 100644
--- a/tests/test_zookeeper.py
+++ b/tests/test_zookeeper.py
@@ -17,7 +17,7 @@ class MockKazooClient(Mock):
 
     def __init__(self, *args, **kwargs):
         super(MockKazooClient, self).__init__()
-        self._session_timeout = 30
+        self._session_timeout = 30000
 
     @property
     def client_id(self):