Merge pull request #274 from zalando/feature/disable-automatic-failover

Feature/disable automatic failover

features/patroni_api.feature

@@ -34,20 +34,22 @@ Scenario: check local configuration reload
 	Then I receive a response code 202
 
 Scenario: check dynamic configuration change via DCS
-	Given I issue a PATCH request to http://127.0.0.1:8008/config with {"ttl": 20, "loop_wait": 1, "postgresql": {"parameters": {"max_connections": 101}}}
+	Given I issue a PATCH request to http://127.0.0.1:8008/config with {"ttl": 20, "loop_wait": 2, "postgresql": {"parameters": {"max_connections": 101}}}
 	Then I receive a response code 200
-	And I receive a response loop_wait 1
+	And I receive a response loop_wait 2
 	And Response on GET http://127.0.0.1:8008/patroni contains pending_restart after 11 seconds
 	When I issue a GET request to http://127.0.0.1:8008/config
 	Then I receive a response code 200
-	And I receive a response loop_wait 1
+	And I receive a response loop_wait 2
 	When I issue a GET request to http://127.0.0.1:8008/patroni
 	Then I receive a response code 200
 	And I receive a response tags {'tag': 'new_value'}
 
-Scenario: check API requests for the primary-replica pair
-	Given I start postgres1
-	And replication works from postgres0 to postgres1 after 20 seconds
+Scenario: check API requests for the primary-replica pair in the pause mode
+	Given I run patronictl.py pause batman
+	Then I receive a response returncode 0
+	When I start postgres1
+	Then replication works from postgres0 to postgres1 after 20 seconds
 	When I issue a GET request to http://127.0.0.1:8009/replica
 	Then I receive a response code 200
 	And I receive a response state running
@@ -62,19 +64,24 @@ Scenario: check API requests for the primary-replica pair
 	When I sleep for 10 seconds
 	Then postgres1 role is the secondary after 15 seconds
 
-Scenario: check the failover via the API
+Scenario: check the failover via the API in the pause mode
 	Given I run patronictl.py failover batman --master postgres0 --candidate postgres1 --force
 	Then I receive a response returncode 0
 	And postgres1 is a leader after 5 seconds
-	And postgres1 role is the primary after 5 seconds
+	And postgres1 role is the primary after 10 seconds
 	And postgres0 role is the secondary after 10 seconds
 	And replication works from postgres1 to postgres0 after 20 seconds
 
 Scenario: check the scheduled failover
+	Given I issue a scheduled failover from postgres1 to postgres0 in 1 seconds
+	Then I receive a response returncode 1
+	And I receive a response output "Can't schedule failover in the paused state"
+	When I run patronictl.py resume batman
+	Then I receive a response returncode 0
 	Given I issue a scheduled failover from postgres1 to postgres0 in 1 seconds
 	Then I receive a response returncode 0
 	And postgres0 is a leader after 20 seconds
-	And postgres0 role is the primary after 5 seconds
+	And postgres0 role is the primary after 10 seconds
 	And postgres1 role is the secondary after 10 seconds
 	And replication works from postgres0 to postgres1 after 25 seconds
 
@@ -84,7 +91,7 @@ Scenario: check the scheduled restart
 		And Response on GET http://127.0.0.1:8008/patroni contains pending_restart after 5 seconds
 	Given I issue a scheduled restart at http://127.0.0.1:8008 in 1 seconds with {"role": "replica"}
 	Then I receive a response code 202
-		And I sleep for 10 seconds
+		And I sleep for 2 seconds
 		And Response on GET http://127.0.0.1:8008/patroni contains pending_restart after 10 seconds
 	Given I issue a scheduled restart at http://127.0.0.1:8008 in 1 seconds with {"restart_pending": "True"}
 	Then I receive a response code 202

patroni/__init__.py

@@ -38,9 +38,11 @@ class Patroni(object):
             try:
                 cluster = self.dcs.get_cluster()
                 if cluster and cluster.config:
-                    self.config.set_dynamic_configuration(cluster.config)
+                    if self.config.set_dynamic_configuration(cluster.config):
+                        self.dcs.reload_config(self.config)
                 elif not self.config.dynamic_configuration and 'bootstrap' in self.config:
-                    self.config.set_dynamic_configuration(self.config['bootstrap']['dcs'])
+                    if self.config.set_dynamic_configuration(self.config['bootstrap']['dcs']):
+                        self.dcs.reload_config(self.config)
                 break
             except DCSError:
                 logger.warning('Can not get cluster from dcs')
@@ -51,7 +53,7 @@ class Patroni(object):
 
     @property
     def nofailover(self):
-        return self.tags.get('nofailover', False)
+        return bool(self.tags.get('nofailover', False))
 
     def reload_config(self):
         try:
@@ -76,7 +78,7 @@ class Patroni(object):
 
     @property
     def noloadbalance(self):
-        return self.tags.get('noloadbalance', False)
+        return bool(self.tags.get('noloadbalance', False))
 
     def schedule_next_run(self):
         self.next_run += self.dcs.loop_wait
@@ -132,5 +134,8 @@ def main():
         pass
     finally:
         patroni.api.shutdown()
-        patroni.postgresql.stop(checkpoint=False)
-        patroni.dcs.delete_leader()
+        if patroni.ha.is_paused():
+            logger.info('Leader key is not deleted and Postgresql is not stopped due paused state')
+        else:
+            patroni.postgresql.stop(checkpoint=False)
+            patroni.dcs.delete_leader()

patroni/api.py

@@ -194,11 +194,17 @@ class RestApiHandler(BaseHTTPRequestHandler):
         status_code = 500
         data = 'restart failed'
         request = self._read_json_content(body_is_optional=True)
+        cluster = self.server.patroni.dcs.get_cluster()
         if request is None:
             # failed to parse the json
             return
         if request:
             logger.debug("received restart request: {0}".format(request))
+
+        if cluster.is_paused() and 'schedule' in request:
+            self._write_response(status_code, "Can't schedule restart in the paused state")
+            return
+
         for k in request:
             if k == 'schedule':
                 (_, data, request[k]) = self.parse_schedule(request[k], "restart")
