diff --git a/.changesets/implement-heartbeats.md b/.changesets/implement-heartbeats.md new file mode 100644 index 00000000..469b02c8 --- /dev/null +++ b/.changesets/implement-heartbeats.md @@ -0,0 +1,28 @@ +--- +bump: "patch" +type: "add" +--- + +Add heartbeats support. You can use the `heartbeat` function to send heartbeats directly from your code, to track the execution of certain processes: + +```python +from appsignal import heartbeat + +def send_invoices(): + # ... your code here ... + heartbeat("send_invoices") +``` + +It is also possible to pass a defined function as an argument to the `heartbeat` +function: + +```python +def send_invoices(): + # ... your code here ... + +heartbeat("send_invoices", send_invoices) +``` + +If an exception is raised within the function, the finish event will not be +reported to AppSignal, triggering a notification about the missing heartbeat. +The exception will be raised outside of the heartbeat function. diff --git a/conftest.py b/conftest.py index 1d63b110..81730ebf 100644 --- a/conftest.py +++ b/conftest.py @@ -17,8 +17,8 @@ from appsignal import probes from appsignal.agent import agent -from appsignal.opentelemetry import METRICS_PREFERRED_TEMPORALITY from appsignal.client import _reset_client +from appsignal.opentelemetry import METRICS_PREFERRED_TEMPORALITY @pytest.fixture(scope="function", autouse=True) @@ -100,10 +100,12 @@ def stop_and_clear_probes_after_tests(): def reset_agent_active_state(): agent.active = False + @pytest.fixture(scope="function", autouse=True) def reset_global_client(): _reset_client() + @pytest.fixture(scope="function", autouse=True) def stop_agent(): tmp_path = "/tmp" if platform.system() == "Darwin" else tempfile.gettempdir() diff --git a/pyproject.toml b/pyproject.toml index d69035f5..647372fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -91,6 +91,7 @@ dependencies = [ "types-deprecated", "types-requests", "django-stubs", + "pytest", ] [tool.hatch.envs.lint.scripts] diff --git a/src/appsignal/__init__.py b/src/appsignal/__init__.py index b0af7ab2..37c206d0 100644 --- a/src/appsignal/__init__.py +++ b/src/appsignal/__init__.py @@ -4,6 +4,7 @@ from runpy import run_path from .client import Client as Appsignal +from .heartbeat import Heartbeat, heartbeat from .metrics import add_distribution_value, increment_counter, set_gauge from .tracing import ( send_error, @@ -42,6 +43,8 @@ "increment_counter", "set_gauge", "add_distribution_value", + "Heartbeat", + "heartbeat", "start", ] diff --git a/src/appsignal/client.py b/src/appsignal/client.py index 371b7bc0..62f048b3 100644 --- a/src/appsignal/client.py +++ b/src/appsignal/client.py @@ -17,6 +17,7 @@ _client: Client | None = None + def _reset_client() -> None: global _client _client = None diff --git a/src/appsignal/config.py b/src/appsignal/config.py index 37e20c17..3ba57acd 100644 --- a/src/appsignal/config.py +++ b/src/appsignal/config.py @@ -37,6 +37,7 @@ class Options(TypedDict, total=False): log: str | None log_level: str | None log_path: str | None + logging_endpoint: str | None opentelemetry_port: str | int | None name: str | None push_api_key: str | None @@ -77,6 +78,7 @@ class Config: files_world_accessible=True, log="file", log_level="info", + logging_endpoint="https://appsignal-endpoint.net", opentelemetry_port=8099, send_environment_metadata=True, send_params=True, @@ -172,6 +174,7 @@ def load_from_environment() -> Options: log=os.environ.get("APPSIGNAL_LOG"), log_level=os.environ.get("APPSIGNAL_LOG_LEVEL"), log_path=os.environ.get("APPSIGNAL_LOG_PATH"), + logging_endpoint=os.environ.get("APPSIGNAL_LOGGING_ENDPOINT"), opentelemetry_port=os.environ.get("APPSIGNAL_OPENTELEMETRY_PORT"), name=os.environ.get("APPSIGNAL_APP_NAME"), push_api_key=os.environ.get("APPSIGNAL_PUSH_API_KEY"), @@ -239,6 +242,7 @@ def set_private_environ(self) -> None: "_APPSIGNAL_LOG": options.get("log"), "_APPSIGNAL_LOG_LEVEL": options.get("log_level"), "_APPSIGNAL_LOG_FILE_PATH": self.log_file_path(), + "_APPSIGNAL_LOGGING_ENDPOINT": options.get("logging_endpoint"), "_APPSIGNAL_OPENTELEMETRY_PORT": options.get("opentelemetry_port"), "_APPSIGNAL_PUSH_API_KEY": options.get("push_api_key"), "_APPSIGNAL_PUSH_API_ENDPOINT": options.get("endpoint"), diff --git a/src/appsignal/heartbeat.py b/src/appsignal/heartbeat.py new file mode 100644 index 00000000..0516e2e6 --- /dev/null +++ b/src/appsignal/heartbeat.py @@ -0,0 +1,88 @@ +from __future__ import annotations + +import logging +from binascii import hexlify +from os import urandom +from time import time +from typing import Any, Callable, Literal, TypedDict, TypeVar, Union + +from .client import Client +from .config import Config +from .transmitter import transmit + + +T = TypeVar("T") + +EventKind = Union[Literal["start"], Literal["finish"]] + + +class Event(TypedDict): + name: str + id: str + kind: EventKind + timestamp: int + + +class Heartbeat: + name: str + id: str + + def __init__(self, name: str) -> None: + self.name = name + self.id = hexlify(urandom(8)).decode("utf-8") + + def _event(self, kind: EventKind) -> Event: + return Event(name=self.name, id=self.id, kind=kind, timestamp=int(time())) + + def _transmit(self, event: Event) -> None: + config = Client.config() or Config() + logger = logging.getLogger("appsignal") + + if not config.is_active(): + logger.debug("AppSignal not active, not transmitting heartbeat event") + return + + url = f"{config.option('logging_endpoint')}/heartbeats/json" + try: + response = transmit(url, json=event) + if 200 <= response.status_code <= 299: + logger.debug( + f"Transmitted heartbeat `{event['name']}` ({event['id']}) " + f"{event['kind']} event" + ) + else: + logger.error( + "Failed to transmit heartbeat event: " + f"status code was {response.status_code}" + ) + except Exception as e: + logger.error(f"Failed to transmit heartbeat event: {e}") + + def start(self) -> None: + self._transmit(self._event("start")) + + def finish(self) -> None: + self._transmit(self._event("finish")) + + def __enter__(self) -> None: + self.start() + + def __exit__( + self, exc_type: Any = None, exc_value: Any = None, traceback: Any = None + ) -> Literal[False]: + if exc_type is None: + self.finish() + + return False + + +def heartbeat(name: str, fn: Callable[[], T] | None = None) -> None | T: + heartbeat = Heartbeat(name) + output = None + + if fn is not None: + heartbeat.start() + output = fn() + + heartbeat.finish() + return output diff --git a/tests/diagnose b/tests/diagnose index 31e13d70..152d1e03 160000 --- a/tests/diagnose +++ b/tests/diagnose @@ -1 +1 @@ -Subproject commit 31e13d704577a138cfaba76976f843603d070b5a +Subproject commit 152d1e033c239742f9d58cefd4a830c8cdd778a6 diff --git a/tests/test_heartbeat.py b/tests/test_heartbeat.py new file mode 100644 index 00000000..d747ea56 --- /dev/null +++ b/tests/test_heartbeat.py @@ -0,0 +1,175 @@ +from time import sleep + +from pytest import raises + +from appsignal.client import Client +from appsignal.heartbeat import Heartbeat, heartbeat + + +def init_client(active=True): + Client( + push_api_key="some-push-api-key", + name="some-app", + environment="staging", + hostname="beepboop.local", + active=active, + ) + + +def test_start_finish_when_appsignal_is_not_active_sends_nothing(mocker): + mock_request = mocker.patch("requests.post") + init_client(active=False) + + heartbeat = Heartbeat("some-heartbeat") + heartbeat.start() + heartbeat.finish() + + assert not mock_request.called + + +def test_start_sends_heartbeat_start_event(mocker): + mock_request = mocker.patch("requests.post") + init_client() + + heartbeat = Heartbeat("some-heartbeat") + heartbeat.start() + + assert mock_request.called + assert ( + "https://appsignal-endpoint.net/heartbeats/json?" + in mock_request.call_args[0][0] + ) + # The ordering of query parameters is not guaranteed. + assert "api_key=some-push-api-key" in mock_request.call_args[0][0] + assert "environment=staging" in mock_request.call_args[0][0] + assert "hostname=beepboop.local" in mock_request.call_args[0][0] + assert "name=some-app" in mock_request.call_args[0][0] + + assert mock_request.call_args[1]["json"]["name"] == "some-heartbeat" + assert mock_request.call_args[1]["json"]["kind"] == "start" + assert isinstance(mock_request.call_args[1]["json"]["timestamp"], int) + assert isinstance(mock_request.call_args[1]["json"]["id"], str) + + +def test_finish_sends_heartbeat_finish_event(mocker): + mock_request = mocker.patch("requests.post") + init_client() + + heartbeat = Heartbeat("some-heartbeat") + heartbeat.finish() + + assert mock_request.called + assert ( + "https://appsignal-endpoint.net/heartbeats/json?" + in mock_request.call_args[0][0] + ) + # The ordering of query parameters is not guaranteed. + assert "api_key=some-push-api-key" in mock_request.call_args[0][0] + assert "environment=staging" in mock_request.call_args[0][0] + assert "hostname=beepboop.local" in mock_request.call_args[0][0] + assert "name=some-app" in mock_request.call_args[0][0] + + assert mock_request.call_args[1]["json"]["name"] == "some-heartbeat" + assert mock_request.call_args[1]["json"]["kind"] == "finish" + assert isinstance(mock_request.call_args[1]["json"]["timestamp"], int) + assert isinstance(mock_request.call_args[1]["json"]["id"], str) + + +def test_heartbeat_sends_heartbeat_finish_event(mocker): + mock_request = mocker.patch("requests.post") + init_client() + + heartbeat("some-heartbeat") + + assert mock_request.called + assert len(mock_request.call_args_list) == 1 + + assert mock_request.call_args[1]["json"]["name"] == "some-heartbeat" + assert mock_request.call_args[1]["json"]["kind"] == "finish" + + +def test_heartbeat_with_function_sends_heartbeat_start_and_finish_event(mocker): + mock_request = mocker.patch("requests.post") + init_client() + + def some_function(): + sleep(1.1) + return "output" + + assert heartbeat("some-heartbeat", some_function) == "output" + + assert mock_request.called + assert len(mock_request.call_args_list) == 2 + + assert mock_request.call_args_list[0][1]["json"]["name"] == "some-heartbeat" + assert mock_request.call_args_list[0][1]["json"]["kind"] == "start" + assert mock_request.call_args_list[1][1]["json"]["name"] == "some-heartbeat" + assert mock_request.call_args_list[1][1]["json"]["kind"] == "finish" + assert ( + mock_request.call_args_list[0][1]["json"]["timestamp"] + < mock_request.call_args_list[1][1]["json"]["timestamp"] + ) + assert ( + mock_request.call_args_list[0][1]["json"]["id"] + == mock_request.call_args_list[1][1]["json"]["id"] + ) + + +def test_heartbeat_with_function_does_not_send_heartbeat_finish_event_on_exception( + mocker, +): + mock_request = mocker.patch("requests.post") + init_client() + + def some_function(): + raise Exception("Whoops!") + + with raises(Exception, match="Whoops!"): + heartbeat("some-heartbeat", some_function) + + assert mock_request.called + assert len(mock_request.call_args_list) == 1 + + assert mock_request.call_args_list[0][1]["json"]["name"] == "some-heartbeat" + assert mock_request.call_args_list[0][1]["json"]["kind"] == "start" + + +def test_heartbeat_context_manager_sends_heartbeat_start_and_finish_event(mocker): + mock_request = mocker.patch("requests.post") + init_client() + + with Heartbeat("some-heartbeat"): + sleep(1.1) + + assert mock_request.called + assert len(mock_request.call_args_list) == 2 + + assert mock_request.call_args_list[0][1]["json"]["name"] == "some-heartbeat" + assert mock_request.call_args_list[0][1]["json"]["kind"] == "start" + assert mock_request.call_args_list[1][1]["json"]["name"] == "some-heartbeat" + assert mock_request.call_args_list[1][1]["json"]["kind"] == "finish" + assert ( + mock_request.call_args_list[0][1]["json"]["timestamp"] + < mock_request.call_args_list[1][1]["json"]["timestamp"] + ) + assert ( + mock_request.call_args_list[0][1]["json"]["id"] + == mock_request.call_args_list[1][1]["json"]["id"] + ) + + +def test_heartbeat_context_manager_does_not_send_heartbeat_finish_event_on_exception( + mocker, +): + mock_request = mocker.patch("requests.post") + init_client() + + with raises(Exception, match="Whoops!"): + with Heartbeat("some-heartbeat"): + raise Exception("Whoops!") + + assert mock_request.called + assert len(mock_request.call_args_list) == 1 + + assert mock_request.call_args_list[0][1]["json"]["name"] == "some-heartbeat" + assert mock_request.call_args_list[0][1]["json"]["kind"] == "start"