Skip to content

Commit

Permalink
Delay initialization of retry hooks
Browse files Browse the repository at this point in the history
Fixes #31
  • Loading branch information
hynek committed Oct 21, 2023
1 parent a7ef5ff commit 476f5d6
Show file tree
Hide file tree
Showing 5 changed files with 114 additions and 54 deletions.
3 changes: 3 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,9 @@ You can find our backwards-compatibility policy [here](https://github.com/hynek/
`stamina.retry_context()` now yields instances of `stamina.Attempt`.
[#22](https://github.com/hynek/stamina/pull/22)

- Initialization of instrumentation is now delayed.
This means that if there's no retries, there's no startup overhead from importing *structlog* and *prometheus_client*.


## [23.1.0](https://github.com/hynek/stamina/compare/22.2.0...23.1.0) - 2023-07-04

Expand Down
6 changes: 3 additions & 3 deletions src/stamina/_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,17 +8,17 @@
from threading import Lock
from typing import Iterable

from ._instrumentation import INSTRUMENTS
from .typing import RetryHook


@dataclass
class Config:
is_active: bool
on_retry: Iterable[RetryHook]
on_retry: Iterable[RetryHook] | None


_CONFIG = Config(is_active=True, on_retry=INSTRUMENTS)
# on_retry is lazily initialized to avoid startup overhead.
_CONFIG = Config(is_active=True, on_retry=None)
_LOCK = Lock()


Expand Down
31 changes: 20 additions & 11 deletions src/stamina/_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,8 @@

from stamina.typing import RetryDetails, RetryHook

from ._config import _CONFIG
from ._instrumentation import guess_name
from ._config import _CONFIG, Config
from ._instrumentation import get_default_hooks, guess_name


if sys.version_info >= (3, 10):
Expand Down Expand Up @@ -181,10 +181,8 @@ def __iter__(self) -> Iterator[Attempt]:

for r in _t.Retrying(
before_sleep=_make_before_sleep(
self._name, _CONFIG.on_retry, self._args, self._kw
)
if _CONFIG.on_retry
else None,
self._name, _CONFIG, self._args, self._kw
),
**self._t_kw,
):
yield Attempt(r)
Expand All @@ -193,10 +191,8 @@ def __aiter__(self) -> AsyncIterator[Attempt]:
if _CONFIG.is_active:
self._t_a_retrying = _t.AsyncRetrying(
before_sleep=_make_before_sleep(
self._name, _CONFIG.on_retry, self._args, self._kw
)
if _CONFIG.on_retry
else None,
self._name, _CONFIG, self._args, self._kw
),
**self._t_kw,
)

Expand All @@ -208,9 +204,19 @@ async def __anext__(self) -> Attempt:
return Attempt(await self._t_a_retrying.__anext__())


def _get_before_retry_hooks(config: Config) -> Iterable[RetryHook]:
"""
Return on_retry hooks if they've been initialized, otherwise initialize.
"""
if config.on_retry is None:
config.on_retry = get_default_hooks()

return config.on_retry


def _make_before_sleep(
name: str,
on_retry: Iterable[RetryHook],
config: Config,
args: tuple[object, ...],
kw: dict[str, object],
) -> Callable[[_t.RetryCallState], None]:
Expand All @@ -220,6 +226,9 @@ def _make_before_sleep(
"""

def before_sleep(rcs: _t.RetryCallState) -> None:
if not (on_retry := _get_before_retry_hooks(config)):
return

details = RetryDetails(
name=name,
attempt=rcs.attempt_number,
Expand Down
106 changes: 67 additions & 39 deletions src/stamina/_instrumentation.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,58 +4,86 @@

from __future__ import annotations

from .typing import RetryDetails
from .typing import RetryDetails, RetryHook


try:
import structlog
def get_default_hooks() -> tuple[RetryHook, ...]:
"""
Return the default hooks according to availability.
"""
hooks = []

logger = structlog.get_logger()
except ImportError:
logger = None
if prom := init_prometheus():
hooks.append(prom)

try:
from prometheus_client import Counter
if sl := init_structlog():
hooks.append(sl)

RETRY_COUNTER = Counter(
"stamina_retries_total",
"Total number of retries.",
("callable", "attempt", "error_type"),
)
except ImportError:
RETRY_COUNTER = None # type: ignore[assignment]
return tuple(hooks)


def count_retries(details: RetryDetails) -> None:
"""
Count and log retries for callable *name*.
"""
RETRY_COUNTER.labels(
callable=details.name,
attempt=details.attempt,
error_type=guess_name(details.exception.__class__),
).inc()
RETRY_COUNTER = None


def log_retries(details: RetryDetails) -> None:
logger.warning(
"stamina.retry_scheduled",
callable=details.name,
attempt=details.attempt,
slept=details.idle_for,
error=repr(details.exception),
args=tuple(repr(a) for a in details.args),
kwargs=dict(details.kwargs.items()),
)
def init_prometheus() -> RetryHook | None:
"""
Try to initialize Prometheus instrumentation.
Return None if it's not available.
"""
try:
from prometheus_client import Counter
except ImportError:
return None

global RETRY_COUNTER # noqa: PLW0603

# Mostly for testing so we can call init_prometheus more than once.
if RETRY_COUNTER is None:
RETRY_COUNTER = Counter(
"stamina_retries_total",
"Total number of retries.",
("callable", "attempt", "error_type"),
)

def count_retries(details: RetryDetails) -> None:
"""
Count and log retries for callable *name*.
"""
RETRY_COUNTER.labels(
callable=details.name,
attempt=details.attempt,
error_type=guess_name(details.exception.__class__),
).inc()

return count_retries


def init_structlog() -> RetryHook | None:
"""
Try to initialize structlog instrumentation.
INSTRUMENTS = []
Return None if it's not available.
"""
try:
import structlog
except ImportError:
return None

if RETRY_COUNTER:
INSTRUMENTS.append(count_retries)
logger = structlog.get_logger()

if logger:
INSTRUMENTS.append(log_retries)
def log_retries(details: RetryDetails) -> None:
logger.warning(
"stamina.retry_scheduled",
callable=details.name,
attempt=details.attempt,
slept=details.idle_for,
error=repr(details.exception),
args=tuple(repr(a) for a in details.args),
kwargs=dict(details.kwargs.items()),
)

return log_retries


def guess_name(obj: object) -> str:
Expand Down
22 changes: 21 additions & 1 deletion tests/test_instrumentation.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,18 @@

import pytest

from stamina._instrumentation import guess_name
from stamina._instrumentation import get_default_hooks, guess_name


try:
import structlog
except ImportError:
structlog = None

try:
import prometheus_client
except ImportError:
prometheus_client = None


def function():
Expand Down Expand Up @@ -65,3 +76,12 @@ async def async_f():
"tests.test_instrumentation.TestGuessName.test_local.<locals>.async_f"
== guess_name(async_f)
)


def test_get_default_hooks():
"""
Both default instrumentations are detected.
"""
assert len([m for m in (structlog, prometheus_client) if m]) == len(
get_default_hooks()
)

0 comments on commit 476f5d6

Please sign in to comment.