Restart and reinitialize via api

POST /restart -- will restart postgres
You you are restartung leader node, lock would be maintained during
restart.

POST /reinitialize -- will reinitialize node from the leader.
It's not possible to reinitialize current leader.
Command will fail when the leader is unknown.

patroni/__init__.py

@@ -48,8 +48,6 @@ class Patroni:
             logger.info('waiting on DCS')
             sleep(5)
 
-        self.postgresql.schedule_load_slots = self.postgresql.is_running() and self.postgresql.use_slots
-
     def schedule_next_run(self):
         self.next_run += self.nap_time
         current_time = time.time()
@@ -75,7 +73,7 @@ class Patroni:
 
 
 def main():
-    logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=logging.INFO)
+    logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=logging.DEBUG)
     logging.getLogger('requests').setLevel(logging.WARNING)
     setup_signal_handlers()
 

patroni/api.py

@@ -44,23 +44,23 @@ class RestApiHandler(BaseHTTPRequestHandler):
     def do_GET(self):
         """Default method for processing all GET requests which can not be routed to other methods"""
 
+        path = '/master' if self.path == '/' else self.path
         response = self.get_postgresql_status()
 
-        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
+        patroni = self.server.patroni
+        if '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'))
 
-    @check_auth
-    def do_GET_sampleauth(self):
-        self.send_response(200)
-        self.send_header('Content-Type', 'text/html')
-        self.end_headers()
-        self.wfile.write(b'Hello!')
-
     def do_GET_patroni(self):
         response = self.get_postgresql_status(True)
 
@@ -69,6 +69,51 @@ class RestApiHandler(BaseHTTPRequestHandler):
         self.end_headers()
         self.wfile.write(json.dumps(response).encode('utf-8'))
 
+    @check_auth
+    def do_POST_restart(self):
+        action = self.server.patroni.ha.schedule_restart()
+        if action is not None:
+            status_code = 503
+            data = (action + ' already in progress').encode('utf-8')
+        else:
+            status_code = 503
+            data = b'restart failed'
+            try:
+                if self.server.patroni.ha.restart():
+                    status_code = 200
+                    data = b'restarted successfully'
+            except:
+                logger.exception('Exception during restart')
+
+        self.send_response(status_code)
+        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
 
@@ -104,9 +149,9 @@ class RestApiHandler(BaseHTTPRequestHandler):
                                        pg_last_xlog_replay_location(),
                                        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],
@@ -116,7 +161,7 @@ class RestApiHandler(BaseHTTPRequestHandler):
             }
         except (psycopg2.Error, RetryFailedError, PostgresConnectionException):
             logger.exception('get_postgresql_status')
-            return {'running': self.server.patroni.postgresql.is_running()}
+            return {'state': self.server.patroni.postgresql.state}
 
 
 class RestApiServer(ThreadingMixIn, HTTPServer, Thread):

patroni/dcs.py

@@ -120,14 +120,17 @@ class AbstractDCS:
             running as a master and exception raised instance would be demoted."""
 
     @abc.abstractmethod
-    def update_leader(self, state_handler):
-        """Update leader key (or session) ttl and `/optime/leader` key in DCS.
+    def write_leader_optime(self, last_operation):
+        """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."""
 

patroni/etcd.py

@@ -222,14 +222,12 @@ class Etcd(AbstractDCS):
         return False
 
     @catch_etcd_errors
-    def write_leader_optime(self, state_handler):
-        return self.client.set(self.leader_optime_path, state_handler.last_operation())
+    def write_leader_optime(self, last_operation):
+        return self.client.set(self.leader_optime_path, last_operation)
 
     @catch_etcd_errors
-    def update_leader(self, state_handler):
-        ret = self.retry(self.client.test_and_set, self.leader_path, self._name, self._name, self.ttl)
-        ret and self.write_leader_optime(state_handler)
-        return ret
+    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):

patroni/ha.py

@@ -2,6 +2,7 @@ import logging
 import psycopg2
 
 from patroni.exceptions import DCSError, PostgresConnectionException
