Merge pull request #20 from zalando/features/refactoring

Try to avoid of double promote

governor.py

@@ -88,6 +88,8 @@ def main():
     try:
         governor.initialize()
         governor.run()
+    except KeyboardInterrupt:
+        pass
     finally:
         governor.touch_member(300)  # schedule member removal
         governor.postgresql.stop()

helpers/api.py

@@ -10,21 +10,16 @@ if sys.hexversion >= 0x03000000:
 else:
     from BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
 
-
 logger = logging.getLogger(__name__)
 
 
 class RestApiHandler(BaseHTTPRequestHandler):
 
     def do_GET(self):
-        try:
-            response = self.get_postgresql_status()
-        except (psycopg2.OperationalError, psycopg2.InterfaceError):
-            logging.exception('get_postgresql_status')
-            response = {'running': False}
+        response = self.get_postgresql_status()
 
         path = '/master' if self.path == '/' else self.path
-        status_code = 200 if response['running'] and response['role'] in path else 503
+        status_code = 200 if response['running'] and 'role' in response and response['role'] in path else 503
 
         self.send_response(status_code)
         self.send_header('Content-Type', 'application/json')
@@ -32,29 +27,31 @@ class RestApiHandler(BaseHTTPRequestHandler):
         self.wfile.write(json.dumps(response).encode('utf-8'))
 
     def get_postgresql_status(self):
-        if not self.server.governor.postgresql.is_running():
-            return {'running': False}
-        cursor = self.server._cursor()
-        cursor.execute("""SELECT to_char(pg_postmaster_start_time(), 'YYYY-MM-DD HH24:MI:SS.MS TZ'),
-                                 pg_is_in_recovery(),
-                                 CASE WHEN pg_is_in_recovery()
-                                      THEN null
-                                      ELSE pg_current_xlog_location() END,
-                                 pg_last_xlog_receive_location(),
-                                 pg_last_xlog_replay_location(),
-                                 pg_is_in_recovery() AND pg_is_xlog_replay_paused()""")
-        row = cursor.fetchone()
-        return {
-            'running': True,
-            'postmaster_start_time': row[0],
-            'role': 'slave' if row[1] else 'master',
-            'xlog': ({
-                'received_location': row[3],
-                'replayed_location': row[4],
-                'paused': row[5]} if row[1] else {
-                'location': row[2]
-            })
-        }
+        try:
+            cursor = self.server._cursor()
+            cursor.execute("""SELECT to_char(pg_postmaster_start_time(), 'YYYY-MM-DD HH24:MI:SS.MS TZ'),
+                                     pg_is_in_recovery(),
+                                     CASE WHEN pg_is_in_recovery()
+                                          THEN null
+                                          ELSE pg_current_xlog_location() END,
+                                     pg_last_xlog_receive_location(),
+                                     pg_last_xlog_replay_location(),
+                                     pg_is_in_recovery() AND pg_is_xlog_replay_paused()""")
+            row = cursor.fetchone()
+            return {
+                'running': True,
+                'postmaster_start_time': row[0],
+                'role': 'slave' if row[1] else 'master',
+                'xlog': ({
+                    'received_location': row[3],
+                    'replayed_location': row[4],
+                    'paused': row[5]} if row[1] else {
+                    'location': row[2]
+                })
+            }
+        except (psycopg2.OperationalError, psycopg2.InterfaceError):
+            logger.exception('get_postgresql_status')
+            return {'running': self.server.governor.postgresql.is_running()}
 
 
 class RestApiServer(HTTPServer, Thread):

helpers/ha.py

@@ -48,10 +48,10 @@ class Ha:
             if self.cluster.is_unlocked():
                 if self.state_handler.is_healthiest_node(self.cluster):
                     if self.acquire_lock():
-                        if not self.state_handler.is_leader():
-                            self.state_handler.promote()
-                            return 'promoted self to leader by acquiring session lock'
-                        return 'acquired session lock as a leader'
+                        if self.state_handler.is_leader() or self.state_handler.is_promoted:
+                            return 'acquired session lock as a leader'
+                        self.state_handler.promote()
+                        return 'promoted self to leader by acquiring session lock'
                     else:
                         self.load_cluster_from_etcd()
                         if self.state_handler.is_leader():
@@ -71,14 +71,13 @@ class Ha:
             else:
                 if self.has_lock() and self.update_lock():
                     try:
-                        if not self.state_handler.is_leader():
-                            self.state_handler.promote()
-                            return 'promoted self to leader because i had the session lock'
-                        else:
+                        if self.state_handler.is_leader() or self.state_handler.is_promoted:
                             return 'no action.  i am the leader with the lock'
