Implement reload of config.yaml on SIGHUP

If some changes require restart of postgres patroni will expose
`restart_pending` flag in DCS and via REST API

patroni/__init__.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import signal
 import sys
 import time
 import yaml
@@ -17,10 +18,12 @@ logger = logging.getLogger(__name__)
 class Patroni(object):
     PATRONI_CONFIG_VARIABLE = 'PATRONI_CONFIGURATION'
 
-    def __init__(self, config):
+    def __init__(self, config_file=None, config_env=None):
+        self._config_file = config_file
+        config = yaml.load(config_env) if config_env else self._load_config()
+
         self.nap_time = config['loop_wait']
-        self.tags = {tag: value for tag, value in config.get('tags', {}).items()
-                     if tag not in ('clonefrom', 'nofailover', 'noloadbalance') or value}
+        self.tags = self.get_tags(config)
         self.postgresql = Postgresql(config['postgresql'])
         self.dcs = self.get_dcs(self.postgresql.name, config)
         self.version = __version__
@@ -28,6 +31,32 @@ class Patroni(object):
         self.ha = Ha(self)
         self.next_run = time.time()
 
+        self._reload_config_scheduled = False
+
+    @staticmethod
+    def get_tags(config):
+        return {tag: value for tag, value in config.get('tags', {}).items()
+                if tag not in ('clonefrom', 'nofailover', 'noloadbalance') or value}
+
+    def _load_config(self, fail=True):
+        with open(self._config_file) as f:
+            return yaml.load(f)
+
+    def reload_config(self):
+        try:
+            config = self._load_config()
+            self.tags = self.get_tags(config)
+            self.nap_time = config['loop_wait']
+            self.dcs.set_ttl(config.get('ttl') or 30)
+            self.api.reload_config(config['restapi'])
+            self.postgresql.reload_config(config['postgresql'])
+        except Exception:
+            logger.exception('Failed to reload config_file=%s', self._config_file)
+        self._reload_config_scheduled = False
+
+    def sighup_handler(self, *args):
+        self._reload_config_scheduled = True
+
     @property
     def noloadbalance(self):
         return self.tags.get('noloadbalance', False)
@@ -64,38 +93,34 @@ class Patroni(object):
 
     def run(self):
         self.api.start()
+        signal.signal(signal.SIGHUP, self.sighup_handler)
         self.next_run = time.time()
 
         while True:
+            if self._reload_config_scheduled:
+                self.reload_config()
             logger.info(self.ha.run_cycle())
             reap_children()
             self.schedule_next_run()
 
 
 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 reads the configuration from the command-line argument if it exists, and from the environment otherwise.
-    use_env = False
-    use_file = (len(sys.argv) >= 2 and os.path.isfile(sys.argv[1]))
-    if not use_file:
+    config_env = False
+    config_file = len(sys.argv) >= 2 and os.path.isfile(sys.argv[1]) and sys.argv[1]
+    if not config_file:
         config_env = os.environ.get(Patroni.PATRONI_CONFIG_VARIABLE)
-        use_env = config_env is not None
-        if not use_env:
+        if config_env is None:
             print('Usage: {0} config.yml'.format(sys.argv[0]))
             print('\tPatroni may also read the configuration from the {} environment variable'.
                   format(Patroni.PATRONI_CONFIG_VARIABLE))
             return
 
-    if use_file:
-        with open(sys.argv[1], 'r') as f:
-            config = yaml.load(f)
-    elif use_env:
-        config = yaml.load(config_env)
-
-    patroni = Patroni(config)
+    patroni = Patroni(config_file, config_env)
     try:
         patroni.run()
     except KeyboardInterrupt:

patroni/api.py

@@ -71,6 +71,8 @@ class RestApiHandler(BaseHTTPRequestHandler):
             response.update({'tags': patroni.tags} if patroni.tags else {})
             if patroni.postgresql.sysid:
                 response['database_system_identifier'] = patroni.postgresql.sysid
+            if patroni.postgresql.restart_pending:
+                response['restart_pending'] = True
             response['patroni'] = {'version': patroni.version, 'scope': patroni.postgresql.scope}
             body = json.dumps(response)
         self._write_response(status_code, body, {'Content-Type': 'application/json'})