+from threading import Lock
 
 logger = logging.getLogger(__name__)
 
@@ -13,6 +14,10 @@ class Ha:
         self.dcs = etcd
         self.cluster = None
         self.old_cluster = None
+        self.scheduled_action = None
+        self.scheduled_action_lock = Lock()
+        self.restart_in_progress = False
+        self.restart_thread_lock = Lock()
 
     def load_cluster_from_dcs(self):
         cluster = self.dcs.get_cluster()
@@ -28,7 +33,13 @@ 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
@@ -37,8 +48,9 @@ class Ha:
 
     def bootstrap(self):
         if not self.cluster.is_unlocked():  # cluster already has leader
-            logger.info('trying to bootstrap from leader', )
+            logger.info('trying to bootstrap from leader')
             if self.state_handler.bootstrap(self.cluster.leader):
+                self.reinitialize_scheduled() and self.reset_scheduled_action()
                 return 'bootstrapped from leader'
             else:
                 self.state_handler.stop('immediate')
@@ -63,15 +75,17 @@ class Ha:
             return 'waiting for leader to bootstrap'
 
     def recover(self):
-        if self.state_handler.is_healthy():
-            return False
         has_lock = self.has_lock()
         self.state_handler.write_recovery_conf(None if has_lock else self.cluster.leader)
-        self.state_handler.start()
-        if has_lock:
-            logger.info('started as readonly because i had the session lock')
-            self.load_cluster_from_dcs()
-        return True
+        if not self.state_handler.start():
+            if not has_lock:
+                return 'failed to start postgres'
+            self.dcs.delete_leader()
+            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()
@@ -112,7 +126,68 @@ class Ha:
         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 run_cycle(self):
+    def schedule_action(self, action):
+        with self.scheduled_action_lock:
+            if self.scheduled_action is not None:
+                return self.scheduled_action
+            self.scheduled_action = action
+        return None
+
+    def get_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 schedule_restart(self):
+        return self.schedule_action('restart')
+
+    def restart_scheduled(self):
+        return self.get_scheduled_action() == 'restart'
+
+    def schedule_reinitialize(self):
+        return self.schedule_action('reinitialize')
+
+    def reinitialize_scheduled(self):
+        return self.get_scheduled_action() == 'reinitialize'
+
+    def restart(self):
+        with self.restart_thread_lock:
+            self.restart_in_progress = True
+        try:
+            return self.state_handler.restart()
+        finally:
+            with self.restart_thread_lock:
+                self.restart_in_progress = False
+            self.reset_scheduled_action()
+
+    def process_scheduled_action(self):
+        if self.reinitialize_scheduled():
+            if self.cluster.is_unlocked():
+                logger.error('Cluster has no leader, can not reinitialize')
+                self.reset_scheduled_action()
+            elif self.has_lock():
+                logger.error('I am the leader, can not reinitialize')
+                self.reset_scheduled_action()
+            else:
+                self.state_handler.stop('immediate')
+                self.state_handler.remove_data_directory()
+                self.load_cluster_from_dcs()
+
+    def handle_restart_in_progress(self):
+        if self.has_lock():
+            if self.update_lock():
+                return 'updated leader lock during restart'
+            else:
+                return 'failed to update leader lock during restart'
+        elif self.cluster.is_unlocked():
+            return 'not healthy enough for leader race'
+        else:
+            return 'restart in progress'
+
+    def _run_cycle(self):
         try:
             self.load_cluster_from_dcs()
 
@@ -120,6 +195,9 @@ class Ha:
             if not self.cluster.is_unlocked() and not self.cluster.initialize:
                 self.dcs.initialize()  # fix it
 
+            # currently it can trigger only reinitialize
+            self.process_scheduled_action()
+
             # is data directory empty?
             if self.state_handler.data_directory_empty():
                 return self.bootstrap()  # new node
@@ -127,10 +205,14 @@ class Ha:
             elif not self.cluster.initialize and self.cluster.is_unlocked():
                 self.dcs.initialize()
 
