mirror of
https://github.com/outbackdingo/patroni.git
synced 2026-09-01 09:09:21 +00:00
Alexander KukushkinandGitHub
76b3b99de2
- added pyrightconfig.json with typeCheckingMode=strict - added type hints to all files except api.py - added type stubs for dns, etcd, consul, kazoo, pysyncobj and other modules - added type stubs for psycopg2 and urllib3 with some little fixes - fixes most of the issues reported by pyright - remaining issues will be addressed later, along with enabling CI linting task
136 lines
4.5 KiB
Python
136 lines
4.5 KiB
Python
import logging
|
|
import psutil
|
|
import subprocess
|
|
|
|
from patroni.exceptions import PostgresException
|
|
from patroni.utils import polling_loop
|
|
from threading import Lock
|
|
from typing import Any, Dict, List, Optional, Union
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class CancellableExecutor(object):
|
|
|
|
"""
|
|
There must be only one such process so that AsyncExecutor can easily cancel it.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._process = None
|
|
self._process_cmd = None
|
|
self._process_children: List[psutil.Process] = []
|
|
self._lock = Lock()
|
|
|
|
def _start_process(self, cmd: List[str], *args: Any, **kwargs: Any) -> Optional[bool]:
|
|
"""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) -> None:
|
|
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 children
|
|
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) -> None:
|
|
waitlist: List[psutil.Process] = []
|
|
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) -> None:
|
|
super(CancellableSubprocess, self).__init__()
|
|
self._is_cancelled = False
|
|
|
|
def call(self, *args: Any, **kwargs: Union[Any, Dict[str, str]]) -> Optional[int]:
|
|
for s in ('stdin', 'stdout', 'stderr'):
|
|
kwargs.pop(s, None)
|
|
|
|
communicate: Optional[Dict[str, str]] = kwargs.pop('communicate', None)
|
|
input_data = None
|
|
if isinstance(communicate, dict):
|
|
input_data = communicate.get('input')
|
|
if input_data:
|
|
if input_data[-1] != '\n':
|
|
input_data += '\n'
|
|
input_data = input_data.encode('utf-8')
|
|
kwargs['stdin'] = subprocess.PIPE
|
|
kwargs['stdout'] = subprocess.PIPE
|
|
kwargs['stderr'] = subprocess.PIPE
|
|
|
|
try:
|
|
with self._lock:
|
|
if self._is_cancelled:
|
|
raise PostgresException('cancelled')
|
|
|
|
self._is_cancelled = False
|
|
started = self._start_process(*args, **kwargs)
|
|
|
|
if started and self._process is not None:
|
|
if isinstance(communicate, dict):
|
|
communicate['stdout'], communicate['stderr'] = self._process.communicate(input_data)
|
|
return self._process.wait()
|
|
finally:
|
|
with self._lock:
|
|
self._process = None
|
|
self._kill_children()
|
|
|
|
def reset_is_cancelled(self) -> None:
|
|
with self._lock:
|
|
self._is_cancelled = False
|
|
|
|
@property
|
|
def is_cancelled(self) -> bool:
|
|
with self._lock:
|
|
return self._is_cancelled
|
|
|
|
def cancel(self, kill: bool = False) -> None:
|
|
with self._lock:
|
|
self._is_cancelled = True
|
|
if self._process is None or not self._process.is_running():
|
|
return
|
|
|
|
logger.info('Terminating %s', self._process_cmd)
|
|
self._process.terminate()
|
|
|
|
for _ in polling_loop(10):
|
|
with self._lock:
|
|
if self._process is None or not self._process.is_running():
|
|
return
|
|
if kill:
|
|
break
|
|
|
|
self._kill_process()
|