|
6 | 6 | import asyncio |
7 | 7 | import gc |
8 | 8 | import json |
| 9 | +import os |
| 10 | +import subprocess |
| 11 | +import sys |
9 | 12 | from contextvars import ContextVar |
10 | 13 | from dataclasses import dataclass |
11 | 14 | from datetime import timedelta |
| 15 | +from functools import partial |
12 | 16 | from typing import Any |
13 | 17 | from weakref import ReferenceType, ref |
14 | 18 |
|
@@ -463,6 +467,114 @@ async def branch(index: int) -> None: |
463 | 467 | assert _execute(orchestrator, [], old).actions == expected |
464 | 468 |
|
465 | 469 |
|
| 470 | +async def _shared_future_cancellation_orchestrator( |
| 471 | + ctx: task.OrchestrationContext, _: Any, *, direct_future: bool, gather_wait: bool, |
| 472 | +) -> str: |
| 473 | + loop = asyncio.get_running_loop() |
| 474 | + futures: list[asyncio.Future[None]] = [ |
| 475 | + asyncio.Future[None]() if direct_future else loop.create_future() for _ in range(2)] |
| 476 | + |
| 477 | + def on_cancelled(index: int, future: asyncio.Future[None]) -> None: |
| 478 | + assert future.cancelled() |
| 479 | + ctx.call_activity(f"cancelled-{index}") |
| 480 | + |
| 481 | + for index, future in enumerate(futures): |
| 482 | + future.add_done_callback(partial(on_cancelled, index)) |
| 483 | + |
| 484 | + async def wait(future: asyncio.Future[None]) -> None: |
| 485 | + if gather_wait: |
| 486 | + await asyncio.gather(future) |
| 487 | + else: |
| 488 | + await future |
| 489 | + |
| 490 | + async def fail() -> None: |
| 491 | + await asyncio.sleep(0) |
| 492 | + raise ValueError("cancel siblings") |
| 493 | + |
| 494 | + try: |
| 495 | + async with asyncio.TaskGroup() as group: |
| 496 | + for future in futures * 2: |
| 497 | + group.create_task(wait(future)) |
| 498 | + group.create_task(fail()) |
| 499 | + except* ValueError: |
| 500 | + pass |
| 501 | + return await ctx.call_activity("join", return_type=str) |
| 502 | + |
| 503 | + |
| 504 | +@pytest.mark.parametrize("direct_future", [False, True]) |
| 505 | +@pytest.mark.parametrize("gather_wait", [False, True]) |
| 506 | +def test_shared_future_cancellation_callbacks_replay_on_fresh_executors( |
| 507 | + direct_future: bool, gather_wait: bool, |
| 508 | +) -> None: |
| 509 | + orchestrator = partial( |
| 510 | + _shared_future_cancellation_orchestrator, direct_future=direct_future, gather_wait=gather_wait) |
| 511 | + start = _start() |
| 512 | + initial = _execute(orchestrator, [], start) |
| 513 | + assert sorted(a.scheduleTask.name for a in initial.actions[:-1]) == ["cancelled-0", "cancelled-1"] |
| 514 | + assert initial.actions[-1].scheduleTask.name == "join" |
| 515 | + recorded = start + [ |
| 516 | + helpers.new_task_scheduled_event(action.id, action.scheduleTask.name) |
| 517 | + for action in initial.actions |
| 518 | + ] |
| 519 | + # Activity workers may complete in any order. Each dispatch reconstructs |
| 520 | + # all local tasks and futures from history, as on a different worker. |
| 521 | + completed = [ |
| 522 | + helpers.new_task_completed_event(initial.actions[1].id, "null"), |
| 523 | + helpers.new_task_completed_event(initial.actions[0].id, "null"), |
| 524 | + ] |
| 525 | + for _ in range(30): |
| 526 | + assert _execute(orchestrator, [], start).actions == initial.actions |
| 527 | + assert _execute(orchestrator, recorded, completed).actions == [] |
| 528 | + complete = _completion(_execute( |
| 529 | + orchestrator, recorded + completed, |
| 530 | + [helpers.new_task_completed_event(initial.actions[-1].id, '"joined"')])) |
| 531 | + assert complete.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED, complete.failureDetails |
| 532 | + assert json.loads(complete.result.value) == "joined" |
| 533 | + |
| 534 | + |
| 535 | +def test_shared_future_cancellation_history_replays_in_other_worker_processes() -> None: |
| 536 | + orchestrator = partial( |
| 537 | + _shared_future_cancellation_orchestrator, direct_future=True, gather_wait=False) |
| 538 | + start = _start() |
| 539 | + initial = _execute(orchestrator, [], start) |
| 540 | + history = start + [ |
| 541 | + helpers.new_task_scheduled_event(action.id, action.scheduleTask.name) |
| 542 | + for action in initial.actions |
| 543 | + ] |
| 544 | + history += [ |
| 545 | + helpers.new_task_completed_event(initial.actions[1].id, "null"), |
| 546 | + helpers.new_task_completed_event(initial.actions[0].id, "null"), |
| 547 | + ] |
| 548 | + payload = json.dumps({ |
| 549 | + "initial": [action.SerializeToString().hex() for action in initial.actions], |
| 550 | + "history": [event.SerializeToString().hex() for event in history], |
| 551 | + "join_id": initial.actions[-1].id, |
| 552 | + }) |
| 553 | + code = """ |
| 554 | +import json |
| 555 | +import sys |
| 556 | +from functools import partial |
| 557 | +from durabletask.internal import helpers, orchestrator_service_pb2 as pb |
| 558 | +from tests.durabletask.test_asyncio_orchestrators import _shared_future_cancellation_orchestrator |
| 559 | +from tests.durabletask.test_coroutine_orchestrators import _completion, _execute, _start |
| 560 | +
|
| 561 | +data = json.load(sys.stdin) |
| 562 | +orchestrator = partial(_shared_future_cancellation_orchestrator, direct_future=True, gather_wait=False) |
| 563 | +initial = _execute(orchestrator, [], _start()) |
| 564 | +assert [action.SerializeToString().hex() for action in initial.actions] == data["initial"] |
| 565 | +history = [pb.HistoryEvent.FromString(bytes.fromhex(event)) for event in data["history"]] |
| 566 | +result = _completion(_execute( |
| 567 | + orchestrator, history, [helpers.new_task_completed_event(data["join_id"], '"joined"')])) |
| 568 | +assert result.orchestrationStatus == pb.ORCHESTRATION_STATUS_COMPLETED, result.failureDetails |
| 569 | +assert json.loads(result.result.value) == "joined" |
| 570 | +""" |
| 571 | + for seed in ("0", "999"): |
| 572 | + process = subprocess.run( |
| 573 | + [sys.executable, "-c", code], input=payload, capture_output=True, text=True, |
| 574 | + env={**os.environ, "PYTHONHASHSEED": seed}, timeout=30) |
| 575 | + assert process.returncode == 0, process.stdout + process.stderr |
| 576 | + |
| 577 | + |
466 | 578 | def test_cancelling_last_loop_timer_cannot_hide_foreign_wait() -> None: |
467 | 579 | async def orchestrator(ctx: task.OrchestrationContext, _: Any) -> None: |
468 | 580 | handle = asyncio.get_running_loop().call_later(1, lambda: None) |
|
0 commit comments