Run long time operations asynchronously

i.e. restart, reinitialize, demote

patroni/__init__.py

@@ -56,12 +56,8 @@ class Patroni:
         self.next_run = time.time()
 
         while True:
-            logger.info(self.ha.run_cycle())
-            try:
-                self.ha.cluster and self.ha.state_handler.sync_replication_slots(self.ha.cluster)
-            except:
-                logger.exception('Exception when changing replication slots')
             self.touch_member()
+            logger.info(self.ha.run_cycle())
             reap_children()
             self.schedule_next_run()
 

patroni/api.py

@@ -71,19 +71,14 @@ class RestApiHandler(BaseHTTPRequestHandler):
 
     @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')
+        status_code = 503
+        data = b'restart failed'
+        try:
+            status, msg = self.server.patroni.ha.restart()
+            status_code = 200 if status else 503
+            data = msg.encode('utf-8')
+        except:
+            logger.exception('Exception during restart')
 
         self.send_response(status_code)
         self.send_header('Content-Type', 'text/html')
@@ -135,7 +130,7 @@ class RestApiHandler(BaseHTTPRequestHandler):
     def query(self, sql, *params, **kwargs):
         if not kwargs.get('retry', False):
             return self.server.query(sql, *params)
-        retry = Retry(delay=2, retry_exceptions=PostgresConnectionException)
+        retry = Retry(delay=1, retry_exceptions=PostgresConnectionException)
         return retry(self.server.query, sql, *params)
 
     def get_postgresql_status(self, retry=False):
@@ -164,6 +159,7 @@ class RestApiHandler(BaseHTTPRequestHandler):
             state = self.server.patroni.postgresql.state
             if state in ['stopped', 'starting', 'stopping', 'restarting', 'running']:
                 logger.exception('get_postgresql_status')
+                state = 'unknown' if state == 'running' else state
             return {'state': state}
 
 

patroni/ha.py

@@ -4,7 +4,7 @@ import requests
 
 from patroni.exceptions import DCSError, PostgresConnectionException
 from multiprocessing.pool import ThreadPool
-from threading import Lock
+from threading import Lock, Thread
 
 logger = logging.getLogger(__name__)
 
@@ -16,10 +16,10 @@ class Ha:
         self.dcs = dcs
         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()
+        self._scheduled_action = None
+        self._scheduled_action_lock = Lock()
+        self._long_action_in_progress = False
+        self._long_action_thread_lock = Lock()
 
     def load_cluster_from_dcs(self):
         cluster = self.dcs.get_cluster()
@@ -48,16 +48,37 @@ class Ha:
         logger.info('Lock owner: %s; I am %s', lock_owner, self.state_handler.name)
         return lock_owner == self.state_handler.name
 