@@ -249,22 +255,12 @@ class RestApiHandler(BaseHTTPRequestHandler):
 
     @check_auth
     def do_POST_reinitialize(self):
-        patroni = self.server.patroni
-        cluster = patroni.dcs.get_cluster()
-        if cluster.is_unlocked():
-            status_code = 503
-            data = 'Cluster has no leader, can not reinitialize'
-        elif cluster.leader.name == patroni.ha.state_handler.name:
-            status_code = 503
-            data = 'I am the leader, can not reinitialize'
+        data = self.server.patroni.ha.reinitialize()
+        if data is None:
+            status_code = 200
+            data = 'reinitialize started'
         else:
-            action = patroni.ha.schedule_reinitialize()
-            if action is not None:
-                status_code = 503
-                data = action + ' already in progress'
-            else:
-                status_code = 200
-                data = 'reinitialize scheduled'
+            status_code = 503
         self._write_response(status_code, data)
 
     def poll_failover_result(self, leader, candidate):
@@ -285,7 +281,7 @@ class RestApiHandler(BaseHTTPRequestHandler):
         return 503, 'Failover status unknown'
 
     def is_failover_possible(self, cluster, leader, candidate):
-        if leader and not cluster.leader or cluster.leader.name != leader:
+        if leader and (not cluster.leader or cluster.leader.name != leader):
             return 'leader name does not match'
         if candidate:
             members = [m for m in cluster.members if m.name == candidate]
@@ -303,6 +299,7 @@ class RestApiHandler(BaseHTTPRequestHandler):
     @check_auth
     def do_POST_failover(self):
         request = self._read_json_content()
+        status_code = 500
         if not request:
             return
 
@@ -310,7 +307,9 @@ class RestApiHandler(BaseHTTPRequestHandler):
         candidate = request.get('candidate') or request.get('member')
         scheduled_at = request.get('scheduled_at')
         cluster = self.server.patroni.dcs.get_cluster()
-        status_code = 500
+
+        if scheduled_at and cluster.is_paused():
+            self._write_response(status_code, "Can't schedule failover in the paused state")
 
         logger.info("received failover request with leader=%s candidate=%s scheduled_at=%s",
                     leader, candidate, scheduled_at)

patroni/async_executor.py

@@ -7,21 +7,19 @@ logger = logging.getLogger(__name__)
 class AsyncExecutor(object):
 
     def __init__(self):
-        self._busy = False
         self._thread_lock = RLock()
         self._scheduled_action = None
         self._scheduled_action_lock = RLock()
 
     @property
     def busy(self):
-        return self._busy
+        return self.scheduled_action is not None
 
     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
@@ -32,7 +30,6 @@ class AsyncExecutor(object):
     def reset_scheduled_action(self):
         with self._scheduled_action_lock:
             self._scheduled_action = None
-            self._busy = False
 
     def run(self, func, args=()):
         try:
@@ -41,11 +38,9 @@ class AsyncExecutor(object):
             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):

patroni/ctl.py

