diff --git a/patroni/postgresql/callback_executor.py b/patroni/postgresql/callback_executor.py index 2d00b64c..40ebf5a1 100644 --- a/patroni/postgresql/callback_executor.py +++ b/patroni/postgresql/callback_executor.py @@ -1,40 +1,36 @@ import logging -import subprocess -from threading import Event, Lock, Thread + +from patroni.postgresql.cancellable import CancellableExecutor +from threading import Condition, Thread logger = logging.getLogger(__name__) -class CallbackExecutor(Thread): +class CallbackExecutor(CancellableExecutor, Thread): def __init__(self): - super(CallbackExecutor, self).__init__() + CancellableExecutor.__init__(self) + Thread.__init__(self) self.daemon = True - self._lock = Lock() self._cmd = None - self._process = None - self._callback_event = Event() + self._condition = Condition() self.start() def call(self, cmd): - with self._lock: - if self._process and self._process.poll() is None: - try: - self._process.kill() - logger.warning('Killed the old callback process because it was still running: %s', self._cmd) - except OSError: - logger.exception('Failed to kill the old callback') - self._cmd = cmd - self._callback_event.set() + self._kill_process() + with self._condition: + self._cmd = cmd + self._condition.notify() def run(self): while True: - self._callback_event.wait() - self._callback_event.clear() + with self._condition: + if self._cmd is None: + self._condition.wait() + cmd, self._cmd = self._cmd, None + with self._lock: - try: - self._process = subprocess.Popen(self._cmd, close_fds=True) - except Exception: - logger.exception('Failed to execute %s', self._cmd) + if not self._start_process(cmd, close_fds=True): continue self._process.wait() + self._kill_children() diff --git a/patroni/postgresql/cancellable.py b/patroni/postgresql/cancellable.py index 866fa194..ffb7040e 100644 --- a/patroni/postgresql/cancellable.py +++ b/patroni/postgresql/cancellable.py @@ -1,5 +1,6 @@ import logging import os +import psutil import subprocess from patroni.exceptions import PostgresException @@ -10,13 +11,66 @@ from threading import Lock logger = logging.getLogger(__name__) -class CancellableSubprocess(object): +class CancellableExecutor(object): def __init__(self): - self._is_cancelled = False self._process = None + self._process_cmd = None + self._process_children = [] self._lock = Lock() + def _start_process(self, cmd, *args, **kwargs): + """This method must be executed only when the `_lock` is acquired""" + + try: + self._process_children = [] + self._process_cmd = cmd + self._process = psutil.Popen(cmd, *args, **kwargs) + except Exception: + return logger.exception('Failed to execute %s', cmd) + return True + + def _kill_process(self): + with self._lock: + if self._process is not None and self._process.is_running() and not self._process_children: + try: + self._process.suspend() # Suspend the process before getting list of childrens + except psutil.Error as e: + logger.info('Failed to suspend the process: %s', e.msg) + + try: + self._process_children = self._process.children(recursive=True) + except psutil.Error: + pass + + try: + self._process.kill() + logger.warning('Killed %s because it was still running', self._process_cmd) + except psutil.NoSuchProcess: + pass + except psutil.AccessDenied as e: + logger.warning('Failed to kill the process: %s', e.msg) + + def _kill_children(self): + waitlist = [] + with self._lock: + for child in self._process_children: + try: + child.kill() + except psutil.NoSuchProcess: + continue + except psutil.AccessDenied as e: + logger.info('Failed to kill child process: %s', e.msg) + waitlist.append(child) + psutil.wait_procs(waitlist) + + +class CancellableSubprocess(CancellableExecutor): + + def __init__(self): + super(CancellableSubprocess, self).__init__() + self._is_cancelled = False + def call(self, *args, **kwargs): for s in ('stdin', 'stdout', 'stderr'): kwargs.pop(s, None) @@ -38,17 +92,18 @@ class CancellableSubprocess(object): raise PostgresException('cancelled') self._is_cancelled = False - self._process = subprocess.Popen(*args, **kwargs) + started = self._start_process(*args, **kwargs) - if communicate_input: - if input_data: - self._process.communicate(input_data) - self._process.stdin.close() - - return self._process.wait() + if started: + if communicate_input: + if input_data: + self._process.communicate(input_data) + self._process.stdin.close() + return self._process.wait() finally: with self._lock: self._process = None + self._kill_children() def reset_is_cancelled(self): with self._lock: @@ -62,15 +117,13 @@ class CancellableSubprocess(object): def cancel(self): with self._lock: self._is_cancelled = True - if self._process is None or self._process.returncode is not None: + if self._process is None or not self._process.is_running(): return self._process.terminate() for _ in polling_loop(10): with self._lock: - if self._process is None or self._process.returncode is not None: + if self._process is None or not self._process.is_running(): return - with self._lock: - if self._process is not None and self._process.returncode is None: - self._process.kill() + self._kill_process() diff --git a/tests/test_bootstrap.py b/tests/test_bootstrap.py index f00dbe58..3268c617 100644 --- a/tests/test_bootstrap.py +++ b/tests/test_bootstrap.py @@ -16,6 +16,7 @@ from . import psycopg2_connect, BaseTestPostgresql @patch('os.rename', Mock()) class TestBootstrap(BaseTestPostgresql): + @patch('patroni.postgresql.CallbackExecutor', Mock()) def setUp(self): super(TestBootstrap, self).setUp() self.b = self.p.bootstrap diff --git a/tests/test_callback_executor.py b/tests/test_callback_executor.py index 92f51782..cc0c4184 100644 --- a/tests/test_callback_executor.py +++ b/tests/test_callback_executor.py @@ -1,3 +1,4 @@ +import psutil import unittest from mock import Mock, patch @@ -6,22 +7,28 @@ from patroni.postgresql.callback_executor import CallbackExecutor class TestCallbackExecutor(unittest.TestCase): - @patch('subprocess.Popen') + @patch('psutil.Popen') def test_callback_executor(self, mock_popen): - mock_popen.return_value.wait.side_effect = Exception - mock_popen.return_value.poll.return_value = None + mock_popen.return_value.children.return_value = [] + mock_popen.return_value.is_running.return_value = True ce = CallbackExecutor() + ce._kill_children = Mock(side_effect=Exception) self.assertIsNone(ce.call([])) ce.join() self.assertIsNone(ce.call([])) - mock_popen.return_value.kill.side_effect = OSError + mock_popen.return_value.kill.side_effect = psutil.AccessDenied() + self.assertIsNone(ce.call([])) + + ce._process_children = [] + mock_popen.return_value.children.side_effect = psutil.Error() + mock_popen.return_value.kill.side_effect = psutil.NoSuchProcess(123) self.assertIsNone(ce.call([])) mock_popen.side_effect = Exception ce = CallbackExecutor() - ce._callback_event.wait = Mock(side_effect=[None, Exception]) + ce._condition.wait = Mock(side_effect=[None, Exception]) self.assertIsNone(ce.call([])) ce.join() diff --git a/tests/test_cancellable.py b/tests/test_cancellable.py index 33b72e2a..7a0e24b7 100644 --- a/tests/test_cancellable.py +++ b/tests/test_cancellable.py @@ -1,6 +1,7 @@ +import psutil import unittest -from mock import Mock, PropertyMock, patch +from mock import Mock, patch from patroni.exceptions import PostgresException from patroni.postgresql.cancellable import CancellableSubprocess @@ -14,10 +15,20 @@ class TestCancellableSubprocess(unittest.TestCase): self.c.cancel() self.assertRaises(PostgresException, self.c.call, communicate_input=None) + def test__kill_children(self): + self.c._process_children = [Mock()] + self.c._kill_children() + self.c._process_children[0].kill.side_effect = psutil.AccessDenied() + self.c._kill_children() + self.c._process_children[0].kill.side_effect = psutil.NoSuchProcess(123) + self.c._kill_children() + @patch('patroni.postgresql.cancellable.polling_loop', Mock(return_value=[0, 0])) def test_cancel(self): self.c._process = Mock() - self.c._process.returncode = None + self.c._process.is_running.return_value = True + self.c._process.children.side_effect = psutil.Error() + self.c._process.suspend.side_effect = psutil.Error() self.c.cancel() - type(self.c._process).returncode = PropertyMock(side_effect=[None, -15]) + self.c._process.is_running.side_effect = [True, False] self.c.cancel() diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index 62976faf..be96422f 100644 --- a/tests/test_postgresql.py +++ b/tests/test_postgresql.py @@ -90,6 +90,7 @@ class TestPostgresql(BaseTestPostgresql): @patch('subprocess.call', Mock(return_value=0)) @patch('os.rename', Mock()) + @patch('patroni.postgresql.CallbackExecutor', Mock()) @patch.object(Postgresql, 'get_major_version', Mock(return_value=120000)) @patch.object(Postgresql, 'is_running', Mock(return_value=True)) def setUp(self): @@ -643,7 +644,7 @@ class TestPostgresql(BaseTestPostgresql): data = self.p.read_postmaster_opts() self.assertEqual(data, dict()) - @patch('subprocess.Popen') + @patch('psutil.Popen') def test_single_user_mode(self, subprocess_popen_mock): subprocess_popen_mock.return_value.wait.return_value = 0 self.assertEqual(self.p.single_user_mode('CHECKPOINT', {'archive_mode': 'on'}), 0)