|
2 | 2 | # Licensed under the MIT License. |
3 | 3 |
|
4 | 4 | """Unit tests for the _EntityExecutor class in durabletask.worker.""" |
| 5 | +import inspect |
5 | 6 | import json |
6 | 7 | import logging |
| 8 | +import threading |
7 | 9 |
|
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 |
9 | 14 | from durabletask.internal.entity_state_shim import StateShim |
10 | 15 | 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 | +) |
12 | 22 |
|
13 | 23 |
|
14 | 24 | def _make_executor(*entity_args) -> _EntityExecutor: |
@@ -289,3 +299,227 @@ def test_falsy_serialized_state_is_not_dropped(self): |
289 | 299 | state = StateShim("0", JsonDataConverter(), is_serialized=True) |
290 | 300 | assert state.get_state(int) == 0 |
291 | 301 | 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