Merge pull request #59 from zalando/feature/pg_rewind_unclean_shutdown

Feature/pg rewind unclean shutdown. Allow pg_rewind to rewind the former master if it was uncleanly shutdown or just lost the connection and advanced past the point where the new master had promoted. Add configuration block pg_rewind with username and password as mandatory parameters inside this block.

README.rst

@@ -52,94 +52,62 @@ YAML Configuration
 For an example file, see ``postgres0.yml``. Below is an explanation of
 settings:
 
--  *ttl*: the TTL to acquire the leader lock. Think of it as the length
+-  *ttl*: the TTL to acquire the leader lock. Think of it as the length of time before automatic failover process is initiated.
-   of time before automatic failover process is initiated.
 -  *loop\_wait*: the number of seconds the loop will sleep
 
--  *restapi*
+-  *restapi*:
--  *listen*: ip address + port that Patroni will listen to provide
+    -  *listen*: ip address + port that Patroni will listen to provide health-check information for haproxy.
-   health-check information for haproxy.
+    -  *connect\_address*: ip address + port through which restapi is accessible.
--  *connect\_address*: ip address + port through which restapi is
+    -  *auth*: (optional) 'username:password' to protect some dangerous REST API endpoints.
-   accessible.
+    -  *certfile*: (optional) Specifies a file with the certificate in the PEM format. If certfile is not specified or empty API server will work without SSL.
+    -  *keyfile*: (optional) Specifies a file with the secret key in the PEM format.
 
--  *etcd*
+-  *etcd*:
--  *scope*: the relative path used on etcd's http api for this
+    -  *scope*: the relative path used on etcd's http api for this deployment, thus you can run multiple HA deployments from a single etcd
-   deployment, thus you can run multiple HA deployments from a single
+    -  *ttl*: the TTL to acquire the leader lock. Think of it as the length of time before automatic failover process is initiated.
-   etcd
+    -  *host*: the host:port for the etcd endpoint
--  *ttl*: the TTL to acquire the leader lock. Think of it as the length
-   of time before automatic failover process is initiated.
--  *host*: the host:port for the etcd endpoint
 
--  *zookeeper*
+-  *zookeeper*:
--  *scope*: the relative path used on etcd's http api for this
+    -  *scope*: the relative path used on etcd's http api for this deployment, thus you can run multiple HA deployments from a single etcd
-   deployment, thus you can run multiple HA deployments from a single
+    -  *session\_timeout*: the TTL to acquire the leader lock. Think of it as the length of time before automatic failover process is initiated.
-   etcd
+    -  *reconnect\_timeout*: how long we should try to reconnect to ZooKeeper after connection loss. After this timeout we assume that we don't have lock anymore and will restart in read-only mode.
--  *session\_timeout*: the TTL to acquire the leader lock. Think of it
+    -  *hosts*: list of ZooKeeper cluster members in format: ['host1:port1', 'host2:port2', 'etc...']
-   as the length of time before automatic failover process is initiated.
+    -  *exhibitor*: if you are running ZooKeeper cluster under Exhibitor supervisory the following section could be interesting for you
--  *reconnect\_timeout*: how long we should try to reconnect to
+        -  *poll\_interval*: how often list of ZooKeeper and Exhibitor nodes should be updated from Exhibitor
-   ZooKeeper after connection loss. After this timeout we assume that we
+        -  *port*: Exhibitor port
-   don't have lock anymore and will restart in read-only mode.
+        -  *hosts*: initial list of Exhibitor (ZooKeeper) nodes in format: ['host1', 'host2', 'etc...' ]. This list would be updated automatically when Exhibitor (ZooKeeper) cluster topology changes.
--  *hosts*: list of ZooKeeper cluster members in format: [
-   'host1:port1', 'host2:port2', 'etc...']
--  *exhibitor*: if you are running ZooKeeper cluster under Exhibitor
-   supervisory the following section could be interesting for you
 
-   -  *poll\_interval*: how often list of ZooKeeper and Exhibitor nodes
+-  *postgresql*:
-      should be updated from Exhibitor
+    -  *name*: the name of the Postgres host, must be unique for the cluster
-   -  *port*: Exhibitor port
+    -  *listen*: ip address + port that Postgres listening. Must be accessible from other nodes in the cluster if using streaming replication.
-   -  *hosts*: initial list of Exhibitor (ZooKeeper) nodes in format: [
+    -  *connect\_address*: ip address + port through which Postgres is accessible from other nodes and applications.
-      'host1', 'host2', 'etc...' ]. This list would be updated
+    -  *data\_dir*: file path to initialize and store Postgres data files
-      automatically when Exhibitor (ZooKeeper) cluster topology changes.
+    -  *maximum\_lag\_on\_failover*: the maximum bytes a follower may lag
+    -  *use\_slots*: whether or not to use replication_slots.  Must be False for PostgreSQL 9.3, and you should comment out max_replication_slots. before it is not eligible become leader
+    -  *pg\_hba*: list of lines which should be added to pg\_hba.conf
+        -  *- host all all 0.0.0.0/0 md5*
 
--  *postgresql*
+    -  *replication*:
--  *name*: the name of the Postgres host, must be unique for the cluster
+        -  *username*: replication username, user will be created during initialization
--  *listen*: ip address + port that Postgres listening. Must be
+        -  *password*: replication password, user will be created during initialization
-   accessible from other nodes in the cluster if using streaming
+        -  *network*: network setting for replication in pg\_hba.conf
-   replication.
--  *connect\_address*: ip address + port through which Postgres is
-   accessible from other nodes and applications.
--  *data\_dir*: file path to initialize and store Postgres data files
--  *maximum\_lag\_on\_failover*: the maximum bytes a follower may lag
--  *use\_slots*: whether or not to use replication_slots.  Must be False for PostgreSQL 9.3, and you should comment out max_replication_slots.
-   before it is not eligible become leader
--  *pg\_hba*: list of lines which should be added to pg\_hba.conf
 
-   -  *- host all all 0.0.0.0/0 md5*
+    -  *callbacks* callback scripts to run on certain actions. Patroni will pass current action, role and cluster name. See scripts/aws.py as an example on how to write them.
+        -  *on\_start*: a script to run when the cluster starts
+        -  *on\_stop*: a script to run when the cluster stops
+        -  *on\_restart*: a script to run when the cluster restarts
+        -  *on\_reload*: a script to run when configuration reload is triggered
+        -  *on\_role\_change*: a script to run when the cluster is being promoted or demoted
 
--  *replication*
+    -  *superuser*:
+        -  *password*: password for postgres user. It would be set during initialization
 
-   -  *username*: replication username, user will be created during
+    -  *admin*:
-      initialization
+        -  *username*: admin username, user will be created during initialization. It would have CREATEDB and CREATEROLE privileges
-   -  *password*: replication password, user will be created during
+        -  *password*: admin password, user will be created during initialization.
-      initialization
-   -  *network*: network setting for replication in pg\_hba.conf
 
--  *callbacks* callback scripts to run on certain actions. Patroni will
+    -  *recovery\_conf*: additional configuration settings written to recovery.conf when configuring follower
-   pass current action, role and cluster name. See scripts/aws.py as an
+        -  *parameters*: list of configuration settings for Postgres.  Many of these are required for replication to work.
-   example on how to write them.
-
-   -  *on\_start*: a script to run when the cluster starts
-   -  *on\_stop*: a script to run when the cluster stops
-   -  *on\_restart*: a script to run when the cluster restarts
-   -  *on\_reload*: a script to run when configuration reload is
-      triggered
-   -  *on\_role\_change*: a script to run when the cluster is being
-      promoted or demoted
-
--  *superuser*
-
-   -  *password*: password for postgres user. It would be set during
-      initialization
-
--  *admin*:
-
-   -  *username*: admin username, user will be created during
-      initialization. It would have CREATEDB and CREATEROLE privileges
-   -  *password*: admin password, user will be created during
-      initialization.
-
--  *recovery\_conf*: additional configuration settings written to recovery.conf when configuring follower
--  *parameters*: list of configuration settings for Postgres.  Many of these are required for replication to work.
 
 Replication choices
 -------------------

docker/README.md

@@ -1,7 +1,7 @@
 # Patroni Dockerfile
 You can run Patroni in a docker container using this Dockerfile, or by using one of the Docker image at
 
-    https://os-registry.stups.zalan.do/v1/repositories/acid/patroni/tags
+    https://registry.opensource.zalan.do/v1/repositories/acid/patroni/tags
 
 This Dockerfile is meant in aiding development of Patroni and quick testing of features. It is not a production-worthy
 Dockerfile
@@ -10,7 +10,7 @@ Dockerfile
 
 ## Standalone Patroni
 
-    docker run -d os-registry.stups.zalan.do/acid/patroni:1.0-SNAPSHOT
+    docker run -d registry.opensource.zalan.do/acid/patroni:1.0-SNAPSHOT
 
 ## Multiple Patroni's communicating with a standalone etcd inside Docker
 
@@ -36,12 +36,12 @@ To automate this you can run the following script:
 
 Example session:
 
-    $ ./dev_patroni_cluster.sh --image os-registry.stups.zalan.do/acid/patroni:1.0-SNAPSHOT --members=2 --name=bravo
+    $ ./dev_patroni_cluster.sh --image registry.opensource.zalan.do/acid/patroni:1.0-SNAPSHOT --members=2 --name=bravo
     The etcd container is 6be871a11cb373406ca5ea1c6b39e1.0-SNAPSHOTfdde9fb1d6177212d6ad0c0d1bd9b563, ip=172.17.1.24
     Started Patroni container 67e611f2eca7c40f9e6e0e24a4a8f2cba7e3e56d22a420e15ab9240a37a9d7a4, ip=172.17.1.25
     Started Patroni container 47dd12ae635ab83b039f5889e250048b606ed5e48e3650b69e365e7e1d4acbcf, ip=172.17.1.26
     $ docker ps
     CONTAINER ID        IMAGE                                         COMMAND                CREATED             STATUS              PORTS                          NAMES
-    47dd12ae635a        os-registry.stups.zalan.do/acid/patroni:1.0-SNAPSHOT   "/bin/bash /entrypoi   10 seconds ago      Up 8 seconds        4001/tcp, 5432/tcp, 2380/tcp   bravo_OR64g8bx
+    47dd12ae635a        registry.opensource.zalan.do/acid/patroni:1.0-SNAPSHOT   "/bin/bash /entrypoi   10 seconds ago      Up 8 seconds        4001/tcp, 5432/tcp, 2380/tcp   bravo_OR64g8bx
-    67e611f2eca7        os-registry.stups.zalan.do/acid/patroni:1.0-SNAPSHOT   "/bin/bash /entrypoi   11 seconds ago      Up 10 seconds       2380/tcp, 4001/tcp, 5432/tcp   bravo_si9no8iz
+    67e611f2eca7        registry.opensource.zalan.do/acid/patroni:1.0-SNAPSHOT   "/bin/bash /entrypoi   11 seconds ago      Up 10 seconds       2380/tcp, 4001/tcp, 5432/tcp   bravo_si9no8iz
-    6be871a11cb3        os-registry.stups.zalan.do/acid/patroni:1.0-SNAPSHOT   "/bin/bash /entrypoi   12 seconds ago      Up 10 seconds       4001/tcp, 5432/tcp, 2380/tcp   bravo_etcd
+    6be871a11cb3        registry.opensource.zalan.do/acid/patroni:1.0-SNAPSHOT   "/bin/bash /entrypoi   12 seconds ago      Up 10 seconds       4001/tcp, 5432/tcp, 2380/tcp   bravo_etcd

extras/startup-scripts/README.md

@@ -0,0 +1,10 @@
+# startup scripts for Patroni
+
+This directory contains sample startup scripts for various OSes 
+and management tools for Patroni.
+
+Scripts supplied:
+
+### patroni.upstart.conf
+
+Upstart job for Ubuntu 12.04 or 14.04.  Requires Upstart > 1.4. Intended for systems where Patroni has been installed on a base system, rather than in Docker.

extras/startup-scripts/patroni.upstart.conf

@@ -0,0 +1,36 @@
+# patroni - patroni daemon
+#
+# controls startup/shutdown of postgres
+# you should disable any postgres start jobs
+#
+# assumes that patroni has been installed into the
+# pythonpath by using setup.py install
+
+description    "patroni start daemon"
+
+start on net-device-up
+stop on runlevel [06]
+
+respawn
+respawn limit 5 10
+
+# set location of patroni
+env PATRONI=/usr/local/bin/patroni
+# virtualenv example
+# env PATRONI=/var/lib/postgresql/patronienv/bin/patroni
+
+# set location of config file
+env PATRONICONF=/etc/patroni/patroni.yml
+
+# set log dir for patroni logs
+# postgres user must have write permission
+env POSTGRESLOGDIR=/var/log/postgresql
+
+setuid postgres
+setgid postgres
+
+script
+    exec start-stop-daemon --start \
+        --exec $PATRONI -- $PATRONICONF \
+        >> $POSTGRESLOGDIR/patroni.log 2>&1
+end script

patroni/__init__.py

@@ -6,10 +6,9 @@ import yaml
 
 from patroni.api import RestApiServer
 from patroni.etcd import Etcd
-from patroni.exceptions import DCSError
 from patroni.ha import Ha
 from patroni.postgresql import Postgresql
-from patroni.utils import setup_signal_handlers, sleep, reap_children
+from patroni.utils import setup_signal_handlers, reap_children
 from patroni.zookeeper import ZooKeeper
 
 logger = logging.getLogger(__name__)
@@ -20,11 +19,11 @@ class Patroni:
     def __init__(self, config):
         self.nap_time = config['loop_wait']
         self.postgresql = Postgresql(config['postgresql'])
-        self.ha = Ha(self.postgresql, self.get_dcs(self.postgresql.name, config))
+        self.dcs = self.get_dcs(self.postgresql.name, config)
         host, port = config['restapi']['listen'].split(':')
         self.api = RestApiServer(self, config['restapi'])
+        self.ha = Ha(self)
         self.next_run = time.time()
-        self.shutdown_member_ttl = 300
 
     @staticmethod
     def get_dcs(name, config):
@@ -34,80 +33,21 @@ class Patroni:
             return ZooKeeper(name, config['zookeeper'])
         raise Exception('Can not find sutable configuration of distributed configuration store')
 
-    def touch_member(self, ttl=None):
-        connection_string = self.postgresql.connection_string + '?application_name=' + self.api.connection_string
-        if self.ha.cluster:
-            for m in self.ha.cluster.members:
-                # Do not update member TTL when it is far from being expired
-                if m.name == self.postgresql.name and m.real_ttl() > self.shutdown_member_ttl:
-                    return True
-        return self.ha.dcs.touch_member(connection_string, ttl)
-
-    def cleanup_on_failed_initialization(self):
-        """ cleanup the DCS if initialization was not successfull """
-        logger.info("removing initialize key after failed attempt to initialize the cluster")
-        self.ha.dcs.cancel_initialization()
-        self.touch_member(self.shutdown_member_ttl)
-        self.postgresql.stop()
-        self.postgresql.move_data_directory()
-
-    def initialize(self):
-        # wait for etcd to be available
-        while not self.touch_member():
-            logger.info('waiting on DCS')
-            sleep(5)
-
-        # is data directory empty?
-        if self.postgresql.data_directory_empty():
-            while True:
-                try:
-                    cluster = self.ha.dcs.get_cluster()
-                    if not cluster.is_unlocked():  # the leader already exists
-                        if not cluster.initialize:
-                            self.ha.dcs.initialize()
-                        self.postgresql.bootstrap(cluster.leader)
-                        break
-                    # racing to initialize
-                    elif not cluster.initialize and self.ha.dcs.initialize():
-                        try:
-                            self.postgresql.bootstrap()
-                        except:
-                            # bail out and clean the initialize flag.
-                            self.cleanup_on_failed_initialization()
-                            raise
-                        self.ha.dcs.take_leader()
-                        break
-                except DCSError:
-                    logger.info('waiting on DCS')
-                sleep(5)
-        elif self.postgresql.is_running():
-            self.postgresql.load_replication_slots()
-
     def schedule_next_run(self):
-        if self.postgresql.is_promoted:
-            self.next_run = time.time()
         self.next_run += self.nap_time
         current_time = time.time()
         nap_time = self.next_run - current_time
         if nap_time <= 0:
             self.next_run = current_time
-        else:
+        elif self.dcs.watch(nap_time):
-            self.ha.dcs.watch(nap_time)
+            self.next_run = time.time()
 
     def run(self):
         self.api.start()
         self.next_run = time.time()
 
         while True:
-            self.touch_member()
             logger.info(self.ha.run_cycle())
-            try:
-                if self.ha.state_handler.is_leader():
-                    self.ha.cluster and self.ha.state_handler.create_replication_slots(self.ha.cluster)
-                else:
-                    self.ha.state_handler.drop_replication_slots()
-            except:
-                logger.exception('Exception when changing replication slots')
             reap_children()
             self.schedule_next_run()
 
@@ -125,12 +65,11 @@ def main():
         config = yaml.load(f)
 
     patroni = Patroni(config)
-    patroni.initialize()
     try:
         patroni.run()
     except KeyboardInterrupt:
         pass
     finally:
-        patroni.touch_member(patroni.shutdown_member_ttl)  # schedule member removal
+        patroni.api.shutdown()
         patroni.postgresql.stop()
-        patroni.ha.dcs.delete_leader()
+        patroni.dcs.delete_leader()

patroni/api.py

@@ -1,8 +1,11 @@
+import base64
 import fcntl
 import json
 import logging
 import psycopg2
 
+from patroni.exceptions import PostgresConnectionException
+from patroni.utils import Retry, RetryFailedError
 from six.moves.BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
 from six.moves.socketserver import ThreadingMixIn
 from threading import Thread
@@ -10,33 +13,153 @@ from threading import Thread
 logger = logging.getLogger(__name__)
 
 
+def check_auth(func):
+    """Decorator function to check authorization header.
+
+    Usage example:
+    @check_auth
+    def do_PUT_foo():
+        pass
+    """
+    def wrapper(handler):
+        if handler.check_auth_header():
+            return func(handler)
+    return wrapper
+
+
 class RestApiHandler(BaseHTTPRequestHandler):
 
+    def send_auth_request(self, body):
+        self.send_response(401)
+        self.send_header('WWW-Authenticate', 'Basic realm=\"Patroni\"')
+        self.send_header('Content-type', 'text/html')
+        self.end_headers()
+        self.wfile.write(body.encode('utf-8'))
+
+    def check_auth_header(self):
+        auth_header = self.headers.get('Authorization')
+        status = self.server.check_auth_header(auth_header)
+        return not status or self.send_auth_request(status)
+
     def do_GET(self):
-        response = self.get_postgresql_status()
+        """Default method for processing all GET requests which can not be routed to other methods"""
 
         path = '/master' if self.path == '/' else self.path
-        status_code = 200 if response['running'] and 'role' in response and response['role'] in path else 503
+        response = self.get_postgresql_status()
+
+        patroni = self.server.patroni
+        cluster = patroni.dcs.cluster
+        if cluster:  # dcs available
+            if cluster.leader and cluster.leader.name == patroni.postgresql.name:  # is_leader
+                status_code = 200 if 'master' in path else 503
+            elif 'role' not in response:
+                status_code = 503
+            elif response['role'] == 'master':  # running as master but without leader lock!!!!
+                status_code = 503
+            elif response['role'] in path:
+                status_code = 200
+            else:
+                status_code = 503
+        elif 'role' in response and response['role'] in path:
+            status_code = 200
+        elif patroni.ha.restart_scheduled() and patroni.postgresql.role == 'master' and 'master' in path:
+            # exceptional case for master node when the postgres is being restarted via API
+            status_code = 200
+        else:
+            status_code = 503
 
         self.send_response(status_code)
         self.send_header('Content-Type', 'application/json')
         self.end_headers()
         self.wfile.write(json.dumps(response).encode('utf-8'))
 
-    def get_postgresql_status(self):
+    def do_GET_patroni(self):
+        response = self.get_postgresql_status(True)
+
+        self.send_response(200)
+        self.send_header('Content-Type', 'application/json')
+        self.end_headers()
+        self.wfile.write(json.dumps(response).encode('utf-8'))
+
+    @check_auth
+    def do_POST_restart(self):
+        status_code = 503
+        data = b'restart failed'
         try:
-            row = self.server.query("""SELECT to_char(pg_postmaster_start_time(), 'YYYY-MM-DD HH24:MI:SS.MS TZ'),
+            status, msg = self.server.patroni.ha.restart()
-                                              pg_is_in_recovery(),
+            status_code = 200 if status else 503
-                                              CASE WHEN pg_is_in_recovery()
+            data = msg.encode('utf-8')
-                                                   THEN null
+        except:
-                                                   ELSE pg_current_xlog_location() END,
+            logger.exception('Exception during restart')
-                                              pg_last_xlog_receive_location(),
+
-                                              pg_last_xlog_replay_location(),
+        self.send_response(status_code)
-                                              pg_is_in_recovery() AND pg_is_xlog_replay_paused()""")[0]
+        self.send_header('Content-Type', 'text/html')
+        self.end_headers()
+        self.wfile.write(data)
+
+    @check_auth
+    def do_POST_reinitialize(self):
+        ha = self.server.patroni.ha
+        cluster = ha.dcs.get_cluster()
+        if cluster.is_unlocked():
+            status_code = 503
+            data = b'Cluster has no leader, can not reinitialize'
+        elif cluster.leader.name == ha.state_handler.name:
+            status_code = 503
+            data = b'I am the leader, can not reinitialize'
+        else:
+            action = ha.schedule_reinitialize()
+            if action is not None:
+                status_code = 503
+                data = (action + ' already in progress').encode('utf-8')
+            else:
+                status_code = 200
+                data = b'reinitialize scheduled'
+
+        self.send_response(status_code)
+        self.send_header('Content-Type', 'text/html')
+        self.end_headers()
+        self.wfile.write(data)
+
+    def parse_request(self):
+        """Override parse_request method to enrich basic functionality of `BaseHTTPRequestHandler` class
+
+        Original class can only invoke do_GET, do_POST, do_PUT, etc method implementations if they are defined.
+        But we would like to have at least some simple routing mechanism, i.e.:
+        GET /uri1/part2 request should invoke `do_GET_uri1()`
+        POST /other should invoke `do_POST_other()`
+
+        If the `do_<REQUEST_METHOD>_<first_part_url>` method does not exists we'll fallback to original behavior."""
+
+        ret = BaseHTTPRequestHandler.parse_request(self)
+        if ret:
+            mname = self.path.lstrip('/').split('/')[0]
+            mname = self.command + ('_' + mname if mname else '')
+            if hasattr(self, 'do_' + mname):
+                self.command = mname
+        return ret
+
+    def query(self, sql, *params, **kwargs):
+        if not kwargs.get('retry', False):
+            return self.server.query(sql, *params)
+        retry = Retry(delay=1, retry_exceptions=PostgresConnectionException)
+        return retry(self.server.query, sql, *params)
+
+    def get_postgresql_status(self, retry=False):
+        try:
+            row = self.query("""SELECT to_char(pg_postmaster_start_time(), 'YYYY-MM-DD HH24:MI:SS.MS TZ'),
+                                       pg_is_in_recovery(),
+                                       CASE WHEN pg_is_in_recovery()
+                                            THEN 0
+                                            ELSE pg_xlog_location_diff(pg_current_xlog_location(), '0/0')::bigint
+                                       END,
+                                       pg_xlog_location_diff(pg_last_xlog_receive_location(), '0/0')::bigint,
+                                       pg_xlog_location_diff(pg_last_xlog_replay_location(), '0/0')::bigint,
+                                       pg_is_in_recovery() AND pg_is_xlog_replay_paused()""", retry=retry)[0]
             return {
-                'running': True,
+                'state': self.server.patroni.postgresql.state,
                 'postmaster_start_time': row[0],
-                'role': 'slave' if row[1] else 'master',
+                'role': 'replica' if row[1] else 'master',
                 'xlog': ({
                     'received_location': row[3],
                     'replayed_location': row[4],
@@ -44,30 +167,60 @@ class RestApiHandler(BaseHTTPRequestHandler):
                     'location': row[2]
                 })
             }
-        except (psycopg2.OperationalError, psycopg2.InterfaceError):
+        except (psycopg2.Error, RetryFailedError, PostgresConnectionException):
-            logger.exception('get_postgresql_status')
+            state = self.server.patroni.postgresql.state
-            return {'running': self.server.patroni.postgresql.is_running()}
+            if state in ['stopped', 'starting', 'stopping', 'restarting', 'running']:
+                logger.exception('get_postgresql_status')
+                state = 'unknown' if state == 'running' else state
+            return {'state': state}
 
 
 class RestApiServer(ThreadingMixIn, HTTPServer, Thread):
 
     def __init__(self, patroni, config):
-        self.connection_string = 'http://{}/patroni'.format(config.get('connect_address', None) or config['listen'])
+        self._auth_key = base64.b64encode(config['auth'].encode('utf-8')).decode('utf-8') if 'auth' in config else None
         host, port = config['listen'].split(':')
         HTTPServer.__init__(self, (host, int(port)), RestApiHandler)
         Thread.__init__(self, target=self.serve_forever)
         self._set_fd_cloexec(self.socket)
+
+        protocol = 'http'
+
+        # wrap socket with ssl if 'certfile' is defined in a config.yaml
+        # Sometime it's also needed to pass reference to a 'keyfile'.
+        options = {option: config[option] for option in ['certfile', 'keyfile'] if option in config}
+        if options.get('certfile', None):
+            import ssl
+            self.socket = ssl.wrap_socket(self.socket, server_side=True, **options)
+            protocol = 'https'
+
+        self.connection_string = '{}://{}/patroni'.format(protocol, config.get('connect_address', config['listen']))
+
         self.patroni = patroni
         self.daemon = True
 
     def query(self, sql, *params):
-        cursor = self.patroni.postgresql.connection().cursor()
+        cursor = None
-        cursor.execute(sql, params)
+        try:
-        ret = [r for r in cursor]
+            with self.patroni.postgresql.connection().cursor() as cursor:
-        cursor.close()
+                cursor.execute(sql, params)
-        return ret
+                return [r for r in cursor]
+        except psycopg2.Error as e:
+            if cursor and cursor.connection.closed == 0:
+                raise e
+            raise PostgresConnectionException('connection problems')
 
     @staticmethod
     def _set_fd_cloexec(fd):
         flags = fcntl.fcntl(fd, fcntl.F_GETFD)
         fcntl.fcntl(fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC)
+
+    def check_basic_auth_key(self, key):
+        return self._auth_key == key
+
+    def check_auth_header(self, auth_header):
+        if self._auth_key:
+            if auth_header is None:
+                return 'no auth header received'
+            if not auth_header.startswith('Basic ') or not self.check_basic_auth_key(auth_header[6:]):
+                return 'not authenticated'

patroni/async_executor.py

@@ -0,0 +1,55 @@
+import logging
+from threading import Lock, Thread
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncExecutor:
+
+    def __init__(self):
+        Lock.__init__(self)
+        self._busy = False
+        self._thread_lock = Lock()
+        self._scheduled_action = None
+        self._scheduled_action_lock = Lock()
+
+    @property
+    def busy(self):
+        return self._busy
+
+    def schedule(self, action, immediately=False):
+        with self._scheduled_action_lock:
+            if self._scheduled_action is not None:
+                return self._scheduled_action
+            self._scheduled_action = action
+            self._busy = immediately
+        return None
+
+    @property
+    def scheduled_action(self):
+        with self._scheduled_action_lock:
+            return self._scheduled_action
+
+    def reset_scheduled_action(self):
+        with self._scheduled_action_lock:
+            self._scheduled_action = None
+
+    def run(self, func, args=()):
+        try:
+            return func(*args) if args else func()
+        except:
+            logger.exception('Exception during execution of long running task %s', self.scheduled_action)
+        finally:
+            with self:
+                self._busy = False
+                self.reset_scheduled_action()
+
+    def run_async(self, func, args=()):
+        self._busy = True
+        Thread(target=self.run, args=(func, args)).start()
+
+    def __enter__(self):
+        self._thread_lock.acquire()
+
+    def __exit__(self, type, value, traceback):
+        self._thread_lock.release()

patroni/dcs.py

@@ -1,9 +1,10 @@
 import abc
+import json
 
 from collections import namedtuple
 from patroni.exceptions import DCSError
-from patroni.utils import calculate_ttl, sleep
 from six.moves.urllib_parse import urlparse, urlunparse, parse_qsl
+from threading import Event, Lock
 
 
 def parse_connection_string(value):
@@ -23,28 +24,52 @@ def parse_connection_string(value):
     return conn_url, api_url
 
 
-class Member(namedtuple('Member', 'index,name,conn_url,api_url,expiration,ttl')):
+class Member(namedtuple('Member', 'index,name,session,data')):
 
     """Immutable object (namedtuple) which represents single member of PostgreSQL cluster.
     Consists of the following fields:
     :param index: modification index of a given member key in a Configuration Store
     :param name: name of PostgreSQL cluster member
-    :param conn_url: connection string containing host, user and password which could be used to access this member.
+    :param session: either session id or just ttl in seconds
-    :param api_url: REST API url of patroni instance
+    :param data: arbitrary data i.e. conn_url, api_url, xlog location, state, role, tags, etc...
-    :param expiration: expiration time of given member key
-    :param ttl: ttl of given member key in seconds"""
 
-    def real_ttl(self):
+    There are two mandatory keys in a data:
-        return calculate_ttl(self.expiration) or -1
+    conn_url: connection string containing host, user and password which could be used to access this member.
+    api_url: REST API url of patroni instance"""
+
+    @staticmethod
+    def from_node(index, name, session, data):
+        """
+        >>> Member.from_node(-1, '', '', '{"conn_url": "postgres://foo@bar/postgres"}') is not None
+        True
+        >>> Member.from_node(-1, '', '', '{')
+        Member(index=-1, name='', session='', data={})
+        """
+        if data.startswith('postgres'):
+            conn_url, api_url = parse_connection_string(data)
+            data = {'conn_url': conn_url, 'api_url': api_url}
+        else:
+            try:
+                data = json.loads(data)
+            except:
+                data = {}
+        return Member(index, name, session, data)
+
+    @property
+    def conn_url(self):
+        return self.data.get('conn_url', None)
+
+    @property
+    def api_url(self):
+        return self.data.get('api_url', None)
 
 
-class Leader(namedtuple('Leader', 'index,expiration,ttl,member')):
+class Leader(namedtuple('Leader', 'index,session,member')):
 
     """Immutable object (namedtuple) which represents leader key.
     Consists of the following fields:
     :param index: modification index of a leader key in a Configuration Store
-    :param expiration: expiration time of the leader key
+    :param session: either session id or just ttl in seconds
-    :param ttl: ttl of the leader key
     :param member: reference to a `Member` object which represents current leader (see `Cluster.members`)"""
 
     @property
@@ -56,7 +81,15 @@ class Leader(namedtuple('Leader', 'index,expiration,ttl,member')):
         return self.member.conn_url
 
 
-class Cluster(namedtuple('Cluster', 'initialize,leader,last_leader_operation,members')):
+class Failover(namedtuple('Failover', 'index,leader,member')):
+
+    @staticmethod
+    def from_node(index, value):
+        t = [a.strip() for a in value.split(':')] + ['']
+        return Failover(index, t[0], t[1]) if t[0] or t[1] else None
+
+
+class Cluster(namedtuple('Cluster', 'initialize,leader,last_leader_operation,members,failover')):
 
     """Immutable object (namedtuple) which represents PostgreSQL cluster.
     Consists of the following fields:
@@ -64,7 +97,8 @@ class Cluster(namedtuple('Cluster', 'initialize,leader,last_leader_operation,mem
     :param leader: `Leader` object which represents current leader of the cluster
     :param last_leader_operation: int or long object containing position of last known leader operation.
         This value is stored in `/optime/leader` key
-    :param members: list of Member object, all PostgreSQL cluster members including leader"""
+    :param members: list of Member object, all PostgreSQL cluster members including leader
+    :param failover: reference to `Failover` object"""
 
     def is_unlocked(self):
         return not (self.leader and self.leader.name)
@@ -76,6 +110,7 @@ class AbstractDCS:
 
     _INITIALIZE = 'initialize'
     _LEADER = 'leader'
+    _FAILOVER = 'failover'
     _MEMBERS = 'members/'
     _OPTIME = 'optime'
     _LEADER_OPTIME = _OPTIME + '/' + _LEADER
@@ -90,6 +125,10 @@ class AbstractDCS:
         self._scope = config['scope']
         self._base_path = '/service/' + self._scope
 
+        self._cluster = None
+        self._cluster_thread_lock = Lock()
+        self.event = Event()
+
     def client_path(self, path):
         return '/'.join([self._base_path, path.lstrip('/')])
 
@@ -109,25 +148,54 @@ class AbstractDCS:
     def leader_path(self):
         return self.client_path(self._LEADER)
 
+    @property
+    def failover_path(self):
+        return self.client_path(self._FAILOVER)
+
     @property
     def leader_optime_path(self):
         return self.client_path(self._LEADER_OPTIME)
 
     @abc.abstractmethod
+    def _load_cluster(self):
+        """Internally this method should build  `Cluster` object which
+           represents current state and topology of the cluster in DCS.
+           this method supposed to be called only by `get_cluster` method.
+
+           raise `~DCSError` in case of communication or other problems with DCS.
+           If the current node was running as a master and exception raised,
+           instance would be demoted."""
+
     def get_cluster(self):
-        """:returns: `Cluster` object which represent current state and topology of the cluster
+        with self._cluster_thread_lock:
-        raise `~DCSError` in case of communication or other problems with DCS. If current instance was
+            try:
-            running as a master and exception raised instance would be demoted."""
+                self._load_cluster()
+            except:
+                self._cluster = None
+                raise
+            return self._cluster
+
+    @property
+    def cluster(self):
+        with self._cluster_thread_lock:
+            return self._cluster
+
+    def reset_cluster(self):
+        with self._cluster_thread_lock:
+            self._cluster = None
 
     @abc.abstractmethod
-    def update_leader(self, state_handler):
+    def write_leader_optime(self, last_operation):
-        """Update leader key (or session) ttl and `/optime/leader` key in DCS.
+        """write current xlog location into `/optime/leader` key in DCS
+        :param last_operation: absolute xlog location in bytes"""
+
+    @abc.abstractmethod
+    def update_leader(self):
+        """Update leader key (or session) ttl
 
-        :param state_handler: reference to `Postgresql` object
         :returns: `!True` if leader key (or session) has been updated successfully.
             If not, `!False` must be returned and current instance would be demoted.
 
-        If you failed to update `/optime/leader` this error is not critical and you can return `!True`
         You have to use CAS (Compare And Swap) operation in order to update leader key,
         for example for etcd `prevValue` parameter must be used."""
 
@@ -140,6 +208,13 @@ class AbstractDCS:
         Key must be created atomically. In case if key already exists it should not be
         overwritten and `!False` must be returned"""
 
+    @abc.abstractmethod
+    def set_failover_value(self, value, index=None):
+        """Create or update `/failover` key"""
+
+    def manual_failover(self, leader, member, index=None):
+        return self.set_failover_value(leader + (':' + member if member else ''), index)
+
     def current_leader(self):
         try:
             cluster = self.get_cluster()
@@ -182,4 +257,11 @@ class AbstractDCS:
         """ Removes the initialize key for a cluster """
 
     def watch(self, timeout):
-        sleep(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 timeout: timeout in seconds
+        :returns: `!True` if you would like to reschedule the next run of ha cycle"""
+
+        self.event.wait(timeout)
+        return self.event.isSet()

patroni/etcd.py

@@ -10,7 +10,8 @@ import urllib3
 
 from dns.exception import DNSException
 from dns import resolver
-from patroni.dcs import AbstractDCS, Cluster, DCSError, Leader, Member, parse_connection_string
+from patroni.dcs import AbstractDCS, Cluster, Failover, Leader, Member
+from patroni.exceptions import DCSError
 from patroni.utils import Retry, RetryFailedError, sleep
 from requests.exceptions import RequestException
 
@@ -80,7 +81,7 @@ class Client(etcd.Client):
         for host, port in self.get_srv_record(discovery_srv):
             url = '{}://{}:{}/members'.format(self._protocol, host, port)
             try:
-                response = requests.get(url)
+                response = requests.get(url, timeout=5)
                 if response.ok:
                     for member in response.json():
                         ret.extend(member['clientURLs'])
@@ -146,14 +147,12 @@ class Etcd(AbstractDCS):
     def __init__(self, name, config):
         super(Etcd, self).__init__(name, config)
         self.ttl = config['ttl']
-        self.member_ttl = config.get('member_ttl', 3600)
         self._retry = Retry(deadline=10, max_delay=1, max_tries=-1,
                             retry_exceptions=(etcd.EtcdConnectionFailed,
                                               etcd.EtcdLeaderElectionInProgress,
                                               etcd.EtcdWatcherCleared,
                                               etcd.EtcdEventIndexCleared))
         self.client = self.get_etcd_client(config)
-        self.cluster = None
 
     def retry(self, *args, **kwargs):
         return self._retry.copy()(*args, **kwargs)
@@ -170,10 +169,9 @@ class Etcd(AbstractDCS):
 
     @staticmethod
     def member(node):
-        conn_url, api_url = parse_connection_string(node.value)
+        return Member.from_node(node.modifiedIndex, os.path.basename(node.key), node.ttl, node.value)
-        return Member(node.modifiedIndex, os.path.basename(node.key), conn_url, api_url, node.expiration, node.ttl)
 
-    def get_cluster(self):
+    def _load_cluster(self):
         try:
             result = self.retry(self.client.read, self.client_path(''), recursive=True)
             nodes = {os.path.relpath(node.key, result.key): node for node in result.leaves}
@@ -191,22 +189,25 @@ class Etcd(AbstractDCS):
             # get leader
             leader = nodes.get(self._LEADER, None)
             if leader:
-                member = Member(-1, leader.value, None, None, None, None)
+                member = Member(-1, leader.value, None, {})
                 member = ([m for m in members if m.name == leader.value] or [member])[0]
-                leader = Leader(leader.modifiedIndex, leader.expiration, leader.ttl, member)
+                leader = Leader(leader.modifiedIndex, leader.ttl, member)
 
-            self.cluster = Cluster(initialize, leader, last_leader_operation, members)
+            # failover key
+            failover = nodes.get(self._FAILOVER, None)
+            if failover:
+                failover = Failover.from_node(failover.modifiedIndex, failover.value)
+
+            self._cluster = Cluster(initialize, leader, last_leader_operation, members, failover)
         except etcd.EtcdKeyNotFound:
-            self.cluster = Cluster(False, None, None, [])
+            self._cluster = Cluster(False, None, None, [], None)
         except:
-            self.cluster = None
             logger.exception('get_cluster')
             raise EtcdError('Etcd is not responding properly')
-        return self.cluster
 
     @catch_etcd_errors
     def touch_member(self, connection_string, ttl=None):
-        return self.retry(self.client.set, self.member_path, connection_string, ttl or self.member_ttl)
+        return self.retry(self.client.set, self.member_path, connection_string, ttl or self.ttl)
 
     @catch_etcd_errors
     def take_leader(self):
@@ -222,18 +223,20 @@ class Etcd(AbstractDCS):
         return False
 
     @catch_etcd_errors
-    def write_leader_optime(self, state_handler):
+    def set_failover_value(self, value, index=None):
-        return self.client.set(self.leader_optime_path, state_handler.last_operation())
+        return self.client.write(self.failover_path, value, prevIndex=index or 0)
 
     @catch_etcd_errors
-    def update_leader(self, state_handler):
+    def write_leader_optime(self, last_operation):
-        ret = self.retry(self.client.test_and_set, self.leader_path, self._name, self._name, self.ttl)
+        return self.client.set(self.leader_optime_path, last_operation)
-        ret and self.write_leader_optime(state_handler)
+
-        return ret
+    @catch_etcd_errors
+    def update_leader(self):
+        return self.retry(self.client.test_and_set, self.leader_path, self._name, self._name, self.ttl)
 
     @catch_etcd_errors
     def initialize(self):
-        return self.client.write(self.initialize_path, self._name, prevExist=False)
+        return self.retry(self.client.write, self.initialize_path, self._name, prevExist=False)
 
     @catch_etcd_errors
     def delete_leader(self):
@@ -241,25 +244,30 @@ class Etcd(AbstractDCS):
 
     @catch_etcd_errors
     def cancel_initialization(self):
-        return self.client.delete(self.initialize_path, prevValue=self._name)
+        return self.retry(self.client.delete, self.initialize_path, prevValue=self._name)
 
     def watch(self, timeout):
+        cluster = self.cluster
         # 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.name != self._name:
+        if cluster and cluster.leader and cluster.leader.name != self._name:
             end_time = time.time() + timeout
-            index = self.cluster.leader.index
+            index = 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.leader_path, index=index + 1, timeout=timeout)
+                    self.client.watch(self.leader_path, index=index + 1, timeout=timeout + 0.5)
-                    if res.action not in ['set', 'compareAndSwap'] or res.value != self.cluster.leader.name:
+                    # Synchronous work of all cluster members with etcd is less expensive
-                        return
+                    # than reestablishing http connection every time from every replica.
-                    index = res.modifiedIndex
+                    return True
                 except urllib3.exceptions.TimeoutError:
                     self.client.http.clear()
-                    return
+                    return False
                 except etcd.EtcdException:
-                    index = None
+                    logging.exception('watch')
 
                 timeout = end_time - time.time()
 
-        timeout > 0 and super(Etcd, self).watch(timeout)
+        try:
+            return super(Etcd, self).watch(timeout)
+        finally:
+            self.event.clear()

patroni/exceptions.py

@@ -7,7 +7,7 @@ class PatroniException(Exception):
 
     def __str__(self):
         """
-        >>> str(DCSError('foo'))
+        >>> str(PatroniException('foo'))
         "'foo'"
         """
         return repr(self.value)
@@ -19,3 +19,7 @@ class PostgresException(PatroniException):
 
 class DCSError(PatroniException):
     pass
+
+
+class PostgresConnectionException(PostgresException):
+    pass

patroni/ha.py

@@ -1,26 +1,30 @@
+import json
 import logging
+import psycopg2
+import requests
 
-from patroni.dcs import DCSError
+from patroni.async_executor import AsyncExecutor
-from psycopg2 import InterfaceError, OperationalError
+from patroni.exceptions import DCSError, PostgresConnectionException
+from multiprocessing.pool import ThreadPool
 
 logger = logging.getLogger(__name__)
 
 
 class Ha:
 
-    def __init__(self, state_handler, etcd):
+    def __init__(self, patroni):
-        self.state_handler = state_handler
+        self.patroni = patroni
-        self.dcs = etcd
+        self.state_handler = patroni.postgresql
+        self.dcs = patroni.dcs
         self.cluster = None
         self.old_cluster = None
+        self._async_executor = AsyncExecutor()
 
     def load_cluster_from_dcs(self):
         cluster = self.dcs.get_cluster()
 
         # We want to keep the state of cluster when it was healhy
-        if cluster.is_unlocked() and self.cluster and not self.cluster.is_unlocked():
+        if not cluster.is_unlocked() or not self.old_cluster:
-            self.old_cluster = self.cluster
-        if not self.old_cluster:
             self.old_cluster = cluster
         self.cluster = cluster
 
@@ -28,74 +32,370 @@ class Ha:
         return self.dcs.attempt_to_acquire_leader()
 
     def update_lock(self):
-        return self.dcs.update_leader(self.state_handler)
+        ret = self.dcs.update_leader()
+        if ret:
+            try:
+                self.dcs.write_leader_optime(self.state_handler.last_operation())
+            except:
+                pass
+        return ret
 
     def has_lock(self):
         lock_owner = self.cluster.leader and self.cluster.leader.name
         logger.info('Lock owner: %s; I am %s', lock_owner, self.state_handler.name)
         return lock_owner == self.state_handler.name
 
-    def demote(self):
+    def touch_member(self):
-        return self.state_handler.demote(self.cluster.leader)
+        data = {
+            'conn_url': self.state_handler.connection_string,
+            'api_url': self.patroni.api.connection_string,
+            'state': self.state_handler.state,
+            'role': self.state_handler.role
+        }
+        if data['state'] in ['running', 'restarting', 'starting']:
+            try:
+                data['xlog_location'] = self.state_handler.xlog_position()
+            except:
+                pass
+        self.dcs.touch_member(json.dumps(data, separators=(',', ':')))
 
-    def follow_the_leader(self):
+    def copy_backup_from_leader(self, leader):
-        return self.state_handler.follow_the_leader(self.cluster.leader)
+        if self.state_handler.bootstrap(leader):
+            logger.info('bootstrapped from leader')
+        else:
+            self.state_handler.stop('immediate')
+            self.state_handler.remove_data_directory()
+            logger.error('failed to bootstrap from leader')
 
-    def run_cycle(self):
+    def bootstrap(self):
+        if not self.cluster.is_unlocked():  # cluster already has leader
+            self._async_executor.schedule('bootstrap from leader')
+            self._async_executor.run_async(self.copy_backup_from_leader, args=(self.cluster.leader, ))
+            return 'trying to bootstrap from leader'
+        elif not self.cluster.initialize:  # no initialize key
+            if self.dcs.initialize():  # race for initialization
+                try:
+                    self.state_handler.bootstrap()
+                except:  # initdb or start failed
+                    # remove initialization key and give a chance to other members
+                    logger.info("removing initialize key after failed attempt to initialize the cluster")
+                    self.dcs.cancel_initialization()
+                    self.state_handler.stop('immediate')
+                    self.state_handler.move_data_directory()
+                    raise
+                self.dcs.take_leader()
+                return 'initialized a new cluster'
+            else:
+                return 'failed to acquire initialize lock'
+        else:
+            return 'waiting for leader to bootstrap'
+
+    def recover(self):
+        has_lock = self.has_lock()
+
+        # try to see if we are the former master that crashed. If so - we likely need to run pg_rewind
+        # in order to join the former standby being promoted.
+        pg_controldata = self.state_handler.controldata()
+        if not has_lock and pg_controldata and\
+                pg_controldata.get('Database cluster state', '') == 'in production':  # crashed master
+            self.state_handler.require_rewind()
+
+        # XXX: follow the leader calls stop, which might take quite some time.
+        # perhaps we should run sync asynchronously
+        # (we still need the exit code from follow_the_leader)
+        ret = self.state_handler.follow_the_leader(None if has_lock else self.cluster.leader, recovery=True)
+        if not ret:
+            if not has_lock:
+                return 'failed to start postgres'
+            self.dcs.delete_leader()
+            self.dcs.reset_cluster()
+            return 'removed leader key after trying and failing to start postgres'
+        if not has_lock:
+            return 'started as a secondary'
+        logger.info('started as readonly because i had the session lock')
+        self.load_cluster_from_dcs()
+
+    def follow_the_leader(self, demote_reason, follow_reason, refresh=True):
+        refresh and self.load_cluster_from_dcs()
+        ret = demote_reason if self.state_handler.is_leader() else follow_reason
+        leader = self.cluster.leader
+        leader = None if (leader and leader.name) == self.state_handler.name else leader
+        if not self.state_handler.check_recovery_conf(leader):
+            self._async_executor.schedule('changing primary_conninfo and restarting')
+            self._async_executor.run_async(self.state_handler.follow_the_leader, (leader, ))
+        return ret
+
+    def enforce_master_role(self, message, promote_message):
+        if self.state_handler.is_leader() or self.state_handler.role == 'master':
+            return message
+        else:
+            self.state_handler.promote()
+            return promote_message
+
+    @staticmethod
+    def fetch_node_status(member):
+        """This function perform http get request on member.api_url and fetches its status
+        :returns: tuple(`member`, reachable, in_recovery, xlog_location)
+
+        reachable - `!False` if the node is not reachable or is not responding with correct JSON
+        in_recovery - `!True` if pg_is_in_recovery() == true
+        xlog_location - value of `replayed_location` or `location` from JSON, dependin on its role."""
+
+        try:
+            response = requests.get(member.api_url, timeout=2, verify=False)
+            logger.info('Got response from %s %s: %s', member.name, member.api_url, response.content)
+            json = response.json()
+            is_master = json['role'] == 'master'
+            xlog_location = json['xlog']['location' if is_master else 'replayed_location']
+            return (member, True, not is_master, xlog_location)
+        except:
+            logging.exception('request failed: GET %s', member.api_url)
+        return (member, False, None, 0)
+
+    def fetch_nodes_statuses(self, members):
+        pool = ThreadPool(len(members))
+        results = pool.map(self.fetch_node_status, members)  # Run API calls on members in parallel
+        pool.close()
+        pool.join()
+        return results
+
+    def _is_healthiest_node(self, members, check_replication_lag=True):
+        """This method tries to determine whether I am healthy enough to became a new leader candidate or not."""
+
+        if self.state_handler.is_leader():
+            return True
+
+        if check_replication_lag and not self.state_handler.check_replication_lag(self.cluster.last_leader_operation):
+            return False  # Too far behind last reported xlog location on master
+
+        # Prepare list of nodes to run check against
+        members = [m for m in members if m.name != self.state_handler.name and m.api_url]
+
+        if members:
+            my_xlog_location = self.state_handler.xlog_position()
+            for member, reachable, in_recovery, xlog_location in self.fetch_nodes_statuses(members):
+                if reachable:  # If the node is unreachable it's not healhy
+                    if not in_recovery:
+                        logger.warning('Master (%s) is still alive', member.name)
+                        return False
+                    if my_xlog_location < xlog_location:
+                        return False
+        return True
+
+    def is_failover_possible(self, members):
+        ret = False
+        members = [m for m in members if m.name != self.state_handler.name and m.api_url]
+        if members:
+            for member, reachable, in_recovery, xlog_location in self.fetch_nodes_statuses(members):
+                if reachable:
+                    ret = True  # TODO: check xlog_location
+                else:
+                    logger.info('Member %s is not reachable', member.name)
+        else:
+            logger.warning('manual failover: members list is empty')
+        return ret
+
+    def manual_failover_process_no_leader(self):
+        failover = self.cluster.failover
+        if failover.member:  # manual failover to specific member
+            if failover.member == self.state_handler.name:  # manual failover to me
+                return True
+
+            # find specific node and check that it is healthy
+            members = [m for m in self.cluster.members if m.name == failover.member]
+            if members:
+                member, reachable, in_recovery, xlog_location = self.fetch_node_status(members[0])
+                if reachable:  # node is healthy
+                    logger.info('manual failover: to %s, i am %s', member.name, self.state_handler.name)
+                    return False
+                # we wanted to failover to specific member but it is not healthy
+                logger.warning('manual failover: member %s is unhealthy', member.name)
+
+            # at this point we should consider all members as a candidates for failover
+            # i.e. we assume that failover.member is None
+
+        # try to pick some other members to failover and check that they are healthy
+        if failover.leader:
+            if self.state_handler.name == failover.leader:  # I was the leader
+                # exclude me and desired member which is unhealthy (failover.member can be None)
+                members = [m for m in self.cluster.members if m.name != failover.member]
+                if self.is_failover_possible(members):  # check that there are healthy members
+                    return False
+                else:  # I was the leader and it looks like currently I am the only healthy member
+                    return True
+
+            # at this point we assume that our node is a candidate for a failover among all nodes except former leader
+
+        # exclude former leader from the list (failover.leader can be None)
+        members = [m for m in self.cluster.members if m.name != failover.leader]
+        return self._is_healthiest_node(members, check_replication_lag=False)
+
+    def is_healthiest_node(self):
+        if self.cluster.failover:
+            return self.manual_failover_process_no_leader()
+
+        # run usual health check
+        members = {m.name: m for m in self.cluster.members + self.old_cluster.members}
+        return self._is_healthiest_node(members.values())
+
+    def demote(self, delete_leader=True):
+        if delete_leader:
+            self.state_handler.stop()
+            self.dcs.delete_leader()
+            self.dcs.reset_cluster()
+        self.state_handler.follow_the_leader(None)
+
+    def process_manual_failover_from_leader(self):
+        failover = self.cluster.failover
+        if not failover.leader or failover.leader == self.state_handler.name:
+            if not failover.member or failover.member != self.state_handler.name:
+                members = [m for m in self.cluster.members if not failover.member or m.name == failover.member]
+                if self.is_failover_possible(members):  # check that there are healthy members
+                    self._async_executor.schedule('manual failover: demote')
+                    self._async_executor.run_async(self.demote)
+                    return 'manual failover: demoting myself'
+                else:
+                    logger.warning('manual failover: no healthy members found, failover is not possible')
+            else:
+                logger.warning('manual failover: I am already the leader, no need to failover')
+        else:
+            logger.warning('manual failover: leader name does not match: %s != %s',
+                           self.cluster.failover.leader, self.state_handler.name)
+
+        logger.info('Trying to clean up failover key')
+        self.dcs.manual_failover('', '', self.cluster.failover.index)
+
+    def process_unhealthy_cluster(self):
+        if self.is_healthiest_node():
+            if self.acquire_lock():
+                if self.cluster.failover:
+                    logger.info('Cleanning up failover key after acquiring leader lock...')
+                    self.dcs.manual_failover('', '')
+                return self.enforce_master_role('acquired session lock as a leader',
+                                                'promoted self to leader by acquiring session lock')
+            else:
+                return self.follow_the_leader('demoted self due after trying and failing to obtain lock',
+                                              'following new leader after trying and failing to obtain lock')
+        else:
+            return self.follow_the_leader('demoting self because i am not the healthiest node',
+                                          'following a different leader because i am not the healthiest node')
+
+    def process_healthy_cluster(self):
+        if self.has_lock():
+            if self.cluster.failover:
+                msg = self.process_manual_failover_from_leader()
+                if msg is not None:
+                    return msg
+
+            if self.update_lock():
+                return self.enforce_master_role('no action.  i am the leader with the lock',
+                                                'promoted self to leader because i had the session lock')
+            else:
+                # Either there is no connection to DCS or someone else acquired the lock
+                logger.error('failed to update leader lock')
+                self.load_cluster_from_dcs()
+        else:
+            logger.info('does not have lock')
+        return self.follow_the_leader('demoting self because i do not have the lock and i was a leader',
+                                      'no action.  i am a secondary and i am following a leader', False)
+
+    def schedule(self, action):
+        with self._async_executor:
+            return self._async_executor.schedule(action)
+
+    def restart_scheduled(self):
+        return self._async_executor.scheduled_action == 'restart'
+
+    def schedule_reinitialize(self):
+        return self.schedule('reinitialize')
+
+    def reinitialize_scheduled(self):
+        return self._async_executor.scheduled_action == 'reinitialize'
+
+    def restart(self):
+        with self._async_executor:
+            prev = self._async_executor.schedule('restart', True)
+            if prev is not None:
+                return (False, prev + ' already in progress')
+        if self._async_executor.run(self.state_handler.restart):
+            return (True, 'restarted successfully')
+        else:
+            return (False, 'restart failed')
+
+    def reinitialize(self, cluster):
+        self.state_handler.stop('immediate')
+        self.state_handler.remove_data_directory()
+        self.copy_backup_from_leader(cluster.leader)
+
+    def process_scheduled_action(self):
+        if self.reinitialize_scheduled():
+            if self.cluster.is_unlocked():
+                logger.error('Cluster has no leader, can not reinitialize')
+                self._async_executor.reset_scheduled_action()
+            elif self.has_lock():
+                logger.error('I am the leader, can not reinitialize')
+                self._async_executor.reset_scheduled_action()
+            else:
+                self._async_executor.run_async(self.reinitialize, args=(self.cluster, ))
+                return 'reinitialize started'
+
+    def handle_long_action_in_progress(self):
+        if self.has_lock():
+            if self.update_lock():
+                return 'updated leader lock during ' + self._async_executor.scheduled_action
+            else:
+                return 'failed to update leader lock during ' + self._async_executor.scheduled_action
+        elif self.cluster.is_unlocked():
+            return 'not healthy enough for leader race'
+        else:
+            return self._async_executor.scheduled_action + ' in progress'
+
+    def _run_cycle(self):
         try:
             self.load_cluster_from_dcs()
-            if not self.state_handler.is_healthy():
-                has_lock = self.has_lock()
-                self.state_handler.write_recovery_conf(None if has_lock else self.cluster.leader)
-                self.state_handler.start()
-                if not has_lock:
-                    return 'started as a secondary'
-                logger.info('started as readonly because i had the session lock')
-                self.load_cluster_from_dcs()
 
-            if self.cluster.is_unlocked():
+            self.touch_member()
-                if self.state_handler.is_healthiest_node(self.old_cluster):
+
-                    if self.acquire_lock():
+            # cluster has leader key but not initialize key
-                        if self.state_handler.is_leader() or self.state_handler.is_promoted:
+            if not self.cluster.is_unlocked() and not self.cluster.initialize:
-                            return 'acquired session lock as a leader'
+                self.dcs.initialize()  # fix it
-                        else:
+
-                            self.state_handler.promote()
+            if self._async_executor.busy:
-                            return 'promoted self to leader by acquiring session lock'
+                return self.handle_long_action_in_progress()
-                    else:
+
-                        self.load_cluster_from_dcs()
+            # currently it can trigger only reinitialize
-                        if self.state_handler.is_leader():
+            msg = self.process_scheduled_action()
-                            self.demote()
+            if msg is not None:
-                            return 'demoted self due after trying and failing to obtain lock'
+                return msg
-                        else:
+
-                            self.follow_the_leader()
+            # is data directory empty?
-                            return 'following new leader after trying and failing to obtain lock'
+            if self.state_handler.data_directory_empty():
+                return self.bootstrap()  # new node
+            # "bootstrap", but data directory is not empty
+            elif not self.cluster.initialize and self.cluster.is_unlocked():
+                self.dcs.initialize()
+
+            # try to start dead postgres
+            if not self.state_handler.is_healthy():
+                msg = self.recover()
+                if msg is not None:
+                    return msg
+
+            try:
+                if self.cluster.is_unlocked():
+                    return self.process_unhealthy_cluster()
                 else:
-                    self.load_cluster_from_dcs()
+                    return self.process_healthy_cluster()
-                    if self.state_handler.is_leader():
+            finally:
-                        self.demote()
+                self.state_handler.sync_replication_slots(self.cluster)
-                        return 'demoting self because i am not the healthiest node'
-                    else:
-                        self.follow_the_leader()
-                        return 'following a different leader because i am not the healthiest node'
-            else:
-                if self.has_lock() and self.update_lock():
-                    if self.state_handler.is_leader() or self.state_handler.is_promoted:
-                        return 'no action.  i am the leader with the lock'
-                    else:
-                        self.state_handler.promote()
-                        return 'promoted self to leader because i had the session lock'
-                else:
-                    logger.info('does not have lock')
-                    if self.state_handler.is_leader():
-                        self.demote()
-                        return 'demoting self because i do not have the lock and i was a leader'
-                    else:
-                        self.follow_the_leader()
-                        return 'no action.  i am a secondary and i am following a leader'
         except DCSError:
             logger.error('Error communicating with DCS')
-            if self.state_handler.is_leader():
+            if self.state_handler.is_running() and self.state_handler.is_leader():
-                self.state_handler.demote(None)
+                self.demote(delete_leader=False)
                 return 'demoted self because DCS is not accessible and i was a leader'
-        except (InterfaceError, OperationalError):
+        except (psycopg2.Error, PostgresConnectionException):
-            logger.error('Error communicating with Postgresql.  Will try again')
+            logger.exception('Error communicating with Postgresql. Will try again later')
+
+    def run_cycle(self):
+        with self._async_executor:
+            return self._run_cycle()

patroni/postgresql.py

@@ -6,9 +6,10 @@ import shutil
 import subprocess
 import time
 
-from patroni.exceptions import PostgresException
+from patroni.exceptions import PostgresConnectionException, PostgresException
-from patroni.utils import sleep
+from patroni.utils import Retry, RetryFailedError
 from six.moves.urllib_parse import urlparse
+from threading import Lock
 
 logger = logging.getLogger(__name__)
 
@@ -47,15 +48,16 @@ class Postgresql:
         self.replication = config['replication']
         self.superuser = config['superuser']
         self.admin = config['admin']
+        self.pg_rewind = config.get('pg_rewind', {})
         self.callback = config.get('callbacks', {})
         self.use_slots = config.get('use_slots', True)
+        self.schedule_load_slots = self.use_slots
         self.recovery_conf = os.path.join(self.data_dir, 'recovery.conf')
         self.configuration_to_save = (os.path.join(self.data_dir, 'pg_hba.conf'),
                                       os.path.join(self.data_dir, 'postgresql.conf'))
         self.postmaster_pid = os.path.join(self.data_dir, 'postmaster.pid')
         self.trigger_file = config.get('recovery_conf', {}).get('trigger_file', None) or 'promote'
         self.trigger_file = os.path.abspath(os.path.join(self.data_dir, self.trigger_file))
-        self.is_promoted = False
 
         self._pg_ctl = ['pg_ctl', '-w', '-D', self.data_dir]
 
@@ -66,7 +68,45 @@ class Postgresql:
 
         self._connection = None
         self._cursor_holder = None
-        self.members = []  # list of already existing replication slots
+        self._need_rewind = False
+        self.replication_slots = []  # list of already existing replication slots
+        self.retry = Retry(max_tries=-1, deadline=5, max_delay=1, retry_exceptions=PostgresConnectionException)
+
+        self._state = 'stopped'
+        self._state_lock = Lock()
+        self._role = 'replica'
+        self._role_lock = Lock()
+
+        if self.is_running():
+            self._state = 'running'
+            self._role = 'master' if self.is_leader() else 'replica'
+
+    @property
+    def can_rewind(self):
+        """ check if pg_rewind executable is there and that pg_controldata indicates
+            we have either wal_log_hints or checksums turned on
+        """
+        # low-hanging fruit: check if pg_rewind configuration is there
+        if not self.pg_rewind or\
+           not (self.pg_rewind.get('username', '') and self.pg_rewind.get('password', '')):
+            return False
+
+        cmd = ['pg_rewind', '--help']
+        try:
+            ret = subprocess.call(cmd, stdout=open(os.devnull, 'w'), stderr=subprocess.STDOUT)
+            if ret != 0:  # pg_rewind is not there, close up the shop and go home
+                return False
+        except OSError:
+            return False
+        # check if the cluster's configuration permits pg_rewind
+        data = self.controldata()
+        if data:
+            return data.get('wal_log_hints setting', 'off') == 'on' or\
+                data.get('Data page checksum version', '0') != '0'
+        return False
+
+    def require_rewind(self):
+        self._need_rewind = True
 
     def get_local_address(self):
         listen_addresses = self.listen_addresses.split(',')
@@ -86,56 +126,57 @@ class Postgresql:
         return self._connection
 
     def _cursor(self):
-        if not self._cursor_holder or self._cursor_holder.closed:
+        if not self._cursor_holder or self._cursor_holder.closed or self._cursor_holder.connection.closed != 0:
             self._cursor_holder = self.connection().cursor()
         return self._cursor_holder
 
-    def disconnect(self):
+    def _query(self, sql, *params):
-        self._connection and self._connection.close()
+        cursor = None
-        self._connection = self._cursor_holder = None
+        try:
+            cursor = self._cursor()
+            cursor.execute(sql, params)
+            return cursor
+        except psycopg2.Error as e:
+            if cursor and cursor.connection.closed == 0:
+                raise e
+            if self.state == 'restarting':
+                raise RetryFailedError('cluster is being restarted')
+            raise PostgresConnectionException('connection problems')
 
     def query(self, sql, *params):
-        max_attempts = 0
+        try:
-        while True:
+            return self.retry(self._query, sql, *params)
-            ex = None
+        except RetryFailedError as e:
-            try:
+            raise PostgresConnectionException(str(e))
-                cursor = self._cursor()
-                cursor.execute(sql, params)
-                return cursor
-            except psycopg2.InterfaceError as e:
-                ex = e
-            except psycopg2.OperationalError as e:
-                if self._connection and self._connection.closed == 0:
-                    raise e
-                ex = e
-            if ex:
-                self.disconnect()
-                max_attempts += 1
-                if max_attempts >= 3:
-                    raise ex
-                sleep(5)
 
     def data_directory_empty(self):
         return not os.path.exists(self.data_dir) or os.listdir(self.data_dir) == []
 
     def initialize(self):
+        self.set_state('initalizing new cluster')
         ret = subprocess.call(self._pg_ctl + ['initdb', '-o', '--encoding=UTF8']) == 0
-        ret and self.write_pg_hba()
+        if ret:
+            self.write_pg_hba()
+        else:
+            self.set_state('initdb failed')
         return ret
 
     def delete_trigger_file(self):
         os.path.exists(self.trigger_file) and os.unlink(self.trigger_file)
 
-    def sync_from_leader(self, leader):
+    def write_pgpass(self, record):
-        r = parseurl(leader.conn_url)
-
         pgpass = 'pgpass'
         with open(pgpass, 'w') as f:
             os.fchmod(f.fileno(), 0o600)
-            f.write('{host}:{port}:*:{user}:{password}\n'.format(**r))
+            f.write('{host}:{port}:*:{user}:{password}\n'.format(**record))
-
         env = os.environ.copy()
         env['PGPASSFILE'] = pgpass
+        return env
+
+    def sync_from_leader(self, leader):
+        r = parseurl(leader.conn_url)
+
+        env = self.write_pgpass(r)
         return self.create_replica(r, env) == 0
 
     @staticmethod
@@ -143,6 +184,7 @@ class Postgresql:
         return "host={host} port={port} user={user}".format(**conn)
 
     def create_replica(self, master_connection, env):
+        self.set_state('building replica from {host}:{port}'.format(**master_connection))
         connstring = self.build_connstring(master_connection)
         cmd = self.config['restore']
         try:
@@ -150,82 +192,116 @@ class Postgresql:
             self.delete_trigger_file()
         except:
             logger.exception('Error when creating replica')
-            return 1
+            ret = 1
+        if ret != 0:
+            self.set_state('failed to build replica from {host}:{port}'.format(**master_connection))
         return ret
 
-    def is_leader(self, check_only=False):
+    def is_leader(self):
-        ret = not self.query('SELECT pg_is_in_recovery()').fetchone()[0]
+        return not self.query('SELECT pg_is_in_recovery()').fetchone()[0]
-        if ret and self.is_promoted and not check_only:
-            self.delete_trigger_file()
-            self.is_promoted = False
-        return ret
 
     def is_running(self):
         return subprocess.call(' '.join(self._pg_ctl) + ' status > /dev/null 2>&1', shell=True) == 0
 
-    def call_nowait(self, cb_name, is_leader=None):
+    def call_nowait(self, cb_name):
         """ pick a callback command and call it without waiting for it to finish """
         if not self.callback or cb_name not in self.callback:
             return False
         cmd = self.callback[cb_name]
-        if is_leader is None:
-            try:
-                is_leader = self.is_leader(check_only=True)
-            except psycopg2.OperationalError as e:
-                logger.warning("unable to perform {0} action, cannot obtain the cluster role: {1}".format(cb_name, e))
-                return False
         try:
-            role = "master" if is_leader else "replica"
+            subprocess.Popen(shlex.split(cmd) + [cb_name, self.role, self.scope])
-            subprocess.Popen(shlex.split(cmd) + [cb_name, role, self.scope])
         except:
-            logger.exception('callback %s %s %s %s failed', cmd, cb_name, role, self.scope)
+            logger.exception('callback %s %s %s %s failed', cmd, cb_name, self.role, self.scope)
             return False
         return True
 
-    def start(self):
+    @property
-        if self.is_running():
+    def role(self):
-            self.load_replication_slots()
+        with self._role_lock:
-            logger.error('Cannot start PostgreSQL because one is already running.')
+            return self._role
-            return False
 
+    def set_role(self, value):
+        with self._role_lock:
+            self._role = value
+
+    @property
+    def state(self):
+        with self._state_lock:
+            return self._state
+
+    def set_state(self, value):
+        with self._state_lock:
+            self._state = value
+
+    def start(self, block_callbacks=False):
+        if self.is_running():
+            logger.error('Cannot start PostgreSQL because one is already running.')
+            return True
+
+        self.set_role('replica' if os.path.exists(self.recovery_conf) else 'master')
         if os.path.exists(self.postmaster_pid):
             os.remove(self.postmaster_pid)
             logger.info('Removed %s', self.postmaster_pid)
 
+        if not block_callbacks:
+            self.set_state('starting')
+
         ret = subprocess.call(self._pg_ctl + ['start', '-o', self.server_options()]) == 0
-        ret and self.load_replication_slots()
+
+        self.set_state('running' if ret else 'start failed')
+
+        self.schedule_load_slots = ret and self.use_slots
         self.save_configuration_files()
-        if ret and ACTION_ON_START in self.callback:
+        # block_callbacks is used during restart to avoid
-            self.call_nowait(ACTION_ON_START)
+        # running start/stop callbacks in addition to restart ones
+        ret and not block_callbacks and self.call_nowait(ACTION_ON_START)
         return ret
 
-    def stop(self):
+    def checkpoint(self):
         try:
-            is_leader = self.is_leader(check_only=True)
+            r = parseurl('postgres://{}/postgres'.format(self.local_address))
+            r['options'] = '-c statement_timeout=0'
+            with psycopg2.connect(**r) as conn:
+                conn.autocommit = True
+                with conn.cursor() as cur:
+                    cur.execute('CHECKPOINT')
         except:
-            is_leader = None
+            logging.exception('Exception during CHECKPOINT')
-            pass
+
-        ret = subprocess.call(self._pg_ctl + ['stop', '-m', 'fast'])
+    def stop(self, mode='fast', block_callbacks=False):
-        if ret == 0 and ACTION_ON_STOP in self.callback:
+        if not self.is_running():
-            self.call_nowait(ACTION_ON_STOP, is_leader=is_leader)
+            if not block_callbacks:
-        return ret == 0
+                self.set_state('stopped')
+            return True
+
+        if block_callbacks:
+            self.checkpoint()
+        else:
+            self.set_state('stopping')
+
+        ret = subprocess.call(self._pg_ctl + ['stop', '-m', mode]) == 0
+        # block_callbacks is used during restart to avoid
+        # running start/stop callbacks in addition to restart ones
+        if not ret:
+            self.set_state('stop failed')
+        elif not block_callbacks:
+            self.set_state('stopped')
+            self.call_nowait(ACTION_ON_STOP)
+        return ret
 
     def reload(self):
-        ret = subprocess.call(self._pg_ctl + ['reload'])
+        ret = subprocess.call(self._pg_ctl + ['reload']) == 0
-        if ret == 0 and ACTION_ON_RELOAD in self.callback:
+        ret and self.call_nowait(ACTION_ON_RELOAD)
-            self.call_nowait(ACTION_ON_RELOAD)
+        return ret
-        return ret == 0
 
     def restart(self):
-        try:
+        self.set_state('restarting')
-            is_leader = self.is_leader(check_only=True)
+        ret = self.stop(block_callbacks=True) and self.start(block_callbacks=True)
-        except:
+        if ret:
-            is_leader = None
+            self.call_nowait(ACTION_ON_RESTART)
-            pass
+        else:
-        ret = subprocess.call(self._pg_ctl + ['restart', '-m', 'fast'])
+            self.set_state('restart failed ({})'.format(self.state))
-        if ret == 0 and ACTION_ON_RESTART in self.callback:
+        return ret
-            self.call_nowait(ACTION_ON_RESTART, is_leader=is_leader)
-        return ret == 0
 
     def server_options(self):
         options = "--listen_addresses='{}' --port={}".format(self.listen_addresses, self.port)
@@ -239,36 +315,8 @@ class Postgresql:
             return False
         return True
 
-    def is_healthiest_node(self, cluster):
+    def check_replication_lag(self, last_leader_operation):
-        if self.is_leader():
+        return last_leader_operation - self.xlog_position() <= self.config.get('maximum_lag_on_failover', 0)
-            return True
-
-        if cluster.last_leader_operation - self.xlog_position() > self.config.get('maximum_lag_on_failover', 0):
-            return False
-
-        for member in cluster.members:
-            if member.name == self.name:
-                continue
-            try:
-                r = parseurl(member.conn_url)
-                member_conn = psycopg2.connect(**r)
-                member_conn.autocommit = True
-                member_cursor = member_conn.cursor()
-                member_cursor.execute(
-                    "SELECT pg_is_in_recovery(), %s - pg_xlog_location_diff(pg_last_xlog_replay_location(), '0/0')",
-                    (self.xlog_position(),))
-                row = member_cursor.fetchone()
-                member_cursor.close()
-                member_conn.close()
-                logger.error([self.name, member.name, row])
-                if not row[0]:
-                    logger.warning('Master (%s) is still alive', member.name)
-                    return False
-                if row[1] < 0:
-                    return False
-            except psycopg2.Error:
-                continue
-        return True
 
     def write_pg_hba(self):
         with open(os.path.join(self.data_dir, 'pg_hba.conf'), 'a') as f:
@@ -292,10 +340,7 @@ class Postgresql:
         with open(self.recovery_conf, 'r') as f:
             for line in f:
                 if line.startswith('primary_conninfo'):
-                    if not pattern:
+                    return pattern and (pattern in line)
-                        return False
-                    return pattern in line
-
         return not pattern
 
     def write_recovery_conf(self, leader):
@@ -310,12 +355,120 @@ recovery_target_timeline = 'latest'
                 for name, value in self.config.get('recovery_conf', {}).items():
                     f.write("{} = '{}'\n".format(name, value))
 
-    def follow_the_leader(self, leader):
+    def rewind(self, leader):
-        if not self.check_recovery_conf(leader):
+        # prepare pg_rewind connection
+        r = parseurl(leader.conn_url)
+        r.update(self.pg_rewind)
+        r['user'] = r['username']
+        env = self.write_pgpass(r)
+        pc = "user={user} host={host} port={port} dbname=postgres sslmode=prefer sslcompression=1".format(**r)
+        logger.info("running pg_rewind from {}".format(pc))
+        pg_rewind = ['pg_rewind', '-D', self.data_dir, '--source-server', pc]
+        try:
+            ret = (subprocess.call(pg_rewind, env=env) == 0)
+        except:
+            ret = False
+        if ret:
             self.write_recovery_conf(leader)
-            self.restart()
+        return ret
-            if ACTION_ON_ROLE_CHANGE in self.callback:
+
-                self.call_nowait(ACTION_ON_ROLE_CHANGE)
+    def controldata(self):
+        """ return the contents of pg_controldata, or non-True value if pg_controldata call failed """
+        result = {}
+        try:
+            data = subprocess.check_output(['pg_controldata', self.data_dir])
+            if data:
+                data = data.splitlines()
+                result = {l.split(':')[0]: l.split(':')[1].strip() for l in data if l}
+        except subprocess.CalledProcessError:
+            logger.exception("Error when calling pg_controldata")
+        finally:
+            return result
+
+    def read_postmaster_opts(self):
+        """ returns the list of option names/values from postgres.opts, Empty dict if read failed or no file """
+        result = {}
+        try:
+            with open(os.path.join(self.data_dir, "postmaster.opts")) as f:
+                data = f.read()
+                opts = [opt.strip('"\n') for opt in data.split(' "')]
+                for opt in opts:
+                    if '=' in opt and opt.startswith('--'):
+                        name, val = opt.split('=', 1)
+                        name = name.strip('-')
+                        result[name] = val
+        except IOError:
+            logger.exception('Error when reading postmaster.opts')
+        finally:
+            return result
+
+    def single_user_mode(self, command=None, options={}):
+        """ run a given command in a single-user mode. If the command is empty - then just start and stop """
+        cmd = ['postgres', '--single', '-D', self.data_dir]
+        for opt in sorted(options):
+            cmd.extend(['-c', '{0}={1}'.format(opt, options[opt])])
+        # need a database name to connect
+        cmd.append('postgres')
+        p = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=open(os.devnull, 'w'), stderr=subprocess.STDOUT)
+        if p:
+            command and p.communicate('{}\n'.format(command))
+            p.stdin.close()
+            return p.wait()
+        return 1
+
+    def cleanup_archive_status(self):
+        status_dir = os.path.join(self.data_dir, 'pg_xlog', 'archive_status')
+        if os.path.isdir(status_dir):
+            for f in os.listdir(status_dir):
+                path = os.path.join(status_dir, f)
+                try:
+                    if os.path.islink(path):
+                        os.unlink(path)
+                    elif os.path.isfile(path):
+                        os.remove(path)
+                except:
+                    logger.exception("Unable to remove {}".format(path))
+
+    def follow_the_leader(self, leader, recovery=False):
+        if not self.check_recovery_conf(leader) or recovery:
+            change_role = (self.role == 'master')
+
+            self._need_rewind = (self._need_rewind or change_role) and self.can_rewind
+            if self._need_rewind:
+                logger.info("set the rewind flag after demote")
+            self.write_recovery_conf(leader)
+            if not leader or not self._need_rewind:  # do not rewind until the leader becomes available
+                ret = self.restart()
+            else:  # we have a leader and need to rewind
+                if self.is_running():
+                    self.stop()
+                # at present, pg_rewind only runs when the cluster is shut down cleanly
+                # and not shutdown in recovery. We have to remove the recovery.conf if present
+                # and start/shutdown in a single user mode to emulate this.
+                # XXX: if recovery.conf is linked, it will be written anew as a normal file.
+                if os.path.islink(self.recovery_conf):
+                    os.unlink(self.recovery_conf)
+                else:
+                    os.remove(self.recovery_conf)
+                # Archived segments might be useful to pg_rewind,
+                # clean the flags that tell we should remove them.
+                self.cleanup_archive_status()
+                # Start in a single user mode and stop to produce a clean shutdown
+                opts = self.read_postmaster_opts()
+                opts['archive_mode'] = 'on'
+                opts['archive_command'] = 'false'
+                self.single_user_mode(options=opts)
+                if self.rewind(leader):
+                    ret = self.start()
+                else:
+                    logger.error("unable to rewind the former master")
+                    self.remove_data_directory()
+                    ret = True
+                self._need_rewind = False
+            change_role and ret and self.call_nowait(ACTION_ON_ROLE_CHANGE)
+            return ret
+        else:
+            return True
 
     def save_configuration_files(self):
         """
@@ -334,13 +487,18 @@ recovery_target_timeline = 'latest'
             logger.exception('unable to restore configuration from WAL-E backup')
 
     def promote(self):
-        self.is_promoted = subprocess.call(self._pg_ctl + ['promote']) == 0
+        if self.role == 'master':
-        if self.is_promoted and ACTION_ON_ROLE_CHANGE in self.callback:
+            return True
+        ret = subprocess.call(self._pg_ctl + ['promote']) == 0
+        if ret:
+            self.set_role('master')
+            logger.info("cleared rewind flag after becoming the leader")
+            self._need_rewind = False
             self.call_nowait(ACTION_ON_ROLE_CHANGE)
-        return self.is_promoted
+        return ret
 
-    def demote(self, leader):
+    def demote(self):
-        self.follow_the_leader(leader)
+        self.follow_the_leader(None)
 
     def create_replication_user(self):
         self.query('CREATE USER "{}" WITH REPLICATION ENCRYPTED PASSWORD %s'.format(
@@ -362,34 +520,34 @@ recovery_target_timeline = 'latest'
         return self.query("""SELECT pg_xlog_location_diff(CASE WHEN pg_is_in_recovery()
                                                                THEN pg_last_xlog_replay_location()
                                                                ELSE pg_current_xlog_location()
-                                                          END, '0/0')""").fetchone()[0]
+                                                          END, '0/0')::bigint""").fetchone()[0]
 
     def load_replication_slots(self):
-        if self.use_slots:
+        if self.use_slots and self.schedule_load_slots:
             cursor = self.query("SELECT slot_name FROM pg_replication_slots WHERE slot_type='physical'")
-            self.members = [r[0] for r in cursor]
+            self.replication_slots = [r[0] for r in cursor]
+            self.schedule_load_slots = False
 
-    def sync_replication_slots(self, members):
+    def sync_replication_slots(self, cluster):
         if self.use_slots:
-            # drop unused slots
+            try:
-            for slot in set(self.members) - set(members):
+                self.load_replication_slots()
-                self.query("""SELECT pg_drop_replication_slot(%s)
+                slots = [m.name for m in cluster.members if m.name != self.name] if self.role == 'master' else []
-                               WHERE EXISTS(SELECT 1 FROM pg_replication_slots
+                # drop unused slots
-                               WHERE slot_name = %s)""", slot, slot)
+                for slot in set(self.replication_slots) - set(slots):
+                    self.query("""SELECT pg_drop_replication_slot(%s)
+                                   WHERE EXISTS(SELECT 1 FROM pg_replication_slots
+                                   WHERE slot_name = %s)""", slot, slot)
 
-            # create new slots
+                # create new slots
-            for slot in set(members) - set(self.members):
+                for slot in set(slots) - set(self.replication_slots):
-                self.query("""SELECT pg_create_physical_replication_slot(%s)
+                    self.query("""SELECT pg_create_physical_replication_slot(%s)
-                               WHERE NOT EXISTS (SELECT 1 FROM pg_replication_slots
+                                   WHERE NOT EXISTS (SELECT 1 FROM pg_replication_slots
-                               WHERE slot_name = %s)""", slot, slot)
+                                   WHERE slot_name = %s)""", slot, slot)
 
-        self.members = members
+                self.replication_slots = slots
-
+            except:
-    def create_replication_slots(self, cluster):
+                logger.exception('Exception when changing replication slots')
-        self.sync_replication_slots([m.name for m in cluster.members if m.name != self.name])
-
-    def drop_replication_slots(self):
-        self.sync_replication_slots([])
 
     def last_operation(self):
         return str(self.xlog_position())
@@ -421,6 +579,23 @@ recovery_target_timeline = 'latest'
     def move_data_directory(self):
         if os.path.isdir(self.data_dir) and not self.is_running():
             try:
-                os.rename(self.data_dir, '{0}_{1}'.format(self.data_dir, time.strftime('%Y-%m-%d-%H-%M-%S')))
+                new_name = '{0}_{1}'.format(self.data_dir, time.strftime('%Y-%m-%d-%H-%M-%S'))
+                logger.info('renaming data directory to %s', new_name)
+                os.rename(self.data_dir, new_name)
             except:
-                logger.exception("Could not rename data directory {0}".format(self.data_dir))
+                logger.exception("Could not rename data directory %s", self.data_dir)
+
+    def remove_data_directory(self):
+        logger.info('Removing data directory: %s', self.data_dir)
+        try:
+            if os.path.islink(self.data_dir):
+                os.unlink(self.data_dir)
+            elif not os.path.exists(self.data_dir):
+                return
+            elif os.path.isfile(self.data_dir):
+                os.remove(self.data_dir)
+            elif os.path.isdir(self.data_dir):
+                shutil.rmtree(self.data_dir)
+        except:
+            logger.exception('Could not remove data directory %s', self.data_dir)
+            self.move_data_directory()

patroni/scripts/aws.py

@@ -1,4 +1,4 @@
-#!/usr/bin/python
+#!/usr/bin/env python
 
 import logging
 import requests

patroni/scripts/restore.py

@@ -1,4 +1,4 @@
-#!/usr/bin/python
+#!/usr/bin/env python
 # arguments are:
 #   - cluster scope
 #   - cluster role

patroni/utils.py

@@ -6,8 +6,9 @@ import signal
 import sys
 import time
 
-from patroni.exceptions import DCSError
+from patroni.exceptions import PatroniException
 
+ignore_sigterm = False
 interrupted_sleep = False
 reap_children = False
 
@@ -35,6 +36,8 @@ def calculate_ttl(expiration):
     """
     >>> calculate_ttl(None)
     >>> calculate_ttl('2015-06-10 12:56:30.552539016Z')
+    >>> calculate_ttl('2015-06-10T12:56:30.552539016Z') < 0
+    True
     """
     if not expiration:
         return None
@@ -46,7 +49,10 @@ def calculate_ttl(expiration):
 
 
 def sigterm_handler(signo, stack_frame):
-    sys.exit()
+    global ignore_sigterm
+    if not ignore_sigterm:
+        ignore_sigterm = True
+        sys.exit()
 
 
 def sigchld_handler(signo, stack_frame):
@@ -86,7 +92,7 @@ def reap_children():
             reap_children = False
 
 
-class RetryFailedError(DCSError):
+class RetryFailedError(PatroniException):
 
     """Raised when retrying an operation ultimately failed, after retrying the maximum number of attempts."""
 
@@ -96,7 +102,7 @@ class Retry:
     """Helper for retrying a method in the face of retry-able exceptions"""
 
     def __init__(self, max_tries=1, delay=0.1, backoff=2, max_jitter=0.8, max_delay=3600,
-                 sleep_func=time.sleep, deadline=None, retry_exceptions=DCSError):
+                 sleep_func=sleep, deadline=None, retry_exceptions=PatroniException):
         """Create a :class:`Retry` instance for retrying function calls
 
         :param max_tries: How many times to retry the command. -1 means infinite tries.
@@ -150,13 +156,10 @@ class Retry:
                 if self._attempts == self.max_tries:
                     raise RetryFailedError("Too many retry attempts")
                 self._attempts += 1
-                sleeptime = self._cur_delay + (
+                sleeptime = self._cur_delay + (random.randint(0, self.max_jitter) / 100.0)
-                    random.randint(0, self.max_jitter) / 100.0)
 
-                if self._cur_stoptime is not None and \
+                if self._cur_stoptime is not None and time.time() + sleeptime >= self._cur_stoptime:
-                   time.time() + sleeptime >= self._cur_stoptime:
                     raise RetryFailedError("Exceeded retry deadline")
                 else:
                     self.sleep_func(sleeptime)
-                self._cur_delay = min(self._cur_delay * self.backoff,
+                self._cur_delay = min(self._cur_delay * self.backoff, self.max_delay)
-                                      self.max_delay)

patroni/zookeeper.py

@@ -5,7 +5,8 @@ import time
 
 from kazoo.client import KazooClient, KazooState
 from kazoo.exceptions import NoNodeError, NodeExistsError
-from patroni.dcs import AbstractDCS, Cluster, DCSError, Leader, Member, parse_connection_string
+from patroni.dcs import AbstractDCS, Cluster, Failover, Leader, Member
+from patroni.exceptions import DCSError
 from patroni.utils import sleep
 from requests.exceptions import RequestException
 
@@ -90,9 +91,8 @@ class ZooKeeper(AbstractDCS):
                                       'max_tries': -1},
                                   connection_retry={'max_delay': 1, 'max_tries': -1})
         self.client.add_listener(self.session_listener)
-        self.cluster_event = self.client.handler.event_object()
 
-        self.cluster = None
+        self._my_member_data = None
         self.fetch_cluster = True
         self.last_leader_operation = 0
 
@@ -104,18 +104,18 @@ class ZooKeeper(AbstractDCS):
 
     def cluster_watcher(self, event):
         self.fetch_cluster = True
-        self.cluster_event.set()
+        self.event.set()
 
     def get_node(self, key, watch=None):
         try:
-            return self.client.get(key, watch)
+            ret = self.client.get(key, watch)
+            return (ret[0].decode('utf-8'), ret[1])
         except NoNodeError:
             return None
 
     @staticmethod
     def member(name, value, znode):
-        conn_url, api_url = parse_connection_string(value)
+        return Member.from_node(znode.version, name, znode.ephemeralOwner, value)
-        return Member(znode.mzxid, name, conn_url, api_url, None, None)
 
     def get_children(self, key, watch=None):
         try:
@@ -132,8 +132,11 @@ class ZooKeeper(AbstractDCS):
         return members
 
     def _inner_load_cluster(self):
-        self.cluster_event.clear()
+        self.fetch_cluster = False
-        nodes = set(self.get_children(self.client_path('')))
+        self.event.clear()
+        nodes = set(self.get_children(self.client_path(''), self.cluster_watcher))
+        if not nodes:
+            self.fetch_cluster = True
 
         # get initialize flag
         initialize = self._INITIALIZE in nodes
@@ -142,7 +145,7 @@ class ZooKeeper(AbstractDCS):
         members = self.load_members() if self._MEMBERS[:-1] in nodes else []
 
         # get leader
-        leader = self.get_node(self.leader_path, self.cluster_watcher) if self._LEADER in nodes else None
+        leader = self.get_node(self.leader_path) if self._LEADER in nodes else None
         if leader:
             client_id = self.client.client_id
             if leader[0] == self._name and client_id is not None and client_id[0] != leader[1].ephemeralOwner:
@@ -151,17 +154,22 @@ class ZooKeeper(AbstractDCS):
                 leader = None
 
             if leader:
-                member = Member(-1, leader[0], None, None, None, None)
+                member = Member(-1, leader[0], None, {})
                 member = ([m for m in members if m.name == leader[0]] or [member])[0]
-                leader = Leader(leader[1].mzxid, None, None, member)
+                leader = Leader(leader[1].version, leader[1].ephemeralOwner, member)
                 self.fetch_cluster = member.index == -1
 
-        # get last leader operation
+        # failover key
-        self.last_leader_operation = self.get_node(self.leader_optime_path) if self.fetch_cluster else None
+        failover = self.get_node(self.failover_path, watch=self.cluster_watcher) if self._FAILOVER in nodes else None
-        self.last_leader_operation = 0 if self.last_leader_operation is None else int(self.last_leader_operation[0])
+        if failover:
-        self.cluster = Cluster(initialize, leader, self.last_leader_operation, members)
+            failover = Failover.from_node(failover[1].version, failover[0])
 
-    def get_cluster(self):
+        # get last leader operation
+        optime = self.get_node(self.leader_optime_path) if self._OPTIME in nodes and self.fetch_cluster else None
+        self.last_leader_operation = 0 if optime is None else int(optime[0])
+        self._cluster = Cluster(initialize, leader, self.last_leader_operation, members, failover)
+
+    def _load_cluster(self):
         if self.exhibitor and self.exhibitor.poll():
             self.client.set_hosts(self.exhibitor.zookeeper_hosts)
 
@@ -172,11 +180,10 @@ class ZooKeeper(AbstractDCS):
                 logger.exception('get_cluster')
                 self.session_listener(KazooState.LOST)
                 raise ZooKeeperError('ZooKeeper in not responding properly')
-        return self.cluster
 
     def _create(self, path, value, **kwargs):
         try:
-            self.client.retry(self.client.create, path, value, **kwargs)
+            self.client.retry(self.client.create, path, value.encode('utf-8'), **kwargs)
             return True
         except:
             return False
@@ -186,30 +193,60 @@ class ZooKeeper(AbstractDCS):
         ret or logger.info('Could not take out TTL lock')
         return ret
 
+    def set_failover_value(self, value, index=None):
+        try:
+            self.client.retry(self.client.set, self.failover_path, value.encode('utf-8'), version=index or -1)
+            return True
+        except NoNodeError:
+            return value == '' or (not index and self._create(self.failover_path, value.encode('utf-8')))
+        except:
+            logging.exception('set_failover_value')
+            return False
+
     def initialize(self):
         return self._create(self.initialize_path, self._name, makepath=True)
 
-    def touch_member(self, connection_string, ttl=None):
+    def touch_member(self, data, ttl=None):
-        if self.cluster and any(m.name == self._name for m in self.cluster.members):
+        cluster = self.cluster
-            return True
+        me = cluster and ([m for m in cluster.members if m.name == self._name] or [None])[0]
         path = self.member_path
+        data = data.encode('utf-8')
+        create = not me
+        if me and self.client.client_id is not None and me.session != self.client.client_id[0]:
+            try:
+                self.client.retry(self.client.delete, path)
+            except NoNodeError:
+                pass
+            except:
+                return False
+            create = True
+
+        if not create and data == self._my_member_data:
+            return True
+
         try:
-            self.client.retry(self.client.create, path, connection_string, makepath=True, ephemeral=True)
+            if create:
+                self.client.retry(self.client.create, path, data, makepath=True, ephemeral=True)
+            else:
+                self.client.retry(self.client.set, path, data)
+            self._my_member_data = data
             return True
         except NodeExistsError:
             try:
-                self.client.retry(self.client.delete, path)
+                self.client.retry(self.client.set, path, data)
-                self.client.retry(self.client.create, path, connection_string, makepath=True, ephemeral=True)
+                self._my_member_data = data
                 return True
             except:
                 logger.exception('touch_member')
+        except:
+            logger.exception('touch_member')
         return False
 
     def take_leader(self):
         return self.attempt_to_acquire_leader()
 
-    def update_leader(self, state_handler):
+    def write_leader_optime(self, last_operation):
-        last_operation = state_handler.last_operation()
+        last_operation = last_operation.encode('utf-8')
         if last_operation != self.last_leader_operation:
             self.last_leader_operation = last_operation
             path = self.leader_optime_path
@@ -222,16 +259,19 @@ class ZooKeeper(AbstractDCS):
                     logger.exception('Failed to create %s', path)
             except:
                 logger.exception('Failed to update %s', path)
+
+    def update_leader(self):
         return True
 
     def delete_leader(self):
-        if isinstance(self.cluster, Cluster) and self.cluster.leader.name == self._name:
+        self.client.restart()
-            self.client.delete(self.leader_path, version=self.cluster.leader.index)
+        self._my_member_data = None
+        return True
 
     def _cancel_initialization(self):
         node = self.get_node(self.initialize_path)
         if node and node[0] == self._name:
-            self.client.delete(self.initialize_path, version=node[1].mzxid)
+            self.client.delete(self.initialize_path, version=node[1].version)
 
     def cancel_initialization(self):
         try:
@@ -240,6 +280,6 @@ class ZooKeeper(AbstractDCS):
             logger.exception("Unable to delete initialize key")
 
     def watch(self, timeout):
-        self.cluster_event.wait(timeout)
+        if super(ZooKeeper, self).watch(timeout):
-        if self.cluster_event.isSet():
             self.fetch_cluster = True
+        return self.fetch_cluster

postgres0.yml

@@ -4,6 +4,9 @@ scope: &scope batman
 restapi:
   listen: 127.0.0.1:8008
   connect_address: 127.0.0.1:8008
+  auth: 'username:password'
+#  certfile: /etc/ssl/certs/ssl-cert-snakeoil.pem
+#  keyfile: /etc/ssl/private/ssl-cert-snakeoil.key
 etcd:
   scope: *scope
   ttl: *ttl
@@ -31,6 +34,9 @@ postgresql:
   data_dir: data/postgresql0
   maximum_lag_on_failover: 1048576 # 1 megabyte in bytes
   use_slots: True
+  pg_rewind:
+    username: postgres
+    password: zalando
   pg_hba:
   - host all all 0.0.0.0/0 md5
   - hostssl all all 0.0.0.0/0 md5
@@ -39,6 +45,7 @@ postgresql:
     password: rep-pass
     network:  127.0.0.1/32
   superuser:
+    username: postgres
     password: zalando
   admin:
     username: admin
@@ -59,3 +66,4 @@ postgresql:
     archive_timeout: 1800s
     max_replication_slots: 5
     hot_standby: "on"
+    wal_log_hints: "on"

postgres1.yml

@@ -4,6 +4,9 @@ scope: &scope batman
 restapi:
   listen: 127.0.0.1:8009
   connect_address: 127.0.0.1:8009
+  auth: 'username:password'
+#  certfile: /etc/ssl/certs/ssl-cert-snakeoil.pem
+#  keyfile: /etc/ssl/private/ssl-cert-snakeoil.key
 etcd:
   scope: *scope
   ttl: *ttl
@@ -31,6 +34,9 @@ postgresql:
   data_dir: data/postgresql1
   maximum_lag_on_failover: 1048576 # 1 megabyte in bytes
   use_slots: True
+  pg_rewind:
+    username: postgres
+    password: zalando
   pg_hba:
   - host all all 0.0.0.0/0 md5
   - hostssl all all 0.0.0.0/0 md5
@@ -39,6 +45,7 @@ postgresql:
     password: rep-pass
     network: 127.0.0.1/32
   superuser:
+    user: postgres
     password: zalando
   admin:
     username: admin
@@ -59,3 +66,4 @@ postgresql:
     archive_timeout: 1800s
     max_replication_slots: 5
     hot_standby: "on"
+    wal_log_hints: "on"

setup.py

@@ -89,7 +89,6 @@ class PyTest(TestCommand):
         params = {'args': self.test_args}
         if self.cov:
             params['args'] += self.cov
-            params['plugins'] = ['cov']
         if self.junitxml:
             params['args'] += self.junitxml
         params['args'] += ['--doctest-modules', MAIN_PACKAGE, '-s', '-vv']

tests/test_api.py

@@ -1,16 +1,18 @@
 import psycopg2
 import unittest
 
+from mock import Mock, patch
 from patroni.api import RestApiHandler, RestApiServer
 from six import BytesIO as IO
-from test_postgresql import psycopg2_connect
+from six.moves import BaseHTTPServer
+from test_postgresql import psycopg2_connect, MockCursor
 
 
-def throws(*args, **kwargs):
+class MockPostgresql(Mock):
-    raise psycopg2.OperationalError()
 
-
+    name = 'test'
-class MockPostgresql:
+    state = 'running'
+    role = 'master'
 
     def connection(self):
         return psycopg2_connect()
@@ -19,10 +21,29 @@ class MockPostgresql:
         return True
 
 
+class MockHa(Mock):
+
+    dcs = Mock()
+    state_handler = MockPostgresql()
+
+    def schedule_restart(self):
+        return 'restart'
+
+    def schedule_reinitialize(self):
+        return 'reinitialize'
+
+    def restart(self):
+        return (True, '')
+
+    def restart_scheduled(self):
+        return False
+
+
 class MockPatroni:
 
-    def __init__(self):
+    postgresql = MockPostgresql()
-        self.postgresql = MockPostgresql()
+    ha = MockHa()
+    dcs = Mock()
 
 
 class MockRequest:
@@ -36,18 +57,63 @@ class MockRequest:
 
 class MockRestApiServer(RestApiServer):
 
-    def __init__(self, Handler, path, *args):
+    def __init__(self, Handler, path):
-        self.patroni = MockPatroni()
+        self.socket = 0
-        if len(args) > 0:
+        BaseHTTPServer.HTTPServer.__init__ = Mock()
-            self.query = args[0]
+        MockRestApiServer._BaseServer__is_shut_down = Mock()
+        MockRestApiServer._BaseServer__shutdown_request = True
+        config = {'listen': '127.0.0.1:8008', 'auth': 'test:test', 'certfile': 'dumb'}
+        super(MockRestApiServer, self).__init__(MockPatroni(), config)
         Handler(MockRequest(path), ('0.0.0.0', 8080), self)
 
 
+@patch('ssl.wrap_socket', Mock(return_value=0))
 class TestRestApiHandler(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        super(TestRestApiHandler, self).__init__(method_name)
-
     def test_do_GET(self):
-        MockRestApiServer(RestApiHandler, b'GET /')
+        MockRestApiServer(RestApiHandler, b'GET /replica')
-        MockRestApiServer(RestApiHandler, b'GET /', throws)
+        with patch.object(RestApiHandler, 'get_postgresql_status', Mock(return_value={})):
+            MockRestApiServer(RestApiHandler, b'GET /replica')
+        with patch.object(RestApiHandler, 'get_postgresql_status', Mock(return_value={'role': 'master'})):
+            MockRestApiServer(RestApiHandler, b'GET /replica')
+        MockRestApiServer(RestApiHandler, b'GET /master')
+        MockPatroni.dcs.cluster.leader.name = MockPostgresql.name
+        MockRestApiServer(RestApiHandler, b'GET /replica')
+        MockPatroni.dcs.cluster = None
+        with patch.object(RestApiHandler, 'get_postgresql_status', Mock(return_value={'role': 'master'})):
+            MockRestApiServer(RestApiHandler, b'GET /master')
+        with patch.object(MockHa, 'restart_scheduled', Mock(return_value=True)):
+            MockRestApiServer(RestApiHandler, b'GET /master')
+        MockRestApiServer(RestApiHandler, b'GET /master')
+
+    def test_do_GET_patroni(self):
+        MockRestApiServer(RestApiHandler, b'GET /patroni')
+
+    def test_basicauth(self):
+        MockRestApiServer(RestApiHandler, b'POST /restart HTTP/1.0')
+        MockRestApiServer(RestApiHandler, b'POST /restart HTTP/1.0\nAuthorization:')
+
+    def test_do_POST_restart(self):
+        request = b'POST /restart HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0'
+        MockRestApiServer(RestApiHandler, request)
+        with patch.object(MockHa, 'restart', Mock(side_effect=Exception)):
+            MockRestApiServer(RestApiHandler, request)
+
+    @patch.object(MockHa, 'dcs')
+    def test_do_POST_reinitialize(self, dcs):
+        cluster = dcs.get_cluster.return_value
+        request = b'POST /reinitialize HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0'
+        MockRestApiServer(RestApiHandler, request)
+        cluster.is_unlocked.return_value = False
+        MockRestApiServer(RestApiHandler, request)
+        with patch.object(MockHa, 'schedule_reinitialize', Mock(return_value=None)):
+            MockRestApiServer(RestApiHandler, request)
+        cluster.leader.name = 'test'
+        MockRestApiServer(RestApiHandler, request)
+
+    @patch('time.sleep', Mock())
+    def test_RestApiServer_query(self):
+        with patch.object(MockCursor, 'execute', Mock(side_effect=psycopg2.OperationalError)):
+            MockRestApiServer(RestApiHandler, b'GET /patroni')
+        with patch.object(MockPostgresql, 'connection', Mock(side_effect=psycopg2.OperationalError)):
+            MockRestApiServer(RestApiHandler, b'GET /patroni')

tests/test_async_executor.py

@@ -0,0 +1,18 @@
+import unittest
+
+from mock import Mock, patch
+from patroni.async_executor import AsyncExecutor
+from threading import Thread
+
+
+class TestAsyncExecutor(unittest.TestCase):
+
+    def setUp(self):
+        self.a = AsyncExecutor()
+
+    @patch.object(Thread, 'start', Mock())
+    def test_run_async(self):
+        self.a.run_async(Mock(return_value=True))
+
+    def test_run(self):
+        self.a.run(Mock(side_effect=Exception()))

tests/test_etcd.py

@@ -1,16 +1,13 @@
-import datetime
-import dns.resolver
 import etcd
 import json
 import requests
 import urllib3
 import socket
-import time
 import unittest
 
 from dns.exception import DNSException
 from mock import Mock, patch
-from patroni.dcs import Cluster, DCSError, Leader, Member
+from patroni.dcs import Cluster, DCSError, Leader
 from patroni.etcd import Client, Etcd
 
 
@@ -40,11 +37,10 @@ class MockResponse:
         return ''
 
 
-class MockPostgresql:
+class MockPostgresql(Mock):
-    name = ''
 
     def last_operation(self):
-        return 0
+        return '0'
 
 
 def requests_get(url, **kwargs):
@@ -53,6 +49,8 @@ def requests_get(url, **kwargs):
     response = MockResponse()
     if url.startswith('http://local'):
         raise requests.exceptions.RequestException()
+    elif ':8011/patroni' in url:
+        response.content = '{"role": "replica", "xlog": {"replayed_location": 0}}'
     elif url.endswith('/members'):
         if url.startswith('http://error'):
             response.content = '[{}]'
@@ -67,11 +65,11 @@ def requests_get(url, **kwargs):
 
 
 def etcd_watch(key, index=None, timeout=None, recursive=None):
-    if timeout == 1:
+    if timeout == 2.0:
         raise urllib3.exceptions.TimeoutError
-    elif timeout == 5:
+    elif timeout == 5.0:
         return etcd.EtcdResult('delete', {})
-    elif timeout == 10:
+    elif timeout == 10.0:
         raise etcd.EtcdException
     elif index == 20729:
         return etcd.EtcdResult('set', {'value': 'postgresql1', 'modifiedIndex': index + 1})
@@ -88,10 +86,6 @@ def etcd_write(key, value, **kwargs):
     raise etcd.EtcdException
 
 
-def etcd_delete(key, **kwargs):
-    raise etcd.EtcdException
-
-
 def etcd_read(key, **kwargs):
     if key == '/service/noleader/':
         raise DCSError('noleader')
@@ -99,6 +93,8 @@ def etcd_read(key, **kwargs):
         raise etcd.EtcdKeyNotFound
 
     response = {"action": "get", "node": {"key": "/service/batman5", "dir": True, "nodes": [
+                {"key": "/service/batman5/failover", "value": "",
+                 "modifiedIndex": 1582, "createdIndex": 1582},
                 {"key": "/service/batman5/initialize", "value": "postgresql0",
                  "modifiedIndex": 1582, "createdIndex": 1582},
                 {"key": "/service/batman5/leader", "value": "postgresql1",
@@ -123,18 +119,10 @@ def etcd_read(key, **kwargs):
     return etcd.EtcdResult(**response)
 
 
-def time_sleep(_):
-    pass
-
-
 class SleepException(Exception):
     pass
 
 
-def time_sleep_exception(_):
-    raise SleepException()
-
-
 class MockSRV:
     port = 2380
     target = '127.0.0.1'
@@ -151,7 +139,7 @@ def dns_query(name, type):
 def socket_getaddrinfo(*args):
     if args[0] == 'ok':
         return [(2, 1, 6, '', ('127.0.0.1', 2379)), (2, 1, 6, '', ('127.0.0.1', 2379))]
-    raise socket.error()
+    raise socket.error
 
 
 def http_request(method, url, **kwargs):
@@ -160,28 +148,14 @@ def http_request(method, url, **kwargs):
     raise socket.error
 
 
-class TestMember(unittest.TestCase):
+@patch('dns.resolver.query', dns_query)
-
+@patch('socket.getaddrinfo', socket_getaddrinfo)
-    def __init__(self, method_name='runTest'):
+@patch('requests.get', requests_get)
-        super(TestMember, self).__init__(method_name)
-
-    def test_real_ttl(self):
-        now = datetime.datetime.utcnow()
-        member = Member(0, 'a', 'b', 'c', (now + datetime.timedelta(seconds=2)).strftime('%Y-%m-%dT%H:%M:%S.%fZ'), None)
-        self.assertLess(member.real_ttl(), 2)
-        self.assertEquals(Member(0, 'a', 'b', 'c', '', None).real_ttl(), -1)
-
-
 class TestClient(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
+    @patch('dns.resolver.query', dns_query)
-        self.setUp = self.set_up
+    @patch('requests.get', requests_get)
-        super(TestClient, self).__init__(method_name)
+    def setUp(self):
-
-    def set_up(self):
-        socket.getaddrinfo = socket_getaddrinfo
-        requests.get = requests_get
-        dns.resolver.query = dns_query
         with patch.object(etcd.Client, 'machines') as mock_machines:
             mock_machines.__get__ = Mock(return_value=['http://localhost:2379', 'http://localhost:4001'])
             self.client = Client({'discovery_srv': 'test'})
@@ -206,11 +180,11 @@ class TestClient(unittest.TestCase):
         self.assertRaises(etcd.EtcdException, self.client._result_from_response, response)
 
     def test__get_machines_cache_from_srv(self):
-        self.client.get_srv_record = lambda e: [('localhost', 2380)]
+        self.client.get_srv_record = Mock(return_value=[('localhost', 2380)])
         self.client._get_machines_cache_from_srv('blabla')
 
     def test__get_machines_cache_from_dns(self):
-        self.client._get_machines_cache_from_dns('ok:2379')
+        self.client._get_machines_cache_from_dns('error:2379')
 
     def test__load_machines_cache(self):
         self.client._config = {}
@@ -219,25 +193,23 @@ class TestClient(unittest.TestCase):
         self.assertRaises(etcd.EtcdException, self.client._load_machines_cache)
 
 
+@patch('requests.get', requests_get)
 class TestEtcd(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
+    def setUp(self):
-        self.setUp = self.set_up
-        super(TestEtcd, self).__init__(method_name)
-
-    def set_up(self):
-        time.sleep = time_sleep
         with patch.object(Client, 'machines') as mock_machines:
             mock_machines.__get__ = Mock(return_value=['http://localhost:2379', 'http://localhost:4001'])
             self.etcd = Etcd('foo', {'ttl': 30, 'host': 'localhost:2379', 'scope': 'test'})
             self.etcd.client.write = etcd_write
             self.etcd.client.read = etcd_read
+            self.etcd.client.delete = Mock(side_effect=etcd.EtcdException())
 
+    @patch('dns.resolver.query', dns_query)
     def test_get_etcd_client(self):
-        time.sleep = time_sleep_exception
         with patch.object(etcd.Client, 'machines') as mock_machines:
             mock_machines.__get__ = Mock(side_effect=etcd.EtcdException)
-            self.assertRaises(SleepException, self.etcd.get_etcd_client, {'discovery_srv': 'test'})
+            with patch('time.sleep', Mock(side_effect=SleepException())):
+                self.assertRaises(SleepException, self.etcd.get_etcd_client, {'discovery_srv': 'test'})
 
     def test_get_cluster(self):
         self.assertIsInstance(self.etcd.get_cluster(), Cluster)
@@ -257,31 +229,32 @@ class TestEtcd(unittest.TestCase):
     def test_take_leader(self):
         self.assertFalse(self.etcd.take_leader())
 
-    def testattempt_to_acquire_leader(self):
+    def test_attempt_to_acquire_leader(self):
         self.etcd._base_path = '/service/exists'
         self.assertFalse(self.etcd.attempt_to_acquire_leader())
         self.etcd._base_path = '/service/failed'
         self.assertFalse(self.etcd.attempt_to_acquire_leader())
 
+    def test_write_leader_optime(self):
+        self.etcd.write_leader_optime('0')
+
     def test_update_leader(self):
-        self.assertTrue(self.etcd.update_leader(MockPostgresql()))
+        self.assertTrue(self.etcd.update_leader())
 
     def test_initialize(self):
         self.assertFalse(self.etcd.initialize())
 
     def test_cancel_initializion(self):
-        self.etcd.client.delete = etcd_delete
         self.assertFalse(self.etcd.cancel_initialization())
 
     def test_delete_leader(self):
-        self.etcd.client.delete = etcd_delete
         self.assertFalse(self.etcd.delete_leader())
 
     def test_watch(self):
         self.etcd.client.watch = etcd_watch
-        self.etcd.watch(100)
+        self.etcd.watch(0)
         self.etcd.get_cluster()
-        self.etcd.watch(1)
+        self.etcd.watch(1.5)
-        self.etcd.watch(5)
+        self.etcd.watch(4.5)
-        self.etcd.watch(10)
+        self.etcd.watch(9.5)
         self.etcd.watch(100)

tests/test_ha.py

@@ -1,10 +1,11 @@
 import unittest
 
 from mock import Mock, patch
-from patroni.dcs import Cluster, DCSError
+from patroni.dcs import Cluster, Failover, Leader, Member
 from patroni.etcd import Client, Etcd
+from patroni.exceptions import DCSError, PostgresException
 from patroni.ha import Ha
-from test_etcd import etcd_read, etcd_write
+from test_etcd import socket_getaddrinfo, etcd_read, etcd_write, requests_get
 
 
 def true(*args, **kwargs):
@@ -15,18 +16,37 @@ def false(*args, **kwargs):
     return False
 
 
-class MockPostgresql:
+def get_cluster(initialize, leader, members, failover):
+    return Cluster(initialize, leader, None, members, failover)
 
-    def __init__(self):
+
-        self.name = 'postgresql0'
+def get_cluster_not_initialized_without_leader():
-        self.is_promoted = False
+    return get_cluster(None, None, [], None)
+
+
+def get_cluster_initialized_without_leader(leader=False, failover=None):
+    m = Member(0, 'leader', 28, {'conn_url': 'postgres://replicator:[email protected]:5435/postgres',
+                                 'api_url': 'http://127.0.0.1:8008/patroni'})
+    l = Leader(0, 0, m) if leader else None
+    o = Member(0, 'other', 28, {'conn_url': 'postgres://replicator:[email protected]:5436/postgres',
+                                'api_url': 'http://127.0.0.1:8011/patroni'})
+    return get_cluster(True, l, [m, o], failover)
+
+
+def get_cluster_initialized_with_leader(failover=None):
+    return get_cluster_initialized_without_leader(leader=True, failover=failover)
+
+
+class MockPostgresql(Mock):
+
+    name = 'postgresql0'
+    role = 'replica'
+    state = 'running'
+    connection_string = 'postgres://foo@bar/postgres'
 
     def is_healthy(self):
         return True
 
-    def write_recovery_conf(self, _):
-        return True
-
     def start(self):
         return True
 
@@ -36,64 +56,80 @@ class MockPostgresql:
     def is_leader(self):
         return True
 
-    def promote(self):
+    def xlog_position(self):
-        return True
+        return 0
-
-    def demote(self, _):
-        return True
-
-    def follow_the_leader(self, _):
-        return True
-
-    def create_replication_slots(self, _):
-        return True
 
     def last_operation(self):
         return 0
 
+    def data_directory_empty(self):
+        return False
 
-def nop(*args, **kwargs):
+    def bootstrap(self, *args, **kwargs):
-    pass
+        return True
+
+    def check_replication_lag(self, last_leader_operation):
+        return True
+
+    def check_recovery_conf(self, leader):
+        return False
 
 
-def dead_etcd():
+class MockPatroni:
-    raise DCSError('Etcd is not responding properly')
+
+    def __init__(self, p, d):
+        self.postgresql = p
+        self.dcs = d
+        self.api = Mock()
+        self.api.connection_string = 'http://127.0.0.1:8008'
 
 
-def get_unlocked_cluster():
+def run_async(func, args=()):
-    return Cluster(False, None, None, [])
+    func(*args) if args else func()
 
 
 class TestHa(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
+    @patch('socket.getaddrinfo', socket_getaddrinfo)
-        self.setUp = self.set_up
+    @patch.object(Client, 'machines')
-        super(TestHa, self).__init__(method_name)
+    def setUp(self, mock_machines):
-
+        mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
-    def set_up(self):
         self.p = MockPostgresql()
-        with patch.object(Client, 'machines') as mock_machines:
+        self.e = Etcd('foo', {'ttl': 30, 'host': 'ok:2379', 'scope': 'test'})
-            mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
+        self.e.client.read = etcd_read
-            self.e = Etcd('foo', {'ttl': 30, 'host': 'remotehost:2379', 'scope': 'test'})
+        self.e.client.write = etcd_write
-            self.e.client.read = etcd_read
+        self.ha = Ha(MockPatroni(self.p, self.e))
-            self.e.client.write = etcd_write
+        self.ha._async_executor.run_async = run_async
-            self.ha = Ha(self.p, self.e)
+        self.ha.old_cluster = self.e.get_cluster()
-            self.ha.load_cluster_from_dcs()
+        self.ha.cluster = get_cluster_not_initialized_without_leader()
-            self.ha.cluster = get_unlocked_cluster()
+        self.ha.load_cluster_from_dcs = Mock()
-            self.ha.load_cluster_from_dcs = nop
 
-    def test_load_cluster_from_dcs(self):
+    def test_update_lock(self):
-        ha = Ha(self.p, self.e)
+        self.p.last_operation = Mock(side_effect=PostgresException(''))
-        ha.load_cluster_from_dcs()
+        self.assertTrue(self.ha.update_lock())
-        self.e.get_cluster = get_unlocked_cluster
-        ha.load_cluster_from_dcs()
 
-    def test_start_as_slave(self):
+    def test_touch_member(self):
+        self.p.xlog_position = Mock(side_effect=Exception)
+        self.ha.touch_member()
+
+    def test_start_as_replica(self):
         self.p.is_healthy = false
         self.assertEquals(self.ha.run_cycle(), 'started as a secondary')
 
+    def test_recover_replica_failed(self):
+        self.p.controldata = lambda: {'Database cluster state': 'in production'}
+        self.p.is_healthy = false
+        self.p.follow_the_leader = false
+        self.assertEquals(self.ha.run_cycle(), 'failed to start postgres')
+
+    def test_recover_master_failed(self):
+        self.p.follow_the_leader = false
+        self.p.is_healthy = false
+        self.ha.has_lock = true
+        self.assertEquals(self.ha.run_cycle(), 'removed leader key after trying and failing to start postgres')
+
+    @patch.object(Cluster, 'is_unlocked', Mock(return_value=False))
     def test_start_as_readonly(self):
-        self.ha.cluster.is_unlocked = false
         self.p.is_leader = self.p.is_healthy = false
         self.ha.has_lock = true
         self.assertEquals(self.ha.run_cycle(), 'promoted self to leader because i had the session lock')
@@ -102,6 +138,7 @@ class TestHa(unittest.TestCase):
         self.assertEquals(self.ha.run_cycle(), 'acquired session lock as a leader')
 
     def test_promoted_by_acquiring_lock(self):
+        self.ha.is_healthiest_node = true
         self.p.is_leader = false
         self.assertEquals(self.ha.run_cycle(), 'promoted self to leader by acquiring session lock')
 
@@ -110,16 +147,17 @@ class TestHa(unittest.TestCase):
         self.assertEquals(self.ha.run_cycle(), 'demoted self due after trying and failing to obtain lock')
 
     def test_follow_new_leader_after_failing_to_obtain_lock(self):
+        self.ha.is_healthiest_node = true
         self.ha.acquire_lock = false
         self.p.is_leader = false
         self.assertEquals(self.ha.run_cycle(), 'following new leader after trying and failing to obtain lock')
 
     def test_demote_because_not_healthiest(self):
-        self.p.is_healthiest_node = false
+        self.ha.is_healthiest_node = false
         self.assertEquals(self.ha.run_cycle(), 'demoting self because i am not the healthiest node')
 
     def test_follow_new_leader_because_not_healthiest(self):
-        self.p.is_healthiest_node = false
+        self.ha.is_healthiest_node = false
         self.p.is_leader = false
         self.assertEquals(self.ha.run_cycle(), 'following a different leader because i am not the healthiest node')
 
@@ -138,11 +176,124 @@ class TestHa(unittest.TestCase):
         self.ha.cluster.is_unlocked = false
         self.assertEquals(self.ha.run_cycle(), 'demoting self because i do not have the lock and i was a leader')
 
+    def test_demote_because_update_lock_failed(self):
+        self.ha.cluster.is_unlocked = false
+        self.ha.has_lock = true
+        self.ha.update_lock = false
+        self.assertEquals(self.ha.run_cycle(), 'demoting self because i do not have the lock and i was a leader')
+
     def test_follow_the_leader(self):
         self.ha.cluster.is_unlocked = false
         self.p.is_leader = false
         self.assertEquals(self.ha.run_cycle(), 'no action.  i am a secondary and i am following a leader')
 
     def test_no_etcd_connection_master_demote(self):
-        self.ha.load_cluster_from_dcs = dead_etcd
+        self.ha.load_cluster_from_dcs = Mock(side_effect=DCSError('Etcd is not responding properly'))
         self.assertEquals(self.ha.run_cycle(), 'demoted self because DCS is not accessible and i was a leader')
+
+    def test_bootstrap_from_leader(self):
+        self.ha.cluster = get_cluster_initialized_with_leader()
+        self.p.bootstrap = false
+        self.assertEquals(self.ha.bootstrap(), 'trying to bootstrap from leader')
+
+    def test_bootstrap_waiting_for_leader(self):
+        self.ha.cluster = get_cluster_initialized_without_leader()
+        self.assertEquals(self.ha.bootstrap(), 'waiting for leader to bootstrap')
+
+    def test_bootstrap_initialize_lock_failed(self):
+        self.ha.cluster = get_cluster_not_initialized_without_leader()
+        self.assertEquals(self.ha.bootstrap(), 'failed to acquire initialize lock')
+
+    def test_bootstrap_initialized_new_cluster(self):
+        self.ha.cluster = get_cluster_not_initialized_without_leader()
+        self.e.initialize = true
+        self.assertEquals(self.ha.bootstrap(), 'initialized a new cluster')
+
+    def test_bootstrap_release_initialize_key_on_failure(self):
+        self.ha.cluster = get_cluster_not_initialized_without_leader()
+        self.e.initialize = true
+        self.p.bootstrap = Mock(side_effect=PostgresException("Could not bootstrap master PostgreSQL"))
+        self.assertRaises(PostgresException, self.ha.bootstrap)
+
+    def test_reinitialize(self):
+        self.ha.schedule_reinitialize()
+        self.ha.schedule_reinitialize()
+        self.ha.run_cycle()
+        self.assertIsNone(self.ha._async_executor.scheduled_action)
+
+        self.ha.cluster = get_cluster_initialized_with_leader()
+        self.ha.has_lock = true
+        self.ha.schedule_reinitialize()
+        self.ha.run_cycle()
+        self.assertIsNone(self.ha._async_executor.scheduled_action)
+
+        self.ha.has_lock = false
+        self.ha.schedule_reinitialize()
+        self.ha.run_cycle()
+
+    def test_restart(self):
+        self.assertEquals(self.ha.restart(), (True, 'restarted successfully'))
+        self.p.restart = false
+        self.assertEquals(self.ha.restart(), (False, 'restart failed'))
+        self.ha.schedule_reinitialize()
+        self.assertEquals(self.ha.restart(), (False, 'reinitialize already in progress'))
+
+    def test_restart_in_progress(self):
+        self.ha._async_executor.schedule('restart', True)
+        self.assertTrue(self.ha.restart_scheduled())
+        self.assertEquals(self.ha.run_cycle(), 'not healthy enough for leader race')
+
+        self.ha.cluster = get_cluster_initialized_with_leader()
+        self.assertEquals(self.ha.run_cycle(), 'restart in progress')
+
+        self.ha.has_lock = true
+        self.assertEquals(self.ha.run_cycle(), 'updated leader lock during restart')
+
+        self.ha.update_lock = false
+        self.assertEquals(self.ha.run_cycle(), 'failed to update leader lock during restart')
+
+    @patch('requests.get', requests_get)
+    def test_manual_failover_from_leader(self):
+        self.ha.has_lock = true
+        self.ha.cluster = get_cluster_initialized_with_leader(Failover(0, 'blabla', ''))
+        self.assertEquals(self.ha.run_cycle(), 'no action.  i am the leader with the lock')
+        self.ha.cluster = get_cluster_initialized_with_leader(Failover(0, '', MockPostgresql.name))
+        self.assertEquals(self.ha.run_cycle(), 'no action.  i am the leader with the lock')
+        self.ha.cluster = get_cluster_initialized_with_leader(Failover(0, '', 'blabla'))
+        self.assertEquals(self.ha.run_cycle(), 'no action.  i am the leader with the lock')
+        f = Failover(0, MockPostgresql.name, '')
+        self.ha.cluster = get_cluster_initialized_with_leader(f)
+        self.assertEquals(self.ha.run_cycle(), 'manual failover: demoting myself')
+
+    @patch('requests.get', requests_get)
+    def test_manual_failover_process_no_leader(self):
+        self.p.is_leader = false
+        self.ha.cluster = get_cluster_initialized_without_leader(failover=Failover(0, '', MockPostgresql.name))
+        self.assertEquals(self.ha.run_cycle(), 'promoted self to leader by acquiring session lock')
+        self.ha.cluster = get_cluster_initialized_without_leader(failover=Failover(0, '', 'leader'))
+        self.assertEquals(self.ha.run_cycle(), 'promoted self to leader by acquiring session lock')
+        self.ha.fetch_node_status = lambda e: (e, True, True, 0)  # accessible, in_recovery
+        self.assertEquals(self.ha.run_cycle(), 'following a different leader because i am not the healthiest node')
+        self.ha.cluster = get_cluster_initialized_without_leader(failover=Failover(0, MockPostgresql.name, ''))
+        self.assertEquals(self.ha.run_cycle(), 'following a different leader because i am not the healthiest node')
+        self.ha.fetch_node_status = lambda e: (e, False, True, 0)  # accessible, in_recovery
+        self.assertEquals(self.ha.run_cycle(), 'promoted self to leader by acquiring session lock')
+
+    def test__is_healthiest_node(self):
+        self.assertTrue(self.ha._is_healthiest_node(self.ha.old_cluster.members))
+        self.p.is_leader = false
+        self.ha.fetch_node_status = lambda e: (e, True, True, 0)  # accessible, in_recovery
+        self.assertTrue(self.ha._is_healthiest_node(self.ha.old_cluster.members))
+        self.ha.fetch_node_status = lambda e: (e, True, False, 0)  # accessible, not in_recovery
+        self.assertFalse(self.ha._is_healthiest_node(self.ha.old_cluster.members))
+        self.ha.fetch_node_status = lambda e: (e, True, True, 1)  # accessible, in_recovery, xlog location ahead
+        self.assertFalse(self.ha._is_healthiest_node(self.ha.old_cluster.members))
+        self.p.check_replication_lag = false
+        self.assertFalse(self.ha._is_healthiest_node(self.ha.old_cluster.members))
+
+    @patch('requests.get', requests_get)
+    def test_fetch_node_status(self):
+        member = Member(0, 'test', 1, {'api_url': 'http://127.0.0.1:8011/patroni'})
+        self.ha.fetch_node_status(member)
+        member = Member(0, 'test', 1, {'api_url': 'http://localhost:8011/patroni'})
+        self.ha.fetch_node_status(member)

tests/test_patroni.py

@@ -1,7 +1,3 @@
-import datetime
-import patroni.zookeeper
-import psycopg2
-import subprocess
 import sys
 import time
 import unittest
@@ -9,192 +5,72 @@ import yaml
 
 from mock import Mock, patch
 from patroni.api import RestApiServer
-from patroni.dcs import Cluster, Member, Leader
+from patroni.async_executor import AsyncExecutor
 from patroni.etcd import Etcd
-from patroni.exceptions import PostgresException
 from patroni import Patroni, main
 from patroni.zookeeper import ZooKeeper
 from six.moves import BaseHTTPServer
-from test_etcd import Client, etcd_read, etcd_write
+from test_etcd import Client, SleepException, etcd_read, etcd_write
-from test_ha import true, false
+from test_postgresql import Postgresql, psycopg2_connect
-from test_postgresql import Postgresql, subprocess_call, psycopg2_connect
 from test_zookeeper import MockKazooClient
 
 
-def nop(*args, **kwargs):
-    pass
-
-
-class SleepException(Exception):
-    pass
-
-
 def time_sleep(*args):
     raise SleepException()
 
 
-class Mock_BaseServer__is_shut_down:
+@patch('time.sleep', Mock())
-
+@patch('subprocess.call', Mock(return_value=0))
-    def set(self):
+@patch('psycopg2.connect', psycopg2_connect)
-        pass
+@patch.object(Postgresql, 'write_pg_hba', Mock())
-
+@patch.object(Postgresql, 'write_recovery_conf', Mock())
-    def clear(self):
+@patch.object(BaseHTTPServer.HTTPServer, '__init__', Mock())
-        pass
+@patch.object(AsyncExecutor, 'run', Mock())
-
-
-def get_cluster(initialize, leader):
-    return Cluster(initialize, leader, None, None)
-
-
-def get_cluster_not_initialized_without_leader():
-    return get_cluster(None, None)
-
-
-def get_cluster_initialized_without_leader():
-    return get_cluster(True, None)
-
-
-def get_cluster_not_initialized_with_leader():
-    return get_cluster(False, Leader(0, 0, 0,
-                       Member(0, 'leader', 'postgres://replicator:[email protected]:5435/postgres',
-                              None, None, 28)))
-
-
-def get_cluster_initialized_with_leader():
-    return get_cluster(True, Leader(0, 0, 0, 
-                       Member(0, 'leader', 'postgres://replicator:[email protected]:5435/postgres',
-                              None, None, 28)))
-
-
 class TestPatroni(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
+    @patch.object(Client, 'machines')
-        self.setUp = self.set_up
+    def setUp(self, mock_machines):
-        self.tearDown = self.tear_down
+        mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
-        super(TestPatroni, self).__init__(method_name)
-
-    def set_up(self):
         self.touched = False
         self.init_cancelled = False
-        subprocess.call = subprocess_call
+        RestApiServer._BaseServer__is_shut_down = Mock()
-        psycopg2.connect = psycopg2_connect
-        self.time_sleep = time.sleep
-        time.sleep = nop
-        self.write_pg_hba = Postgresql.write_pg_hba
-        self.write_recovery_conf = Postgresql.write_recovery_conf
-        Postgresql.write_pg_hba = nop
-        Postgresql.write_recovery_conf = nop
-        BaseHTTPServer.HTTPServer.__init__ = nop
-        RestApiServer._BaseServer__is_shut_down = Mock_BaseServer__is_shut_down()
         RestApiServer._BaseServer__shutdown_request = True
         RestApiServer.socket = 0
         with open('postgres0.yml', 'r') as f:
             config = yaml.load(f)
-            with patch.object(Client, 'machines') as mock_machines:
+            self.p = Patroni(config)
-                mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
+            self.p.ha.dcs.client.write = etcd_write
-                self.p = Patroni(config)
+            self.p.ha.dcs.client.read = etcd_read
-
-    def tear_down(self):
-        time.sleep = self.time_sleep
-        Postgresql.write_pg_hba = self.write_pg_hba
-        Postgresql.write_recovery_conf = self.write_recovery_conf
 
+    @patch('patroni.zookeeper.KazooClient', MockKazooClient())
     def test_get_dcs(self):
-        patroni.zookeeper.KazooClient = MockKazooClient
         self.assertIsInstance(self.p.get_dcs('', {'zookeeper': {'scope': '', 'hosts': ''}}), ZooKeeper)
         self.assertRaises(Exception, self.p.get_dcs, '', {})
 
-    def test_patroni_main(self):
+    @patch('time.sleep', Mock(side_effect=SleepException()))
+    @patch.object(Etcd, 'delete_leader', Mock())
+    @patch.object(Client, 'machines')
+    def test_patroni_main(self, mock_machines):
         main()
         sys.argv = ['patroni.py', 'postgres0.yml']
-        time.sleep = time_sleep
-
-        with patch.object(Client, 'machines') as mock_machines:
-            mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
-            Patroni.initialize = nop
-            touch_member = Patroni.touch_member
-            run = Patroni.run
-
-            Patroni.touch_member = self.touch_member
-            Patroni.run = time_sleep
-
-            Etcd.delete_leader = nop
 
+        mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
+        with patch.object(Patroni, 'run', Mock(side_effect=SleepException())):
             self.assertRaises(SleepException, main)
+        with patch.object(Patroni, 'run', Mock(side_effect=KeyboardInterrupt())):
+            main()
 
-            Patroni.run = run
+    @patch('time.sleep', Mock(side_effect=SleepException()))
-            Patroni.touch_member = touch_member
+    def test_run(self):
-
-    def test_patroni_run(self):
-        time.sleep = time_sleep
-        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.watch = time_sleep
         self.assertRaises(SleepException, self.p.run)
-        self.p.ha.state_handler.is_leader = lambda: False
+
-        self.p.api.start = nop
+        self.p.ha.state_handler.is_leader = Mock(return_value=False)
+        self.p.api.start = Mock()
         self.assertRaises(SleepException, self.p.run)
 
-    def touch_member(self, ttl=None):
-        if not self.touched:
-            self.touched = True
-            return False
-        return True
-
-    def test_touch_member(self):
-        self.p.ha.dcs.client.write = etcd_write
-        self.p.touch_member()
-        now = datetime.datetime.utcnow()
-        member = Member(0, self.p.postgresql.name, 'b', 'c', (now + datetime.timedelta(
-            seconds=self.p.shutdown_member_ttl + 10)).strftime('%Y-%m-%dT%H:%M:%S.%fZ'), None)
-        self.p.ha.cluster = Cluster(True, member, 0, [member])
-        self.p.touch_member()
-
-    def test_patroni_initialize(self):
-        self.p.ha.dcs.client.write = etcd_write
-        self.p.ha.dcs.client.read = etcd_read
-        self.p.touch_member = self.touch_member
-        self.p.postgresql.data_directory_empty = true
-        self.p.ha.dcs.initialize = true
-        self.p.postgresql.initialize = true
-        self.p.postgresql.start = true
-        self.p.ha.dcs.get_cluster = get_cluster_not_initialized_without_leader
-        self.p.initialize()
-
-        self.p.ha.dcs.initialize = false
-        self.p.ha.dcs.get_cluster = get_cluster_initialized_with_leader
-        time.sleep = time_sleep
-        self.p.ha.dcs.client.read = etcd_read
-        self.p.initialize()
-
-        self.p.ha.dcs.get_cluster = get_cluster_initialized_without_leader
-        self.assertRaises(SleepException, self.p.initialize)
-
-        self.p.postgresql.data_directory_empty = false
-        self.p.initialize()
-
-        self.p.ha.dcs.get_cluster = get_cluster_not_initialized_with_leader
-        self.p.postgresql.data_directory_empty = true
-        self.p.initialize()
-
     def test_schedule_next_run(self):
+        self.p.ha.dcs.watch = Mock(return_value=True)
+        self.p.schedule_next_run()
         self.p.next_run = time.time() - self.p.nap_time - 1
         self.p.schedule_next_run()
-
-    def cancel_initialization(self):
-        self.init_cancelled = True
-
-    def test_cleanup_on_initialization(self):
-        self.p.ha.dcs.client.write = etcd_write
-        self.p.ha.dcs.client.read = etcd_read
-        self.p.ha.dcs.get_cluster = get_cluster_not_initialized_without_leader
-        self.p.touch_member = self.touch_member
-        self.p.postgresql.data_directory_empty = true
-        self.p.ha.dcs.initialize = true
-        self.p.postgresql.initialize = true
-        self.p.postgresql.start = false
-
-        self.p.ha.dcs.cancel_initialization = self.cancel_initialization
-        self.assertRaises(PostgresException, self.p.initialize)
-        self.assertTrue(self.init_cancelled)

tests/test_postgresql.py

@@ -1,51 +1,41 @@
+import mock  # for the mock.call method, importing it without a namespace breaks python3
 import os
 import psycopg2
 import shutil
-import subprocess
 import unittest
 
+from sys import version_info
+if version_info.major == 2:
+    import __builtin__ as builtins
+else:
+    import builtins
+
+from mock import Mock, MagicMock, PropertyMock, patch, mock_open
 from patroni.dcs import Cluster, Leader, Member
+from patroni.exceptions import PostgresException, PostgresConnectionException
 from patroni.postgresql import Postgresql
-from test_ha import true, false
+from patroni.utils import RetryFailedError
-
+from test_ha import false
-
+import subprocess
-def nop(*args, **kwargs):
-    pass
-
-
-def subprocess_call(cmd, shell=False, env=None):
-    return 0
-
-
-def false(*args, **kwargs):
-    return False
 
 
 class MockCursor:
 
-    def __init__(self):
+    def __init__(self, connection):
+        self.connection = connection
         self.closed = False
         self.results = []
 
     def execute(self, sql, *params):
-        if sql.startswith('blabla'):
+        if sql.startswith('blabla') or sql == 'CHECKPOINT':
             raise psycopg2.OperationalError()
-        elif sql.startswith('InterfaceError'):
+        elif sql.startswith('RetryFailedError'):
-            raise psycopg2.InterfaceError()
+            raise RetryFailedError('retry')
         elif sql.startswith('SELECT slot_name'):
             self.results = [('blabla',), ('foobar',)]
-        elif sql.startswith('SELECT pg_current_xlog_location()'):
-            self.results = [(0,)]
-        elif sql.startswith('SELECT pg_is_in_recovery(), %s'):
-            if params[0][0] == 1:
-                raise psycopg2.OperationalError()
-            elif params[0][0] == 2:
-                self.results = [(True, -1)]
-            else:
-                self.results = [(False, 0)]
         elif sql.startswith('SELECT pg_xlog_location_diff'):
             self.results = [(0,)]
-        elif sql.startswith('SELECT pg_is_in_recovery()'):
+        elif sql == 'SELECT pg_is_in_recovery()':
             self.results = [(False, )]
         elif sql.startswith('SELECT to_char(pg_postmaster_start_time'):
             self.results = [('', True, '', '', '', False)]
@@ -73,44 +63,108 @@ class MockCursor:
         for i in self.results:
             yield i
 
+    def __enter__(self):
+        return self
 
-class MockConnect:
+    def __exit__(self, *args):
-
-    def __init__(self):
-        self.autocommit = False
-        self.closed = 0
-
-    def cursor(self):
-        return MockCursor()
-
-    def close(self):
         pass
 
 
-def psycopg2_connect(*args, **kwargs):
+class MockConnect(Mock):
 
+    autocommit = False
+    closed = 0
+
+    def cursor(self):
+        return MockCursor(self)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+
+def pg_controldata_string(*args, **kwargs):
+    return """
+pg_control version number:            942
+Catalog version number:               201509161
+Database system identifier:           6200971513092291716
+Database cluster state:               shut down in recovery
+pg_control last modified:             Fri Oct  2 10:57:06 2015
+Latest checkpoint location:           0/30000C8
+Prior checkpoint location:            0/2000060
+Latest checkpoint's REDO location:    0/3000090
+Latest checkpoint's REDO WAL file:    000000020000000000000003
+Latest checkpoint's TimeLineID:       2
+Latest checkpoint's PrevTimeLineID:   2
+Latest checkpoint's full_page_writes: on
+Latest checkpoint's NextXID:          0/943
+Latest checkpoint's NextOID:          24576
+Latest checkpoint's NextMultiXactId:  1
+Latest checkpoint's NextMultiOffset:  0
+Latest checkpoint's oldestXID:        931
+Latest checkpoint's oldestXID's DB:   1
+Latest checkpoint's oldestActiveXID:  943
+Latest checkpoint's oldestMultiXid:   1
+Latest checkpoint's oldestMulti's DB: 1
+Latest checkpoint's oldestCommitTs:   0
+Latest checkpoint's newestCommitTs:   0
+Time of latest checkpoint:            Fri Oct  2 10:56:54 2015
+Fake LSN counter for unlogged rels:   0/1
+Minimum recovery ending location:     0/30241F8
+Min recovery ending loc's timeline:   2
+Backup start location:                0/0
+Backup end location:                  0/0
+End-of-backup record required:        no
+wal_level setting:                    hot_standby
+wal_log_hints setting:                on
+max_connections setting:              100
+max_worker_processes setting:         8
+max_prepared_xacts setting:           0
+max_locks_per_xact setting:           64
+track_commit_timestamp setting:       off
+Maximum data alignment:               8
+Database block size:                  8192
+Blocks per segment of large relation: 131072
+WAL block size:                       8192
+Bytes per WAL segment:                16777216
+Maximum length of identifiers:        64
+Maximum columns in an index:          32
+Maximum size of a TOAST chunk:        1996
+Size of a large-object chunk:         2048
+Date/time type storage:               64-bit integers
+Float4 argument passing:              by value
+Float8 argument passing:              by value
+Data page checksum version:           0
+"""
+
+
+def postmaster_opts_string(*args, **kwargs):
+    return '/usr/local/pgsql/bin/postgres "-D" "data/postgresql0" "--listen_addresses=127.0.0.1" "--port=5432"'\
+        ' "--hot_standby=on" "--wal_keep_segments=8" "--wal_level=hot_standby" "--archive_command=mkdir -p ../wal_archive \n'\
+        '&& cp %p ../wal_archive/%f" "--wal_log_hints=on" "--max_wal_senders=5" "--archive_timeout=1800s" "--archive_mode=on"'\
+        ' "--max_replication_slots=5"\n'
+
+
+def psycopg2_connect(*args, **kwargs):
     return MockConnect()
 
 
-def is_running():
+@patch('subprocess.call', Mock(return_value=0))
-    return False
+@patch('psycopg2.connect', psycopg2_connect)
-
+@patch('shutil.copy', Mock())
-
 class TestPostgresql(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
+    @patch('subprocess.call', Mock(return_value=0))
-        self.setUp = self.set_up
+    @patch('psycopg2.connect', psycopg2_connect)
-        self.tearDown = self.tear_down
+    def setUp(self):
-        super(TestPostgresql, self).__init__(method_name)
-
-    def set_up(self):
-        subprocess.call = subprocess_call
-        shutil.copy = nop
         self.p = Postgresql({'name': 'test0', 'scope': 'batman', 'data_dir': 'data/test0',
                              'listen': '127.0.0.1, *:5432', 'connect_address': '127.0.0.2:5432',
                              'pg_hba': ['hostssl all all 0.0.0.0/0 md5', 'host all all 0.0.0.0/0 md5'],
                              'superuser': {'password': ''},
                              'admin': {'username': 'admin', 'password': 'admin'},
+                             'pg_rewind': {'username': 'admin', 'password': 'admin'},
                              'replication': {'username': 'replicator',
                                              'password': 'rep-pass',
                                              'network': '127.0.0.1/32'},
@@ -120,20 +174,16 @@ class TestPostgresql(unittest.TestCase):
                                            'on_reload': 'true'
                                            },
                              'restore': 'true'})
-        psycopg2.connect = psycopg2_connect
         if not os.path.exists(self.p.data_dir):
             os.makedirs(self.p.data_dir)
-        self.leadermem = Member(0, 'leader', 'postgres://replicator:[email protected]:5435/postgres', None, None, 28)
+        self.leadermem = Member(0, 'leader', 28, {'conn_url': 'postgres://replicator:[email protected]:5435/postgres'})
-        self.leader = Leader(-1, None, 28, self.leadermem)
+        self.leader = Leader(-1, 28, self.leadermem)
-        self.other = Member(0, 'test1', 'postgres://replicator:[email protected]:5433/postgres', None, None, 28)
+        self.other = Member(0, 'test1', 28, {'conn_url': 'postgres://replicator:[email protected]:5433/postgres'})
-        self.me = Member(0, 'test0', 'postgres://replicator:[email protected]:5434/postgres', None, None, 28)
+        self.me = Member(0, 'test0', 28, {'conn_url': 'postgres://replicator:[email protected]:5434/postgres'})
 
-    def tear_down(self):
+    def tearDown(self):
         shutil.rmtree('data')
 
-    def mock_query(self, p):
-        raise psycopg2.OperationalError("not supported")
-
     def test_data_directory_empty(self):
         self.assertTrue(self.p.data_directory_empty())
 
@@ -141,23 +191,74 @@ class TestPostgresql(unittest.TestCase):
         self.assertTrue(self.p.initialize())
         self.assertTrue(os.path.exists(os.path.join(self.p.data_dir, 'pg_hba.conf')))
 
-    def test_start_stop(self):
+    def test_start(self):
-        self.assertFalse(self.p.start())
-        self.p.is_running = is_running
-        with open(os.path.join(self.p.data_dir, 'postmaster.pid'), 'w'):
-            pass
         self.assertTrue(self.p.start())
+        self.p.is_running = false
+        open(os.path.join(self.p.data_dir, 'postmaster.pid'), 'w').close()
+        self.assertTrue(self.p.start())
+
+    def test_stop(self):
         self.assertTrue(self.p.stop())
+        with patch('subprocess.call', Mock(return_value=1)):
+            self.assertTrue(self.p.stop())
+            self.p.is_running = Mock(return_value=True)
+            self.assertFalse(self.p.stop())
+
+    def test_restart(self):
+        self.p.start = false
+        self.p.is_running = false
+        self.assertFalse(self.p.restart())
+        self.assertEquals(self.p.state, 'restart failed (restarting)')
 
     def test_sync_from_leader(self):
         self.assertTrue(self.p.sync_from_leader(self.leader))
 
-    def test_follow_the_leader(self):
+    @patch('subprocess.call', side_effect=Exception("Test"))
-        self.p.demote(self.leader)
+    def test_pg_rewind(self, mock_call):
+        self.assertTrue(self.p.rewind(self.leader))
+        subprocess.call = mock_call
+        self.assertFalse(self.p.rewind(self.leader))
+
+    @patch('patroni.postgresql.Postgresql.rewind', return_value=False)
+    @patch('patroni.postgresql.Postgresql.remove_data_directory', MagicMock(return_value=True))
+    @patch('patroni.postgresql.Postgresql.single_user_mode', MagicMock(return_value=1))
+    def test_follow_the_leader(self, mock_pg_rewind):
+        self.p.demote()
         self.p.follow_the_leader(None)
-        self.p.demote(self.leader)
+        self.p.demote()
         self.p.follow_the_leader(self.leader)
-        self.p.follow_the_leader(Leader(-1, None, 28, self.other))
+        self.p.follow_the_leader(Leader(-1, 28, self.other))
+        self.p.rewind = mock_pg_rewind
+        self.p.follow_the_leader(self.leader)
+        self.p.require_rewind()
+        with mock.patch('os.path.islink', MagicMock(return_value=True)):
+            with mock.patch('os.unlink', MagicMock(return_value=True)):
+                with mock.patch('patroni.postgresql.Postgresql.can_rewind', new_callable=PropertyMock(return_value=True)):
+                    self.p.follow_the_leader(self.leader, recovery=True)
+        self.p.require_rewind()
+        with mock.patch('patroni.postgresql.Postgresql.can_rewind', new_callable=PropertyMock(return_value=True)):
+            self.p.rewind.return_value = True
+            self.p.follow_the_leader(self.leader, recovery=True)
+            self.p.rewind.return_value = False
+            self.p.follow_the_leader(self.leader, recovery=True)
+
+    def test_can_rewind(self):
+        tmp = self.p.pg_rewind
+        self.p.pg_rewind = None
+        self.assertFalse(self.p.can_rewind)
+        self.p.pg_rewind = tmp
+        with mock.patch('subprocess.call', MagicMock(return_value=1)):
+            self.assertFalse(self.p.can_rewind)
+        with mock.patch('subprocess.call', side_effect=OSError("foo")):
+            self.assertFalse(self.p.can_rewind)
+        tmp = self.p.controldata()
+        self.p.controldata = lambda: {'wal_log_hints setting': 'on'}
+        self.assertTrue(self.p.can_rewind)
+        self.p.controldata = tmp
+
+    def test_create_replica(self):
+        self.p.delete_trigger_file = Mock(side_effect=OSError())
+        self.assertEquals(self.p.create_replica({'host': '', 'port': '', 'user': ''}, ''), 1)
 
     def test_create_connection_users(self):
         cfg = self.p.config
@@ -165,62 +266,166 @@ class TestPostgresql(unittest.TestCase):
         p = Postgresql(cfg)
         p.create_connection_users()
 
-    def test_create_replication_slots(self):
+    def test_sync_replication_slots(self):
         self.p.start()
-        cluster = Cluster(True, self.leader, 0, [self.me, self.other, self.leadermem])
+        cluster = Cluster(True, self.leader, 0, [self.me, self.other, self.leadermem], None)
-        self.p.create_replication_slots(cluster)
+        self.p.sync_replication_slots(cluster)
+        self.p.query = Mock(side_effect=psycopg2.OperationalError)
+        self.p.schedule_load_slots = True
+        self.p.sync_replication_slots(cluster)
+
+    @patch.object(MockConnect, 'closed', 2)
+    def test__query(self):
+        self.assertRaises(PostgresConnectionException, self.p._query, 'blabla')
+        self.p._state = 'restarting'
+        self.assertRaises(RetryFailedError, self.p._query, 'blabla')
 
     def test_query(self):
         self.p.query('select 1')
-        self.assertRaises(psycopg2.InterfaceError, self.p.query, 'InterfaceError')
+        self.assertRaises(PostgresConnectionException, self.p.query, 'RetryFailedError')
         self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla')
-        self.p._connection.closed = 2
-        self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla')
-        self.p._connection.closed = 2
-        self.p.disconnect = false
-        self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla')
-
-    def test_is_healthiest_node(self):
-        cluster = Cluster(True, self.leader, 0, [self.me, self.other, self.leadermem])
-        self.assertTrue(self.p.is_healthiest_node(cluster))
-        self.p.is_leader = false
-        self.assertFalse(self.p.is_healthiest_node(cluster))
-        self.p.xlog_position = lambda: 1
-        self.assertTrue(self.p.is_healthiest_node(cluster))
-        self.p.xlog_position = lambda: 2
-        self.assertFalse(self.p.is_healthiest_node(cluster))
-        self.p.config['maximum_lag_on_failover'] = -3
-        self.assertFalse(self.p.is_healthiest_node(cluster))
 
     def test_is_leader(self):
-        self.p.is_promoted = True
         self.assertTrue(self.p.is_leader())
-        self.assertFalse(self.p.is_promoted)
 
     def test_reload(self):
         self.assertTrue(self.p.reload())
 
     def test_is_healthy(self):
         self.assertTrue(self.p.is_healthy())
-        self.p.is_running = is_running
+        self.p.is_running = false
         self.assertFalse(self.p.is_healthy())
 
     def test_promote(self):
+        self.p._role = 'replica'
+        self.assertTrue(self.p.promote())
         self.assertTrue(self.p.promote())
 
     def test_last_operation(self):
         self.assertEquals(self.p.last_operation(), '0')
 
+    @patch('subprocess.Popen', Mock(side_effect=OSError()))
+    def test_call_nowait(self):
+        self.assertFalse(self.p.call_nowait('on_start'))
+
     def test_non_existing_callback(self):
         self.assertFalse(self.p.call_nowait('foobar'))
 
     def test_is_leader_exception(self):
         self.p.start()
-        self.p.query = self.mock_query
+        self.p.query = Mock(side_effect=psycopg2.OperationalError("not supported"))
         self.assertTrue(self.p.stop())
 
+    def test_check_replication_lag(self):
+        self.assertTrue(self.p.check_replication_lag(0))
+
+    @patch('os.rename', Mock())
+    @patch('os.path.isdir', Mock(return_value=True))
     def test_move_data_directory(self):
-        self.p.is_running = is_running
+        self.p.is_running = false
-        os.rename = nop
-        os.path.isdir = true
         self.p.move_data_directory()
+        with patch('os.rename', Mock(side_effect=OSError())):
+            self.p.move_data_directory()
+
+    def test_bootstrap(self):
+        with patch('subprocess.call', Mock(return_value=1)):
+            self.assertRaises(PostgresException, self.p.bootstrap)
+        self.p.bootstrap()
+        self.p.bootstrap(self.leader)
+
+    def test_remove_data_directory(self):
+        self.p.data_dir = 'data_dir'
+        self.p.remove_data_directory()
+        os.mkdir(self.p.data_dir)
+        self.p.remove_data_directory()
+        open(self.p.data_dir, 'w').close()
+        self.p.remove_data_directory()
+        os.symlink('unexisting', self.p.data_dir)
+        with patch('os.unlink', Mock(side_effect=Exception)):
+            self.p.remove_data_directory()
+        self.p.remove_data_directory()
+
+    @patch('subprocess.check_output', MagicMock(return_value=0, side_effect=pg_controldata_string))
+    @patch('subprocess.check_output', side_effect=subprocess.CalledProcessError)
+    @patch('subprocess.check_output', side_effect=Exception('Failed'))
+    def test_controldata(self, check_output_call_error, check_output_generic_exception):
+        data = self.p.controldata()
+        self.assertEquals(len(data), 50)
+        self.assertEquals(data['Database cluster state'], 'shut down in recovery')
+        self.assertEquals(data['wal_log_hints setting'], 'on')
+        self.assertEquals(int(data['Database block size']), 8192)
+
+        subprocess.check_output = check_output_call_error
+        data = self.p.controldata()
+        self.assertEquals(data, dict())
+
+        subprocess.check_output = check_output_generic_exception
+        self.assertRaises(Exception, self.p.controldata())
+
+    def test_read_postmaster_opts(self):
+        m = mock_open(read_data=postmaster_opts_string())
+        with patch.object(builtins, 'open', m):
+            data = self.p.read_postmaster_opts()
+            self.assertEquals(data['wal_level'], 'hot_standby')
+            self.assertEquals(int(data['max_replication_slots']), 5)
+            self.assertEqual(data.get('D'), None)
+
+            m.side_effect = IOError("foo")
+            data = self.p.read_postmaster_opts()
+            self.assertEqual(data, dict())
+
+            m.side_effect = Exception("foo")
+            self.assertRaises(Exception, self.p.read_postmaster_opts())
+
+    @patch('subprocess.Popen')
+    @patch.object(builtins, 'open', MagicMock(return_value=42))
+    def test_single_user_mode(self, subprocess_popen_mock):
+        subprocess_popen_mock.return_value.wait.return_value = 0
+        self.assertEquals(self.p.single_user_mode(options=dict(archive_mode='on', archive_command='false')), 0)
+        subprocess_popen_mock.assert_called_once_with(['postgres', '--single', '-D', self.p.data_dir,
+                                                      '-c', 'archive_command=false', '-c', 'archive_mode=on',
+                                                       'postgres'], stdin=subprocess.PIPE,
+                                                      stdout=42,
+                                                      stderr=subprocess.STDOUT)
+        subprocess_popen_mock.reset_mock()
+        self.assertEquals(self.p.single_user_mode(command="CHECKPOINT"), 0)
+        subprocess_popen_mock.assert_called_once_with(['postgres', '--single', '-D', self.p.data_dir,
+                                                      'postgres'], stdin=subprocess.PIPE,
+                                                      stdout=42,
+                                                      stderr=subprocess.STDOUT)
+        subprocess_popen_mock.return_value = None
+        self.assertEquals(self.p.single_user_mode(), 1)
+
+    def fake_listdir(path):
+        if path.endswith(os.path.join('pg_xlog', 'archive_status')):
+            return ["a", "b", "c"]
+        return []
+
+    @patch('os.listdir', MagicMock(side_effect=fake_listdir))
+    @patch('os.path.isdir', MagicMock(return_value=True))
+    @patch('os.unlink', return_value=True)
+    @patch('os.remove', return_value=True)
+    @patch('os.path.islink', return_value=False)
+    @patch('os.path.isfile', return_value=True)
+    def test_cleanup_archive_status(self, mock_file, mock_link, mock_remove, mock_unlink):
+        ap = os.path.join(self.p.data_dir, 'pg_xlog', 'archive_status/')
+        self.p.cleanup_archive_status()
+        mock_remove.assert_has_calls([mock.call(ap+'a'), mock.call(ap+'b'), mock.call(ap+'c')])
+        mock_unlink.assert_not_called()
+
+        mock_remove.reset_mock()
+
+        mock_file.return_value = False
+        mock_link.return_value = True
+        self.p.cleanup_archive_status()
+        mock_unlink.assert_has_calls([mock.call(ap+'a'), mock.call(ap+'b'), mock.call(ap+'c')])
+        mock_remove.assert_not_called()
+
+        mock_unlink.reset_mock()
+        mock_remove.reset_mock()
+
+        mock_file.side_effect = Exception("foo")
+        mock_link.side_effect = Exception("foo")
+        self.p.cleanup_archive_status()
+        mock_unlink.assert_not_called()
+        mock_remove.assert_not_called()

tests/test_utils.py

@@ -1,56 +1,34 @@
-import os
-import time
 import unittest
 
-from patroni.exceptions import DCSError
+from mock import Mock, patch
+from patroni.exceptions import PatroniException
 from patroni.utils import Retry, RetryFailedError, reap_children, sigchld_handler, sigterm_handler, sleep
 
 
-def nop(*args, **kwargs):
-    pass
-
-
-def os_waitpid(a, b):
-    return (0, 0)
-
-
 def time_sleep(_):
     sigchld_handler(None, None)
 
 
 class TestUtils(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        self.tearDown = self.tear_down
-        super(TestUtils, self).__init__(method_name)
-
-    def set_up(self):
-        self.time_sleep = time.sleep
-        time.sleep = nop
-
-    def tear_down(self):
-        time.sleep = self.time_sleep
-
     def test_sigterm_handler(self):
         self.assertRaises(SystemExit, sigterm_handler, None, None)
 
+    @patch('time.sleep', Mock())
     def test_reap_children(self):
         reap_children()
-        os.waitpid = os_waitpid
+        with patch('os.waitpid', Mock(return_value=(0, 0))):
-        sigchld_handler(None, None)
+            sigchld_handler(None, None)
-        reap_children()
+            reap_children()
 
+    @patch('time.sleep', time_sleep)
     def test_sleep(self):
-        time.sleep = time_sleep
         sleep(0.01)
 
 
+@patch('time.sleep', Mock())
 class TestRetrySleeper(unittest.TestCase):
 
-    def _pass(self):
-        pass
-
     def _fail(self, times=1):
         scope = dict(times=0)
 
@@ -59,7 +37,7 @@ class TestRetrySleeper(unittest.TestCase):
                 pass
             else:
                 scope['times'] += 1
-                raise DCSError('Failed!')
+                raise PatroniException('Failed!')
         return inner
 
     def _makeOne(self, *args, **kwargs):
@@ -78,20 +56,14 @@ class TestRetrySleeper(unittest.TestCase):
         self.assertEquals(retry._attempts, 1)
 
     def test_maximum_delay(self):
-        def sleep_func(_time):
+        retry = self._makeOne(delay=10, max_tries=100)
-            pass
-
-        retry = self._makeOne(delay=10, max_tries=100, sleep_func=sleep_func)
         retry(self._fail(times=10))
         self.assertTrue(retry._cur_delay < 4000, retry._cur_delay)
         # gevent's sleep function is picky about the type
         self.assertEquals(type(retry._cur_delay), float)
 
     def test_deadline(self):
-        def sleep_func(_time):
+        retry = self._makeOne(deadline=0.0001)
-            pass
-
-        retry = self._makeOne(deadline=0.0001, sleep_func=sleep_func)
         self.assertRaises(RetryFailedError, retry, self._fail(times=100))
 
     def test_copy(self):

tests/test_zookeeper.py

@@ -1,136 +1,112 @@
-import patroni.zookeeper
+import six
-import requests
 import unittest
 
+from mock import Mock, patch
 from patroni.dcs import Leader
 from patroni.zookeeper import ExhibitorEnsembleProvider, ZooKeeper, ZooKeeperError
 from kazoo.client import KazooState
 from kazoo.exceptions import NoNodeError, NodeExistsError
 from kazoo.protocol.states import ZnodeStat
-from test_etcd import MockPostgresql, requests_get
+from test_etcd import SleepException, requests_get
 
 
-class MockEvent:
+class MockKazooClient(Mock):
 
-    def clear(self):
+    leader = False
-        pass
+    exists = True
-
+    handler = Mock()
-    def set(self):
-        pass
-
-    def wait(self, timeout):
-        pass
-
-    def isSet(self):
-        return True
-
-
-class MockEventHandler:
-
-    def event_object(self):
-        return MockEvent()
-
-
-class SleepException(Exception):
-    pass
-
-
-class MockKazooClient:
-
-    def __init__(self, **kwargs):
-        self.handler = MockEventHandler()
-        self.leader = False
-        self.exists = True
-
-    def start(self, timeout):
-        pass
 
     @property
     def client_id(self):
         return (-1, '')
 
-    def add_listener(self, cb):
-        pass
-
     def retry(self, func, *args, **kwargs):
         func(*args, **kwargs)
 
     def get(self, path, watch=None):
+        if not isinstance(path, six.string_types):
+            raise TypeError("Invalid type for 'path' (string expected)")
         if path == '/no_node':
             raise NoNodeError
         elif '/members/' in path:
             return (
-                'postgres://repuser:rep-pass@localhost:5434/postgres?application_name=http://127.0.0.1:8009/patroni',
+                b'postgres://repuser:rep-pass@localhost:5434/postgres?application_name=http://127.0.0.1:8009/patroni',
-                ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
+                ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0 if self.exists else -1, 0, 0, 0)
             )
         elif path.endswith('/optime/leader'):
-            return '1'
+            return (b'1', ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))
         elif path.endswith('/leader'):
             if self.leader:
-                return ('foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, -1, 0, 0, 0))
+                return (b'foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, -1, 0, 0, 0))
-            return ('foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))
+            return (b'foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))
         elif path.endswith('/initialize'):
-            return ('foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))
+            return (b'foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))
+        return (b'', ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))
 
     def get_children(self, path, watch=None, include_data=False):
-        if path == '/no_node':
+        if not isinstance(path, six.string_types):
+            raise TypeError("Invalid type for 'path' (string expected)")
+        if path.startswith('/no_node'):
             raise NoNodeError
         elif path in ['/service/bla/', '/service/test/']:
-            return ['initialize', 'leader', 'members', 'optime']
+            return ['initialize', 'leader', 'members', 'optime', 'failover']
         return ['foo', 'bar', 'buzz']
 
-    def create(self, path, value="", acl=None, ephemeral=False, sequence=False, makepath=False):
+    def create(self, path, value=b"", acl=None, ephemeral=False, sequence=False, makepath=False):
+        if not isinstance(path, six.string_types):
+            raise TypeError("Invalid type for 'path' (string expected)")
+        if not isinstance(value, (six.binary_type,)):
+            raise TypeError("Invalid type for 'value' (must be a byte string)")
         if path.endswith('/initialize') or path == '/service/test/optime/leader':
             raise Exception
-        elif value == 'retry' or (value == 'exists' and self.exists):
+        elif value == b'retry' or (value == b'exists' and self.exists):
             raise NodeExistsError
 
     def set(self, path, value, version=-1):
+        if not isinstance(path, six.string_types):
+            raise TypeError("Invalid type for 'path' (string expected)")
+        if not isinstance(value, (six.binary_type,)):
+            raise TypeError("Invalid type for 'value' (must be a byte string)")
         if path == '/service/bla/optime/leader':
             raise Exception
+        if path == '/service/test/members/bar':
+            if value == b'retry':
+                return
+        if path == '/service/test/failover':
+            if value == b'Exception':
+                raise Exception
+            elif value == b'ok':
+                return
         raise NoNodeError
 
     def delete(self, path, version=-1, recursive=False):
+        if not isinstance(path, six.string_types):
+            raise TypeError("Invalid type for 'path' (string expected)")
         self.exists = False
         if path == '/service/test/leader':
             if self.leader:
                 return
             self.leader = True
             raise Exception
-        elif path.endswith('/initialize'):
+        elif path == '/service/test/members/buzz':
+            raise Exception
+        elif path.endswith('/initialize') or path == '/service/test/members/bar':
             raise NoNodeError
 
-    def set_hosts(self, hosts, randomize_hosts=None):
-        pass
-
-
-def exhibitor_sleep(_):
-    raise SleepException
-
 
+@patch('requests.get', requests_get)
+@patch('patroni.zookeeper.sleep', Mock(side_effect=SleepException()))
 class TestExhibitorEnsembleProvider(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        super(TestExhibitorEnsembleProvider, self).__init__(method_name)
-
-    def set_up(self):
-        requests.get = requests_get
-        patroni.zookeeper.sleep = exhibitor_sleep
-
     def test_init(self):
         self.assertRaises(SleepException, ExhibitorEnsembleProvider, ['localhost'], 8181)
 
 
 class TestZooKeeper(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
+    @patch('requests.get', requests_get)
-        self.setUp = self.set_up
+    @patch('patroni.zookeeper.KazooClient', MockKazooClient)
-        super(TestZooKeeper, self).__init__(method_name)
+    def setUp(self):
-
-    def set_up(self):
-        requests.get = requests_get
-        patroni.zookeeper.KazooClient = MockKazooClient
         self.zk = ZooKeeper('foo', {'exhibitor': {'hosts': ['localhost', 'exhibitor'], 'port': 8181}, 'scope': 'test'})
 
     def test_session_listener(self):
@@ -145,6 +121,8 @@ class TestZooKeeper(unittest.TestCase):
     def test__inner_load_cluster(self):
         self.zk._base_path = self.zk._base_path.replace('test', 'bla')
         self.zk._inner_load_cluster()
+        self.zk._base_path = self.zk._base_path = '/no_node'
+        self.zk._inner_load_cluster()
 
     def test_get_cluster(self):
         self.assertRaises(ZooKeeperError, self.zk.get_cluster)
@@ -154,6 +132,11 @@ class TestZooKeeper(unittest.TestCase):
         self.zk.touch_member('foo')
         self.zk.delete_leader()
 
+    def test_set_failover_value(self):
+        self.zk.set_failover_value('')
+        self.zk.set_failover_value('ok')
+        self.zk.set_failover_value('Exception')
+
     def test_initialize(self):
         self.assertFalse(self.zk.initialize())
 
@@ -161,19 +144,33 @@ class TestZooKeeper(unittest.TestCase):
         self.zk.cancel_initialization()
 
     def test_touch_member(self):
+        self.zk._name = 'buzz'
+        self.zk.get_cluster()
         self.zk.touch_member('new')
+        self.zk._name = 'bar'
+        self.zk.touch_member('new')
+        self.zk._name = 'na'
+        self.zk.client.exists = 1
         self.zk.touch_member('exists')
+        self.zk._name = 'bar'
+        self.zk.touch_member('retry')
+        self.zk.fetch_cluster = True
+        self.zk.get_cluster()
         self.zk.touch_member('retry')
 
     def test_take_leader(self):
         self.zk.take_leader()
 
     def test_update_leader(self):
-        self.zk.last_leader_operation = -1
+        self.assertTrue(self.zk.update_leader())
-        self.assertTrue(self.zk.update_leader(MockPostgresql()))
+
+    def test_write_leader_optime(self):
+        self.zk.last_leader_operation = '0'
+        self.zk.write_leader_optime('1')
         self.zk._base_path = self.zk._base_path.replace('test', 'bla')
-        self.zk.last_leader_operation = -1
+        self.zk.write_leader_optime('2')
-        self.assertTrue(self.zk.update_leader(MockPostgresql()))
 
     def test_watch(self):
         self.zk.watch(0)
+        self.zk.event.isSet = lambda: True
+        self.zk.watch(0)