@@ -294,25 +296,9 @@ class RestApiHandler(BaseHTTPRequestHandler):
 class RestApiServer(ThreadingMixIn, HTTPServer, Thread):
 
     def __init__(self, patroni, config):
-        self._auth_key = base64.b64encode(config['auth'].encode('utf-8')).decode('utf-8') if 'auth' in config else None
-        host, port = config['listen'].split(':')
-        HTTPServer.__init__(self, (host, int(port)), RestApiHandler)
-        Thread.__init__(self, target=self.serve_forever)
-        self._set_fd_cloexec(self.socket)
-
-        protocol = 'http'
-
-        # wrap socket with ssl if 'certfile' is defined in a config.yaml
-        # Sometime it's also needed to pass reference to a 'keyfile'.
-        options = {option: config[option] for option in ['certfile', 'keyfile'] if option in config}
-        if options.get('certfile'):
-            import ssl
-            self.socket = ssl.wrap_socket(self.socket, server_side=True, **options)
-            protocol = 'https'
-
-        self.connection_string = '{0}://{1}/patroni'.format(protocol, config.get('connect_address', config['listen']))
-
         self.patroni = patroni
+        self.__initialize(config)
+        self.__set_config_parameters(config)
         self.daemon = True
 
     def query(self, sql, *params):
@@ -332,11 +318,47 @@ class RestApiServer(ThreadingMixIn, HTTPServer, Thread):
         fcntl.fcntl(fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC)
 
     def check_basic_auth_key(self, key):
-        return self._auth_key == key
+        return self.__auth_key == key
 
     def check_auth_header(self, auth_header):
-        if self._auth_key:
+        if self.__auth_key:
             if auth_header is None:
                 return 'no auth header received'
             if not auth_header.startswith('Basic ') or not self.check_basic_auth_key(auth_header[6:]):
                 return 'not authenticated'
+
+    @staticmethod
+    def __get_ssl_options(config):
+        return {option: config[option] for option in ['certfile', 'keyfile'] if option in config}
+
+    def __set_connection_string(self, connect_address):
+        self.connection_string = '{0}://{1}/patroni'.format(self.__protocol, connect_address or self.__listen)
+
+    def __set_config_parameters(self, config):
+        self.__auth_key = base64.b64encode(config['auth'].encode('utf-8')).decode('utf-8') if 'auth' in config else None
+        self.__set_connection_string(config.get('connect_address'))
+
+    def __initialize(self, config):
+        self.__ssl_options = self.__get_ssl_options(config)
+        self.__listen = config['listen']
+        host, port = config['listen'].split(':')
+        HTTPServer.__init__(self, (host, int(port)), RestApiHandler)
+        Thread.__init__(self, target=self.serve_forever)
+        self._set_fd_cloexec(self.socket)
+
+        self.__protocol = 'http'
+
+        # wrap socket with ssl if 'certfile' is defined in a config.yaml
+        # Sometime it's also needed to pass reference to a 'keyfile'.
+        if self.__ssl_options.get('certfile'):
+            import ssl
+            self.socket = ssl.wrap_socket(self.socket, server_side=True, **self.__ssl_options)
+            self.__protocol = 'https'
+        self.__set_connection_string(config.get('connect_address'))
+
+    def reload_config(self, config):
+        self.__set_config_parameters(config)
+        if self.__listen != config['listen'] or self.__ssl_options != self.__get_ssl_options(config):
+            self.shutdown()
+            self.__initialize(config)
+            self.start()

patroni/consul.py

@@ -72,12 +72,13 @@ class Consul(AbstractDCS):
 
     def __init__(self, name, config):
         super(Consul, self).__init__(name, config)
-        self.ttl = int((config.get('ttl') or 30)/2)  # My experiments have shown that session expires after 2*ttl time
+        self._ttl = None
+        self._session = None
+        self._my_member_data = None
+        self.set_ttl(config.get('ttl') or 30)
         host, port = config.get('host', '127.0.0.1:8500').split(':')
         self._client = ConsulClient(host=host, port=port)
         self._scope = config['scope']