+    def _run_async(self, func, args=()):
+        try:
+            return func(*args) if args else func()
+        except:
+            logger.exception('Exception during execution of long running task %s', self.get_scheduled_action())
+        finally:
+            with self._long_action_thread_lock:
+                self._long_action_in_progress = False
+                self._reset_scheduled_action()
+
+    def run_async(self, func, args=()):
+        self._long_action_in_progress = True
+        Thread(target=self._run_async, args=(func, args)).start()
+
+    def copy_backup_from_leader(self):
+        if self.state_handler.bootstrap(self.cluster.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 bootstrap(self):
         if not self.cluster.is_unlocked():  # cluster already has 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'
+            if self._long_action_in_progress:
+                self.copy_backup_from_leader()
             else:
-                self.state_handler.stop('immediate')
-                self.state_handler.remove_data_directory()
-                return 'failed to bootstrap from leader'
+                with self._scheduled_action_lock:
+                    self._scheduled_action = 'bootstrap from leader'
+                self.run_async(self.copy_backup_from_leader)
+            return 'trying to bootstrap from leader'
         elif not self.cluster.initialize:  # no initialize key
             if self.dcs.initialize():  # race for initialization
                 try:
@@ -92,7 +113,10 @@ class Ha:
     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
-        self.state_handler.follow_the_leader(self.cluster.leader)
+        if not self.state_handler.check_recovery_conf(self.cluster.leader):
+            with self._scheduled_action_lock:
+                self._scheduled_action = 'changing primary_conninfo and restarting'
+            self.run_async(self.state_handler.follow_the_leader, (self.cluster.leader, ))
         return ret
 
     def enforce_master_role(self, message, promote_message):
@@ -208,20 +232,22 @@ class Ha:
         members = {m.name: m for m in self.old_cluster.members + self.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.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.state_handler.stop()
-                    if self.dcs.delete_leader():
-                        ret = 'manual failover: demoted self and released leader lock'
-                    else:
-                        ret = 'manual failover: demoted self but failed to release leader lock'
-                    self.state_handler.follow_the_leader(None)
-                    self.cluster = None
-                    return ret
+                    with self._scheduled_action_lock:
+                        self._scheduled_action = 'manual failover: demote'
+                    self.run_async(self.demote)
+                    return 'manual failover: demoting myself'
                 else:
                     logger.warning('manual failover: no healthy members found, failover is not possible')
             else:
@@ -268,22 +294,20 @@ class Ha:
                                       'no action.  i am a secondary and i am following a leader', False)
 
     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
+        with self._long_action_thread_lock:
+            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
+        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 _reset_scheduled_action(self):
+        with self._scheduled_action_lock:
+            self._scheduled_action = None
 
     def restart_scheduled(self):
         return self.get_scheduled_action() == 'restart'
@@ -295,38 +319,45 @@ class Ha:
         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()
+        with self._long_action_thread_lock:
+            with self._scheduled_action_lock:
+                if self._scheduled_action is not None:
+                    return False, self._scheduled_action + ' already in progress'
+                self._scheduled_action = 'restart'
+            self._long_action_in_progress = True
+        if self._run_async(self.state_handler.restart):
+            return True, 'restarted successfully'
+        else:
+            return False, 'restart failed'
+
+    def reinitialize(self):
+        self.state_handler.stop('immediate')
+        self.state_handler.remove_data_directory()
+        self.load_cluster_from_dcs()
+        self.bootstrap()
 
     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()
+                self._reset_scheduled_action()
             elif self.has_lock():
                 logger.error('I am the leader, can not reinitialize')
-                self.reset_scheduled_action()
+                self._reset_scheduled_action()
             else:
-                self.state_handler.stop('immediate')
-                self.state_handler.remove_data_directory()
-                self.load_cluster_from_dcs()
+                self.run_async(self.reinitialize)
+                return True
 
-    def handle_restart_in_progress(self):
+    def handle_long_action_in_progress(self):
         if self.has_lock():
             if self.update_lock():
-                return 'updated leader lock during restart'
+                return 'updated leader lock during ' + self.get_scheduled_action()
             else:
-                return 'failed to update leader lock during restart'
+                return 'failed to update leader lock during ' + self.get_scheduled_action()
         elif self.cluster.is_unlocked():
             return 'not healthy enough for leader race'
         else:
-            return 'restart in progress'
+            return self.get_scheduled_action() + ' in progress'
 
     def _run_cycle(self):
         try:
@@ -336,8 +367,12 @@ class Ha:
             if not self.cluster.is_unlocked() and not self.cluster.initialize:
                 self.dcs.initialize()  # fix it
 
+            if self._long_action_in_progress:
+                return self.handle_long_action_in_progress()
+
             # currently it can trigger only reinitialize
-            self.process_scheduled_action()
+            if self.process_scheduled_action():
+                return 'reinitialize started'
 
             # is data directory empty?
             if self.state_handler.data_directory_empty():
@@ -346,27 +381,27 @@ 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 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()
-            else:
-                return self.process_healthy_cluster()
+            try:
+                if self.cluster.is_unlocked():
+                    return self.process_unhealthy_cluster()
+                else:
+                    return self.process_healthy_cluster()
+            finally:
+                self.state_handler.sync_replication_slots(self.cluster)
         except DCSError:
             logger.error('Error communicating with DCS')
             if self.state_handler.is_running() and self.state_handler.is_leader():
-                self.state_handler.demote()
+                self.demote(delete_leader=False)
                 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')
+            logger.exception('Error communicating with Postgresql. Will try again later')
 
     def run_cycle(self):
-        with self.restart_thread_lock:
+        with self._long_action_thread_lock:
             return self._run_cycle()

patroni/postgresql.py

@@ -227,10 +227,14 @@ class Postgresql:
 
     def checkpoint(self):
         try:
-            self.query('SET statement_timeout TO 0')
-            self.query('CHECKPOINT')
+            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:
-            logging.exception('Exception diring CHECKPOINT')
+            logging.exception('Exception during CHECKPOINT')
 
     def stop(self, mode='fast', block_callbacks=False):
         if not self.is_running():

patroni/zookeeper.py

@@ -225,6 +225,8 @@ class ZooKeeper(AbstractDCS):
                 return True
             except:
                 logger.exception('touch_member')
+        except:
+            logger.exception('touch_member')
         return False
 
     def take_leader(self):

tests/test_api.py

@@ -33,7 +33,7 @@ class MockHa(Mock):
         return 'reinitialize'
 
     def restart(self):
-        return True
+        return (True, '')
 
     def restart_scheduled(self):
         return False
@@ -85,10 +85,8 @@ class TestRestApiHandler(unittest.TestCase):
     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)):
