From 5a9e6497f711734592131561ad43f4c6f3cd56db Mon Sep 17 00:00:00 2001 From: Noemi <45180344+unflxw@users.noreply.github.com> Date: Thu, 14 Mar 2024 18:48:16 +0100 Subject: [PATCH] Implement minutely probes Implement a minutely probes system, heavily borrowing from the one in our Ruby integration. --- conftest.py | 10 ++++ src/appsignal/client.py | 2 + src/appsignal/probes.py | 82 ++++++++++++++++++++++++++++++ tests/test_probes.py | 109 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 203 insertions(+) create mode 100644 src/appsignal/probes.py create mode 100644 tests/test_probes.py diff --git a/conftest.py b/conftest.py index c3c4ea4..dde403a 100644 --- a/conftest.py +++ b/conftest.py @@ -86,6 +86,16 @@ def remove_logging_handlers_after_tests(): logger.removeHandler(handler) +@pytest.fixture(scope="function", autouse=True) +def remove_probes_after_tests(): + yield + + from appsignal.probes import _probes, _probe_states + + _probes.clear() + _probe_states.clear() + + @pytest.fixture(scope="function", autouse=True) def reset_agent_active_state(): agent.active = False diff --git a/src/appsignal/client.py b/src/appsignal/client.py index 8ef96de..dcebe17 100644 --- a/src/appsignal/client.py +++ b/src/appsignal/client.py @@ -8,6 +8,7 @@ from .agent import agent from .config import Config, Options from .opentelemetry import start_opentelemetry +from .probes import start_probes if TYPE_CHECKING: @@ -38,6 +39,7 @@ def start(self) -> None: self._logger.info("Starting AppSignal") agent.start(self._config) start_opentelemetry(self._config) + start_probes() def start_logger(self) -> None: self._logger = logging.getLogger("appsignal") diff --git a/src/appsignal/probes.py b/src/appsignal/probes.py new file mode 100644 index 0000000..a0a4169 --- /dev/null +++ b/src/appsignal/probes.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +import logging +from inspect import signature +from threading import Thread +from time import gmtime, sleep +from typing import Any, Callable, NoReturn, TypeVar, cast + + +T = TypeVar("T") + +Probe = Callable[[], None] | Callable[[T | None], T | None] + +_probes: dict[str, Probe] = {} +_probe_states: dict[str, Any] = {} + + +def start_probes() -> None: + Thread(target=_start_probes, daemon=True).start() + + +def _start_probes() -> NoReturn: + sleep(_initial_wait_time()) + + while True: + _run_probes() + sleep(_wait_time()) + + +def _run_probes() -> None: + for name in _probes: + _run_probe(name) + + +def _run_probe(name: str) -> None: + logger = logging.getLogger("appsignal") + logger.debug(f"Gathering minutely metrics with `{name}` probe") + + try: + probe = _probes[name] + + if len(signature(probe).parameters) > 0: + probe = cast(Callable[[Any], Any], probe) + state = _probe_states.get(name) + result = probe(state) + _probe_states[name] = result + else: + probe = cast(Callable[[], None], probe) + probe() + + except Exception as e: + logger.debug(f"Error in minutely probe `{name}`: {e}") + + +def _wait_time() -> int: + return 60 - gmtime().tm_sec + + +def _initial_wait_time() -> int: + remaining_seconds = _wait_time() + if remaining_seconds > 30: + return remaining_seconds + + return remaining_seconds + 60 + + +def register(name: str, probe: Probe) -> None: + if name in _probes: + logger = logging.getLogger("appsignal") + logger.debug( + f"A probe with the name `{name}` is already " + "registered. Overwriting the entry with the new probe." + ) + + _probes[name] = probe + + +def unregister(name: str) -> None: + if name in _probes: + del _probes[name] + if name in _probe_states: + del _probe_states[name] diff --git a/tests/test_probes.py b/tests/test_probes.py new file mode 100644 index 0000000..94ec4a3 --- /dev/null +++ b/tests/test_probes.py @@ -0,0 +1,109 @@ +from typing import Any, Callable, cast + +from appsignal.probes import _probe_states, _probes, _run_probes, register, unregister + + +def test_register(mocker): + probe = mocker.Mock() + register("probe_name", probe) + + assert "probe_name" in _probes + + _run_probes() + + probe.assert_called_once() + + +def test_register_with_state(mocker): + probe = mocker.Mock() + probe.return_value = "state" + + register("probe_name", probe) + + assert "probe_name" in _probes + + _run_probes() + + probe.assert_called_once_with(None) + assert _probe_states["probe_name"] == "state" + + _run_probes() + + probe.assert_called_with("state") + + +def test_register_signatures(mocker): + # `mocker.Mock` is not used here because we want to test against + # specific function signatures. + + no_args_called: Any = False + # An explicit "not called" string is used here because the first + # call to the probe will pass `None`. + one_arg_called_with: Any = "not called" + optional_arg_called_with: Any = "not called" + many_optional_args_called_with: Any = "not called" + too_many_args_called: Any = False + + def no_args(): + nonlocal no_args_called + no_args_called = True + + def one_arg(state): + nonlocal one_arg_called_with + one_arg_called_with = state + return "state" + + def optional_arg(state=None): + nonlocal optional_arg_called_with + optional_arg_called_with = state + return "state" + + def many_optional_args(state, wut=None): + nonlocal many_optional_args_called_with + many_optional_args_called_with = state + return "state" + + # Calling this should fail and log an error, but other probes should + # still run. + def too_many_args(state, wut): + nonlocal too_many_args_called + # This should not happen. + too_many_args_called = True + + register("no_args", no_args) + register("one_arg", one_arg) + register("optional_arg", optional_arg) + register("many_optional_args", many_optional_args) + register("too_many_args", cast(Callable[[], None], too_many_args)) + + _run_probes() + + assert no_args_called + assert one_arg_called_with is None + assert optional_arg_called_with is None + assert many_optional_args_called_with is None + assert not too_many_args_called + + no_args_called = False + _run_probes() + + assert no_args_called + assert one_arg_called_with == "state" + assert optional_arg_called_with == "state" + assert many_optional_args_called_with == "state" + assert not too_many_args_called + + +def test_unregister(mocker): + probe = mocker.Mock() + probe.return_value = "state" + + register("probe_name", probe) + _run_probes() + + unregister("probe_name") + _run_probes() + + assert "probe_name" not in _probes + assert "probe_name" not in _probe_states + probe.assert_called_once_with(None)