+                        self.state_handler.promote()
+                        return 'promoted self to leader because i had the session lock'
                     finally:
                         # create replication slots
-                        self.state_handler.create_replication_slots([m.hostname for m in self.cluster.members])
+                        self.state_handler.create_replication_slots(self.cluster)
                 else:
                     logger.info('does not have lock')
                     if self.state_handler.is_leader():

helpers/postgresql.py

@@ -46,6 +46,10 @@ class Postgresql:
         self.configuration_to_save = (os.path.join(self.data_dir, 'pg_hba.conf'),
                                       os.path.join(self.data_dir, 'postgresql.conf'))
         self.pid_path = os.path.join(self.data_dir, 'postmaster.pid')
+        self.trigger_file = config.get('recovery_conf', {}).get('trigger_file', None) or 'promote'
+        self.trigger_file = os.path.abspath(os.path.join(self.data_dir, self.trigger_file))
+        self.is_promoted = False
+
         self._pg_ctl = ['pg_ctl', '-w', '-D', self.data_dir]
         self.wal_e = config.get('wal_e', None)
         if self.wal_e:
@@ -110,6 +114,9 @@ class Postgresql:
         ret and self.write_pg_hba()
         return ret
 
+    def delete_trigger_file(self):
+        os.path.exists(self.trigger_file) and os.unlink(self.trigger_file)
+
     def sync_from_leader(self, leader):
         r = parseurl(leader.address)
 
@@ -118,24 +125,24 @@ class Postgresql:
             os.fchmod(f.fileno(), 0o600)
             f.write('{host}:{port}:*:{user}:{password}\n'.format(**r))
 
-        try:
-            os.environ['PGPASSFILE'] = pgpass
-            return self.create_replica(r) == 0
-        finally:
-            os.environ.pop('PGPASSFILE')
+        env = os.environ.copy()
+        env['PGPASSFILE'] = pgpass
+        return self.create_replica(r, env) == 0
 
-    def create_replica(self, master_connection):
+    def create_replica(self, master_connection, env):
         """ creates a new replica using either pg_basebackup or WAL-E """
         if self.should_use_s3_to_create_replica(master_connection):
             result = self.create_replica_with_s3()
             # if restore from the backup on S3 failed - try with the pg_basebackup
             if result == 0:
                 return result
-        return self.create_replica_with_pg_basebackup(master_connection)
+        return self.create_replica_with_pg_basebackup(master_connection, env)
 
-    def create_replica_with_pg_basebackup(self, master_connection):
-        return subprocess.call(['pg_basebackup', '-R', '-D', self.data_dir, '--host=' + master_connection['host'],
-                                '--port=' + str(master_connection['port']), '-U', master_connection['user']])
+    def create_replica_with_pg_basebackup(self, master_connection, env):
+        ret = subprocess.call(['pg_basebackup', '-R', '-D', self.data_dir, '--host=' + master_connection['host'],
+                               '--port=' + str(master_connection['port']), '-U', master_connection['user']], env=env)
+        self.delete_trigger_file()
+        return ret
 
     def create_replica_with_s3(self):
         if not self.wal_e or not self.wal_e_path:
@@ -215,7 +222,11 @@ class Postgresql:
                (diff_in_bytes < long(backup_size) * float(threshold_backup_size_percentage) / 100)
 
     def is_leader(self):
-        return not self.query('SELECT pg_is_in_recovery()').fetchone()[0]
+        ret = not self.query('SELECT pg_is_in_recovery()').fetchone()[0]
+        if ret and self.is_promoted:
+            self.delete_trigger_file()
+            self.is_promoted = False
+        return ret
 
     def is_running(self):
         return subprocess.call(' '.join(self._pg_ctl) + ' status > /dev/null', shell=True) == 0
@@ -328,10 +339,9 @@ primary_conninfo = '{}'
                     f.write("{} = '{}'\n".format(name, value))
 
     def follow_the_leader(self, leader):