@@ -499,6 +499,8 @@ def restart(cluster_name, member_names, config_file, dcs, force, role, p_any, sc
 
     scheduled_at = parse_scheduled(scheduled)
     if scheduled_at:
+        if cluster.is_paused():
+            raise PatroniCtlException("Can't schedule restart in the paused state")
         content['schedule'] = scheduled_at.isoformat()
 
     for member in members:
@@ -554,16 +556,16 @@ def failover(config_file, cluster_name, master, candidate, force, dcs, scheduled
 
     config, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
 
-    if cluster.leader is None:
+    if cluster.leader is None and not cluster.is_paused():
         raise PatroniCtlException('This cluster has no master')
 
-    if master is None:
+    if master is None and (not cluster.is_paused() or cluster.leader):
         if force:
             master = cluster.leader.member.name
         else:
             master = click.prompt('Master', type=str, default=cluster.leader.member.name)
 
-    if cluster.leader.member.name != master:
+    if master is not None and cluster.leader and cluster.leader.member.name != master:
         raise PatroniCtlException('Member {0} is not the leader of cluster {1}'.format(master, cluster_name))
 
     candidate_names = [str(m.name) for m in cluster.members if m.name != master]
@@ -589,9 +591,12 @@ def failover(config_file, cluster_name, master, candidate, force, dcs, scheduled
     scheduled_at = parse_scheduled(scheduled)
 
     if scheduled_at:
+        if cluster.is_paused():
+            raise PatroniCtlException("Can't schedule failover in the paused state")
         scheduled_at = scheduled_at.isoformat()
 
     failover_value = {'leader': master, 'candidate': candidate, 'scheduled_at': scheduled_at}
+
     logging.debug(failover_value)
 
     # By now we have established that the leader exists and the candidate exists
@@ -607,7 +612,9 @@ def failover(config_file, cluster_name, master, candidate, force, dcs, scheduled
 
     r = None
     try:
-        r = request_patroni(cluster.leader.member, 'post', 'failover', failover_value, auth_header(config))
+        member = cluster.leader.member if cluster.leader else [m for m in cluster.members if m.name == candidate][0]
+
+        r = request_patroni(member, 'post', 'failover', failover_value, auth_header(config))
         if r.status_code in (200, 202):
             logging.debug(r)
             cluster = dcs.get_cluster()
@@ -748,7 +755,7 @@ def touch_member(config, dcs):
 
 
 def set_defaults(config, cluster_name):
-    ''' fill-in some basic configuration parameters if config file is not set '''
+    """fill-in some basic configuration parameters if config file is not set """
     config['postgresql'].setdefault('name', cluster_name)
     config['postgresql'].setdefault('scope', cluster_name)
     config['postgresql'].setdefault('listen', '127.0.0.1')
@@ -800,3 +807,33 @@ def flush(cluster_name, member_names, config_file, dcs, force, role, target):
                 check_response(r, member.name, 'flush scheduled restart')
             else:
                 click.echo('No scheduled restart for member {0}'.format(member.name))
+
+
+def toggle_pause(config_file, cluster_name, dcs, paused):
+    config, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+    if cluster.is_paused() == paused:
+        raise PatroniCtlException('Cluster is {0} paused'.format(paused and 'already' or 'not'))
+
+    r = request_patroni(cluster.leader.member, 'patch', 'config', {'pause': paused or None}, auth_header(config))
+
+    if r.status_code == 200:
+        click.echo('Success: cluster management is {0}'.format(paused and 'paused' or 'resumed'))
+    else:
+        click.echo('Failed: {0} cluster management status code={1}, ({2})'.format(
+                   paused and 'pause' or 'resume', r.status_code, r.text))
+
+
+@ctl.command('pause', help='Disable auto failover')
+@click.argument('cluster_name')
+@option_config_file
+@option_dcs
+def pause(config_file, cluster_name, dcs):
+    return toggle_pause(config_file, cluster_name, dcs, True)
+
+
+@ctl.command('resume', help='Resume auto failover')
+@click.argument('cluster_name')
+@option_config_file
+@option_dcs
+def resume(config_file, cluster_name, dcs):
+    return toggle_pause(config_file, cluster_name, dcs, False)

patroni/dcs/__init__.py

@@ -202,6 +202,9 @@ class Failover(namedtuple('Failover', 'index,leader,candidate,scheduled_at')):
 
         return Failover(index, data.get('leader'), data.get('member'), data.get('scheduled_at'))
 
+    def __len__(self):
+        return int(bool(self.leader)) + int(bool(self.candidate))
+
 
 class ClusterConfig(namedtuple('ClusterConfig', 'index,data,modify_index')):
 
@@ -244,6 +247,9 @@ class Cluster(namedtuple('Cluster', 'initialize,config,leader,last_leader_operat
         candidates = [m for m in self.members if m.clonefrom and (not self.leader or m.name != self.leader.name)]
         return candidates[randint(0, len(candidates) - 1)] if candidates else self.leader
 
+    def is_paused(self):
+        return self.config and self.config.data.get('pause', False) or False
+
 
 @six.add_metaclass(abc.ABCMeta)
 class AbstractDCS(object):

patroni/ha.py

@@ -26,6 +26,9 @@ class Ha(object):
         self.recovering = False
         self._async_executor = AsyncExecutor()
 
+    def is_paused(self):
+        return self.cluster and self.cluster.is_paused()
+
     def load_cluster_from_dcs(self):
         cluster = self.dcs.get_cluster()
 
@@ -131,23 +134,34 @@ class Ha(object):
 
         return node_to_follow if node_to_follow and node_to_follow.name != self.state_handler.name else None
 
-    def follow(self, demote_reason, follow_reason, refresh=True, recovery=False):
+    def follow(self, demote_reason, follow_reason, refresh=True, recovery=False, need_rewind=None):
         if refresh:
             self.load_cluster_from_dcs()
 
         if recovery:
             ret = demote_reason if self.has_lock() else follow_reason
         else:
-            ret = demote_reason if self.state_handler.is_leader() else follow_reason
+            is_leader = self.state_handler.is_leader()
+            ret = demote_reason if is_leader else follow_reason
 
         node_to_follow = self._get_node_to_follow(self.cluster)
 
-        self.state_handler.follow(node_to_follow, self.cluster.leader, recovery, self._async_executor)
+        if self.is_paused() and not self.state_handler.need_rewind:
+            self.state_handler.set_role('master' if is_leader else 'replica')
+            if is_leader:
+                return 'continue to run as master without lock'
+            elif not node_to_follow:
+                return 'no action'
+
+        self.state_handler.follow(node_to_follow, self.cluster.leader, recovery, self._async_executor, need_rewind)
 
         return ret
 
     def enforce_master_role(self, message, promote_message):
         if self.state_handler.is_leader() or self.state_handler.role == 'master':
+            # Inform the state handler about its master role.
+            # It may be unaware of it if postgres is promoted manually.
+            self.state_handler.set_role('master')
             return message
         else:
             self.state_handler.promote()
@@ -223,6 +237,15 @@ class Ha(object):
         if failover.candidate:  # manual failover to specific member
             if failover.candidate == self.state_handler.name:  # manual failover to me
                 return True
+            elif self.is_paused():
+                # Remove failover key if the node to failover has terminated to avoid waiting for it indefinitely
+                # In order to avoid attempts to delete this key from all nodes only the master is allowed to do it.
+                if (not self.cluster.get_member(failover.candidate, fallback_to_leader=False) and
+                   self.state_handler.is_leader()):
+                    logger.warning("manual failover: removing failover key because failover candidate is not running")
+                    self.dcs.manual_failover('', '', index=self.cluster.failover.index)
+                    return None
+                return False
 
             # find specific node and check that it is healthy
             member = self.cluster.get_member(failover.candidate, fallback_to_leader=False)
@@ -239,6 +262,8 @@ class Ha(object):
 
             # at this point we should consider all members as a candidates for failover
             # i.e. we assume that failover.candidate is None
+        elif self.is_paused():
+            return False
 
         # try to pick some other members to failover and check that they are healthy
         if failover.leader:
@@ -257,9 +282,18 @@ class Ha(object):
         return self._is_healthiest_node(members, check_replication_lag=False)
 
     def is_healthiest_node(self):
+        if self.is_paused() and not self.patroni.nofailover and \
+                self.cluster.failover and not self.cluster.failover.scheduled_at:
+            ret = self.manual_failover_process_no_leader()
+            if ret is not None:  # continue if we just deleted the stale failover key as a master
+                return ret
+
         if self.state_handler.is_leader():  # leader is always the healthiest
             return True
 
+        if self.is_paused():
+            return False
+
         if self.patroni.nofailover:  # nofailover tag makes node always unhealthy
             return False
 
@@ -273,19 +307,19 @@ class Ha(object):
     def demote(self, delete_leader=True):
         if delete_leader:
             self.state_handler.stop()
-            self.state_handler.set_role('unknown')
+            self.state_handler.set_role('demoted')
             self.dcs.delete_leader()
             self.touch_member()
             self.dcs.reset_cluster()
-            sleep(2)  # Give a time to somebody to promote
+            sleep(2)  # Give a time to somebody to take the leader lock
             cluster = self.dcs.get_cluster()
             node_to_follow = self._get_node_to_follow(cluster)
-            self.state_handler.follow(node_to_follow, cluster.leader, True)
+            self.state_handler.follow(node_to_follow, cluster.leader, recovery=True, need_rewind=True)
         else:
             self.state_handler.follow(None, None)
 
     def should_run_scheduled_action(self, action_name, scheduled_at, cleanup_fn):
-        if scheduled_at:
+        if scheduled_at and not self.is_paused():
             # If the scheduled action is in the far future, we shouldn't do anything and just return.
             # If the scheduled action is in the past, we consider the value to be stale and we remove
             # the value.
@@ -304,7 +338,6 @@ class Ha(object):
                     logger.warning('Found a stale %s value, cleaning up: %s',
                                    action_name, scheduled_at.isoformat())
                     cleanup_fn()
-                    self.dcs.manual_failover('', '', index=self.cluster.failover.index)
                     return False
 
                 # The value is very close to now
@@ -321,33 +354,44 @@ class Ha(object):
 
         if (failover.scheduled_at and not
             self.should_run_scheduled_action("failover", failover.scheduled_at, lambda:
-                                             self.dcs.manual_failover('', '', index=self.cluster.failover.index))):
+                                             self.dcs.manual_failover('', '', index=failover.index))):
             return
 
         if not failover.leader or failover.leader == self.state_handler.name:
             if not failover.candidate or failover.candidate != self.state_handler.name:
-                members = [m for m in self.cluster.members if not failover.candidate or m.name == failover.candidate]
-                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'
+                if not failover.candidate and self.is_paused():
+                    logger.warning('Failover is possible only to a specific candidate in a paused state')
                 else:
-                    logger.warning('manual failover: no healthy members found, failover is not possible')
+                    members = [m for m in self.cluster.members
+                               if not failover.candidate or m.name == failover.candidate]
+                    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)
+                           failover.leader, self.state_handler.name)
 
-        logger.info('Trying to clean up failover key')
-        self.dcs.manual_failover('', '', index=self.cluster.failover.index)
+        logger.info('Cleaning up failover key')
+        self.dcs.manual_failover('', '', index=failover.index)
 
     def process_unhealthy_cluster(self):
+        """Cluster has no leader key"""
+
         if self.is_healthiest_node():
             if self.acquire_lock():
-                if self.cluster.failover:
-                    logger.info('Cleaning up failover key after acquiring leader lock...')
-                    self.dcs.manual_failover('', '')
+                failover = self.cluster.failover
+                if failover:
+                    if self.is_paused() and failover.leader and failover.candidate:
+                        logger.info('Updating failover key after acquiring leader lock...')
+                        self.dcs.manual_failover('', failover.candidate, failover.scheduled_at, failover.index)
+                    else:
+                        logger.info('Cleaning up failover key after acquiring leader lock...')
+                        self.dcs.manual_failover('', '')
                 self.load_cluster_from_dcs()
                 return self.enforce_master_role('acquired session lock as a leader',
                                                 'promoted self to leader by acquiring session lock')
@@ -355,19 +399,35 @@ class Ha(object):
                 return self.follow('demoted self after trying and failing to obtain lock',
                                    'following new leader after trying and failing to obtain lock')
         else:
+            # when we are doing manual failover there is no guaranty that new leader is ahead of any other node
+            # node tagged as nofailover can be ahead of the new leader either, but it is always excluded from elections
+            need_rewind = bool(self.cluster.failover) or self.patroni.nofailover
+            if need_rewind:
+                sleep(2)  # Give a time to somebody to take the leader lock
+
             if self.patroni.nofailover:
                 return self.follow('demoting self because I am not allowed to become master',
-                                   'following a different leader because I am not allowed to promote')
+                                   'following a different leader because I am not allowed to promote',
+                                   need_rewind=need_rewind)
             return self.follow('demoting self because i am not the healthiest node',
-                               'following a different leader because i am not the healthiest node')
+                               'following a different leader because i am not the healthiest node',
+                               need_rewind=need_rewind)
 
     def process_healthy_cluster(self):
         if self.has_lock():
