Factor out postmaster process (#561)

Introduces a PostmasterProcess object that identifies a running process via pid and start time.
When pid file is parsed and the correct process identified this object is passed around.
When the process goes away we try to find a new one in case somebody restarted postgres behind our back.

patroni/ha.py

@@ -880,7 +880,7 @@ class Ha(object):
                     # background thread and has not even written a pid file yet.
                     with self._async_executor.critical_task as task:
                         if not task.cancel():
-                            self.state_handler.terminate_starting_postmaster(pid=task.result)
+                            self.state_handler.terminate_starting_postmaster(postmaster=task.result)
                 self.demote('immediate-nolock')
                 return 'lost leader lock during ' + self._async_executor.scheduled_action
 
@@ -1025,8 +1025,8 @@ class Ha(object):
             # is data directory empty?
             if self.state_handler.data_directory_empty():
                 self.state_handler.set_role('uninitialized')
-                self.state_handler.stop()
-                # In case datadir went away while we were master. TODO: check for this and try to stop postgresql.
+                self.state_handler.stop('immediate')
+                # In case datadir went away while we were master.
                 self.watchdog.disable()
 
                 # is this instance the leader?

patroni/postgresql.py

@@ -1,12 +1,9 @@
 import logging
-import errno
 import os
 import psycopg2
-import psutil
 import re
 import shlex
 import shutil
-import signal
 import socket
 import subprocess
 import tempfile
@@ -14,10 +11,10 @@ import time
 
 from collections import defaultdict
 from contextlib import contextmanager
-from patroni import call_self
 from patroni.callback_executor import CallbackExecutor
 from patroni.exceptions import PostgresConnectionException
-from patroni.utils import compare_values, parse_bool, parse_int, Retry, RetryFailedError, polling_loop, int_or_none
+from patroni.utils import compare_values, parse_bool, parse_int, Retry, RetryFailedError, polling_loop
+from patroni.postmaster import PostmasterProcess
 from six import string_types
 from six.moves.urllib.parse import quote_plus
 from threading import current_thread, Lock
@@ -35,11 +32,6 @@ STATE_REJECT = 'rejecting connections'
 STATE_NO_RESPONSE = 'not responding'
 STATE_UNKNOWN = 'unknown'
 
-STOP_SIGNALS = {
-    'smart': signal.SIGTERM,
-    'fast': signal.SIGINT,
-    'immediate': signal.SIGQUIT,
-}
 STOP_POLLING_INTERVAL = 1
 REWIND_STATUS = type('Enum', (), {'INITIAL': 0, 'CHECK': 1, 'NEED': 2, 'NOT_NEED': 3, 'SUCCESS': 4, 'FAILED': 5})
 sync_standby_name_re = re.compile('^[A-Za-z_][A-Za-z_0-9\$]*$')
@@ -71,25 +63,6 @@ def null_context():
     yield
 
 
-def _update_postmaster_cached_info(func):
-    def wrapper(self):
-        ret = func(self)
-        if ret and 'pid' in ret and 'start_time' in ret:
-            old_pid = self._postmaster_cached_info.get('pid', 0)
-            old_start_time = self._postmaster_cached_info.get('start_time', 0)
-            try:
-                pmpid = int(ret['pid'])
-                pmstart = int(ret['start_time'])
-                if pmpid != old_pid or pmstart != old_start_time:  # this check removes repeating messages from logs
-                    self._postmaster_cached_info = {'pid': pmpid, 'start_time': pmstart}
-                    logger.info("Updated postmaster info: %s .", self._postmaster_cached_info)
-            except ValueError:
-                logger.warning('Cannot update postmaster info with data due garbage in pid file: %s', ret)
-        return ret
-
-    return wrapper
-
-
 class Postgresql(object):
 
     # List of parameters which must be always passed to postmaster as command line options
@@ -161,7 +134,6 @@ class Postgresql(object):
         self._pg_hba_conf = os.path.join(self._config_dir, 'pg_hba.conf')
         self._recovery_conf = os.path.join(self._data_dir, 'recovery.conf')
         self._postmaster_pid = os.path.join(self._data_dir, 'postmaster.pid')
-        self._postmaster_cached_info = {'pid': 0, 'start_time': 0}
         self._trigger_file = config.get('recovery_conf', {}).get('trigger_file') or 'promote'
         self._trigger_file = os.path.abspath(os.path.join(self._data_dir, self._trigger_file))
 
@@ -184,6 +156,9 @@ class Postgresql(object):
 
         self._state_entry_timestamp = None
 
+        # Last known running process
+        self._postmaster_proc = None
+
         if self.is_running():
             self.set_state('running')
             self.set_role('master' if self.is_leader() else 'replica')
@@ -724,16 +699,17 @@ class Postgresql(object):
             raise PostgresConnectionException(str(e))
 
     def is_running(self):
-        if not (self._version_file_exists() and os.path.isfile(self._postmaster_pid)):
-            # XXX: This is dangerous in case somebody deletes the data directory while PostgreSQL is still running.
-            return False
+        """Returns PostmasterProcess if one is running on the data directory or None. If most recently seen process
+        is running udpates the cached process based on pid file."""
+        if self._postmaster_proc:
+            if self._postmaster_proc.is_running():
+                return self._postmaster_proc
+            self._postmaster_proc = None
 
-        pidfile = self.read_pid_file()
-        return self._is_postmaster_pid_running(int_or_none(pidfile.get('pid')),
-                                               start_time=int_or_none(pidfile.get('start_time')))
+        self._postmaster_proc = PostmasterProcess.from_pidfile(self._read_pid_file())
+        return self._postmaster_proc
 
-    @_update_postmaster_cached_info
-    def read_pid_file(self):
+    def _read_pid_file(self):
         """Reads and parses postmaster.pid from the data directory
 
         :returns dictionary of values if successful, empty dictionary otherwise
@@ -745,61 +721,6 @@ class Postgresql(object):
         except IOError:
             return {}
 
-    def get_pid(self):
-        """Fetches pid value from postmaster.pid using read_pid_file
-
-        :returns pid if successful, 0 if pid file is not present"""
-        # TODO: figure out what to do on permission errors
-        pid = self.read_pid_file().get('pid', 0)
-        try:
-            return int(pid)
-        except ValueError:
-            logger.warning("Garbage pid in postmaster.pid: {0!r}".format(pid))
-            return 0
-
-    def get_pid_with_lost_data_dir(self):
-        logger.info("Trying to check if process running without directory "
-                    "with cached postmaster info: %s .", self._postmaster_cached_info)
-        try:
-            process = psutil.Process(self._postmaster_cached_info['pid'])
-            # check difference instead of values because of rounding issues
-            if abs(self._postmaster_cached_info["start_time"] - process.create_time()) < 2:
-                return process.pid
-            else:
-                logger.info("Process with pid %s was started at different time %s .",
-                            process.pid, process.create_time())
-        except psutil.NoSuchProcess:
-            logger.info("Cannot find process %s .", self._postmaster_cached_info['pid'])
-        return 0
-
-    def clean_postmaster_cached_info(self):
-        self._postmaster_cached_info = {'pid': 0, 'start_time': 0}
-        logger.info("postmaster info was cleaned.")
-
-    @staticmethod
-    def _is_postmaster_pid_running(pid, start_time=None):
-        # Normalize pid handling missing values and negative pids from postmaster.pid
-        if not pid:
-            return False
-        if pid < 0:
-            pid = -pid
-
-        try:
-            proc = psutil.Process(pid)
-        except psutil.NoSuchProcess:
-            return False
-
-        # If the process is Patroni or Patronis host process or Patronis child process then it's a false positive
-        my_pid = os.getpid()
-        if pid == my_pid or pid == os.getppid() or proc.parent() == my_pid:
-            return False
-
-        # If process start time differs by more than 3 seconds it's a false positive
-        if start_time is not None and abs(proc.create_time() - start_time) > 3:
-            return False
-
-        return True
-
     @property
     def cb_called(self):
         return self.__cb_called
@@ -844,30 +765,20 @@ class Postgresql(object):
     def is_starting(self):
         return self.state == 'starting'
 
-    def wait_for_port_open(self, pid, initiated, timeout):
+    def wait_for_port_open(self, postmaster, timeout):
         """Waits until PostgreSQL opens ports."""
         for _ in polling_loop(timeout):
-            pid_file = self.read_pid_file()
-            if len(pid_file) > 5:
-                try:
-                    pmpid = int(pid_file['pid'])
-                    pmstart = int(pid_file['start_time'])
-
-                    if pmstart >= initiated - 2 and pmpid == pid:
-                        isready = self.pg_isready()
-                        if isready != STATE_NO_RESPONSE:
-                            if isready not in [STATE_REJECT, STATE_RUNNING]:
-                                logger.warning("Can't determine PostgreSQL startup status, assuming running")
-                            return True
-                except ValueError:
-                    # Garbage in the pid file
-                    pass
-
-            if not self._is_postmaster_pid_running(pid, start_time=initiated):
+            if not postmaster.is_running():
                 logger.error('postmaster is not running')
                 self.set_state('start failed')
                 return False
 
+            isready = self.pg_isready()
+            if isready != STATE_NO_RESPONSE:
+                if isready not in [STATE_REJECT, STATE_RUNNING]:
+                    logger.warning("Can't determine PostgreSQL startup status, assuming running")
+                return True
+
         logger.warning("Timed out waiting for PostgreSQL to start")
         return False
 
@@ -903,33 +814,17 @@ class Postgresql(object):
         options = ['--{0}={1}'.format(p, self._server_parameters[p]) for p in self.CMDLINE_OPTIONS
                    if p in self._server_parameters and p != 'wal_keep_segments']
 
-        # Unfortunately `pg_ctl start` does not return postmaster pid to us. Without this information
-        # it is hard to know the current state of postgres startup, so we had to reimplement pg_ctl start
-        # in python. It will start postgres, wait for port to be open and wait until postgres will start
-        # accepting connections.
-        # Important!!! We can't just start postgres using subprocess.Popen, because in this case it
-        # will be our child for the rest of our live and we will have to take care of it (`waitpid`).
-        # So we will use the same approach as pg_ctl uses: start a new process, which will start postgres.
-        # This process will write postmaster pid to stdout and exit immediately. Now it's responsibility
-        # of init process to take care about postmaster.
-        # In order to make everything portable we can't use fork&exec approach here, so  we will call
-        # ourselves and pass list of arguments which must be used to start postgres.
         with task or null_context():
             if task and task.is_cancelled:
                 logger.info("PostgreSQL start cancelled.")
                 return False
 
-            start_initiated = time.time()
-            proc = call_self(['pg_ctl_start', self._pgcommand('postgres'), '-D', self._data_dir,
-                             '--config-file={}'.format(self._postgresql_conf)] + options, close_fds=True,
-                             preexec_fn=os.setsid, stdout=subprocess.PIPE,
-                             env={p: os.environ[p] for p in ('PATH', 'LC_ALL', 'LANG') if p in os.environ})
-            pid = int(proc.stdout.readline().strip())
-            proc.wait()
-            logger.info('postmaster pid=%s', pid)
-
+            self._postmaster_proc = PostmasterProcess.start(self._pgcommand('postgres'),
+                                                            self._data_dir,
+                                                            self._postgresql_conf,
+                                                            options)
             if task:
-                task.complete(pid)
+                task.complete(self._postmaster_proc)
 
         start_timeout = timeout
         if not start_timeout:
@@ -939,7 +834,7 @@ class Postgresql(object):
                 start_timeout = 60
 
         # We want postmaster to open ports before we continue
-        if not self.wait_for_port_open(pid, start_initiated, start_timeout):
+        if not self._postmaster_proc or not self.wait_for_port_open(self._postmaster_proc, start_timeout):
             return False
 
         ret = self.wait_for_startup(start_timeout)
@@ -967,7 +862,7 @@ class Postgresql(object):
             logging.exception('Exception during CHECKPOINT')
             return 'not accessible or not healty'
 
-    def stop(self, mode='fast', block_callbacks=False, checkpoint=True, on_safepoint=None):
+    def stop(self, mode='fast', block_callbacks=False, checkpoint=None, on_safepoint=None):
         """Stop PostgreSQL
 
         Supports a callback when a safepoint is reached. A safepoint is when no user backend can return a successful
@@ -976,6 +871,9 @@ class Postgresql(object):
 
         :param on_safepoint: This callback is called when no user backends are running.
         """
+        if checkpoint is None:
+            checkpoint = False if mode == 'immediate' else True
+
         success, pg_signaled = self._do_stop(mode, block_callbacks, checkpoint, on_safepoint)
         if success:
             # block_callbacks is used during restart to avoid
@@ -990,13 +888,8 @@ class Postgresql(object):
         return success
 
     def _do_stop(self, mode, block_callbacks, checkpoint, on_safepoint):
-        if not self.is_running():
-            if self.data_directory_empty() and self._postmaster_cached_info['pid']:
-                pid = self.get_pid_with_lost_data_dir()
-                if pid > 0:
-                    self.terminate_starting_postmaster(pid)
-                    self.clean_postmaster_cached_info()
-                    return True, True
+        postmaster = self.is_running()
+        if not postmaster:
             if on_safepoint:
                 on_safepoint()
             return True, False
@@ -1008,86 +901,38 @@ class Postgresql(object):
             self.set_state('stopping')
 
         # Send signal to postmaster to stop
-        pid, result = self._signal_postmaster_stop(mode)
-        if result is not None:
-            if result and on_safepoint:
+        success = postmaster.signal_stop(mode)
+        if success is not None:
+            if success and on_safepoint:
                 on_safepoint()
-            return result, True
+            return success, True
 
         # We can skip safepoint detection if we don't have a callback
         if on_safepoint:
             # Wait for our connection to terminate so we can be sure that no new connections are being initiated
-            self._wait_for_connection_close(pid)
-            self._wait_for_user_backends_to_close(pid)
+            self._wait_for_connection_close(postmaster)
+            postmaster.wait_for_user_backends_to_close()
             on_safepoint()
 
-        self._wait_for_postmaster_stop(pid)
-        self.clean_postmaster_cached_info()
+        postmaster.wait()
 
         return True, True
 
-    def _wait_for_postmaster_stop(self, pid):
-        # This wait loop differs subtly from pg_ctl as we check for both the pid file going
-        # away and if the pid is running. This seems safer.
-        while pid == self.get_pid() and self._is_postmaster_pid_running(pid):
-            time.sleep(STOP_POLLING_INTERVAL)
-
-    def _signal_postmaster_stop(self, mode):
-        pid = self.get_pid()
-        if pid == 0:
-            return None, True
-        elif pid < 0:
-            logger.warning("Cannot stop server; single-user server is running (PID: {0})".format(-pid))
-            return None, False
-        try:
-            os.kill(pid, STOP_SIGNALS[mode])
-        except OSError as e:
-            if e.errno == errno.ESRCH:
-                return None, True
-            else:
-                logger.warning("Could not send stop signal to PostgreSQL (error: {0})".format(e.errno))
-                return None, False
-        return pid, None
-
-    def terminate_starting_postmaster(self, pid):
+    def terminate_starting_postmaster(self, postmaster):
         """Terminates a postmaster that has not yet opened ports or possibly even written a pid file. Blocks
         until the process goes away."""
-        try:
-            os.kill(pid, STOP_SIGNALS['immediate'])
-        except OSError as e:
-            if e.errno == errno.ESRCH:
-                return
-            logger.warning("Could not send stop signal to PostgreSQL (error: {0})".format(e.errno))
+        postmaster.signal_stop('immediate')
+        postmaster.wait()
 
-        while self._is_postmaster_pid_running(pid):
-            time.sleep(STOP_POLLING_INTERVAL)
-
-    def _wait_for_connection_close(self, pid):
+    def _wait_for_connection_close(self, postmaster):
         try:
             with self.connection().cursor() as cur:
-                while pid == self.get_pid() and self._is_postmaster_pid_running(pid):  # Need a timeout here?
+                while postmaster.is_running():  # Need a timeout here?
                     cur.execute("SELECT 1")
                     time.sleep(STOP_POLLING_INTERVAL)
         except psycopg2.Error:
             pass
 
-    @staticmethod
-    def _wait_for_user_backends_to_close(postmaster_pid):
-        # These regexps are cross checked against versions PostgreSQL 9.1 .. 9.6
-        aux_proc_re = re.compile("(?:postgres:)( .*:)? (?:""(?:startup|logger|checkpointer|writer|wal writer|"
-                                 "autovacuum launcher|autovacuum worker|stats collector|wal receiver|archiver|"
-                                 "wal sender) process|bgworker: )")
-
-        try:
-            postmaster = psutil.Process(postmaster_pid)
-            user_backends = [p for p in postmaster.children() if not aux_proc_re.match(p.cmdline()[0])]
-            logger.debug("Waiting for user backends {0} to close".format(
-                ",".join(p.cmdline()[0] for p in user_backends)))
-            psutil.wait_procs(user_backends)
-            logger.debug("Backends closed")
-        except psutil.NoSuchProcess:
-            return
-
     def reload(self):
         ret = self.pg_ctl('reload')
         if ret:

patroni/postmaster.py

@@ -0,0 +1,115 @@
+import logging
+import os
+import psutil
+import re
+import signal
+import subprocess
+
+from patroni import call_self
+
+logger = logging.getLogger(__name__)
+
+STOP_SIGNALS = {
+    'smart': signal.SIGTERM,
+    'fast': signal.SIGINT,
+    'immediate': signal.SIGQUIT,
+}
+
+
+class PostmasterProcess(psutil.Process):
+    def __init__(self, pid):
+        self.is_single_user = False
+        if pid < 0:
+            pid = -pid
+            self.is_single_user = True
+        super(PostmasterProcess, self).__init__(pid)
+
+    @classmethod
+    def from_pidfile(cls, pidfile):
+        try:
+            pid = int(pidfile.get('pid', 0))
+            if not pid:
+                return None
+        except ValueError:
+            return None
+
+        try:
+            proc = cls(pid)
+        except psutil.NoSuchProcess:
+            return None
+
+        try:
+            start_time = int(pidfile.get('start_time', 0))
+            if start_time and abs(proc.create_time() - start_time) > 3:
+                return None
+        except ValueError:
+            logger.warning("Garbage start time value in pid file: %r", pidfile.get('start_time'))
+
+        # Extra safety check. The process can't be ourselves, our parent or our direct child.
+        if proc.pid == os.getpid() or proc.pid == os.getppid() or proc.parent() == os.getpid():
+            return None
+
+        return proc
+
+    @classmethod
+    def from_pid(cls, pid):
+        try:
+            return cls(pid)
+        except psutil.NoSuchProcess:
+            return None
+
+    def signal_stop(self, mode):
+        """Signal postmaster process to stop
+
+        :returns None if signaled, True if process is already gone, False if error
+        """
+        if self.is_single_user:
+            logger.warning("Cannot stop server; single-user server is running (PID: {0})".format(self.pid))
+            return False
+        try:
+            self.send_signal(STOP_SIGNALS[mode])
+        except psutil.NoSuchProcess:
+            return True
+        except psutil.AccessDenied as e:
+            logger.warning("Could not send stop signal to PostgreSQL (error: {0})".format(e))
+            return False
+
+        return None
+
+    def wait_for_user_backends_to_close(self):
+        # These regexps are cross checked against versions PostgreSQL 9.1 .. 9.6
+        aux_proc_re = re.compile("(?:postgres:)( .*:)? (?:""(?:startup|logger|checkpointer|writer|wal writer|"
+                                 "autovacuum launcher|autovacuum worker|stats collector|wal receiver|archiver|"
+                                 "wal sender) process|bgworker: )")
+
+        user_backends = [p for p in self.children() if not aux_proc_re.match(p.cmdline()[0])]
+        logger.debug("Waiting for user backends {0} to close".format(
+            ",".join(p.cmdline()[0] for p in user_backends)))
+        psutil.wait_procs(user_backends)
+        logger.debug("Backends closed")
+
+    @classmethod
+    def start(cls, pgcommand, data_dir, conf, options):
+        # Unfortunately `pg_ctl start` does not return postmaster pid to us. Without this information
+        # it is hard to know the current state of postgres startup, so we had to reimplement pg_ctl start
+        # in python. It will start postgres, wait for port to be open and wait until postgres will start
+        # accepting connections.
+        # Important!!! We can't just start postgres using subprocess.Popen, because in this case it
+        # will be our child for the rest of our live and we will have to take care of it (`waitpid`).
+        # So we will use the same approach as pg_ctl uses: start a new process, which will start postgres.
+        # This process will write postmaster pid to stdout and exit immediately. Now it's responsibility
+        # of init process to take care about postmaster.
+        # In order to make everything portable we can't use fork&exec approach here, so  we will call
+        # ourselves and pass list of arguments which must be used to start postgres.
+        proc = call_self(['pg_ctl_start', pgcommand, '-D', data_dir,
+                          '--config-file={}'.format(conf)] + options, close_fds=True,
+                         preexec_fn=os.setsid, stdout=subprocess.PIPE,
+                         env={p: os.environ[p] for p in ('PATH', 'LC_ALL', 'LANG') if p in os.environ})
+        pid = int(proc.stdout.readline().strip())
+        proc.wait()
+        logger.info('postmaster pid=%s', pid)
+
+        # TODO: In an extremely unlikely case, the process could have exited and the pid reassigned. The start
+        # initiation time is not accurate enough to compare to create time as start time would also likely
+        # be relatively close. We need the subprocess extract pid+start_time in a race free manner.
+        return PostmasterProcess.from_pid(pid)

patroni/utils.py

@@ -280,11 +280,3 @@ def polling_loop(timeout, interval=1):
         yield iteration
         iteration += 1
         time.sleep(interval)
-
-
-def int_or_none(val):
-    """Returns integer value of the parameter if convertible to int, None otherwise."""
-    try:
-        return int(val)
-    except (ValueError, TypeError):
-        return None

tests/test_ha.py

@@ -13,7 +13,7 @@ from patroni.postgresql import Postgresql
 from patroni.watchdog import Watchdog
 from patroni.utils import tzutc
 from test_etcd import socket_getaddrinfo, etcd_read, etcd_write, requests_get
-from test_postgresql import psycopg2_connect
+from test_postgresql import psycopg2_connect, MockPostmaster
 
 
 def true(*args, **kwargs):
@@ -112,7 +112,7 @@ def run_async(self, func, args=()):
     return func(*args) if args else func()
 
 
-@patch.object(Postgresql, 'is_running', Mock(return_value=True))
+@patch.object(Postgresql, 'is_running', Mock(return_value=MockPostmaster()))
 @patch.object(Postgresql, 'is_leader', Mock(return_value=True))
 @patch.object(Postgresql, 'wal_position', Mock(return_value=10))
 @patch.object(Postgresql, 'call_nowait', Mock(return_value=True))
@@ -345,7 +345,7 @@ class TestHa(unittest.TestCase):
         self.ha.cluster = get_cluster_not_initialized_without_leader()
         self.e.initialize = true
         self.ha.bootstrap()
-        self.p.is_running = true
+        self.p.is_running.return_value = MockPostmaster()
         self.p.is_leader = true
         with patch.object(Watchdog, 'activate', Mock(return_value=False)):
             self.assertEquals(self.ha.post_bootstrap(), 'running post_bootstrap')
@@ -392,9 +392,9 @@ class TestHa(unittest.TestCase):
             self.ha.update_lock = false
             self.p.set_role('master')
             with patch('patroni.async_executor.CriticalTask.cancel', Mock(return_value=False)):
-                with patch('patroni.postgresql.Postgresql.stop') as stop_mock:
+                with patch('patroni.postgresql.Postgresql.terminate_starting_postmaster') as mock_terminate:
                     self.assertEquals(self.ha.run_cycle(), 'lost leader lock during restart')
-                    stop_mock.assert_called()
+                    mock_terminate.assert_called()
 
     @patch('requests.get', requests_get)
     def test_manual_failover_from_leader(self):

tests/test_patroni.py

@@ -12,7 +12,7 @@ from patroni.exceptions import DCSError
 from patroni import Patroni, main as _main, patroni_main
 from six.moves import BaseHTTPServer
 from test_etcd import SleepException, etcd_read, etcd_write
-from test_postgresql import Postgresql, psycopg2_connect
+from test_postgresql import Postgresql, psycopg2_connect, MockPostmaster
 
 
 class MockFrozenImporter(object):
@@ -26,7 +26,7 @@ class MockFrozenImporter(object):
 @patch.object(Postgresql, 'write_pg_hba', Mock())
 @patch.object(Postgresql, '_write_postgresql_conf', Mock())
 @patch.object(Postgresql, 'write_recovery_conf', Mock())
-@patch.object(Postgresql, 'is_running', Mock(return_value=True))
+@patch.object(Postgresql, 'is_running', Mock(return_value=MockPostmaster()))
 @patch.object(Postgresql, 'call_nowait', Mock())
 @patch.object(BaseHTTPServer.HTTPServer, '__init__', Mock())
 @patch.object(AsyncExecutor, 'run', Mock())

tests/test_postgresql.py

@@ -12,6 +12,7 @@ from patroni.async_executor import CriticalTask
 from patroni.dcs import Cluster, Leader, Member, SyncState
 from patroni.exceptions import PostgresConnectionException
 from patroni.postgresql import Postgresql, STATE_REJECT, STATE_NO_RESPONSE
+from patroni.postmaster import PostmasterProcess
 from patroni.utils import RetryFailedError
 from six.moves import builtins
 from threading import Thread
@@ -96,6 +97,14 @@ class MockConnect(object):
         pass
 
 
+class MockPostmaster(object):
+    def __init__(self, is_running=True, is_single_master=False):
+        self.is_running = Mock(return_value=is_running)
+        self.is_single_master = Mock(return_value=is_single_master)
+        self.wait_for_user_backends_to_close = Mock()
+        self.signal_stop = Mock(return_value=None)
+        self.wait = Mock()
+
 def pg_controldata_string(*args, **kwargs):
     return b"""
 pg_control version number:            942
@@ -213,100 +222,77 @@ class TestPostgresql(unittest.TestCase):
     @patch.object(Postgresql, 'wait_for_port_open')
     @patch.object(Postgresql, 'is_running')
     def test_start(self, mock_is_running, mock_wait_for_port_open, mock_wait_for_startup, mock_popen):
-        mock_is_running.return_value = True
+        mock_is_running.return_value = MockPostmaster()
         mock_wait_for_port_open.return_value = True
         mock_wait_for_startup.return_value = False
         mock_popen.return_value.stdout.readline.return_value = '123'
         self.assertTrue(self.p.start())
-        mock_is_running.return_value = False
-        open(os.path.join(self.data_dir, 'postmaster.pid'), 'w').close()
-        pg_conf = os.path.join(self.data_dir, 'postgresql.conf')
-        open(pg_conf, 'w').close()
-        self.assertFalse(self.p.start(task=CriticalTask()))
-        with open(pg_conf) as f:
-            lines = f.readlines()
-            self.assertTrue("f.oo = 'bar'\n" in lines)
+        mock_is_running.return_value = None
 
-        mock_wait_for_startup.return_value = None
-        self.assertFalse(self.p.start(10))
-        self.assertIsNone(self.p.start())
+        mock_postmaster = MockPostmaster()
+        with patch.object(PostmasterProcess, 'start', return_value=mock_postmaster):
+            pg_conf = os.path.join(self.data_dir, 'postgresql.conf')
+            open(pg_conf, 'w').close()
+            self.assertFalse(self.p.start(task=CriticalTask()))
 
-        mock_wait_for_port_open.return_value = False
-        self.assertFalse(self.p.start())
-        task = CriticalTask()
-        task.cancel()
-        self.assertFalse(self.p.start(task=task))
+            with open(pg_conf) as f:
+                lines = f.readlines()
+                self.assertTrue("f.oo = 'bar'\n" in lines)
+
+            mock_wait_for_startup.return_value = None
+            self.assertFalse(self.p.start(10))
+            self.assertIsNone(self.p.start())
+
+            mock_wait_for_port_open.return_value = False
+            self.assertFalse(self.p.start())
+            task = CriticalTask()
+            task.cancel()
+            self.assertFalse(self.p.start(task=task))
 
     @patch.object(Postgresql, 'pg_isready')
-    @patch.object(Postgresql, 'read_pid_file')
-    @patch.object(Postgresql, '_is_postmaster_pid_running')
     @patch('patroni.postgresql.polling_loop', Mock(return_value=range(1)))
-    def test_wait_for_port_open(self, mock_is_postmaster_pid_running, mock_read_pid_file, mock_pg_isready):
-        mock_is_postmaster_pid_running.return_value = False
+    def test_wait_for_port_open(self, mock_pg_isready):
         mock_pg_isready.return_value = STATE_NO_RESPONSE
+        mock_postmaster = MockPostmaster(is_running=False)
 
         # No pid file and postmaster death
-        mock_read_pid_file.return_value = {}
-        self.assertFalse(self.p.wait_for_port_open(42, 100., 1))
+        self.assertFalse(self.p.wait_for_port_open(mock_postmaster, 1))
 
-        mock_is_postmaster_pid_running.return_value = True
+        mock_postmaster.is_running.return_value = True
 
         # timeout
-        mock_read_pid_file.return_value = {'pid', 1}
-        self.assertFalse(self.p.wait_for_port_open(42, 100., 1))
-
-        # Garbage pid
-        mock_read_pid_file.return_value = {'pid': 'garbage', 'start_time': '101', 'data_dir': '',
-                                           'socket_dir': '', 'port': '', 'listen_addr': ''}
-        self.assertFalse(self.p.wait_for_port_open(42, 100., 1))
-
-        # Not ready
-        mock_read_pid_file.return_value = {'pid': '42', 'start_time': '101', 'data_dir': '',
-                                           'socket_dir': '', 'port': '', 'listen_addr': ''}
-        self.assertFalse(self.p.wait_for_port_open(42, 100., 1))
+        self.assertFalse(self.p.wait_for_port_open(mock_postmaster, 1))
 
         # pg_isready failure
         mock_pg_isready.return_value = 'garbage'
-        self.assertTrue(self.p.wait_for_port_open(42, 100., 1))
+        self.assertTrue(self.p.wait_for_port_open(mock_postmaster, 1))
 
     @patch('time.sleep', Mock())
     @patch.object(Postgresql, 'is_running')
-    @patch.object(Postgresql, 'get_pid')
-    def test_stop(self, mock_get_pid, mock_is_running):
+    @patch.object(Postgresql, '_wait_for_connection_close', Mock())
+    def test_stop(self, mock_is_running):
+        # Postmaster is not running
         mock_callback = Mock()
-        mock_is_running.return_value = False
+        mock_is_running.return_value = None
         self.assertTrue(self.p.stop(on_safepoint=mock_callback))
         mock_callback.assert_called()
 
-        with patch.object(Postgresql, '_is_postmaster_pid_running', Mock(return_value=False)), \
-                patch.object(Postgresql, 'data_directory_empty', Mock(return_value=True)):
-            with patch('psutil.Process') as mock_psutil:
-                self.p._postmaster_cached_info = {'pid': 1, 'start_time': 1}
-                mock_psutil.return_value.pid = 1
-                mock_psutil.return_value.create_time.return_value = 1
-                self.assertTrue(self.p.stop())
-                self.p._postmaster_cached_info = {'pid': 1, 'start_time': 1}
-                mock_psutil.return_value.create_time.return_value = 100
-                self.assertTrue(self.p.stop())
-                self.p._postmaster_cached_info = {'pid': 1, 'start_time': 1}
-                mock_psutil.side_effect = psutil.NoSuchProcess('')
-                self.assertTrue(self.p.stop())
-
-        mock_is_running.return_value = True
-        mock_get_pid.return_value = 0
+        # Is running, stopped successfully
+        mock_is_running.return_value = mock_postmaster = MockPostmaster()
         mock_callback.reset_mock()
         self.assertTrue(self.p.stop(on_safepoint=mock_callback))
         mock_callback.assert_called()
-        mock_get_pid.return_value = -1
+        mock_postmaster.signal_stop.assert_called()
+
+        # Stop signal failed
+        mock_postmaster.signal_stop.return_value = False
         self.assertFalse(self.p.stop())
-        mock_get_pid.return_value = 123
-        with patch('os.kill', Mock(side_effect=[OSError(errno.ESRCH, ''), OSError, None])):
-            self.assertTrue(self.p.stop())
-            self.assertFalse(self.p.stop())
-            self.assertTrue(self.p.stop())
-        with patch.object(Postgresql, '_signal_postmaster_stop', Mock(return_value=(123, None))):
-            with patch.object(Postgresql, '_is_postmaster_pid_running', Mock(side_effect=[True, False, False])):
-                self.assertTrue(self.p.stop())
+
+        # Stop signal failed to find process
+        mock_postmaster.signal_stop.return_value = True
+        mock_callback.reset_mock()
+        self.assertTrue(self.p.stop(on_safepoint=mock_callback))
+        mock_callback.assert_called()
 
     def test_restart(self):
         self.p.start = Mock(return_value=False)
@@ -496,14 +482,22 @@ class TestPostgresql(unittest.TestCase):
         self.assertEquals(self.p.last_operation(), '2')
         Thread(target=self.p.last_operation).start()
 
-    @patch('os.path.isfile', Mock(return_value=True))
-    @patch('os.kill', Mock(side_effect=Exception))
-    @patch('os.getpid', Mock(return_value=2))
-    @patch('os.getppid', Mock(return_value=2))
-    @patch.object(builtins, 'open', mock_open(read_data='-1'))
-    @patch.object(Postgresql, '_version_file_exists', Mock(return_value=True))
-    def test_is_running(self):
-        self.assertFalse(self.p.is_running())
+    @patch.object(PostmasterProcess, 'from_pidfile')
+    def test_is_running(self, mock_frompidfile):
+        # Cached postmaster running
+        mock_postmaster = self.p._postmaster_proc = MockPostmaster()
+        self.assertEquals(self.p.is_running(), mock_postmaster)
+
+        # Cached postmaster not running, no postmaster running
+        mock_postmaster.is_running.return_value = False
+        mock_frompidfile.return_value = None
+        self.assertEquals(self.p.is_running(), None)
+        self.assertEquals(self.p._postmaster_proc, None)
+
+        # No cached postmaster, postmaster running
+        mock_frompidfile.return_value = mock_postmaster2 = MockPostmaster()
+        self.assertEquals(self.p.is_running(), mock_postmaster2)
+        self.assertEquals(self.p._postmaster_proc, mock_postmaster2)
 
     @patch('shlex.split', Mock(side_effect=OSError))
     def test_call_nowait(self):
@@ -515,7 +509,7 @@ class TestPostgresql(unittest.TestCase):
     def test_non_existing_callback(self):
         self.assertFalse(self.p.call_nowait('foobar'))
 
-    @patch.object(Postgresql, 'is_running', Mock(return_value=True))
+    @patch.object(Postgresql, 'is_running', Mock(return_value=MockPostmaster()))
     def test_is_leader_exception(self):
         self.p.start()
         self.p.query = Mock(side_effect=psycopg2.OperationalError("not supported"))
@@ -786,21 +780,11 @@ class TestPostgresql(unittest.TestCase):
         pidfile = os.path.join(self.data_dir, 'postmaster.pid')
         if os.path.exists(pidfile):
             os.remove(pidfile)
-        self.assertEquals(self.p.read_pid_file(), {})
-
-    @patch.object(Postgresql, '_version_file_exists', Mock(return_value=True))
-    @patch('os.path.isfile', Mock(return_value=True))
-    @patch.object(Postgresql, 'read_pid_file')
-    @patch('psutil.Process')
-    def test_is_postmaster_pid_running(self, mock_psutil, mock_read_pid_file):
-        mock_psutil.return_value.create_time.return_value = 1
-        mock_read_pid_file.return_value = {'pid': -100, 'start_time': 1}
-        self.assertTrue(self.p.is_running())
-        with patch('os.getpid', Mock(return_value=100)):
-            mock_read_pid_file.return_value = {'pid': 100, 'start_time': 1}
-            self.assertFalse(self.p.is_running())
-        mock_read_pid_file.return_value = {'pid': 100, 'start_time': 100}
-        self.assertFalse(self.p.is_running())
+        self.assertEquals(self.p._read_pid_file(), {})
+        with open(pidfile, 'w') as fd:
+            fd.write("123\n/foo/bar\n123456789\n5432")
+        self.assertEquals(self.p._read_pid_file(), {"pid": "123", "data_dir": "/foo/bar",
+                                                    "start_time": "123456789", "port": "5432"})
 
     def test_pick_sync_standby(self):
         cluster = Cluster(True, None, self.leader, 0, [self.me, self.other, self.leadermem], None,
@@ -870,43 +854,23 @@ class TestPostgresql(unittest.TestCase):
         self.p.set_synchronous_standby('foo')
         self.p.get_server_parameters(config)
 
-    @patch.object(Postgresql, 'read_pid_file', Mock(return_value={'pid': 'z'}))
-    def test_get_pid(self):
-        self.p.get_pid()
-
-    @patch.object(Postgresql, 'is_running', Mock(return_value=True))
-    @patch.object(Postgresql, '_signal_postmaster_stop', Mock(return_value=(123, None)))
-    @patch.object(Postgresql, 'get_pid', Mock(return_value=123))
     @patch('time.sleep', Mock())
-    @patch.object(Postgresql, '_is_postmaster_pid_running')
-    def test__wait_for_connection_close(self, mock_is_postmaster_pid_running):
-        mock_is_postmaster_pid_running.side_effect = [True, False, False]
-        mock_callback = Mock()
-        self.p.stop(on_safepoint=mock_callback)
-
-        mock_is_postmaster_pid_running.side_effect = [True, False, False]
-        with patch.object(MockCursor, "execute", Mock(side_effect=psycopg2.Error)):
+    def test__wait_for_connection_close(self):
+        mock_postmaster = MockPostmaster()
+        with patch.object(Postgresql, 'is_running', Mock(return_value=mock_postmaster)):
+            mock_postmaster.is_running.side_effect = [True, False, False]
+            mock_callback = Mock()
             self.p.stop(on_safepoint=mock_callback)
 
-    @patch.object(Postgresql, 'is_running', Mock(return_value=True))
-    @patch.object(Postgresql, '_signal_postmaster_stop', Mock(return_value=(123, None)))
-    @patch.object(Postgresql, 'get_pid', Mock(return_value=123))
-    @patch.object(Postgresql, '_is_postmaster_pid_running', Mock(return_value=False))
-    @patch('psutil.Process')
-    def test__wait_for_user_backends_to_close(self, mock_psutil):
-        child = Mock()
-        child.cmdline.return_value = ['foo']
-        mock_psutil.return_value.children.return_value = [child]
-        mock_callback = Mock()
-        self.p.stop(on_safepoint=mock_callback)
+            mock_postmaster.is_running.side_effect = [True, False, False]
+            with patch.object(MockCursor, "execute", Mock(side_effect=psycopg2.Error)):
+                self.p.stop(on_safepoint=mock_callback)
 
-    @patch('os.kill', Mock(side_effect=[OSError(errno.ESRCH, ''), OSError]))
-    @patch('psutil.Process', Mock(side_effect=[psutil.NoSuchProcess]))
-    @patch('time.sleep', Mock())
-    @patch.object(Postgresql, '_is_postmaster_pid_running', Mock(side_effect=[True, False]))
     def test_terminate_starting_postmaster(self):
-        self.p.terminate_starting_postmaster(123)
-        self.p.terminate_starting_postmaster(123)
+        mock_postmaster = MockPostmaster()
+        self.p.terminate_starting_postmaster(mock_postmaster)
+        mock_postmaster.signal_stop.assert_called()
+        mock_postmaster.wait.assert_called()
 
     def test_read_postmaster_opts(self):
         m = mock_open(read_data='/usr/lib/postgres/9.6/bin/postgres "-D" "data/postgresql0" \
@@ -945,11 +909,3 @@ class TestPostgresql(unittest.TestCase):
     @patch.object(Postgresql, 'single_user_mode', Mock(return_value=0))
     def test_fix_cluster_state(self):
         self.assertTrue(self.p.fix_cluster_state())
-
-    def test__update_postmaster_cached_info(self):
-        with open(os.path.join(self.data_dir, 'postmaster.pid'), 'w') as f:
-            f.write('1\n\n1\n')
-        self.p.read_pid_file()
-        with open(os.path.join(self.data_dir, 'postmaster.pid'), 'w') as f:
-            f.write('a\n\n1\n')
-        self.p.read_pid_file()

tests/test_postmaster.py

@@ -0,0 +1,77 @@
+import unittest
+
+from mock import Mock, patch
+from patroni.postmaster import PostmasterProcess
+import psutil
+
+class TestPostmasterProcess(unittest.TestCase):
+    @patch('psutil.Process.__init__', Mock())
+    def test_init(self):
+        proc = PostmasterProcess(-123)
+        self.assertTrue(proc.is_single_user)
+
+    @patch('psutil.Process.create_time')
+    @patch('psutil.Process.__init__')
+    def test_from_pidfile(self, mock_init, mock_create_time):
+        mock_init.side_effect = psutil.NoSuchProcess(123)
+        self.assertEquals(PostmasterProcess.from_pidfile({}), None)
+        self.assertEquals(PostmasterProcess.from_pidfile({"pid": "foo"}), None)
+        self.assertEquals(PostmasterProcess.from_pidfile({"pid": "123"}), None)
+
+        mock_init.side_effect = None
+        with patch.object(psutil.Process, 'pid', 123), \
+             patch.object(psutil.Process, 'parent', return_value=124), \
+             patch('os.getpid', return_value=125) as mock_ospid, \
+             patch('os.getppid', return_value=126):
+
+            self.assertNotEquals(PostmasterProcess.from_pidfile({"pid": "123"}), None)
+
+            mock_create_time.return_value = 100000
+            self.assertEquals(PostmasterProcess.from_pidfile({"pid": "123", "start_time": "200000"}), None)
+            self.assertNotEquals(PostmasterProcess.from_pidfile({"pid": "123", "start_time": "foobar"}), None)
+
+            mock_ospid.return_value = 123
+            self.assertEquals(PostmasterProcess.from_pidfile({"pid": "123", "start_time": "100000"}), None)
+
+    @patch('psutil.Process.__init__')
+    def test_from_pid(self, mock_init):
+        mock_init.side_effect = psutil.NoSuchProcess(123)
+        self.assertEquals(PostmasterProcess.from_pid(123), None)
+        mock_init.side_effect = None
+        self.assertNotEquals(PostmasterProcess.from_pid(123), None)
+
+    @patch('psutil.Process.__init__', Mock())
+    @patch('psutil.Process.send_signal')
+    @patch('psutil.Process.pid', Mock(return_value=123))
+    def test_signal_stop(self, mock_send_signal):
+        proc = PostmasterProcess(-123)
+        self.assertEquals(proc.signal_stop('immediate'), False)
+
+        mock_send_signal.side_effect = [None, psutil.NoSuchProcess(123), psutil.AccessDenied()]
+        proc = PostmasterProcess(123)
+        self.assertEquals(proc.signal_stop('immediate'), None)
+        self.assertEquals(proc.signal_stop('immediate'), True)
+        self.assertEquals(proc.signal_stop('immediate'), False)
+
+    @patch('psutil.Process.__init__', Mock())
+    @patch('psutil.wait_procs')
+    def test_wait_for_user_backends_to_close(self, mock_wait):
+        c1 = Mock()
+        c1.cmdline = Mock(return_value=["postgres: startup process"])
+        c2 = Mock()
+        c2.cmdline = Mock(return_value=["postgres: postgres postgres [local] idle"])
+        with patch('psutil.Process.children', Mock(return_value=[c1, c2])):
+            proc = PostmasterProcess(123)
+            proc.wait_for_user_backends_to_close()
+            mock_wait.assert_called_with([c2])
+
+    @patch('subprocess.Popen')
+    @patch.object(PostmasterProcess, 'from_pid')
+    def test_start(self, mock_frompid, mock_popen):
+        mock_frompid.return_value = "proc 123"
+        mock_popen.return_value.stdout.readline.return_value = '123'
+        self.assertEquals(
+            PostmasterProcess.start('/bin/true', '/tmp/', '/tmp/test.conf', ['--foo=bar', '--bar=baz']),
+            "proc 123"
+        )
+        mock_frompid.assert_called_with(123)