Skip to content

Commit 9a1f101

Browse files
chrisguidryclaude
andcommitted
Add cache size monitoring to detect potential memory leaks
Added CACHE_SIZE gauge metric to track the size of global caches that could grow unbounded during infinite retry scenarios. Instrumented three key caches: - Function signature cache (_signature_cache) - Dependency parameter cache (_parameter_cache) - Annotation cache (_cache) 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
1 parent 819bfc9 commit 9a1f101

File tree

5 files changed

+450
-417
lines changed

5 files changed

+450
-417
lines changed

src/docket/annotations.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22
import inspect
33
from typing import Any, Iterable, Mapping, Self
44

5+
from .instrumentation import CACHE_SIZE
6+
57

68
class Annotation(abc.ABC):
79
_cache: dict[tuple[type[Self], inspect.Signature], Mapping[str, Self]] = {}
@@ -10,6 +12,7 @@ class Annotation(abc.ABC):
1012
def annotated_parameters(cls, signature: inspect.Signature) -> Mapping[str, Self]:
1113
key = (cls, signature)
1214
if key in cls._cache:
15+
CACHE_SIZE.set(len(cls._cache), {"cache": "annotation"})
1316
return cls._cache[key]
1417

1518
annotated: dict[str, Self] = {}
@@ -30,6 +33,7 @@ def annotated_parameters(cls, signature: inspect.Signature) -> Mapping[str, Self
3033
annotated[param_name] = arg_type()
3134

3235
cls._cache[key] = annotated
36+
CACHE_SIZE.set(len(cls._cache), {"cache": "annotation"})
3337
return annotated
3438

3539

src/docket/dependencies.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
from .docket import Docket
2323
from .execution import Execution, TaskFunction, get_signature
24+
from .instrumentation import CACHE_SIZE
2425

2526
if TYPE_CHECKING: # pragma: no cover
2627
from .worker import Worker
@@ -415,6 +416,7 @@ def get_dependency_parameters(
415416
function: TaskFunction | DependencyFunction[Any],
416417
) -> dict[str, Dependency]:
417418
if function in _parameter_cache:
419+
CACHE_SIZE.set(len(_parameter_cache), {"cache": "parameter"})
418420
return _parameter_cache[function]
419421

420422
dependencies: dict[str, Dependency] = {}
@@ -428,6 +430,7 @@ def get_dependency_parameters(
428430
dependencies[parameter] = param.default
429431

430432
_parameter_cache[function] = dependencies
433+
CACHE_SIZE.set(len(_parameter_cache), {"cache": "parameter"})
431434
return dependencies
432435

433436

src/docket/execution.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
from opentelemetry import propagate, trace
2020

2121
from .annotations import Logged
22-
from .instrumentation import message_getter
22+
from .instrumentation import CACHE_SIZE, message_getter
2323

2424
logger: logging.Logger = logging.getLogger(__name__)
2525

@@ -32,10 +32,12 @@
3232

3333
def get_signature(function: Callable[..., Any]) -> inspect.Signature:
3434
if function in _signature_cache:
35+
CACHE_SIZE.set(len(_signature_cache), {"cache": "signature"})
3536
return _signature_cache[function]
3637

3738
signature = inspect.signature(function)
3839
_signature_cache[function] = signature
40+
CACHE_SIZE.set(len(_signature_cache), {"cache": "signature"})
3941
return signature
4042

4143

src/docket/instrumentation.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,12 @@
117117
unit="1",
118118
)
119119

120+
CACHE_SIZE = meter.create_gauge(
121+
"docket_cache_size",
122+
description="Size of internal docket caches",
123+
unit="1",
124+
)
125+
120126
Message = dict[bytes, bytes]
121127

122128

0 commit comments

Comments
 (0)