Kill callback child processes when it is necessary (#1242)

Not doing so makes it hard to implement callbacks in bash and eventually can lead to the situation when two callbacks are running at the same time. In case if we failed to kill the child process we will still wait for it to finish.

The same problem could happen with custom bootstrap, therefore if we happen to kill the custom bootstrap process we also kill all child subprocesses.

Closes https://github.com/zalando/patroni/issues/1238

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._kill_process()
+        with self._condition:
             self._cmd = cmd
-        self._callback_event.set()
+            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()

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 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()

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

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()

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()

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)