+            if self.restart_in_progress:
+                return self.handle_restart_in_progress()
+
             # try to start dead postgres
-            if self.recover() and not self.has_lock():
-                # no lock, do not try to promote immediately
-                return 'started as a secondary'
+            if not self.state_handler.is_healthy():
+                msg = self.recover()
+                if msg is not None:
+                    return msg
 
             if self.cluster.is_unlocked():
                 return self.process_unhealthy_cluster()
@@ -143,3 +225,7 @@ class Ha:
                 return 'demoted self because DCS is not accessible and i was a leader'
         except (psycopg2.Error, PostgresConnectionException):
             logger.exception('Error communicating with Postgresql.  Will try again')
+
+    def run_cycle(self):
+        with self.restart_thread_lock:
+            return self._run_cycle()

patroni/postgresql.py

@@ -56,7 +56,6 @@ class Postgresql:
         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._role = 'replica'
 
         self._pg_ctl = ['pg_ctl', '-w', '-D', self.data_dir]
 
@@ -67,9 +66,16 @@ class Postgresql:
 
         self._connection = None
         self._cursor_holder = None
-        self.members = []  # list of already existing replication slots
+        self.replication_slots = []  # list of already existing replication slots
         self.retry = Retry(max_tries=-1, deadline=10, max_delay=1, retry_exceptions=PostgresConnectionException)
 
+        self._state = 'stopped'
+        self._role = 'replica'
+
+        if self.is_running():
+            self._state = 'running'
+            self._role = 'master' if self.is_leader() else 'replica'
+
     def get_local_address(self):
         listen_addresses = self.listen_addresses.split(',')
         local_address = listen_addresses[0].strip()  # take first address from listen_addresses
@@ -101,6 +107,8 @@ class Postgresql:
         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):
@@ -113,8 +121,12 @@ class Postgresql:
         return not os.path.exists(self.data_dir) or os.listdir(self.data_dir) == []
 
     def initialize(self):
+        self._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._state = 'initdb failed'
         return ret
 
     def delete_trigger_file(self):
@@ -137,6 +149,7 @@ class Postgresql:
         return "host={host} port={port} user={user}".format(**conn)
 
     def create_replica(self, master_connection, env):
+        self._state = 'building replica from {host}:{port}'.format(**master_connection)
         connstring = self.build_connstring(master_connection)
         cmd = self.config['restore']
         try:
@@ -144,7 +157,9 @@ class Postgresql:
             self.delete_trigger_file()
         except:
             logger.exception('Error when creating replica')
-            return 1
+            ret = 1
+        if ret != 0:
+            self._state = 'failed to build replica from {host}:{port}'.format(**master_connection)
         return ret
 
     def is_leader(self):
@@ -169,38 +184,60 @@ class Postgresql:
     def role(self):
         return self._role
 
+    @property
+    def state(self):
+        return self._state
+
     def start(self, block_callbacks=False):
         if self.is_running():
-            self._role = 'master' if self.is_leader() else 'replica'
-            self.schedule_load_slots = self.use_slots
             logger.error('Cannot start PostgreSQL because one is already running.')
-            return False
+            return True
 
         self._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._state = 'starting'
+
         ret = subprocess.call(self._pg_ctl + ['start', '-o', self.server_options()]) == 0
+
+        self._state = 'running' if ret else 'start failed'
+
         self.schedule_load_slots = ret and self.use_slots
         self.save_configuration_files()
         # block_callbacks is used during restart to avoid
         # running start/stop callbacks in addition to restart ones
-        ret and not block_callbacks and ret and self.call_nowait(ACTION_ON_START)
+        ret and not block_callbacks and self.call_nowait(ACTION_ON_START)
         return ret
 
+    def checkpoint(self):
+        try:
+            self.query('SET statement_timeout TO 0')
+            self.query('CHECKPOINT')
+        except:
+            logging.exception('Exception diring CHECKPOINT')
+
     def stop(self, mode='fast', block_callbacks=False):
+        if not self.is_running():
+            if not block_callbacks:
+                self._state = 'stopped'
+            return True
+
         if block_callbacks:
