|
5 | 5 |
|
6 | 6 | import asyncio |
7 | 7 | import gzip |
| 8 | +import json |
8 | 9 | import threading |
| 10 | +from typing import Any, Generator |
9 | 11 | from unittest.mock import AsyncMock, MagicMock |
10 | 12 | from urllib.parse import quote, urlsplit |
11 | 13 |
|
12 | 14 | import pytest |
13 | 15 | from google.protobuf import wrappers_pb2 |
14 | 16 |
|
| 17 | +from durabletask import client, task, worker |
15 | 18 | import durabletask.internal.orchestrator_service_pb2 as pb |
16 | 19 | from durabletask.internal.exceptions import PayloadSizeExceededError |
17 | 20 | from durabletask.payload.helpers import ( |
|
25 | 28 | LargePayloadStorageOptions, |
26 | 29 | PayloadStore, |
27 | 30 | ) |
| 31 | +from durabletask.testing import create_test_backend |
| 32 | +from tests.durabletask._port_utils import find_free_port |
28 | 33 |
|
29 | 34 |
|
30 | 35 | # ------------------------------------------------------------------ |
@@ -661,6 +666,128 @@ def test_is_known_token(self): |
661 | 666 | assert store.is_known_token("blob:v3:https://example.blob.core.windows.net/c/b") is False |
662 | 667 |
|
663 | 668 |
|
| 669 | +# ------------------------------------------------------------------ |
| 670 | +# Tests: worker isolation |
| 671 | +# ------------------------------------------------------------------ |
| 672 | + |
| 673 | + |
| 674 | +@pytest.mark.parametrize("target", ["orchestration", "activity"]) |
| 675 | +def test_missing_payload_download_does_not_interrupt_other_work( |
| 676 | + monkeypatch: pytest.MonkeyPatch, target: str) -> None: |
| 677 | + missing_value = "missing" * 100 |
| 678 | + missing_blobs: dict[str, bytes] = {} |
| 679 | + download_attempted = threading.Event() |
| 680 | + abandoned = threading.Event() |
| 681 | + entered = [threading.Event(), threading.Event()] |
| 682 | + release = threading.Event() |
| 683 | + |
| 684 | + class MissingPayloadStore(FakePayloadStore): |
| 685 | + def upload(self, data: bytes, *, instance_id: str | None = None) -> str: |
| 686 | + token = super().upload(data, instance_id=instance_id) |
| 687 | + if instance_id == "missing-payload" and not missing_blobs: |
| 688 | + missing_blobs[token] = self._blobs.pop(token) |
| 689 | + return token |
| 690 | + |
| 691 | + def download(self, token: str) -> bytes: |
| 692 | + if token in missing_blobs and token not in self._blobs: |
| 693 | + download_attempted.set() |
| 694 | + return super().download(token) |
| 695 | + |
| 696 | + def restore(self) -> None: |
| 697 | + self._blobs.update(missing_blobs) |
| 698 | + |
| 699 | + store = MissingPayloadStore() |
| 700 | + |
| 701 | + def echo(ctx: task.ActivityContext, value: Any) -> Any: |
| 702 | + return value |
| 703 | + |
| 704 | + def blocked(ctx: task.ActivityContext, value: dict[str, Any]) -> dict[str, Any]: |
| 705 | + entered[int(value["index"])].set() |
| 706 | + release.wait() |
| 707 | + return value |
| 708 | + |
| 709 | + def healthy(ctx: task.OrchestrationContext, value: dict[str, Any]) -> Generator[task.Task[Any], Any, Any]: |
| 710 | + return (yield ctx.call_activity(blocked, input=value)) |
| 711 | + |
| 712 | + def roundtrip(ctx: task.OrchestrationContext, value: Any) -> Generator[task.Task[Any], Any, Any]: |
| 713 | + return (yield ctx.call_activity(echo, input=value)) |
| 714 | + |
| 715 | + def missing_activity(ctx: task.OrchestrationContext, _: Any) -> Generator[task.Task[Any], Any, Any]: |
| 716 | + return (yield ctx.call_activity(echo, input=missing_value)) |
| 717 | + |
| 718 | + def passthrough(ctx: task.OrchestrationContext, value: Any) -> Any: |
| 719 | + return value |
| 720 | + |
| 721 | + port = find_free_port() |
| 722 | + address = f"localhost:{port}" |
| 723 | + backend = create_test_backend(port=port) |
| 724 | + try: |
| 725 | + with worker.TaskHubGrpcWorker( |
| 726 | + host_address=address, payload_store=store, |
| 727 | + concurrency_options=worker.ConcurrencyOptions( |
| 728 | + maximum_concurrent_orchestration_work_items=1, |
| 729 | + maximum_concurrent_activity_work_items=3, |
| 730 | + maximum_thread_pool_workers=4, |
| 731 | + )) as host: |
| 732 | + host.add_orchestrator(healthy) |
| 733 | + host.add_orchestrator(roundtrip) |
| 734 | + host.add_orchestrator(missing_activity) |
| 735 | + host.add_orchestrator(passthrough) |
| 736 | + host.add_activity(echo) |
| 737 | + host.add_activity(blocked) |
| 738 | + cancel_name = "_cancel_orchestrator" if target == "orchestration" else "_cancel_activity" |
| 739 | + original_cancel = getattr(host, cancel_name) |
| 740 | + |
| 741 | + def observe_abandon(request: Any, stub: Any, token: Any) -> None: |
| 742 | + original_cancel(request, stub, token) |
| 743 | + instance_id = request.instanceId if target == "orchestration" else request.orchestrationInstance.instanceId |
| 744 | + assert instance_id == "missing-payload" |
| 745 | + abandoned.set() |
| 746 | + |
| 747 | + monkeypatch.setattr(host, cancel_name, observe_abandon) |
| 748 | + host.start() |
| 749 | + try: |
| 750 | + with client.TaskHubGrpcClient(host_address=address, payload_store=store) as connection: |
| 751 | + inputs = [{"index": index, "data": "healthy" * 100} for index in range(len(entered))] |
| 752 | + instance_ids = [ |
| 753 | + connection.schedule_new_orchestration(healthy, input=value) for value in inputs |
| 754 | + ] |
| 755 | + for event in entered: |
| 756 | + assert event.wait(10), "Healthy activity did not start" |
| 757 | + |
| 758 | + connection.schedule_new_orchestration( |
| 759 | + passthrough if target == "orchestration" else missing_activity, |
| 760 | + instance_id="missing-payload", |
| 761 | + input=missing_value if target == "orchestration" else None, |
| 762 | + ) |
| 763 | + assert abandoned.wait(10), "Missing payload work item was not abandoned" |
| 764 | + assert download_attempted.is_set() |
| 765 | + # Preserve abandonment policy, but stop immediate orchestration |
| 766 | + # redelivery from repeatedly reading the deliberately missing blob. |
| 767 | + store.restore() |
| 768 | + |
| 769 | + # New work must complete while the original healthy activities |
| 770 | + # are still blocked, exercising continued polling and free permits. |
| 771 | + fresh_payload = "fresh" * 100 |
| 772 | + fresh_id = connection.schedule_new_orchestration(roundtrip, input=fresh_payload) |
| 773 | + completed = connection.wait_for_orchestration_completion(fresh_id, timeout=10) |
| 774 | + assert completed is not None |
| 775 | + assert completed.runtime_status == client.OrchestrationStatus.COMPLETED |
| 776 | + assert json.loads(completed.serialized_output or "null") == fresh_payload |
| 777 | + |
| 778 | + release.set() |
| 779 | + for instance_id, value in zip(instance_ids, inputs): |
| 780 | + completed = connection.wait_for_orchestration_completion(instance_id, timeout=10) |
| 781 | + assert completed is not None |
| 782 | + assert completed.runtime_status == client.OrchestrationStatus.COMPLETED |
| 783 | + assert json.loads(completed.serialized_output or "null") == value |
| 784 | + finally: |
| 785 | + store.restore() |
| 786 | + release.set() |
| 787 | + finally: |
| 788 | + backend.stop() |
| 789 | + |
| 790 | + |
664 | 791 | # ------------------------------------------------------------------ |
665 | 792 | # Tests: BlobPayloadStore construction and defaults |
666 | 793 | # ------------------------------------------------------------------ |
|
0 commit comments