-        self._session = None
-        self._my_member_data = None
         self.create_or_restore_session()
 
     def create_or_restore_session(self):
@@ -91,6 +92,12 @@ class Consul(AbstractDCS):
                 logger.info('waiting on consul')
                 sleep(5)
 
+    def set_ttl(self, ttl):
+        ttl = int(ttl/2)  # My experiments have shown that session expires after 2*ttl time
+        if self._ttl != ttl:
+            self._session = None
+        self._ttl = ttl
+
     def refresh_session(self):
         """:returns: `!True` if it had to create new session"""
         if self._session:
@@ -101,7 +108,7 @@ class Consul(AbstractDCS):
         if not self._session:
             name = self._scope + '-' + self._name
             try:
-                self._session = self._client.session.create(name=name, lock_delay=0, behavior='delete', ttl=self.ttl)
+                self._session = self._client.session.create(name=name, lock_delay=0, behavior='delete', ttl=self._ttl)
             except (ConsulException, RequestException):
                 logger.exception('session.create')
         if not self._session:

patroni/dcs.py

@@ -215,6 +215,10 @@ class AbstractDCS(object):
     def leader_optime_path(self):
         return self.client_path(self._LEADER_OPTIME)
 
+    @abc.abstractmethod
+    def set_ttl(self, ttl):
+        """Set the new ttl value for leader key"""
+
     @abc.abstractmethod
     def _load_cluster(self):
         """Internally this method should build  `Cluster` object which
@@ -285,12 +289,12 @@ class AbstractDCS(object):
         return self.set_failover_value(json.dumps(failover_value), index)
 
     @abc.abstractmethod
-    def touch_member(self, connection_string, ttl=None):
+    def touch_member(self, data, ttl=None):
         """Update member key in DCS.
         This method should create or update key with the name = '/members/' + `~self._name`
-        and value = connection_string in a given DCS.
+        and value = data in a given DCS.
 