-            if self.cluster.failover:
+            if self.cluster.failover and (not self.is_paused() or self.state_handler.is_leader()):
                 msg = self.process_manual_failover_from_leader()
                 if msg is not None:
                     return msg
 
+            if self.is_paused() and not self.state_handler.is_leader():
+                if self.cluster.failover and self.cluster.failover.candidate == self.state_handler.name:
+                    return 'waiting to become master after promote...'
+
+                self.dcs.delete_leader()
+                self.dcs.reset_cluster()
+                return 'removed leader lock because postgres is not running as master'
+
             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')
@@ -423,10 +483,6 @@ class Ha(object):
             logger.info("not proceeding with the restart: %s", reason_to_cancel)
         return False
 
-    def schedule(self, action, immediate=False):
-        with self._async_executor:
-            return self._async_executor.schedule(action, immediate)
-
     def schedule_future_restart(self, restart_data):
         with self._async_executor:
             if not self.patroni.scheduled_restart:
@@ -448,15 +504,6 @@ class Ha(object):
         return self.patroni.scheduled_restart.copy() if (self.patroni.scheduled_restart and
                                                          isinstance(self.patroni.scheduled_restart, dict)) else None
 
-    def schedule_reinitialize(self):
-        return self.schedule('reinitialize')
-
-    def reinitialize_scheduled(self):
-        return self._async_executor.scheduled_action == 'reinitialize'
-
-    def schedule_restart(self, immediate=False):
-        return self.schedule('restart', immediate)
-
     def restart_scheduled(self):
         return self._async_executor.scheduled_action == 'restart'
 