-            try:
-                self.query('SET statement_timeout TO 0')
-                self.query('CHECKPOINT')
-            except:
-                logging.exception('Exception diring CHECKPOINT')
+            self.checkpoint()
+        else:
+            self._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
-        ret and not block_callbacks and self.call_nowait(ACTION_ON_STOP)
+        if not ret:
+            self._state = 'stop failed'
+        elif not block_callbacks:
+            self._state = 'stopped'
+            self.call_nowait(ACTION_ON_STOP)
         return ret
 
     def reload(self):
@@ -209,8 +246,12 @@ class Postgresql:
         return ret
 
     def restart(self):
+        self._state = 'restarting'
         ret = self.stop(block_callbacks=True) and self.start(block_callbacks=True)
-        ret and self.call_nowait(ACTION_ON_RESTART)
+        if ret:
+            self.call_nowait(ACTION_ON_RESTART)
+        else:
+            self._state = 'restart failed ({})'.format(self._state)
         return ret
 
     def server_options(self):
@@ -356,26 +397,26 @@ recovery_target_timeline = 'latest'
     def load_replication_slots(self):
         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, cluster):
         if self.use_slots:
             self.load_replication_slots()
-            members = [m.name for m in cluster.members if m.name != self.name] if self.role == 'master' else []
+            slots = [m.name for m in cluster.members if m.name != self.name] if self.role == 'master' else []
             # drop unused slots