-        :param connection_string: how this instance can be accessed by other instances
+        :param data: json serialized information about instance (including connection strings)
         :param ttl: ttl for member key, optional parameter. If it is None `~self.member_ttl will be used`
         :returns: `!True` on success otherwise `!False`
         """

patroni/etcd.py

@@ -193,7 +193,7 @@ class Etcd(AbstractDCS):
 
     def __init__(self, name, config):
         super(Etcd, self).__init__(name, config)
-        self.ttl = config.get('ttl', 30)
+        self.set_ttl(config.get('ttl', 30))
         self._retry = Retry(deadline=10, max_delay=1, max_tries=-1,
                             retry_exceptions=(etcd.EtcdConnectionFailed,
                                               etcd.EtcdLeaderElectionInProgress,
@@ -215,6 +215,9 @@ class Etcd(AbstractDCS):
                 sleep(5)
         return client
 
+    def set_ttl(self, ttl):
+        self.ttl = int(ttl)
+
     @staticmethod
     def member(node):
         return Member.from_node(node.modifiedIndex, os.path.basename(node.key), node.ttl, node.value)
@@ -255,8 +258,8 @@ class Etcd(AbstractDCS):
             raise EtcdError('Etcd is not responding properly')
 
     @catch_etcd_errors
-    def touch_member(self, connection_string, ttl=None):
-        return self.retry(self._client.set, self.member_path, connection_string, ttl or self.ttl)
+    def touch_member(self, data, ttl=None):
+        return self.retry(self._client.set, self.member_path, data, ttl or self.ttl)
 
     @catch_etcd_errors
     def take_leader(self):

patroni/ha.py

@@ -59,6 +59,8 @@ class Ha(object):
         }
         if self.patroni.tags:
             data['tags'] = self.patroni.tags
+        if self.state_handler.restart_pending:
+            data['restart_pending'] = True
         if not self._async_executor.busy and data['state'] in ['running', 'restarting', 'starting']:
             try:
                 data['xlog_location'] = self.state_handler.xlog_position()

patroni/postgresql.py

@@ -44,12 +44,15 @@ class Postgresql(object):
     def __init__(self, config):
         self.config = config
         self.name = config['name']
+        self._restart_pending = False
         self._server_parameters = self.get_server_parameters(config)
         self._listen_addresses, self._port = (config['listen'] + ':5432').split(':')[:2]
+        self._connect_address = config.get('connect_address')
+        self.replication = config['replication']
+        self.resolve_connection_addresses()
 
         self.scope = config['scope']
         self._data_dir = config['data_dir']
-        self.replication = config['replication']
         self.superuser = config.get('superuser') or {}
         self.admin = config.get('admin') or {}
 
@@ -71,11 +74,6 @@ class Postgresql(object):
 
         self._pg_ctl = ['pg_ctl', '-w', '-D', self._data_dir]
 
-        self.local_address = self.get_local_address()
-        connect_address = config.get('connect_address') or self.local_address
-        self.connection_string = 'postgres://{username}:{password}@{connect_address}/postgres'.format(
-            connect_address=connect_address, **self.replication)
-
         self._connection = None
         self._cursor_holder = None
         self._sysid = None
@@ -96,6 +94,40 @@ class Postgresql(object):
     def get_server_parameters(config):
         return {p: v for p, v in (config.get('parameters') or {}).items() if p not in ('listen_addresses', 'port')}
 
+    def resolve_connection_addresses(self):
+        self.local_address = self.get_local_address()
+        self.connection_string = 'postgres://{username}:{password}@{connect_address}/postgres'.format(
+            connect_address=self._connect_address or self.local_address, **self.replication)
+
+    def reload_config(self, config):
+        server_parameters = self.get_server_parameters(config)
+        self._connect_address = config.get('connect_address')
+        listen_addresses, port = (config['listen'] + ':5432').split(':')[:2]
+        if self._listen_addresses == listen_addresses and self._port == port:
+            self.resolve_connection_addresses()
+
+        self._listen_addresses = listen_addresses
+        self._port = port
+        if self.is_healthy():
+            changes = server_parameters.copy()
+            changes.update({p: None for p, v in self._server_parameters.items() if p not in server_parameters})
+            if changes:
+                for r in self.query("""SELECT name, setting
+                                         FROM pg_settings
+                                        WHERE context in ('internal', 'postmaster')
+                                          AND name IN (""" + ', '.join('%s' for _ in changes.keys()) + ')',
+                                    *(list(changes.keys()))):
+                    if server_parameters[r[0]] is None or str(server_parameters[r[0]]) != str(r[1]):
+                        self._restart_pending = True
+                        break
+        self._server_parameters = server_parameters
+        self._write_postgresql_conf()
+        self.reload()
+
+    @property
+    def restart_pending(self):
+        return self._restart_pending
+
     @property
     def can_rewind(self):
         """ check if pg_rewind executable is there and that pg_controldata indicates
@@ -378,10 +410,13 @@ class Postgresql(object):
         self._write_postgresql_conf()
         server_arguments = ['-o', "--listen_addresses='{0}' --port={1}".format(self._listen_addresses, self._port)]
         ret = subprocess.call(self._pg_ctl + ['start'] + server_arguments, env=env, preexec_fn=os.setsid) == 0
+        self._restart_pending = False
 
         self.set_state('running' if ret else 'start failed')
+        if ret:
+            self.resolve_connection_addresses()
+            self._schedule_load_slots = self.use_slots
 
-        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

patroni/zookeeper.py

@@ -103,6 +103,14 @@ class ZooKeeper(AbstractDCS):
         self._fetch_cluster = True
         self.event.set()
 
+    def set_ttl(self, ttl):
+        ttl = int(ttl * 1000)
+        # I know, it's weird to access private attributes and method
+        # but there is no other way to change session_timeout without losing session
+        if self._client._session_timeout != ttl:
+            self._client._session_timeout = ttl
+            self._client._connection._socket.close()
+
     def get_node(self, key, watch=None):
         try:
             ret = self._client.get(key, watch)

requirements.txt

@@ -3,7 +3,7 @@ psycopg2>=2.6.1
 PyYAML
 requests
 six >= 1.7
-kazoo>=2.2.1
+kazoo==2.2.1
 python-etcd==0.4.3
 python-consul==0.6.0
 click>=4.1

tests/test_api.py

@@ -19,6 +19,7 @@ class MockPostgresql(object):
     server_version = '999999'
     sysid = 'dummysysid'
     scope = 'dummy'
+    restart_pending = True
 
     @staticmethod
     def connection():
@@ -73,8 +74,10 @@ class MockRestApiServer(RestApiServer):
         BaseHTTPServer.HTTPServer.__init__ = Mock()
         MockRestApiServer._BaseServer__is_shut_down = Mock()
         MockRestApiServer._BaseServer__shutdown_request = True
-        config = {'listen': '127.0.0.1:8008', 'auth': 'test:test', 'certfile': 'dumb'}
+        config = {'listen': '127.0.0.1:8008', 'auth': 'test:test'}
         super(MockRestApiServer, self).__init__(MockPatroni(), config)
+        config['certfile'] = 'dumb'
+        self.reload_config(config)
         Handler(MockRequest(request), ('0.0.0.0', 8080), self)
 
 

tests/test_patroni.py

@@ -3,7 +3,6 @@ import os
 import sys
 import time
 import unittest
-import yaml
 
 from mock import Mock, patch
 from patroni.api import RestApiServer
@@ -22,6 +21,7 @@ from test_zookeeper import MockKazooClient
 @patch('subprocess.call', Mock(return_value=0))
 @patch('psycopg2.connect', psycopg2_connect)
 @patch.object(Postgresql, 'write_pg_hba', Mock())
+@patch.object(Postgresql, '_write_postgresql_conf', Mock())
 @patch.object(Postgresql, 'write_recovery_conf', Mock())
 @patch.object(BaseHTTPServer.HTTPServer, '__init__', Mock())
 @patch.object(AsyncExecutor, 'run', Mock())
@@ -37,9 +37,7 @@ class TestPatroni(unittest.TestCase):
             RestApiServer._BaseServer__is_shut_down = Mock()
             RestApiServer._BaseServer__shutdown_request = True
             RestApiServer.socket = 0
-            with open('postgres0.yml', 'r') as f:
-                config = yaml.load(f)
-                self.p = Patroni(config)
+            self.p = Patroni('postgres0.yml')
 
     @patch('patroni.zookeeper.KazooClient', MockKazooClient())
     @patch.object(Consul, 'create_or_restore_session', Mock())
@@ -71,6 +69,7 @@ class TestPatroni(unittest.TestCase):
             del os.environ[Patroni.PATRONI_CONFIG_VARIABLE]
 
     def test_run(self):
+        self.p.sighup_handler()
         self.p.ha.dcs.watch = Mock(side_effect=SleepException)
         self.p.api.start = Mock()
         self.assertRaises(SleepException, self.p.run)
@@ -95,3 +94,8 @@ class TestPatroni(unittest.TestCase):
         self.assertIsNone(self.p.replicatefrom)
         self.p.tags['replicatefrom'] = 'foo'
         self.assertEqual(self.p.replicatefrom, 'foo')
+
+    def test_reload_config(self):
+        self.p.reload_config()
+        with patch('yaml.load', Mock(side_effect=Exception)):
+            self.p.reload_config()

tests/test_postgresql.py

@@ -34,19 +34,10 @@ class MockCursor(object):
             self.results = [(False, )]
         elif sql.startswith('SELECT to_char(pg_postmaster_start_time'):
             self.results = [('', True, '', '', '', '', False)]
+        elif sql.startswith('SELECT name, setting'):
+            self.results = [('archive_mode', 'off')]
         else:
-            self.results = [(
-                None,
-                None,
-                None,
-                None,
-                None,
-                None,
-                None,
-                None,
-                None,
-                None,
-            )]
+            self.results = [(None, None, None, None, None, None, None, None, None, None)]
 
     def fetchone(self):
         return self.results[0]

tests/test_zookeeper.py

@@ -109,6 +109,9 @@ class TestZooKeeper(unittest.TestCase):
     def test_session_listener(self):
         self.zk.session_listener(KazooState.SUSPENDED)
 
+    def test_set_ttl(self):
+        self.zk.set_ttl(20)
+
     def test_get_node(self):
         self.assertIsNone(self.zk.get_node('/no_node'))