Watch for changes of leader key

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)

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)

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