@@ -469,37 +516,41 @@ class Ha(object):
             return (False, "restart conditions are not satisfied")
 
         with self._async_executor:
-            prev = self.schedule_restart(immediate=(not run_async))
+            prev = self._async_executor.schedule('restart')
             if prev is not None:
                 return (False, prev + ' already in progress')
-            if not run_async:
-                if self._async_executor.run(self.state_handler.restart):
-                    return (True, 'restarted successfully')
-                else:
-                    return (False, 'restart failed')
-            else:
-                self._async_executor.run_async(self.state_handler.restart)
-                return (True, "restart initiated")
 
-    def reinitialize(self, cluster):
+        if run_async:
+            self._async_executor.run_async(self.state_handler.restart)
+            return (True, 'restart initiated')
+        elif self._async_executor.run(self.state_handler.restart):
+            return (True, 'restarted successfully')
+        else:
+            return (False, 'restart failed')
+
+    def _do_reinitialize(self, cluster):
         self.state_handler.stop('immediate')
         self.state_handler.remove_data_directory()
 
-        clone_member = cluster.get_clone_member()
-        member_role = 'leader' if clone_member == cluster.leader else 'replica'
+        clone_member = self.cluster.get_clone_member()
+        member_role = 'leader' if clone_member == self.cluster.leader else 'replica'
         self.clone(clone_member, "from {0} '{1}'".format(member_role, clone_member.name))
 
-    def process_scheduled_action(self):
-        if self.reinitialize_scheduled():
+    def reinitialize(self):
+        with self._async_executor:
+            self.load_cluster_from_dcs()
+
             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'
+                return 'Cluster has no leader, can not reinitialize'
+
+            if self.cluster.leader.name == self.state_handler.name:
+                return 'I am the leader, can not reinitialize'
+
+            action = self._async_executor.schedule('reinitialize', immediately=True)
+            if action is not None:
+                return '{0} already in progress'.format(action)
+
+        self._async_executor.run_async(self._do_reinitialize, args=(self.cluster, ))
 
     def handle_long_action_in_progress(self):
         if self.has_lock():
@@ -545,22 +596,17 @@ class Ha(object):
                 return self.handle_long_action_in_progress()
 
             # we've got here, so any async action has finished. Check if we tried to recover and failed
-            if self.recovering:
+            if self.recovering and not self.state_handler.need_rewind:
                 self.recovering = False
                 msg = self.post_recover()
                 if msg is not None:
                     return msg
 
-            # currently it can trigger only reinitialize
-            msg = self.process_scheduled_action()
-            if msg is not None:
-                return msg
-
             # is data directory empty?
             if self.state_handler.data_directory_empty():
                 return self.bootstrap()  # new node
             # "bootstrap", but data directory is not empty
-            elif not self.sysid_valid(self.cluster.initialize) and self.cluster.is_unlocked():
+            elif not self.sysid_valid(self.cluster.initialize) and self.cluster.is_unlocked() and not self.is_paused():
                 self.dcs.initialize(create_new=(self.cluster.initialize is None), sysid=self.state_handler.sysid)
             else:
                 # check if we are allowed to join
@@ -569,8 +615,16 @@ class Ha(object):
                                  self.state_handler.name, self.cluster.initialize, self.state_handler.sysid)
                     sys.exit(1)
 
-            # try to start dead postgres
             if not self.state_handler.is_healthy():
+                if self.is_paused():
+                    if self.has_lock():
+                        self.dcs.delete_leader()
+                        self.dcs.reset_cluster()
+                        return 'removed leader lock because postgres is not running'
+                    elif not self.state_handler.need_rewind:
+                        return 'postgres is not running'
+
+                # try to start dead postgres
                 return self.recover()
 
             try:
@@ -591,12 +645,14 @@ class Ha(object):
                     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():
+            if not self.is_paused() and self.state_handler.is_running() and self.state_handler.is_leader():
                 self.demote(delete_leader=False)
                 return 'demoted self because DCS is not accessible and i was a leader'
+            return 'DCS is not accessible'
         except (psycopg2.Error, PostgresConnectionException):
-            logger.exception('Error communicating with PostgreSQL. Will try again later')
+            return 'Error communicating with PostgreSQL. Will try again later'
 
     def run_cycle(self):
         with self._async_executor:
-            return self._run_cycle()
+            info = self._run_cycle()
+            return (self.is_paused() and 'PAUSE: ' or '') + info

patroni/postgresql.py

@@ -757,7 +757,14 @@ class Postgresql(object):
         except OSError:
             logger.exception("Unable to list %s", status_dir)
 
-    def follow(self, member, leader, recovery=False, async_executor=None):
+    @property
+    def need_rewind(self):
+        return self._need_rewind
+
+    def follow(self, member, leader, recovery=False, async_executor=None, need_rewind=None):
+        if need_rewind is not None:
+            self._need_rewind = need_rewind
+
         primary_conninfo = self.primary_conninfo(member)
 
         if self.check_recovery_conf(primary_conninfo) and not recovery:
@@ -770,31 +777,24 @@ class Postgresql(object):
             self._do_follow(primary_conninfo, leader, recovery)
 
     def _do_follow(self, primary_conninfo, leader, recovery=False):
-        change_role = self.role == 'master'
+        change_role = self.role in ('master', 'demoted')
 
-        if change_role:
-            if leader:
-                if leader.name == self.name:
-                    self._need_rewind = False
-                    primary_conninfo = None
-                    if self.is_running():
-                        return
-                else:
-                    self._need_rewind = bool(leader.conn_url) and self.can_rewind
-            else:
-                self._need_rewind = False
-                primary_conninfo = None
+        if leader and leader.name == self.name:
+            primary_conninfo = None
+            self._need_rewind = False
+            if self.is_running():
+                return
+        elif change_role:
+            self._need_rewind = True
+
+        self._need_rewind &= bool(leader and leader.conn_url) and self.can_rewind
 
         if self._need_rewind:
-            logger.info("set the rewind flag after demote")
+            logger.info("rewind flag is set")
 