+        with patch.object(MockHa, 'restart', Mock(side_effect=Exception)):
             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):

tests/test_ha.py

@@ -6,6 +6,7 @@ from patroni.etcd import Client, Etcd
 from patroni.exceptions import DCSError, PostgresException
 from patroni.ha import Ha
 from test_etcd import socket_getaddrinfo, etcd_read, etcd_write, requests_get
+from threading import Thread
 
 
 def true(*args, **kwargs):
@@ -69,6 +70,13 @@ class MockPostgresql(Mock):
     def check_replication_lag(self, last_leader_operation):
         return True
 
+    def check_recovery_conf(self, leader):
+        return False
+
+
+def run_async(func, args=()):
+    func(args) if args else func()
+
 
 class TestHa(unittest.TestCase):
 
@@ -81,6 +89,7 @@ class TestHa(unittest.TestCase):
         self.e.client.read = etcd_read
         self.e.client.write = etcd_write
         self.ha = Ha(self.p, self.e)
+        self.ha.run_async = run_async
         self.ha.load_cluster_from_dcs()
         self.ha.cluster = get_cluster_not_initialized_without_leader()
         self.ha.load_cluster_from_dcs = Mock()
@@ -173,14 +182,20 @@ class TestHa(unittest.TestCase):
         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__run_async(self):
+        self.ha._run_async(Mock(side_effect=Exception()))        
+
+    @patch.object(Thread, 'start', Mock())
+    def test_run_async(self):
+        ha = Ha(self.p, self.e)
+        ha.run_async(true)
+
     def test_bootstrap_from_leader(self):
-        self.ha.cluster = get_cluster_initialized_with_leader()
-        self.assertEquals(self.ha.bootstrap(), 'bootstrapped from leader')
-
-    def test_bootstrap_from_leader_failed(self):
         self.ha.cluster = get_cluster_initialized_with_leader()
         self.p.bootstrap = false
-        self.assertEquals(self.ha.bootstrap(), 'failed to bootstrap from leader')
+        self.assertEquals(self.ha.bootstrap(), 'trying to bootstrap from leader')
+        self.ha._long_action_in_progress = True
+        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()
@@ -202,26 +217,32 @@ class TestHa(unittest.TestCase):
         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.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())
 
+        self.ha.has_lock = false
+        self.ha.schedule_reinitialize()
+        self.ha.run_cycle()
+
     def test_restart(self):
-        self.ha.schedule_restart()
-        self.assertTrue(self.ha.restart_scheduled())
-        self.ha.restart()
+        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.restart_in_progress = True
+        self.ha._long_action_in_progress = True
+        self.ha._scheduled_action = 'restart'
+        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()
@@ -244,10 +265,7 @@ class TestHa(unittest.TestCase):
         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: demoted self but failed to release leader lock')
-        self.ha.cluster = get_cluster_initialized_with_leader(f)
-        self.e.client.delete = Mock(return_value=True)
-        self.assertEquals(self.ha.run_cycle(), 'manual failover: demoted self and released leader lock')
+        self.assertEquals(self.ha.run_cycle(), 'manual failover: demoting myself')
 
     @patch('requests.get', requests_get)
     def test_manual_failover_process_no_leader(self):

tests/test_patroni.py

@@ -8,6 +8,7 @@ from mock import Mock, patch
 from patroni.api import RestApiServer
 from patroni.dcs import Cluster, Member
 from patroni.etcd import Etcd
+from patroni.ha import Ha
 from patroni import Patroni, main
 from patroni.zookeeper import ZooKeeper
 from six.moves import BaseHTTPServer
@@ -26,6 +27,7 @@ def time_sleep(*args):
 @patch.object(Postgresql, 'write_pg_hba', Mock())
 @patch.object(Postgresql, 'write_recovery_conf', Mock())
 @patch.object(BaseHTTPServer.HTTPServer, '__init__', Mock())
+@patch.object(Ha, 'run_async', Mock())
 class TestPatroni(unittest.TestCase):
 
     @patch.object(Client, 'machines')

tests/test_postgresql.py

@@ -70,6 +70,12 @@ class MockConnect(Mock):
     def cursor(self):
         return MockCursor(self)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
 
 def psycopg2_connect(*args, **kwargs):
     return MockConnect()
@@ -214,6 +220,7 @@ class TestPostgresql(unittest.TestCase):
         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'