Skip to content

Commit 7567860

Browse files
Bernd VerstCopilot
andcommitted
Make entity method signature caching effective
_execute_entity_batch() constructed a new _EntityExecutor inside its per-operation loop, so the executor's _entity_method_cache was always cold and inspect.signature() was re-run for every single entity operation. The executor is now constructed once per batch and shares a bounded, thread-safe signature cache owned by the worker, so a given (entity type, operation) pair is reflected once for the lifetime of the worker rather than once per operation. The cache is bounded (_EntityMethodSignatureCache, 1024 entries, oldest evicted first) because entity types and operation names come from user code and are unbounded in principle. Writes are serialized under a lock so concurrently processed entity work items stay safe. Behavior is unchanged: the dispatch logic in _EntityExecutor.execute() is untouched, so the no-argument versus input-argument distinction and all error messages for unknown or invalid operations are preserved. _EntityExecutor's existing three-argument constructor still works; the shared cache is an optional fourth argument. Fixes #186 Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 2265f7b9-f8ad-4b50-a25d-19bd6f1ac96c
1 parent 4d87430 commit 7567860

2 files changed

Lines changed: 298 additions & 5 deletions

File tree

‎durabletask/worker.py‎

Lines changed: 62 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -534,6 +534,9 @@ def __init__(
534534
data_converter: DataConverter | None = None,
535535
):
536536
self._registry = _Registry()
537+
# Shared by every entity batch this worker processes so that reflected
538+
# entity method signatures are computed once instead of per operation.
539+
self._entity_method_cache = _EntityMethodSignatureCache()
537540
self._host_address = (
538541
host_address if host_address else shared.get_default_host_address()
539542
)
@@ -1370,9 +1373,17 @@ def _execute_entity_batch(
13701373
raise RuntimeError(f"Invalid entity instance ID '{instance_id}' in entity operation request.")
13711374

13721375
results: list[pb.OperationResult] = []
1376+
# A single executor serves the whole batch and shares the worker's
1377+
# bounded signature cache, so entity method reflection is not repeated
1378+
# for every operation.
1379+
executor = _EntityExecutor(
1380+
self._registry,
1381+
self._logger,
1382+
self._data_converter,
1383+
entity_method_cache=self._entity_method_cache,
1384+
)
13731385
for operation in req.operations:
13741386
start_time = datetime.now(timezone.utc)
1375-
executor = _EntityExecutor(self._registry, self._logger, self._data_converter)
13761387

13771388
operation_result = None
13781389

@@ -3101,13 +3112,61 @@ def execute(
31013112
return encoded_output
31023113

31033114

3115+
class _EntityMethodSignatureCache:
3116+
"""Bounded, thread-safe cache of reflected entity method signatures.
3117+
3118+
Dispatching an operation on a class-based entity requires knowing whether
3119+
the target method declares a required parameter, which is answered with an
3120+
``inspect.signature()`` call. The answer depends only on the entity type
3121+
and the operation name, so it is cached here and shared across entity
3122+
batches (and therefore across the worker threads that process them
3123+
concurrently).
3124+
3125+
The cache is bounded because entity types and operation names originate
3126+
from user code and are unbounded in principle; once the limit is reached
3127+
the oldest entries are evicted instead of growing without limit.
3128+
"""
3129+
3130+
DEFAULT_MAX_SIZE = 1024
3131+
3132+
def __init__(self, max_size: int = DEFAULT_MAX_SIZE):
3133+
self._max_size = max(1, max_size)
3134+
self._lock = Lock()
3135+
self._entries: dict[tuple[type, str], bool] = {}
3136+
3137+
def get(self, key: tuple[type, str]) -> bool | None:
3138+
"""Returns the cached result for ``key``, or ``None`` when not cached."""
3139+
return self._entries.get(key)
3140+
3141+
def __setitem__(self, key: tuple[type, str], value: bool) -> None:
3142+
with self._lock:
3143+
self._entries[key] = value
3144+
# Dicts preserve insertion order, so the first key is the oldest.
3145+
while len(self._entries) > self._max_size:
3146+
self._entries.pop(next(iter(self._entries)), None)
3147+
3148+
def __contains__(self, key: object) -> bool:
3149+
return key in self._entries
3150+
3151+
def __len__(self) -> int:
3152+
return len(self._entries)
3153+
3154+
31043155
class _EntityExecutor:
31053156
def __init__(self, registry: _Registry, logger: logging.Logger,
3106-
data_converter: DataConverter):
3157+
data_converter: DataConverter,
3158+
entity_method_cache: _EntityMethodSignatureCache | None = None):
31073159
self._registry = registry
31083160
self._logger = logger
31093161
self._data_converter = data_converter
3110-
self._entity_method_cache: dict[tuple[type, str], bool] = {}
3162+
# When the caller supplies a cache (the worker does), reflected entity
3163+
# method signatures are shared with it and therefore survive beyond the
3164+
# lifetime of this executor.
3165+
self._entity_method_cache: _EntityMethodSignatureCache = (
3166+
entity_method_cache
3167+
if entity_method_cache is not None
3168+
else _EntityMethodSignatureCache()
3169+
)
31113170

31123171
def execute(
31133172
self,

‎tests/durabletask/test_entity_executor.py‎

Lines changed: 236 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,23 @@
22
# Licensed under the MIT License.
33

44
"""Unit tests for the _EntityExecutor class in durabletask.worker."""
5+
import inspect
56
import json
67
import logging
8+
import threading
79

8-
from durabletask import entities
10+
from google.protobuf import wrappers_pb2
11+
12+
import durabletask.internal.orchestrator_service_pb2 as pb
13+
from durabletask import entities, worker
914
from durabletask.internal.entity_state_shim import StateShim
1015
from durabletask.serialization import JsonDataConverter
11-
from durabletask.worker import _EntityExecutor, _Registry
16+
from durabletask.worker import (
17+
TaskHubGrpcWorker,
18+
_EntityExecutor,
19+
_EntityMethodSignatureCache,
20+
_Registry,
21+
)
1222

1323

1424
def _make_executor(*entity_args) -> _EntityExecutor:
@@ -289,3 +299,227 @@ def test_falsy_serialized_state_is_not_dropped(self):
289299
state = StateShim("0", JsonDataConverter(), is_serialized=True)
290300
assert state.get_state(int) == 0
291301
assert state.encode_state() == "0"
302+
303+
304+
class _CountingInspect:
305+
"""Delegates to :mod:`inspect` while counting ``signature()`` calls."""
306+
307+
def __init__(self):
308+
self.signature_calls: list[object] = []
309+
310+
def signature(self, obj, *args, **kwargs):
311+
self.signature_calls.append(obj)
312+
return inspect.signature(obj, *args, **kwargs)
313+
314+
def __getattr__(self, name):
315+
return getattr(inspect, name)
316+
317+
318+
class _RecordingEntityStub:
319+
"""Captures the batch results the worker would send to the sidecar."""
320+
321+
def __init__(self):
322+
self.completed: list[pb.EntityBatchResult] = []
323+
324+
def CompleteEntityTask(self, batch_result):
325+
self.completed.append(batch_result)
326+
327+
328+
def _entity_batch_request(instance_id: str, operations) -> pb.EntityBatchRequest:
329+
"""Builds an entity batch request from (operation, encoded_input) pairs."""
330+
return pb.EntityBatchRequest(
331+
instanceId=instance_id,
332+
operations=[
333+
pb.OperationRequest(
334+
operation=name,
335+
requestId=str(index),
336+
input=(
337+
wrappers_pb2.StringValue(value=encoded_input)
338+
if encoded_input is not None
339+
else None
340+
),
341+
)
342+
for index, (name, encoded_input) in enumerate(operations)
343+
],
344+
)
345+
346+
347+
class Counter(entities.DurableEntity):
348+
"""Entity exercising both input-argument and no-argument dispatch."""
349+
350+
def add(self, value: int):
351+
current = self.get_state(int, 0)
352+
self.set_state(current + value)
353+
return current + value
354+
355+
def get(self):
356+
return self.get_state(int, 0)
357+
358+
359+
class TestEntityMethodSignatureCaching:
360+
"""Reflection of an entity method signature must happen only once."""
361+
362+
def test_repeated_operations_reflect_once_per_executor(self, monkeypatch):
363+
executor = _make_executor(Counter)
364+
counting = _CountingInspect()
365+
monkeypatch.setattr(worker, "inspect", counting)
366+
367+
entity_id = entities.EntityInstanceId("Counter", "test-key")
368+
state = StateShim(None, JsonDataConverter())
369+
for _ in range(5):
370+
executor.execute("test-orch", entity_id, "add", state, "1")
371+
state.commit()
372+
373+
assert len(counting.signature_calls) == 1
374+
assert state.get_state(int) == 5
375+
376+
def test_cache_is_shared_between_executors(self, monkeypatch):
377+
shared_cache = _EntityMethodSignatureCache()
378+
registry = _Registry()
379+
registry.add_entity(Counter)
380+
counting = _CountingInspect()
381+
monkeypatch.setattr(worker, "inspect", counting)
382+
383+
entity_id = entities.EntityInstanceId("Counter", "test-key")
384+
for _ in range(3):
385+
executor = _EntityExecutor(
386+
registry, logging.getLogger("test"), JsonDataConverter(),
387+
entity_method_cache=shared_cache,
388+
)
389+
executor.execute(
390+
"test-orch", entity_id, "add",
391+
StateShim(None, JsonDataConverter()), "1")
392+
393+
assert len(counting.signature_calls) == 1
394+
395+
def test_executor_without_shared_cache_still_caches(self, monkeypatch):
396+
# Constructing the executor with the historical three-argument form
397+
# keeps working and still avoids repeated reflection.
398+
registry = _Registry()
399+
registry.add_entity(Counter)
400+
executor = _EntityExecutor(
401+
registry, logging.getLogger("test"), JsonDataConverter())
402+
counting = _CountingInspect()
403+
monkeypatch.setattr(worker, "inspect", counting)
404+
405+
entity_id = entities.EntityInstanceId("Counter", "test-key")
406+
state = StateShim(None, JsonDataConverter())
407+
executor.execute("test-orch", entity_id, "add", state, "1")
408+
executor.execute("test-orch", entity_id, "add", state, "1")
409+
410+
assert len(counting.signature_calls) == 1
411+
412+
413+
class TestEntityBatchSignatureCaching:
414+
"""The worker must not recompute signatures for every batch operation."""
415+
416+
@staticmethod
417+
def _make_worker() -> TaskHubGrpcWorker:
418+
instance = TaskHubGrpcWorker()
419+
instance.add_entity(Counter)
420+
return instance
421+
422+
def test_batch_reflects_once_for_repeated_operations(self, monkeypatch):
423+
instance = self._make_worker()
424+
stub = _RecordingEntityStub()
425+
counting = _CountingInspect()
426+
monkeypatch.setattr(worker, "inspect", counting)
427+
428+
request = _entity_batch_request("@counter@key", [("add", "1")] * 5)
429+
result = instance._execute_entity_batch(request, stub, b"token")
430+
431+
assert [r.WhichOneof("resultType") for r in result.results] == ["success"] * 5
432+
assert json.loads(result.entityState.value) == 5
433+
assert len(counting.signature_calls) == 1
434+
assert len(stub.completed) == 1
435+
436+
def test_cache_survives_across_batches(self, monkeypatch):
437+
instance = self._make_worker()
438+
stub = _RecordingEntityStub()
439+
counting = _CountingInspect()
440+
monkeypatch.setattr(worker, "inspect", counting)
441+
442+
for _ in range(3):
443+
instance._execute_entity_batch(
444+
_entity_batch_request("@counter@key", [("add", "1")]), stub, b"token")
445+
446+
assert len(counting.signature_calls) == 1
447+
448+
def test_worker_cache_is_populated_by_batch(self):
449+
instance = self._make_worker()
450+
instance._execute_entity_batch(
451+
_entity_batch_request("@counter@key", [("add", "1")]),
452+
_RecordingEntityStub(), b"token")
453+
454+
assert (Counter, "add") in instance._entity_method_cache
455+
456+
def test_no_arg_and_input_arg_dispatch_preserved(self, monkeypatch):
457+
instance = self._make_worker()
458+
stub = _RecordingEntityStub()
459+
counting = _CountingInspect()
460+
monkeypatch.setattr(worker, "inspect", counting)
461+
462+
request = _entity_batch_request(
463+
"@counter@key",
464+
[("add", "2"), ("get", None), ("add", "3"), ("get", None)],
465+
)
466+
result = instance._execute_entity_batch(request, stub, b"token")
467+
468+
assert [r.success.result.value for r in result.results] == ["2", "2", "5", "5"]
469+
# One reflection per distinct operation, regardless of repeat count.
470+
assert len(counting.signature_calls) == 2
471+
472+
def test_unknown_operation_failure_is_unchanged(self):
473+
instance = self._make_worker()
474+
result = instance._execute_entity_batch(
475+
_entity_batch_request("@counter@key", [("missing", None)]),
476+
_RecordingEntityStub(), b"token")
477+
478+
failure = result.results[0]
479+
assert failure.WhichOneof("resultType") == "failure"
480+
details = failure.failure.failureDetails
481+
assert details.errorType == "builtins.AttributeError"
482+
assert "does not have operation 'missing'" in details.errorMessage
483+
484+
485+
class TestEntityMethodSignatureCacheBounds:
486+
"""The shared signature cache must stay bounded and thread-safe."""
487+
488+
def test_get_returns_none_for_unknown_key(self):
489+
cache = _EntityMethodSignatureCache()
490+
assert cache.get((int, "missing")) is None
491+
492+
def test_false_values_are_distinguishable_from_misses(self):
493+
cache = _EntityMethodSignatureCache()
494+
cache[(int, "op")] = False
495+
assert cache.get((int, "op")) is False
496+
assert (int, "op") in cache
497+
498+
def test_evicts_oldest_entries_when_full(self):
499+
cache = _EntityMethodSignatureCache(max_size=2)
500+
cache[(int, "op")] = True
501+
cache[(str, "op")] = True
502+
cache[(float, "op")] = True
503+
504+
assert len(cache) == 2
505+
assert cache.get((int, "op")) is None
506+
assert cache.get((str, "op")) is True
507+
assert cache.get((float, "op")) is True
508+
509+
def test_concurrent_writes_stay_within_bound(self):
510+
cache = _EntityMethodSignatureCache(max_size=8)
511+
entity_types = [type(f"Entity{index}", (), {}) for index in range(64)]
512+
barrier = threading.Barrier(4)
513+
514+
def writer(worker_index: int) -> None:
515+
barrier.wait()
516+
for index, entity_type in enumerate(entity_types):
517+
cache[(entity_type, f"op{worker_index}")] = index % 2 == 0
518+
519+
threads = [threading.Thread(target=writer, args=(index,)) for index in range(4)]
520+
for thread in threads:
521+
thread.start()
522+
for thread in threads:
523+
thread.join()
524+
525+
assert len(cache) == 8

0 commit comments

Comments
 (0)