-            self.set_role('unknown')
             if self.is_running() and not self.stop():
                 return logger.warning('Can not run pg_rewind because postgres is still running')
 
-            if not (leader and leader.conn_url):
-                return logger.info('Leader unknown, can not rewind')
-
             # prepare pg_rewind connection
             r = leader.conn_kwargs(self._superuser)
 
@@ -827,7 +827,6 @@ class Postgresql(object):
             else:
                 logger.error('unable to rewind the former master')
                 self.remove_data_directory()
-                self.set_role('uninitialized')
                 ret = True
             self._need_rewind = False
         else:
@@ -990,6 +989,7 @@ $$""".format(name, ' '.join(options)), name, password, password)
                 logger.exception("Could not rename data directory %s", self._data_dir)
 
     def remove_data_directory(self):
+        self.set_role('uninitialized')
         logger.info('Removing data directory: %s', self._data_dir)
         try:
             if os.path.islink(self._data_dir):

requirements.txt

@@ -9,4 +9,4 @@ python-consul==0.6.0
 click>=4.1
 prettytable>=0.7
 tzlocal
-python-dateutil
+python-dateutil

tests/test_api.py

@@ -40,7 +40,7 @@ class MockHa(object):
     state_handler = MockPostgresql()
 
     @staticmethod
-    def schedule_reinitialize():
+    def reinitialize():
         return 'reinitialize'
 
     @staticmethod
@@ -179,7 +179,9 @@ class TestRestApiHandler(unittest.TestCase):
             MockRestApiServer(RestApiHandler, 'POST /reload HTTP/1.0' + self._authorization)
         self.assertIsNotNone(MockRestApiServer(RestApiHandler, 'POST /reload HTTP/1.0' + self._authorization))
 
-    def test_do_POST_restart(self):
+    @patch.object(MockPatroni, 'dcs')
+    def test_do_POST_restart(self, mock_dcs):
+        mock_dcs.get_cluster.return_value.is_paused.return_value = False
         request = 'POST /restart HTTP/1.0' + self._authorization
         self.assertIsNotNone(MockRestApiServer(RestApiHandler, request))
 
@@ -221,6 +223,9 @@ class TestRestApiHandler(unittest.TestCase):
                 request = make_request(role='master', postgres_version='9.5.2')
                 MockRestApiServer(RestApiHandler, request)
 
+        mock_dcs.get_cluster.return_value.is_paused.return_value = True
+        MockRestApiServer(RestApiHandler, make_request(schedule='2016-08-42 12:45TZ+1', role='master'))
+
     def test_do_DELETE_restart(self):
         for retval in (True, False):
             with patch.object(MockHa, 'delete_future_restart', Mock(return_value=retval)):
@@ -228,16 +233,13 @@ class TestRestApiHandler(unittest.TestCase):
                 self.assertIsNotNone(MockRestApiServer(RestApiHandler, request))
 
     @patch.object(MockPatroni, 'dcs')
-    def test_do_POST_reinitialize(self, dcs):
-        cluster = dcs.get_cluster.return_value
+    def test_do_POST_reinitialize(self, mock_dcs):
+        cluster = mock_dcs.get_cluster.return_value
+        cluster.is_paused.return_value = False
         request = 'POST /reinitialize HTTP/1.0' + self._authorization
         MockRestApiServer(RestApiHandler, request)
-        cluster.is_unlocked.return_value = False
-        MockRestApiServer(RestApiHandler, request)
-        with patch.object(MockHa, 'schedule_reinitialize', Mock(return_value=None)):
+        with patch.object(MockHa, 'reinitialize', Mock(return_value=None)):
             MockRestApiServer(RestApiHandler, request)
-        cluster.leader.name = 'test'
-        self.assertIsNotNone(MockRestApiServer(RestApiHandler, request))
 
     @patch('time.sleep', Mock())
     def test_RestApiServer_query(self):

tests/test_ctl.py

@@ -82,6 +82,11 @@ class TestCtl(unittest.TestCase):
         result = self.runner.invoke(ctl, ['failover', 'dummy'], input='leader\nother\n2030-01-01T12:23:00\ny')
         assert result.exit_code == 0
 
+        with patch('patroni.dcs.Cluster.is_paused', Mock(return_value=True)):
+            result = self.runner.invoke(ctl,
+                                        ['failover', 'dummy', '--force', '--scheduled', '2015-01-01T12:00:00+01:00'])
+            assert result.exit_code == 1
+
         # Aborting failover,as we anser NO to the confirmation
         result = self.runner.invoke(ctl, ['failover', 'dummy'], input='leader\nother\n\nN')
         assert result.exit_code == 1
@@ -241,6 +246,11 @@ class TestCtl(unittest.TestCase):
                                               '--scheduled', '2300-10-01T14:30'])
             assert 'Failed: flush scheduled restart' in result.output
 
+        with patch('patroni.dcs.Cluster.is_paused', Mock(return_value=True)):
+            result = self.runner.invoke(ctl,
+                                        ['restart', 'alpha', 'other', '--force', '--scheduled', '2300-10-01T14:30'])
+            assert result.exit_code == 1
+
         with patch('requests.post', Mock(return_value=MockResponse())):
             # normal restart, the schedule is actually parsed, but not validated in patronictl
             result = self.runner.invoke(ctl, ['restart', 'alpha', '--pg-version', '42.0.0',
@@ -393,3 +403,40 @@ class TestCtl(unittest.TestCase):
         with patch.object(requests, 'delete', return_value=MockResponse(404)):
             result = self.runner.invoke(ctl, ['flush', 'dummy', 'restart', '--force'])
             assert 'Failed: flush scheduled restart' in result.output
+
+    @patch('patroni.ctl.get_dcs')
+    def test_pause_cluster(self, mock_get_dcs):
+        mock_get_dcs.return_value = self.e
+        mock_get_dcs.return_value.get_cluster = get_cluster_initialized_with_leader
+
+        with patch('requests.patch', Mock(return_value=MockResponse(200))):
+            result = self.runner.invoke(ctl, ['pause', 'dummy'])
+            assert 'Success' in result.output
+
+        with patch('requests.patch', Mock(return_value=MockResponse(500))):
+            result = self.runner.invoke(ctl, ['pause', 'dummy'])
+            assert 'Failed' in result.output
+
+        with patch('requests.patch', Mock(return_value=MockResponse(200))),\
+                patch('patroni.dcs.Cluster.is_paused', Mock(return_value=True)):
+            result = self.runner.invoke(ctl, ['pause', 'dummy'])
+            assert 'Cluster is already paused' in result.output
+
+    @patch('patroni.ctl.get_dcs')
+    def test_resume_cluster(self, mock_get_dcs):
+        mock_get_dcs.return_value = self.e
+        mock_get_dcs.return_value.get_cluster = get_cluster_initialized_with_leader
+
+        with patch('patroni.dcs.Cluster.is_paused', Mock(return_value=True)):
+            with patch('requests.patch', Mock(return_value=MockResponse(200))):
+                result = self.runner.invoke(ctl, ['resume', 'dummy'])
+                assert 'Success' in result.output
+
+            with patch('requests.patch', Mock(return_value=MockResponse(500))):
+                result = self.runner.invoke(ctl, ['resume', 'dummy'])
+                assert 'Failed' in result.output
+
+            with patch('requests.patch', Mock(return_value=MockResponse(200))),\
+                    patch('patroni.dcs.Cluster.is_paused', Mock(return_value=False)):
+                result = self.runner.invoke(ctl, ['resume', 'dummy'])
+                assert 'Cluster is not paused' in result.output

tests/test_ha.py

@@ -4,7 +4,7 @@ import os
 import pytz
 import unittest
 
-from mock import Mock, MagicMock, patch
+from mock import Mock, MagicMock, PropertyMock, patch
 from patroni.config import Config
 from patroni.dcs import Cluster, Failover, Leader, Member, get_dcs
 from patroni.dcs.etcd import Client
@@ -90,7 +90,7 @@ zookeeper:
                                   'postmaster_start_time': str(postmaster_start_time)}
 
 
-def run_async(func, args=()):
+def run_async(self, func, args=()):
     return func(*args) if args else func()
 
 
@@ -109,6 +109,8 @@ def run_async(func, args=()):
 @patch.object(etcd.Client, 'write', etcd_write)
 @patch.object(etcd.Client, 'read', etcd_read)
 @patch.object(etcd.Client, 'delete', Mock(side_effect=etcd.EtcdException))
+@patch('patroni.async_executor.AsyncExecutor.busy', PropertyMock(return_value=False))
+@patch('patroni.async_executor.AsyncExecutor.run_async', run_async)
 @patch('subprocess.call', Mock(return_value=0))
 class TestHa(unittest.TestCase):
 
@@ -131,7 +133,6 @@ class TestHa(unittest.TestCase):
             self.e = get_dcs({'etcd': {'ttl': 30, 'host': 'ok:2379', 'scope': 'test',
                                        'name': 'foo', 'retry_timeout': 10}})
             self.ha = Ha(MockPatroni(self.p, self.e))
-            self.ha._async_executor.run_async = run_async
             self.ha.old_cluster = self.e.get_cluster()
             self.ha.cluster = get_cluster_not_initialized_without_leader()
             self.ha.load_cluster_from_dcs = Mock()
@@ -166,6 +167,9 @@ class TestHa(unittest.TestCase):
         self.ha.cluster = get_cluster_initialized_with_leader()
         self.assertEquals(self.ha.run_cycle(), 'starting as readonly because i had the session lock')
 
+    def test_do_not_recover_in_pause(self):
+        pass
+
     @patch('sys.exit', return_value=1)
     @patch('patroni.ha.Ha.sysid_valid', MagicMock(return_value=True))
     def test_sysid_no_match(self, exit_mock):
@@ -234,6 +238,13 @@ class TestHa(unittest.TestCase):
         self.ha.patroni.replicatefrom = "foo"
         self.assertEquals(self.ha.run_cycle(), 'no action.  i am a secondary and i am following a leader')
 
+    def test_follow_in_pause(self):
+        self.ha.cluster.is_unlocked = false
+        self.ha.is_paused = true
+        self.assertEquals(self.ha.run_cycle(), 'PAUSE: continue to run as master without lock')
+        self.p.is_leader = false
+        self.assertEquals(self.ha.run_cycle(), 'PAUSE: no action')
+
     def test_no_etcd_connection_master_demote(self):
         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')
@@ -267,43 +278,40 @@ 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._async_executor.scheduled_action)
+        self.assertIsNotNone(self.ha.reinitialize())
 
         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.assertIsNone(self.ha.reinitialize())
 
-        self.ha.has_lock = false
-        self.ha.schedule_reinitialize()
-        self.ha.run_cycle()
+        self.assertIsNotNone(self.ha.reinitialize())
+
+        self.ha.state_handler.name = self.ha.cluster.leader.name
+        self.assertIsNotNone(self.ha.reinitialize())
 
     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.ha.cluster = get_cluster_initialized_with_leader()
+        self.ha.reinitialize()
         self.assertEquals(self.ha.restart(), (False, 'reinitialize already in progress'))
         with patch.object(self.ha, "restart_matches", return_value=False):
             self.assertEquals(self.ha.restart({'foo': 'bar'}), (False, "restart conditions are not satisfied"))
 
     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')
+        with patch('patroni.async_executor.AsyncExecutor.busy', PropertyMock(return_value=True)):
+            self.ha.restart(run_async=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.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.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')
+            self.ha.update_lock = false
+            self.assertEquals(self.ha.run_cycle(), 'failed to update leader lock during restart')
 
     @patch('requests.get', requests_get)
     @patch('time.sleep', Mock())
@@ -346,6 +354,17 @@ class TestHa(unittest.TestCase):
         self.assertEquals('no action.  i am the leader with the lock', self.ha.run_cycle())
 
     @patch('requests.get', requests_get)
+    def test_manual_failover_from_leader_in_pause(self):
+        self.ha.has_lock = true
+        self.ha.is_paused = true
+        scheduled = datetime.datetime.now()
+        self.ha.cluster = get_cluster_initialized_with_leader(Failover(0, 'blabla', self.p.name, scheduled))
+        self.assertEquals('PAUSE: no action.  i am the leader with the lock', self.ha.run_cycle())
+        self.ha.cluster = get_cluster_initialized_with_leader(Failover(0, self.p.name, '', None))
+        self.assertEquals('PAUSE: no action.  i am the leader with the lock', self.ha.run_cycle())
+
+    @patch('requests.get', requests_get)
+    @patch('time.sleep', Mock())
     def test_manual_failover_process_no_leader(self):
         self.p.is_leader = false
         self.ha.cluster = get_cluster_initialized_without_leader(failover=Failover(0, '', self.p.name, None))
@@ -370,11 +389,27 @@ class TestHa(unittest.TestCase):
         self.ha.patroni.nofailover = True
         self.assertEquals(self.ha.run_cycle(), 'following a different leader because I am not allowed to promote')
 
+    @patch('time.sleep', Mock())
+    def test_manual_failover_process_no_leader_in_pause(self):
+        self.ha.is_paused = true
+        self.ha.cluster = get_cluster_initialized_without_leader(failover=Failover(0, '', 'other', None))
+        self.assertEquals(self.ha.run_cycle(), 'PAUSE: continue to run as master without lock')
+        self.ha.cluster = get_cluster_initialized_without_leader(failover=Failover(0, 'leader', '', None))
+        self.assertEquals(self.ha.run_cycle(), 'PAUSE: continue to run as master without lock')
+        self.ha.cluster = get_cluster_initialized_without_leader(failover=Failover(0, 'leader', 'blabla', None))
+        self.assertEquals('PAUSE: acquired session lock as a leader', self.ha.run_cycle())
+        self.p.is_leader = false
+        self.p.set_role('replica')
+        self.ha.cluster = get_cluster_initialized_without_leader(failover=Failover(0, 'leader', self.p.name, None))
+        self.assertEquals(self.ha.run_cycle(), 'PAUSE: promoted self to leader by acquiring session lock')
+
     def test_is_healthiest_node(self):
         self.ha.state_handler.is_leader = false
         self.ha.patroni.nofailover = False
         self.ha.fetch_node_status = lambda e: (e, True, True, 0, {})
         self.assertTrue(self.ha.is_healthiest_node())
+        self.ha.is_paused = true
+        self.assertFalse(self.ha.is_healthiest_node())
 
     def test__is_healthiest_node(self):
         self.assertTrue(self.ha._is_healthiest_node(self.ha.old_cluster.members))
@@ -451,3 +486,24 @@ class TestHa(unittest.TestCase):
         self.p._pending_restart = False
         self.assertFalse(self.ha.restart_matches("replica", "9.5.2", True))
         self.assertTrue(self.ha.restart_matches("replica", "9.5.2", False))
+
+    def test_process_healthy_cluster_in_pause(self):
+        self.p.is_leader = false
+        self.ha.is_paused = true
+        self.p.name = 'leader'
+        self.ha.cluster = get_cluster_initialized_with_leader()
+        self.assertEquals(self.ha.run_cycle(), 'PAUSE: removed leader lock because postgres is not running as master')
+        self.ha.cluster = get_cluster_initialized_with_leader(Failover(0, '', self.p.name, None))
+        self.assertEquals(self.ha.run_cycle(), 'PAUSE: waiting to become master after promote...')
+
+    def test_postgres_unhealthy_in_pause(self):
+        self.ha.is_paused = true
+        self.p.is_healthy = false
+        self.assertEquals(self.ha.run_cycle(), 'PAUSE: postgres is not running')
+        self.ha.has_lock = true
+        self.assertEquals(self.ha.run_cycle(), 'PAUSE: removed leader lock because postgres is not running')
+
+    def test_no_etcd_connection_in_pause(self):
+        self.ha.is_paused = true
+        self.ha.load_cluster_from_dcs = Mock(side_effect=DCSError('Etcd is not responding properly'))
+        self.assertEquals(self.ha.run_cycle(), 'PAUSE: DCS is not accessible')

tests/test_patroni.py

@@ -61,7 +61,8 @@ class TestPatroni(unittest.TestCase):
             with patch.object(Patroni, 'run', Mock(side_effect=SleepException)):
                 self.assertRaises(SleepException, _main)
             with patch.object(Patroni, 'run', Mock(side_effect=KeyboardInterrupt())):
-                _main()
+                with patch('patroni.ha.Ha.is_paused', Mock(return_value=True)):
+                    _main()
 
     @patch('patroni.config.Config.save_cache', Mock())
     @patch('patroni.config.Config.reload_local_configuration', Mock(return_value=True))

tests/test_postgresql.py

@@ -265,20 +265,16 @@ class TestPostgresql(unittest.TestCase):
         with patch.object(Postgresql, 'restart', Mock(return_value=False)):
             self.p.set_role('replica')
             self.p.follow(None, None)  # restart without rewind
-            self.p.set_role('master')
 
         with patch.object(Postgresql, 'stop', Mock(return_value=False)):
-            self.p.follow(self.leader, self.leader)  # failed to stop postgres
-
-        self.p.follow(self.leader, None)  # Leader unknown, can not rewind
+            self.p.follow(self.leader, self.leader, need_rewind=True)  # failed to stop postgres
 
         self.p.follow(self.leader, self.leader)  # "leader" is not accessible or is_in_recovery
 
         with patch.object(Postgresql, 'checkpoint', Mock(return_value=None)):
             self.p.follow(self.leader, self.leader)
-            self.p.set_role('master')
             mock_pg_rewind.return_value = True
-            self.p.follow(self.leader, self.leader)
+            self.p.follow(self.leader, self.leader, need_rewind=True)
 
         self.p.follow(None, None)  # check_recovery_conf...