-        if self.check_recovery_conf(leader):
-            return
-        self.write_recovery_conf(leader)
-        self.restart()
+        if not self.check_recovery_conf(leader):
+            self.write_recovery_conf(leader)
+            self.restart()
 
     def save_configuration_files(self):
         """
@@ -350,7 +360,8 @@ primary_conninfo = '{}'
             logger.error("unable to restore configuration from WAL-E backup: {}".format(e))
 
     def promote(self):
-        return subprocess.call(self._pg_ctl + ['promote']) == 0
+        self.is_promoted = subprocess.call(self._pg_ctl + ['promote']) == 0
+        return self.is_promoted
 
     def demote(self, leader):
         self.follow_the_leader(leader)
@@ -379,7 +390,8 @@ primary_conninfo = '{}'
         cursor = self.query("SELECT slot_name FROM pg_replication_slots WHERE slot_type='physical'")
         self.members = [r[0] for r in cursor]
 
-    def create_replication_slots(self, members):
+    def create_replication_slots(self, cluster):
+        members = [m.hostname for m in cluster.members if m.hostname != self.name]
         # drop unused slots
         for slot in set(self.members) - set(members):
             self.query("""SELECT pg_drop_replication_slot(%s)

tests/test_api.py

@@ -11,10 +11,6 @@ else:
     from StringIO import StringIO as IO
 
 
-def false(*args, **kwargs):
-    return False
-
-
 def throws(*args, **kwargs):
     raise psycopg2.OperationalError()
 
@@ -48,7 +44,7 @@ class MockRestApiServer(RestApiServer):
     def __init__(self, Handler, path, *args):
         self.governor = MockGovernor()
         if len(args) > 0:
-            self.governor.postgresql.is_running = args[0]
+            self._cursor = args[0]
         self._cursor_holder = None
         Handler(MockRequest(path), ('0.0.0.0', 8080), self)
 
@@ -61,6 +57,3 @@ class TestRestApiHandler(unittest.TestCase):
     def test_do_GET(self):
         MockRestApiServer(RestApiHandler, b'GET /')
         MockRestApiServer(RestApiHandler, b'GET /', throws)
-
-    def test_get_postgresql_status(self):
-        MockRestApiServer(RestApiHandler, b'GET /', false)

tests/test_ha.py

@@ -19,6 +19,7 @@ class MockPostgresql:
 
     def __init__(self):
         self.name = 'postgresql0'
+        self.is_promoted = False
 
     def is_healthy(self):
         return True

tests/test_postgresql.py

@@ -15,7 +15,7 @@ def nop(*args, **kwargs):
     pass
 
 
-def subprocess_call(cmd, shell=False):
+def subprocess_call(cmd, shell=False, env=None):
     return 0
 
 
@@ -110,17 +110,13 @@ class TestPostgresql(unittest.TestCase):
     def set_up(self):
         subprocess.call = subprocess_call
         shutil.copy = nop
-        self.p = Postgresql({
-            'name': 'test0',
-            'data_dir': 'data/test0',
-            'listen': '127.0.0.1, 127.0.0.2:5432',
-            'connect_address': '127.0.0.2:5432',
-            'superuser': {'password': ''},
-            'admin': {'username': 'admin', 'password': 'admin'},
-            'replication': {'username': 'replicator', 'password': 'rep-pass', 'network': '127.0.0.1/32'},
-            'parameters': {'foo': 'bar'},
-            'recovery_conf': {'foo': 'bar'},
-        })
+        self.p = Postgresql({'name': 'test0', 'data_dir': 'data/test0', 'listen': '127.0.0.1, 127.0.0.2:5432',
+                             'connect_address': '127.0.0.2:5432', 'superuser': {'password': ''},
+                             'admin': {'username': 'admin', 'password': 'admin'},
+                             'replication': {'username': 'replicator',
+                                             'password': 'rep-pass',
+                                             'network': '127.0.0.1/32'},
+                             'parameters': {'foo': 'bar'}, 'recovery_conf': {'foo': 'bar'}})
         psycopg2.connect = psycopg2_connect
         if not os.path.exists(self.p.data_dir):
             os.makedirs(self.p.data_dir)
@@ -155,7 +151,10 @@ class TestPostgresql(unittest.TestCase):
 
     def test_create_replication_slots(self):
         self.p.start()
-        self.p.create_replication_slots('qaz')
+        me = Member('test0', 'postgres://replicator:[email protected]:5434/postgres', 28)
+        other = Member('test1', 'postgres://replicator:[email protected]:5433/postgres', 28)
+        cluster = Cluster(True, self.leader, 0, [me, other, self.leader])
+        self.p.create_replication_slots(cluster)
 
     def test_query(self):
         self.p.query('select 1')
@@ -181,7 +180,9 @@ class TestPostgresql(unittest.TestCase):
         self.assertFalse(self.p.is_healthiest_node(cluster))
 
     def test_is_leader(self):
+        self.p.is_promoted = True
         self.assertTrue(self.p.is_leader())
+        self.assertFalse(self.p.is_promoted)
 
     def test_reload(self):
         self.assertTrue(self.p.reload())