Delay initialization of retry hooks

CHANGELOG.md

@@ -22,6 +22,10 @@ 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*.
+  [#34](https://github.com/hynek/stamina/pull/34)
+
 
 ## [23.1.0](https://github.com/hynek/stamina/compare/22.2.0...23.1.0) - 2023-07-04
 

src/stamina/_config.py

@@ -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()
 
 

src/stamina/_core.py

@@ -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):
@@ -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)
@@ -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,
             )
 
@@ -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]:
@@ -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,

src/stamina/_instrumentation.py

@@ -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:

tests/test_instrumentation.py

@@ -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():
@@ -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()
+    )