# Copyright 2021-2022 VMware, Inc.
# SPDX-License-Identifier: Apache-2.0
"""
Shelling class implementations.
"""
import atexit
import contextlib
import json
import locale
import logging
import os
import pathlib
import shutil
import subprocess
import sys
from tempfile import SpooledTemporaryFile
from typing import Any
from typing import Callable
from typing import cast
from typing import Dict
from typing import Generator
from typing import List
from typing import Optional
from typing import Tuple
from typing import TYPE_CHECKING
from typing import Union
import attr
import psutil
from pytestskipmarkers.utils import platform
from pytestshellutils.customtypes import Callback
from pytestshellutils.customtypes import EnvironDict
from pytestshellutils.exceptions import CallbackException
from pytestshellutils.exceptions import FactoryNotRunning
from pytestshellutils.exceptions import FactoryNotStarted
from pytestshellutils.exceptions import FactoryTimeout
from pytestshellutils.exceptions import ShellUtilsException
from pytestshellutils.utils import format_callback_to_string
from pytestshellutils.utils import ports
from pytestshellutils.utils import resolved_pathlib_path
from pytestshellutils.utils import time
from pytestshellutils.utils.processes import ProcessResult
from pytestshellutils.utils.processes import terminate_process
from pytestshellutils.utils.processes import terminate_process_list
if TYPE_CHECKING:
from typing import Type
from pytestsysstats.plugin import StatsProcesses
log = logging.getLogger(__name__)
[docs]@attr.s(slots=True, kw_only=True)
class BaseFactory:
"""
Base factory class.
:keyword pathlib.Path cwd:
The path to the desired working directory
:keyword dict environ:
A dictionary of ``key``, ``value`` pairs to add to the environment.
"""
cwd: pathlib.Path = attr.ib(converter=resolved_pathlib_path)
environ: EnvironDict = attr.ib(repr=False)
@cwd.default
def _default_cwd(self) -> pathlib.Path:
"""
Return the default cwd to use.
"""
return pathlib.Path.cwd()
@environ.default
def _default_environ(self) -> EnvironDict:
"""
Return the default ``os.environ`` to use.
"""
return cast(EnvironDict, os.environ.copy())
[docs]@attr.s(slots=True, kw_only=True)
class SubprocessImpl:
"""
Subprocess interaction implementation.
:keyword ~pytestshellutils.shell.Subprocess factory:
The factory instance, either :py:class:`~pytestshellutils.shell.Subprocess` or
a sub-class of it.
"""
factory: "Union[Factory, Subprocess, ScriptSubprocess]" = attr.ib()
_terminal: "Optional[subprocess.Popen[Any]]" = attr.ib(repr=False, init=False, default=None)
_terminal_stdout: "Optional[SpooledTemporaryFile[bytes]]" = attr.ib(
repr=False, init=False, default=None
)
_terminal_stderr: "Optional[SpooledTemporaryFile[bytes]]" = attr.ib(
repr=False, init=False, default=None
)
_terminal_result: Optional[ProcessResult] = attr.ib(repr=False, init=False, default=None)
_terminal_timeout: Union[int, float] = attr.ib(repr=False, init=False, default=None)
_children: List[psutil.Process] = attr.ib(repr=False, init=False, factory=list)
[docs] def cmdline(self, *args: str, **kwargs: Any) -> List[str]:
"""
Construct a list of arguments to use when starting the subprocess.
:param str args:
Additional arguments to use when starting the subprocess
By default, this method will just call it's factory's ``cmdline()``
method, but can be overridden.
"""
return self.factory.cmdline(*args)
[docs] def init_terminal(
self, cmdline: List[str], env: Optional[EnvironDict] = None
) -> "subprocess.Popen[Any]":
"""
Instantiate a terminal with the passed command line(``cmdline``) and return it.
Additionally, it sets a reference to it in ``self._terminal`` and also collects
an initial listing of child processes which will be used when terminating the
terminal
:param list,tuple cmdline:
List of strings to pass as ``args`` to :py:class:`~subprocess.Popen`
:keyword dict environ:
A dictionary of ``key``, ``value`` pairs to add to the
py:attr:`pytestshellutils.shell.Factory.environ`.
"""
environ = self.factory.environ.copy()
if env is not None:
environ.update(env)
self._terminal_stdout = SpooledTemporaryFile(512000, buffering=0)
self._terminal_stderr = SpooledTemporaryFile(512000, buffering=0)
close_fds: bool
if platform.is_windows(): # pragma: is-windows
# Windows does not support closing FDs
close_fds = False
elif platform.is_freebsd() and sys.version_info < (3, 9): # pragma: is-bsd-lt-py39
# Closing FDs in FreeBSD before Py3.9 can be slow
# https://bugs.python.org/issue38061
close_fds = False
else:
close_fds = True
self._terminal = subprocess.Popen(
cmdline,
stdout=self._terminal_stdout,
stderr=self._terminal_stderr,
shell=False,
cwd=str(self.factory.cwd),
universal_newlines=True,
close_fds=close_fds,
env=environ,
bufsize=0,
)
# Reset the previous _terminal_result if set
self._terminal_result = None
try:
# Check if the process starts properly
self._terminal.wait(timeout=0.05)
# If TimeoutExpired is not raised, it means the process failed to start
except subprocess.TimeoutExpired:
# We're good
# Collect any child processes, though, this early there likely is none
with contextlib.suppress(psutil.NoSuchProcess):
for child in psutil.Process(self._terminal.pid).children(
recursive=True
): # pragma: no cover
if child not in self._children:
self._children.append(child)
atexit.register(self.terminate)
return self._terminal
[docs] def is_running(self) -> bool:
"""
Returns true if the sub-process is alive.
:return: Returns true if the sub-process is alive
:rtype: bool
"""
if not self._terminal:
return False
return self._terminal.poll() is None
[docs] def terminate(self) -> ProcessResult:
"""
Terminate the started subprocess.
:rtype: ~pytestshellutils.utils.processes.ProcessResult
"""
return self._terminate()
def _terminate(self) -> ProcessResult:
"""
This method actually terminates the started subprocess.
"""
if self._terminal is None:
if TYPE_CHECKING:
# Make mypy happy
assert self._terminal_result
return self._terminal_result
atexit.unregister(self.terminate)
log.info("Stopping %s", self.factory)
# Collect any child processes information before terminating the process
with contextlib.suppress(psutil.NoSuchProcess):
for child in psutil.Process(self._terminal.pid).children(recursive=True):
if child not in self._children:
self._children.append(child)
with self._terminal:
if self.factory.slow_stop:
self._terminal.terminate()
else:
self._terminal.kill()
try:
# Allow the process to exit by itself in case slow_stop is True
self._terminal.wait(10)
except subprocess.TimeoutExpired: # pragma: no cover
# The process failed to stop, no worries, we'll make sure it exit along with it's
# child processes bellow
pass
# Lets log and kill any child processes left behind, including the main subprocess
# If it failed to properly stop
terminate_process(
pid=self._terminal.pid,
kill_children=True,
children=self._children,
slow_stop=self.factory.slow_stop,
)
# Wait for the process to terminate, to avoid zombies.
self._terminal.wait()
# poll the terminal so the right returncode is set on the popen object
self._terminal.poll()
# This call shouldn't really be necessary
self._terminal.communicate()
if TYPE_CHECKING:
# Make mypy happy
assert self._terminal_stdout
self._terminal_stdout.flush()
self._terminal_stdout.seek(0)
if sys.version_info < (3, 6): # pragma: no cover
stdout = self._terminal._translate_newlines(
self._terminal_stdout.read(),
self.factory.system_encoding,
)
else:
stdout = self._terminal._translate_newlines( # type: ignore[attr-defined]
self._terminal_stdout.read(),
self.factory.system_encoding,
sys.stdout.errors,
)
self._terminal_stdout.close()
if TYPE_CHECKING:
# Make mypy happy
assert self._terminal_stderr
self._terminal_stderr.flush()
self._terminal_stderr.seek(0)
if sys.version_info < (3, 6): # pragma: no cover
stderr = self._terminal._translate_newlines(
self._terminal_stderr.read(),
self.factory.system_encoding,
)
else:
stderr = self._terminal._translate_newlines( # type: ignore[attr-defined]
self._terminal_stderr.read(),
self.factory.system_encoding,
sys.stderr.errors,
)
self._terminal_stderr.close()
try:
self._terminal_result = ProcessResult(
returncode=self._terminal.returncode,
stdout=stdout,
stderr=stderr,
cmdline=cast(List[str], self._terminal.args),
)
log.info("%s %s", self.factory.__class__.__name__, self._terminal_result)
return self._terminal_result
finally:
self._terminal = None
self._terminal_stdout = None
self._terminal_stderr = None
self._children = []
@property
def pid(self) -> Optional[int]:
"""
The pid of the running process. None if not running.
"""
if not self._terminal: # pragma: no cover
return None
return self._terminal.pid
[docs] def run(
self, *args: str, env: Optional[EnvironDict] = None, **kwargs: Any
) -> "subprocess.Popen[Any]":
"""
Run the given command synchronously.
"""
cmdline = self.cmdline(*args, **kwargs)
log.info("%s is running %r in CWD: %s ...", self.factory, cmdline, self.factory.cwd)
return self.init_terminal(cmdline, env=env)
[docs]@attr.s(slots=True, kw_only=True)
class Factory(BaseFactory):
"""
Base shell factory class.
:keyword bool slow_stop:
Whether to terminate the processes by sending a :py:attr:`SIGTERM` signal or by calling
:py:meth:`~subprocess.Popen.terminate` on the sub-process.
When code coverage is enabled, one will want `slow_stop` set to `True` so that coverage data
can be written down to disk.
:keyword str system_encoding:
The system encoding to use when decoding the subprocess output. Defaults to "utf-8".
:keyword int,float timeout:
The default maximum amount of seconds that a script should run.
This value can be overridden when calling :py:meth:`~pytestshellutils.shell.Process.run` through
the ``_timeout`` keyword argument, and, in that case, the timeout value applied would be that
of ``_timeout`` instead of ``self.timeout``.
"""
slow_stop: bool = attr.ib(default=True)
system_encoding: str = attr.ib(repr=False)
timeout: Union[int, float] = attr.ib()
impl: SubprocessImpl = attr.ib(repr=False, init=False)
# Internal attributes
_cmdline: List[Any] = attr.ib(repr=False, init=False, default=None)
@system_encoding.default
def _default_system_encoding(self) -> str:
return self._get_default_system_encoding()
@timeout.default
def _set_timeout(self) -> Optional[int]:
return self._get_default_timeout()
def _get_default_system_encoding(self) -> str:
encoding: Optional[str] = None
if not platform.is_windows() and sys.stdin is not None:
# On Linux we can rely on sys.stdin for the encoding since it
# most commonly matches the file-system encoding. This however
# does not apply to windows
encoding = sys.stdin.encoding
if not encoding:
# If the system is properly configured this should return a valid
# encoding. MS Windows has problems with this and reports the wrong
# encoding
try:
encoding = locale.getdefaultlocale()[-1]
except ValueError:
# A bad locale setting was most likely found:
# https://github.com/saltstack/salt/issues/26063
pass
if not encoding:
# This is most likely ascii which is not the best but we were
# unable to find a better encoding. If this fails, we fall all
# the way back to ascii
encoding = sys.getdefaultencoding()
if not encoding:
if platform.is_darwin():
# Mac OS X uses UTF-8
encoding = "utf-8"
elif platform.is_windows():
# Windows uses a configurable encoding; on Windows, Python uses the name “mbcs”
# to refer to whatever the currently configured encoding is.
encoding = "mbcs"
else:
# On linux default to ascii as a last resort
encoding = "ascii"
if not encoding:
# If we still didn't detect the encoding, default to utf-8
encoding = "utf-8"
return encoding
def _get_default_timeout(self) -> Optional[int]:
return None
def _get_impl_class(self) -> "Type[SubprocessImpl]":
"""
Return the ``impl`` class to use.
"""
return SubprocessImpl
[docs] def __attrs_post_init__(self) -> None:
"""
Post ``attrs`` class initialization routines.
"""
impl_class = self._get_impl_class()
self.impl = impl_class(factory=self)
[docs] def cmdline(self, *args: str) -> List[str]:
"""
Method to construct a command line.
"""
self._cmdline = list(args)
return self._cmdline
[docs] def get_display_name(self) -> str:
"""
Returns a human readable name for the factory.
"""
return "{}({})".format(self.__class__.__name__, self._cmdline or "")
[docs] def is_running(self) -> bool:
"""
Returns true if the sub-process is alive.
"""
return self.impl.is_running()
[docs] def terminate(self) -> ProcessResult:
"""
Terminate the started subprocess.
"""
return self.impl.terminate()
@property
def pid(self) -> Optional[int]:
"""
The pid of the running process. None if not running.
"""
return self.impl.pid
[docs]@attr.s(slots=True, kw_only=True)
class Subprocess(Factory):
"""
Base shell factory class.
"""
[docs] def run(
self,
*args: str,
env: Optional[EnvironDict] = None,
_timeout: Optional[Union[int, float]] = None,
**kwargs: Any
) -> ProcessResult:
"""
Run the given command synchronously.
:keyword list,tuple args:
The list of arguments to pass to :py:meth:`~pytestshellutils.shell.Subprocess.cmdline`
to construct the command to run
:keyword int.float _timeout:
The timeout value for this particular ``run()`` call. If this value is not ``None``,
it will be used instead of :py:attr:`~pytestshellutils.shell.Subprocess.timeout`,
the default timeout.
"""
start_time = time.time()
# Build the cmdline to pass to the terminal
# We set the _terminal_timeout attribute while calling cmdline in case it needs
# access to that information to build the command line
self.impl._terminal_timeout = _timeout or self.timeout
timmed_out = False
try:
self.impl.run(*args, env=env, **kwargs)
if TYPE_CHECKING:
# Make mypy happy
assert self.impl._terminal
self.impl._terminal.communicate(timeout=self.impl._terminal_timeout)
except subprocess.TimeoutExpired:
timmed_out = True
result = self.terminate()
cmdline = result.cmdline
returncode = result.returncode
if timmed_out:
raise FactoryTimeout(
"{} Failed to run: {}; Error: Timed out after {:.2f} seconds!".format(
self, cmdline, time.time() - start_time
),
process_result=result,
)
stdout, stderr, json_out = self.process_output(
result.stdout, result.stderr, cmdline=cmdline
)
log.info(
"%s completed %r in CWD: %s after %.2f seconds",
self,
cmdline,
self.cwd,
time.time() - start_time,
)
return ProcessResult(
returncode=returncode, stdout=stdout, stderr=stderr, data=json_out, cmdline=cmdline
)
[docs] def process_output(
self, stdout: str, stderr: str, cmdline: Optional[List[str]] = None
) -> Tuple[str, str, Optional[Dict[Any, Any]]]:
"""
Process the output. When possible JSON is loaded from the output.
:return:
Returns a tuple in the form of ``(stdout, stderr, loaded_json)``
:rtype: tuple
"""
if stdout:
try:
json_out = json.loads(stdout)
except ValueError:
log.debug("%s failed to load JSON from the following output:\n%r", self, stdout)
json_out = None
else:
json_out = None
return stdout, stderr, json_out
[docs]@attr.s(slots=True, kw_only=True)
class ScriptSubprocess(Subprocess):
"""
Base CLI script/binary class.
:keyword str script_name:
This is the string containing the name of the binary to call on the subprocess, either the
full path to it, or the basename. In case of the basename, the directory containing the
basename must be in your ``$PATH`` variable.
:keyword list,tuple base_script_args:
An list or tuple iterable of the base arguments to use when building the command line to
launch the process
Please look at :py:class:`~pytestshellutils.shell.Factory` for the additional supported keyword
arguments documentation.
"""
script_name: str = attr.ib()
base_script_args: List[str] = attr.ib(factory=list)
[docs] def get_display_name(self) -> str:
"""
Returns a human readable name for the factory.
"""
return "{}({})".format(self.__class__.__name__, pathlib.Path(self.script_name).name)
[docs] def get_script_path(self) -> str:
"""
Returns the path to the script to run.
"""
script_path: Optional[str]
if os.path.isabs(self.script_name):
script_path = self.script_name
else:
script_path = shutil.which(self.script_name)
if not script_path or not os.path.exists(script_path):
raise FileNotFoundError("The CLI script {!r} does not exist".format(self.script_name))
if TYPE_CHECKING:
# Make mypy happy
assert script_path
return script_path
[docs] def get_base_script_args(self) -> List[str]:
"""
Returns any additional arguments to pass to the CLI script.
"""
return list(self.base_script_args)
[docs] def get_script_args(self) -> List[str]: # pylint: disable=no-self-use
"""
Returns any additional arguments to pass to the CLI script.
"""
return []
[docs] def cmdline(self, *args: str) -> List[str]:
"""
Construct a list of arguments to use when starting the subprocess.
:param str args:
Additional arguments to use when starting the subprocess
:rtype: list
"""
return (
[self.get_script_path()]
+ self.get_base_script_args()
+ self.get_script_args()
+ list(args)
)
[docs]@attr.s(kw_only=True, slots=True, frozen=True)
class StartDaemonCallArguments:
"""
This class holds the arguments and keyword arguments used to start a daemon.
It's used when restarting the daemon so that the same call is used.
:keyword list,tuple args:
List of arguments
:keyword dict kwargs:
Dictionary of keyword arguments
"""
args: Tuple[str, ...] = attr.ib()
kwargs: Dict[str, Any] = attr.ib()
[docs]@attr.s(slots=True, kw_only=True)
class DaemonImpl(SubprocessImpl):
"""
Daemon subprocess interaction implementation.
Please look at :py:class:`~pytestshellutils.shell.SubprocessImpl` for the additional supported keyword
arguments documentation.
"""
factory: "Daemon" = attr.ib()
_before_start_callbacks: List[Callback] = attr.ib(repr=False, hash=False, factory=list)
_after_start_callbacks: List[Callback] = attr.ib(repr=False, hash=False, factory=list)
_before_terminate_callbacks: List[Callback] = attr.ib(repr=False, hash=False, factory=list)
_after_terminate_callbacks: List[Callback] = attr.ib(repr=False, hash=False, factory=list)
_start_args_and_kwargs: StartDaemonCallArguments = attr.ib(
init=False, repr=False, hash=False, default=None
)
[docs] def before_start(self, callback: Callable[[], None], *args: Any, **kwargs: Any) -> None:
"""
Register a function callback to run before the daemon starts.
:param ~collections.abc.Callable callback:
The function to call back
:keyword args:
The arguments to pass to the callback
:keyword kwargs:
The keyword arguments to pass to the callback
"""
self._before_start_callbacks.append(Callback(func=callback, args=args, kwargs=kwargs))
[docs] def after_start(self, callback: Callable[[], None], *args: Any, **kwargs: Any) -> None:
"""
Register a function callback to run after the daemon starts.
:param ~collections.abc.Callable callback:
The function to call back
:keyword args:
The arguments to pass to the callback
:keyword kwargs:
The keyword arguments to pass to the callback
"""
self._after_start_callbacks.append(Callback(func=callback, args=args, kwargs=kwargs))
[docs] def before_terminate(self, callback: Callable[[], None], *args: Any, **kwargs: Any) -> None:
"""
Register a function callback to run before the daemon terminates.
:param ~collections.abc.Callable callback:
The function to call back
:keyword args:
The arguments to pass to the callback
:keyword kwargs:
The keyword arguments to pass to the callback
"""
self._before_terminate_callbacks.append(Callback(func=callback, args=args, kwargs=kwargs))
[docs] def after_terminate(self, callback: Callable[[], None], *args: Any, **kwargs: Any) -> None:
"""
Register a function callback to run after the daemon terminates.
:param ~collections.abc.Callable callback:
The function to call back
:keyword args:
The arguments to pass to the callback
:keyword kwargs:
The keyword arguments to pass to the callback
"""
self._after_terminate_callbacks.append(Callback(func=callback, args=args, kwargs=kwargs))
[docs] def start(
self,
*extra_cli_arguments: str,
max_start_attempts: Optional[int] = None,
start_timeout: Optional[Union[int, float]] = None
) -> bool:
"""
Start the daemon.
:keyword tuple extra_cli_arguments:
Extra arguments to pass to the CLI that starts the daemon
:keyword int max_start_attempts:
Maximum number of attempts to try and start the daemon in case of failures
:keyword int,float start_timeout:
The maximum number of seconds to wait before considering that the daemon did not start
"""
if self.is_running(): # pragma: no cover
log.warning("%s is already running.", self)
return True
self._start_args_and_kwargs = StartDaemonCallArguments(
args=extra_cli_arguments,
kwargs={"max_start_attempts": max_start_attempts, "start_timeout": start_timeout},
)
process_running = False
start_time = time.time()
start_attempts = max_start_attempts or self.factory.max_start_attempts
current_attempt = 0
run_arguments = list(extra_cli_arguments)
while True:
if process_running:
break
current_attempt += 1
if current_attempt > start_attempts:
break
log.info(
"Starting %s. Attempt: %d of %d", self.factory, current_attempt, start_attempts
)
for callback in self._before_start_callbacks: # pylint: disable=not-an-iterable
try:
callback()
except CallbackException as exc: # pragma: no cover
log.info(
"Exception raised when running %s: %s",
callback,
exc,
exc_info=True,
)
current_start_time = time.time()
start_running_timeout = current_start_time + (
start_timeout or self.factory.start_timeout
)
if current_attempt > 1 and self.factory.extra_cli_arguments_after_first_start_failure:
run_arguments = list(extra_cli_arguments) + list(
self.factory.extra_cli_arguments_after_first_start_failure
)
self.run(*run_arguments)
if not self.is_running(): # pragma: no cover
# A little breathe time to allow the process to start if not started already
time.sleep(0.5)
while time.time() <= start_running_timeout:
if not self.is_running():
log.warning("%s is no longer running", self.factory)
self.terminate()
break
try:
if (
self.factory.run_start_checks(current_start_time, start_running_timeout)
is False
):
time.sleep(1)
continue
except FactoryNotStarted:
self.terminate()
break
log.info(
"The %s factory is running after %d attempts. Took %1.2f seconds",
self.factory,
current_attempt,
time.time() - start_time,
)
process_running = True
break
else:
# The factory failed to confirm it's running status
self.terminate()
if process_running:
for callback in self._after_start_callbacks: # pylint: disable=not-an-iterable
try:
callback()
except CallbackException as exc: # pragma: no cover
log.info(
"Exception raised when running %s: %s",
callback,
exc,
exc_info=True,
)
return process_running
result = self.terminate()
raise FactoryNotStarted(
"The {} factory has failed to confirm running status after {} attempts, which "
"took {:.2f} seconds".format(
self.factory,
current_attempt - 1,
time.time() - start_time,
),
process_result=result,
)
[docs] def terminate(self) -> ProcessResult:
"""
Terminate the daemon.
"""
if self._terminal_result is not None:
# This factory has already been terminated
return self._terminal_result
for callback in self._before_terminate_callbacks: # pylint: disable=not-an-iterable
try:
callback()
except CallbackException as exc: # pragma: no cover
log.info(
"Exception raised when running %s: %s",
callback,
exc,
exc_info=True,
)
try:
return super().terminate()
finally:
for callback in self._after_terminate_callbacks: # pylint: disable=not-an-iterable
try:
callback()
except CallbackException as exc: # pragma: no cover
log.warning(
"Exception raised when running %s: %s",
callback,
exc,
exc_info=True,
)
[docs] def get_start_arguments(self) -> StartDaemonCallArguments:
"""
Return the arguments and keyword arguments used when starting the daemon.
:rtype: StartDaemonCallArguments
"""
return self._start_args_and_kwargs
[docs]@attr.s(slots=True, kw_only=True)
class Daemon(ScriptSubprocess):
"""
Base daemon factory.
:keyword list check_ports:
List of ports to try and connect to while confirming that the daemon is up and running
:keyword tuple, extra_cli_arguments_after_first_start_failure:
Extra arguments to pass to the CLI that starts the daemon after the first failure
:keyword int max_start_attempts:
Maximum number of attempts to try and start the daemon in case of failures
:keyword int start_timeout:
The maximum number of seconds to wait before considering that the daemon did not start
Please look at :py:class:`~pytestshellutils.shell.Subprocess` for the additional supported keyword
arguments documentation.
"""
impl: DaemonImpl = attr.ib(repr=False, init=False)
script_name: str = attr.ib()
base_script_args: List[str] = attr.ib(factory=list)
check_ports: List[int] = attr.ib(factory=list)
stats_processes: "StatsProcesses" = attr.ib(repr=False, hash=False, default=None)
start_timeout: Union[int, float] = attr.ib(repr=False)
max_start_attempts: int = attr.ib(repr=False, default=3)
extra_cli_arguments_after_first_start_failure: List[str] = attr.ib(hash=False, factory=list)
listen_ports: List[int] = attr.ib(init=False, repr=False, hash=False, factory=list)
def _get_impl_class(self) -> "Type[DaemonImpl]":
"""
Return the ``impl`` class to use.
"""
return DaemonImpl
[docs] def __attrs_post_init__(self) -> None:
"""
Post ``attrs`` class initialization routines.
"""
super().__attrs_post_init__()
if self.check_ports and not isinstance(self.check_ports, (list, tuple)):
self.check_ports = [self.check_ports]
if self.check_ports:
self.listen_ports.extend(self.check_ports)
self.after_start(self._add_factory_to_stats_processes)
self.after_terminate(self._terminate_processes_matching_listen_ports)
self.after_terminate(self._remove_factory_from_stats_processes)
[docs] def before_start(self, callback: Callable[[], None], *args: Any, **kwargs: Any) -> None:
"""
Register a function callback to run before the daemon starts.
:param ~collections.abc.Callable callback:
The function to call back
:keyword args:
The arguments to pass to the callback
:keyword kwargs:
The keyword arguments to pass to the callback
"""
self.impl.before_start(callback, *args, **kwargs)
[docs] def after_start(self, callback: Callable[[], None], *args: Any, **kwargs: Any) -> None:
"""
Register a function callback to run after the daemon starts.
:param ~collections.abc.Callable callback:
The function to call back
:keyword args:
The arguments to pass to the callback
:keyword kwargs:
The keyword arguments to pass to the callback
"""
self.impl.after_start(callback, *args, **kwargs)
[docs] def before_terminate(self, callback: Callable[[], None], *args: Any, **kwargs: Any) -> None:
"""
Register a function callback to run before the daemon terminates.
:param ~collections.abc.Callable callback:
The function to call back
:keyword args:
The arguments to pass to the callback
:keyword kwargs:
The keyword arguments to pass to the callback
"""
self.impl.before_terminate(callback, *args, **kwargs)
[docs] def after_terminate(self, callback: Callable[[], None], *args: Any, **kwargs: Any) -> None:
"""
Register a function callback to run after the daemon terminates.
:param ~collections.abc.Callable callback:
The function to call back
:keyword args:
The arguments to pass to the callback
:keyword kwargs:
The keyword arguments to pass to the callback
"""
self.impl.after_terminate(callback, *args, **kwargs)
[docs] def get_check_ports(self) -> List[int]:
"""
Return a list of ports to check against to ensure the daemon is running.
"""
return self.check_ports or []
[docs] def start(
self,
*extra_cli_arguments: str,
max_start_attempts: Optional[int] = None,
start_timeout: Optional[Union[int, float]] = None
) -> bool:
"""
Start the daemon.
"""
return self.impl.start(
*extra_cli_arguments, max_start_attempts=max_start_attempts, start_timeout=start_timeout
)
[docs] def started(
self,
*extra_cli_arguments: str,
max_start_attempts: Optional[int] = None,
start_timeout: Optional[Union[int, float]] = None
) -> "Daemon":
"""
Start the daemon and return it's instance so it can be used as a context manager.
"""
self.start(
*extra_cli_arguments, max_start_attempts=max_start_attempts, start_timeout=start_timeout
)
return self
[docs] @contextlib.contextmanager
def stopped(
self,
before_stop_callback: Optional[Callable[["Daemon"], None]] = None,
after_stop_callback: Optional[Callable[["Daemon"], None]] = None,
before_start_callback: Optional[Callable[["Daemon"], None]] = None,
after_start_callback: Optional[Callable[["Daemon"], None]] = None,
) -> Generator["Daemon", None, None]:
"""
Stop the daemon and return it's instance so it can be used as a context manager.
:keyword ~collections.abc.Callable before_stop_callback:
A callable to run before stopping the daemon. The callback must accept one argument,
the daemon instance.
:keyword ~collections.abc.Callable after_stop_callback:
A callable to run after stopping the daemon. The callback must accept one argument,
the daemon instance.
:keyword ~collections.abc.Callable before_start_callback:
A callable to run before starting the daemon. The callback must accept one argument,
the daemon instance.
:keyword ~collections.abc.Callable after_start_callback:
A callable to run after starting the daemon. The callback must accept one argument,
the daemon instance.
This context manager will stop the factory while the context is in place, it re-starts it once out of
context.
For example:
.. code-block:: python
assert factory.is_running() is True
with factory.stopped():
assert factory.is_running() is False
assert factory.is_running() is True
"""
if not self.is_running():
raise FactoryNotRunning("{} is not running ".format(self))
start_arguments = self.impl.get_start_arguments()
try:
if before_stop_callback:
try:
before_stop_callback(self)
except CallbackException as exc: # pragma: no cover
log.info(
"Exception raised when running %s: %s",
format_callback_to_string(before_stop_callback),
exc,
exc_info=True,
)
self.terminate()
if after_stop_callback:
try:
after_stop_callback(self)
except CallbackException as exc: # pragma: no cover
log.info(
"Exception raised when running %s: %s",
format_callback_to_string(after_stop_callback),
exc,
exc_info=True,
)
yield self
except ShellUtilsException: # pragma: no cover pylint: disable=try-except-raise
raise
else:
if before_start_callback:
try:
before_start_callback(self)
except CallbackException as exc: # pragma: no cover
log.info(
"Exception raised when running %s: %s",
format_callback_to_string(before_start_callback),
exc,
exc_info=True,
)
_started = self.started(
*start_arguments.args, # pylint: disable=not-an-iterable
**start_arguments.kwargs, # pylint: disable=not-a-mapping
)
if _started:
if after_start_callback:
try:
after_start_callback(self)
except CallbackException as exc: # pragma: no cover
log.info(
"Exception raised when running %s: %s",
format_callback_to_string(after_start_callback),
exc,
exc_info=True,
)
[docs] def run_start_checks(self, started_at: float, timeout_at: float) -> bool:
"""
Run checks to confirm that the daemon has started.
"""
log.debug("%s is running start checks", self)
check_ports = set(self.get_check_ports())
if not check_ports:
log.debug("No ports to check connection to for %s", self)
return True
log.debug("Listening ports to check for %s: %s", self, set(self.get_check_ports()))
checks_start_time = time.time()
while time.time() <= timeout_at:
if not self.is_running():
raise FactoryNotStarted("{} is no longer running".format(self))
if not check_ports:
break
check_ports -= ports.get_connectable_ports(check_ports)
if check_ports:
time.sleep(1.5)
else:
log.error(
"Failed to check ports after %1.2f seconds for %s. Remaining ports to check: %s",
time.time() - checks_start_time,
self,
check_ports,
)
return False
log.debug("All listening ports checked for %s: %s", self, set(self.get_check_ports()))
return True
def _add_factory_to_stats_processes(self) -> None:
if self.stats_processes is not None:
display_name = self.get_display_name()
self.stats_processes.add(display_name, self.pid)
def _remove_factory_from_stats_processes(self) -> None:
if self.stats_processes is not None:
display_name = self.get_display_name()
self.stats_processes.remove(display_name)
def _terminate_processes_matching_listen_ports(self) -> None:
if not self.listen_ports:
return
# If any processes were not terminated and are listening on the ports
# we have set on listen_ports, terminate those processes.
found_processes = []
for process in psutil.process_iter(["connections"]):
try:
for connection in process.connections():
if connection.status != psutil.CONN_LISTEN:
# We only care about listening services
continue
if connection.laddr.port in self.check_ports:
found_processes.append(process)
# We already found one connection, no need to check the others
break
except psutil.AccessDenied: # pragma: no cover
# We've been denied access to this process connections. Carry on.
continue
except psutil.ZombieProcess:
continue
if found_processes:
log.debug(
"The following processes were found listening on ports %s: %s",
", ".join(
[str(port) for port in self.listen_ports], # pylint: disable=not-an-iterable
),
found_processes,
)
terminate_process_list(found_processes, kill=True, slow_stop=False)
else:
log.debug(
"No astray processes were found listening on ports: %s",
", ".join(
[str(port) for port in self.listen_ports], # pylint: disable=not-an-iterable
),
)
[docs] def __enter__(self) -> "Daemon":
"""
Use class as a context manager.
"""
if not self.is_running():
raise RuntimeError(
"Factory not yet started. Perhaps you're after something like:\n\n"
"with {}.started() as factory:\n"
" yield factory".format(self.__class__.__name__)
)
return self
[docs] def __exit__(self, *_: Any) -> None:
"""
Exit the class context manager.
"""
self.terminate()