-            for slot in set(self.members) - set(members):
+            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
-            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)
                                WHERE NOT EXISTS (SELECT 1 FROM pg_replication_slots
                                WHERE slot_name = %s)""", slot, slot)
 
-            self.members = members
+            self.replication_slots = slots
 
     def last_operation(self):
         return str(self.xlog_position())

patroni/zookeeper.py

@@ -211,8 +211,8 @@ class ZooKeeper(AbstractDCS):
     def take_leader(self):
         return self.attempt_to_acquire_leader()
 
-    def update_leader(self, state_handler):
-        last_operation = state_handler.last_operation().encode('utf-8')
+    def write_leader_optime(self, 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
@@ -225,6 +225,8 @@ 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):

tests/test_api.py

@@ -10,6 +10,10 @@ from test_postgresql import psycopg2_connect, MockCursor
 
 class MockPostgresql(Mock):
 
+    name = 'test'
+    state = 'running'
+    role = 'master'
+
     def connection(self):
         return psycopg2_connect()
 
@@ -17,9 +21,28 @@ class MockPostgresql(Mock):
         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:
 
     postgresql = MockPostgresql()
+    ha = MockHa()
 
 
 class MockRequest:
@@ -47,18 +70,38 @@ class MockRestApiServer(RestApiServer):
 class TestRestApiHandler(unittest.TestCase):
 
     def test_do_GET(self):
-        MockRestApiServer(RestApiHandler, b'GET /')
-        with patch.object(RestApiServer, 'query', Mock(side_effect=psycopg2.OperationalError())):
-            MockRestApiServer(RestApiHandler, b'GET /')
-
-    def test_do_GET_sampleauth(self):
-        MockRestApiServer(RestApiHandler, b'GET /sampleauth')
-        MockRestApiServer(RestApiHandler, b'GET /sampleauth\nAuthorization:')
-        MockRestApiServer(RestApiHandler, b'GET /sampleauth\nAuthorization: Basic dGVzdDp0ZXN0')
+        MockRestApiServer(RestApiHandler, b'GET /master')
+        MockRestApiServer(RestApiHandler, b'GET /replica')
+        with patch.object(MockHa, 'restart_scheduled', Mock(return_value=True)):
+            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, 'schedule_restart', Mock(return_value=None)):
+            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)):

tests/test_etcd.py

@@ -242,8 +242,11 @@ class TestEtcd(unittest.TestCase):
         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())

tests/test_ha.py

@@ -82,10 +82,25 @@ class TestHa(unittest.TestCase):
         self.e.get_cluster = get_cluster_not_initialized_without_leader
         ha.load_cluster_from_dcs()
 
-    def test_start_as_slave(self):
+    def test_update_lock(self):
+        self.p.last_operation = Mock(side_effect=PostgresException(''))
+        self.assertTrue(self.ha.update_lock())
+
+    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.is_healthy = false
+        self.p.start = false
+        self.assertEquals(self.ha.run_cycle(), 'failed to start postgres')
+
+    def test_recover_master_failed(self):
+        self.p.is_healthy = false
+        self.p.start = false
+        self.ha.has_lock = true
+        self.assertEquals(self.ha.run_cycle(), 'removed leader key after trying and failing to start postgres')
+
     def test_start_as_readonly(self):
         self.ha.cluster.is_unlocked = false
         self.p.is_leader = self.p.is_healthy = false
@@ -174,3 +189,35 @@ class TestHa(unittest.TestCase):
         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.run_cycle()
+        self.assertIsNone(self.ha.get_scheduled_action())
+
+        self.ha.cluster = get_cluster_initialized_with_leader()
+        self.ha.schedule_reinitialize()
+        self.ha.run_cycle()
+
+        self.ha.has_lock = true
+        self.ha.schedule_reinitialize()
+        self.ha.run_cycle()
+        self.assertIsNone(self.ha.get_scheduled_action())
+
+    def test_restart(self):
+        self.ha.schedule_restart()
+        self.assertTrue(self.ha.restart_scheduled())
+        self.ha.restart()
+
+    def test_restart_in_progress(self):
+        self.ha.restart_in_progress = True
+        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')

tests/test_postgresql.py

@@ -85,10 +85,12 @@ def psycopg2_connect(*args, **kwargs):
 
 
 @patch('subprocess.call', Mock(return_value=0))
-@patch('shutil.copy', Mock())
 @patch('psycopg2.connect', psycopg2_connect)
+@patch('shutil.copy', Mock())
 class TestPostgresql(unittest.TestCase):
 
+    @patch('subprocess.call', Mock(return_value=0))
+    @patch('psycopg2.connect', psycopg2_connect)
     def setUp(self):
         self.p = Postgresql({'name': 'test0', 'scope': 'batman', 'data_dir': 'data/test0',
                              'listen': '127.0.0.1, *:5432', 'connect_address': '127.0.0.2:5432',
@@ -121,13 +123,24 @@ 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):
-        self.assertFalse(self.p.start())
-        self.p.is_running = false
-        with open(os.path.join(self.p.data_dir, 'postmaster.pid'), 'w'):
-            pass
+    def test_start(self):
         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))
@@ -157,6 +170,8 @@ class TestPostgresql(unittest.TestCase):
     @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')
@@ -184,6 +199,7 @@ class TestPostgresql(unittest.TestCase):
         self.assertFalse(self.p.is_healthy())
 
     def test_promote(self):
+        self.p._role = 'replica'
         self.assertTrue(self.p.promote())
         self.assertTrue(self.p.promote())
 
@@ -211,8 +227,8 @@ class TestPostgresql(unittest.TestCase):
             self.p.move_data_directory()
 
     def test_bootstrap(self):
-        self.assertRaises(PostgresException, self.p.bootstrap)
-        self.p.start = Mock(return_value=True)
+        with patch('subprocess.call', Mock(return_value=1)):
+            self.assertRaises(PostgresException, self.p.bootstrap)
         self.p.bootstrap()
 
     def test_remove_data_directory(self):

tests/test_zookeeper.py

@@ -134,11 +134,13 @@ class TestZooKeeper(unittest.TestCase):
         self.zk.take_leader()
 
     def test_update_leader(self):
-        self.zk.last_leader_operation = -1
-        self.assertTrue(self.zk.update_leader(MockPostgresql()))
+        self.assertTrue(self.zk.update_leader())
+
+    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.assertTrue(self.zk.update_leader(MockPostgresql()))
+        self.zk.write_leader_optime('2')
 
     def test_watch(self):
         self.zk.watch(0)