Add cache size monitoring to detect potential memory leaks

src/docket/annotations.py

@@ -2,6 +2,8 @@
 import inspect
 from typing import Any, Iterable, Mapping, Self
 
+from .instrumentation import CACHE_SIZE
+
 
 class Annotation(abc.ABC):
     _cache: dict[tuple[type[Self], inspect.Signature], Mapping[str, Self]] = {}
@@ -10,6 +12,7 @@ class Annotation(abc.ABC):
     def annotated_parameters(cls, signature: inspect.Signature) -> Mapping[str, Self]:
         key = (cls, signature)
         if key in cls._cache:
+            CACHE_SIZE.set(len(cls._cache), {"cache": "annotation"})
             return cls._cache[key]
 
         annotated: dict[str, Self] = {}
@@ -30,6 +33,7 @@ def annotated_parameters(cls, signature: inspect.Signature) -> Mapping[str, Self
                     annotated[param_name] = arg_type()
 
         cls._cache[key] = annotated
+        CACHE_SIZE.set(len(cls._cache), {"cache": "annotation"})
         return annotated
 
 

src/docket/dependencies.py

@@ -21,6 +21,7 @@
 
 from .docket import Docket
 from .execution import Execution, TaskFunction, get_signature
+from .instrumentation import CACHE_SIZE
 
 if TYPE_CHECKING:  # pragma: no cover
     from .worker import Worker
@@ -415,6 +416,7 @@ def get_dependency_parameters(
     function: TaskFunction | DependencyFunction[Any],
 ) -> dict[str, Dependency]:
     if function in _parameter_cache:
+        CACHE_SIZE.set(len(_parameter_cache), {"cache": "parameter"})
         return _parameter_cache[function]
 
     dependencies: dict[str, Dependency] = {}
@@ -428,6 +430,7 @@ def get_dependency_parameters(
         dependencies[parameter] = param.default
 
     _parameter_cache[function] = dependencies
+    CACHE_SIZE.set(len(_parameter_cache), {"cache": "parameter"})
     return dependencies
 
 

src/docket/execution.py

@@ -19,7 +19,7 @@
 from opentelemetry import propagate, trace
 
 from .annotations import Logged
-from .instrumentation import message_getter
+from .instrumentation import CACHE_SIZE, message_getter
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -32,10 +32,12 @@
 
 def get_signature(function: Callable[..., Any]) -> inspect.Signature:
     if function in _signature_cache:
+        CACHE_SIZE.set(len(_signature_cache), {"cache": "signature"})
         return _signature_cache[function]
 
     signature = inspect.signature(function)
     _signature_cache[function] = signature
+    CACHE_SIZE.set(len(_signature_cache), {"cache": "signature"})
     return signature
 
 

src/docket/instrumentation.py

@@ -117,6 +117,12 @@
     unit="1",
 )
 
+CACHE_SIZE = meter.create_gauge(
+    "docket_cache_size",
+    description="Size of internal docket caches",
+    unit="1",
+)
+
 Message = dict[bytes, bytes]