| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176 |
- from __future__ import annotations
- import asyncio
- import json
- import threading
- from collections.abc import Callable
- from typing import Any
- import pytest
- import agent_lab.application.events as events_module
- from agent_lab.application.events import (
- EventDefinition,
- EventKernel,
- EventRegistry,
- EventRequest,
- EventSource,
- EventStatus,
- ResultPolicy,
- )
- from agent_lab.application.tools import ToolDefinition, ToolRegistry
- from agent_lab.domain.events import ToolCallEvent
- def _executor_type():
- executor_type = getattr(events_module, "EventBatchExecutor", None)
- assert executor_type is not None, "EventBatchExecutor is not implemented"
- return executor_type
- def _definition(
- name: str,
- handler: Callable[[EventRequest], Any],
- *,
- conflict_keys: tuple[str, ...] = (),
- timeout_seconds: float | None = None,
- terminal: bool = False,
- normalizer: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
- ) -> EventDefinition:
- return EventDefinition(
- name=name,
- description=f"Execute {name}.",
- parameters={"type": "object"},
- handler=handler,
- conflict_keys=conflict_keys,
- timeout_seconds=timeout_seconds,
- terminal=terminal,
- normalizer=normalizer,
- )
- def _request(event_id: str, name: str, arguments: dict[str, Any] | None = None):
- return EventRequest(
- id=event_id,
- name=name,
- arguments=arguments or {},
- )
- def _executor(
- definitions: list[EventDefinition],
- *,
- max_parallel_events: int = 4,
- batch_timeout_seconds: float = 1.0,
- **options: Any,
- ):
- return _executor_type()(
- EventKernel(EventRegistry(definitions)),
- max_parallel_events=max_parallel_events,
- batch_timeout_seconds=batch_timeout_seconds,
- **options,
- )
- @pytest.mark.asyncio
- async def test_batch_executor_bounds_global_parallelism():
- active = 0
- max_active = 0
- async def handler(request: EventRequest) -> dict[str, Any]:
- nonlocal active, max_active
- active += 1
- max_active = max(max_active, active)
- await asyncio.sleep(0.02)
- active -= 1
- return {"event_id": request.id}
- executor = _executor(
- [_definition("example.work", handler)],
- max_parallel_events=2,
- )
- batch = await executor.execute(
- [
- _request(str(index), "example.work", {"index": index})
- for index in range(4)
- ],
- scope="session-1:turn-1",
- )
- assert max_active == 2
- assert [result.status for result in batch.results] == [
- EventStatus.SUCCESS,
- ] * 4
- @pytest.mark.asyncio
- async def test_batch_executor_preserves_input_order_when_completion_order_differs():
- release_first = asyncio.Event()
- second_finished = asyncio.Event()
- async def handler(request: EventRequest) -> dict[str, Any]:
- if request.id == "first":
- await release_first.wait()
- else:
- second_finished.set()
- return {"event_id": request.id}
- executor = _executor([_definition("example.work", handler)])
- task = asyncio.create_task(
- executor.execute(
- [
- _request("first", "example.work", {"order": 1}),
- _request("second", "example.work", {"order": 2}),
- ],
- scope="session-1:turn-1",
- )
- )
- await asyncio.wait_for(second_finished.wait(), timeout=0.2)
- release_first.set()
- batch = await asyncio.wait_for(task, timeout=0.2)
- assert [result.event_id for result in batch.results] == ["first", "second"]
- @pytest.mark.asyncio
- async def test_batch_executor_serializes_shared_conflict_keys():
- active = 0
- max_active = 0
- async def handler(request: EventRequest) -> dict[str, Any]:
- nonlocal active, max_active
- active += 1
- max_active = max(max_active, active)
- await asyncio.sleep(0.02)
- active -= 1
- return {"event_id": request.id}
- executor = _executor(
- [
- _definition(
- "example.write",
- handler,
- conflict_keys=("shared-resource",),
- )
- ]
- )
- await executor.execute(
- [
- _request("first", "example.write", {"order": 1}),
- _request("second", "example.write", {"order": 2}),
- ],
- scope="session-1:turn-1",
- )
- assert max_active == 1
- @pytest.mark.asyncio
- async def test_batch_executor_isolates_per_event_timeout_from_successful_sibling():
- async def slow_handler(request: EventRequest) -> dict[str, Any]:
- await asyncio.sleep(1)
- return {"event_id": request.id}
- executor = _executor(
- [
- _definition(
- "example.slow",
- slow_handler,
- timeout_seconds=0.01,
- ),
- _definition(
- "example.fast",
- lambda request: {"event_id": request.id},
- ),
- ],
- batch_timeout_seconds=0.2,
- )
- batch = await executor.execute(
- [
- _request("slow", "example.slow"),
- _request("fast", "example.fast"),
- ],
- scope="session-1:turn-1",
- )
- assert [result.status for result in batch.results] == [
- EventStatus.TIMEOUT,
- EventStatus.SUCCESS,
- ]
- assert batch.deadline_exceeded is False
- @pytest.mark.asyncio
- async def test_batch_executor_applies_overall_deadline_to_waiting_work():
- async def blocked_handler(request: EventRequest) -> dict[str, Any]:
- await asyncio.Event().wait()
- return {"event_id": request.id}
- executor = _executor(
- [_definition("example.blocked", blocked_handler)],
- max_parallel_events=1,
- batch_timeout_seconds=0.02,
- )
- batch = await executor.execute(
- [
- _request("first", "example.blocked"),
- _request("second", "example.blocked"),
- ],
- scope="session-1:turn-1",
- )
- assert [result.status for result in batch.results] == [
- EventStatus.TIMEOUT,
- EventStatus.TIMEOUT,
- ]
- assert batch.deadline_exceeded is True
- @pytest.mark.asyncio
- async def test_batch_executor_coalesces_exact_duplicates_inside_one_batch():
- calls = 0
- def handler(request: EventRequest) -> dict[str, Any]:
- nonlocal calls
- calls += 1
- return {"event_id": request.id, "arguments": request.arguments}
- executor = _executor([_definition("example.write", handler)])
- batch = await executor.execute(
- [
- _request("same", "example.write", {"a": 1, "b": 2}),
- _request("same", "example.write", {"b": 2, "a": 1}),
- ],
- scope="session-1:turn-1",
- )
- assert calls == 1
- assert batch.coalesced_count == 1
- assert batch.results[0].event_id == batch.results[1].event_id == "same"
- assert batch.results[0].payload == batch.results[1].payload
- assert batch.results[0].deduplicated_from is None
- assert batch.results[1].deduplicated_from == "same"
- @pytest.mark.asyncio
- async def test_batch_executor_replay_key_includes_scope_event_id_and_request():
- calls = 0
- def handler(request: EventRequest) -> dict[str, Any]:
- nonlocal calls
- calls += 1
- return {"call": calls, "arguments": request.arguments}
- executor = _executor([_definition("example.write", handler)])
- original = _request("same", "example.write", {"a": 1, "b": 2})
- canonical_replay = _request(
- "same",
- "example.write",
- {"b": 2, "a": 1},
- )
- first = await executor.execute([original], scope="session-1:turn-1")
- replay = await executor.execute(
- [canonical_replay],
- scope="session-1:turn-1",
- )
- changed_request = await executor.execute(
- [_request("same", "example.write", {"a": 2})],
- scope="session-1:turn-1",
- )
- changed_id = await executor.execute(
- [_request("other", "example.write", {"a": 1, "b": 2})],
- scope="session-1:turn-1",
- )
- changed_scope = await executor.execute(
- [original],
- scope="session-1:turn-2",
- )
- assert first.results[0] == replay.results[0]
- assert replay.replayed_count == 1
- assert changed_request.replayed_count == 0
- assert changed_id.replayed_count == 0
- assert changed_scope.replayed_count == 0
- assert calls == 4
- @pytest.mark.asyncio
- async def test_batch_executor_runs_terminal_events_after_all_nonterminal_results():
- observed: list[str] = []
- async def nonterminal(request: EventRequest) -> dict[str, Any]:
- observed.append(f"start:{request.id}")
- await asyncio.sleep(0.01)
- observed.append(f"finish:{request.id}")
- return {"event_id": request.id}
- def terminal(request: EventRequest) -> dict[str, Any]:
- observed.append(f"terminal:{request.id}")
- return {"event_id": request.id}
- executor = _executor(
- [
- _definition("example.work", nonterminal),
- _definition("example.terminate", terminal, terminal=True),
- ]
- )
- batch = await executor.execute(
- [
- _request("terminate", "example.terminate"),
- _request("first", "example.work", {"order": 1}),
- _request("second", "example.work", {"order": 2}),
- ],
- scope="session-1:turn-1",
- )
- terminal_index = observed.index("terminal:terminate")
- assert observed.index("finish:first") < terminal_index
- assert observed.index("finish:second") < terminal_index
- assert [result.event_id for result in batch.results] == [
- "terminate",
- "first",
- "second",
- ]
- @pytest.mark.asyncio
- async def test_batch_executor_does_not_swallow_external_cancellation():
- started = asyncio.Event()
- async def handler(request: EventRequest) -> dict[str, Any]:
- started.set()
- await asyncio.Event().wait()
- return {"event_id": request.id}
- executor = _executor([_definition("example.blocked", handler)])
- task = asyncio.create_task(
- executor.execute(
- [_request("blocked", "example.blocked")],
- scope="session-1:turn-1",
- )
- )
- await asyncio.wait_for(started.wait(), timeout=0.2)
- task.cancel()
- with pytest.raises(asyncio.CancelledError):
- await task
- @pytest.mark.asyncio
- async def test_batch_executor_coalesces_normalized_provider_arguments_across_ids():
- calls: list[EventRequest] = []
- def normalize(arguments: dict[str, Any]) -> dict[str, Any]:
- value = arguments["value"]
- return {"value": int(value) if isinstance(value, float) else value}
- executor = _executor(
- [
- _definition(
- "example.normalized",
- lambda request: calls.append(request) or {"value": request.arguments["value"]},
- normalizer=normalize,
- )
- ]
- )
- requests = [
- EventRequest(
- id="integer-id",
- name="example.normalized",
- arguments={"value": 30},
- source=EventSource.PROVIDER_RESOLVED,
- raw_arguments='{"value":30}',
- ),
- EventRequest(
- id="float-id",
- name="example.normalized",
- arguments={"value": 30.0},
- source=EventSource.PROVIDER_RESOLVED,
- raw_arguments='{"value":30.0}',
- ),
- ]
- batch = await executor.execute(requests, scope="normalized-scope")
- assert len(calls) == 1
- assert batch.coalesced_count == 1
- assert [result.event_id for result in batch.results] == [
- "integer-id",
- "float-id",
- ]
- assert batch.results[0].payload == {"value": 30}
- assert batch.results[1].payload == {"value": 30}
- assert getattr(batch.results[0], "deduplicated_from", None) is None
- assert getattr(batch.results[1], "deduplicated_from", None) == "integer-id"
- @pytest.mark.asyncio
- async def test_coalesced_clone_rewrites_only_correlated_payload_event_id():
- correlated = _executor(
- [
- _definition(
- "example.correlated",
- lambda request: {"event_id": request.id, "value": "same"},
- )
- ]
- )
- correlated_batch = await correlated.execute(
- [
- EventRequest(
- id="primary",
- name="example.correlated",
- arguments={},
- raw_arguments='{"call":"primary"}',
- ),
- EventRequest(
- id="duplicate",
- name="example.correlated",
- arguments={},
- raw_arguments='{"call":"duplicate"}',
- ),
- ],
- scope="correlated-payload",
- )
- primary, duplicate = correlated_batch.results
- assert primary.event_id == "primary"
- assert duplicate.event_id == "duplicate"
- assert duplicate.raw_arguments == '{"call":"duplicate"}'
- assert duplicate.payload == {"event_id": "duplicate", "value": "same"}
- assert getattr(duplicate, "deduplicated_from", None) == "primary"
- domain = _executor(
- [
- _definition(
- "example.domain",
- lambda request: {"event_id": "domain-object"},
- )
- ]
- )
- domain_batch = await domain.execute(
- [
- _request("primary", "example.domain"),
- _request("duplicate", "example.domain"),
- ],
- scope="domain-payload",
- )
- assert domain_batch.results[1].payload == {"event_id": "domain-object"}
- assert getattr(domain_batch.results[1], "deduplicated_from", None) == "primary"
- @pytest.mark.asyncio
- async def test_batch_executor_tool_replies_keep_each_coalesced_call_id():
- registry = ToolRegistry(
- [
- ToolDefinition(
- name="example.tool",
- description="Execute a coalesced tool.",
- parameters={"type": "object"},
- handler=lambda event: {"event_id": event.id},
- )
- ]
- )
- executor = _executor_type()(
- EventKernel(registry.event_registry),
- batch_timeout_seconds=1,
- )
- events = [
- ToolCallEvent(id=event_id, name="example.tool", arguments={}, raw_arguments="{}")
- for event_id in ("first-id", "second-id")
- ]
- batch = await executor.execute(
- [
- registry.event_request(event, source=EventSource.PROVIDER_RESOLVED)
- for event in events
- ],
- scope="tool-scope",
- )
- replies = registry.tool_replies(events, batch)
- assert [reply.tool_call_id for reply in replies] == ["first-id", "second-id"]
- assert [json.loads(reply.content) for reply in replies] == [
- {"event_id": "first-id"},
- {"event_id": "second-id"},
- ]
- assert [result.event_id for result in batch.results] == [
- "first-id",
- "second-id",
- ]
- @pytest.mark.asyncio
- async def test_batch_policy_and_compact_summary_skip_logical_duplicates():
- registry = ToolRegistry(
- [
- ToolDefinition(
- name="example.template",
- description="Render one template.",
- parameters={"type": "object"},
- handler=lambda event: {"event_id": event.id},
- result_policy=ResultPolicy.TEMPLATE_FOLLOW_UP,
- result_message_factory=lambda result: "Completed once.",
- )
- ]
- )
- executor = _executor_type()(EventKernel(registry.event_registry))
- events = [
- ToolCallEvent(
- id=event_id,
- name="example.template",
- arguments={},
- raw_arguments="{}",
- )
- for event_id in ("primary", "duplicate")
- ]
- batch = await executor.execute(
- [registry.event_request(event) for event in events],
- scope="policy-dedupe",
- )
- decision = registry.batch_decision(batch)
- summary = registry.compact_results_message(batch)
- assert decision.template_messages == ("Completed once.",)
- assert summary is not None
- assert len(summary.content.splitlines()) == 2
- @pytest.mark.asyncio
- async def test_batch_key_falls_back_to_original_arguments_when_normalizer_fails():
- calls = 0
- def normalize(arguments: dict[str, Any]) -> dict[str, Any]:
- if isinstance(arguments.get("value"), str):
- raise ValueError("unsupported value")
- return {"value": int(arguments["value"])}
- def handler(request: EventRequest) -> dict[str, Any]:
- nonlocal calls
- calls += 1
- return {"value": request.arguments["value"]}
- executor = _executor(
- [_definition("example.normalized", handler, normalizer=normalize)]
- )
- batch = await executor.execute(
- [
- EventRequest(
- id="valid",
- name="example.normalized",
- arguments={"value": 30},
- source=EventSource.PROVIDER_RESOLVED,
- ),
- EventRequest(
- id="invalid",
- name="example.normalized",
- arguments={"value": "30"},
- source=EventSource.PROVIDER_RESOLVED,
- ),
- ],
- scope="normalizer-failure",
- )
- assert batch.coalesced_count == 0
- assert [result.status for result in batch.results] == [
- EventStatus.SUCCESS,
- EventStatus.INVALID_ARGUMENTS,
- ]
- assert calls == 1
- @pytest.mark.asyncio
- @pytest.mark.parametrize(
- ("invalid_arguments", "raw_arguments"),
- [
- ({"value": float("nan")}, '{"value":NaN}'),
- ({"value": object()}, '{"value":"object"}'),
- ],
- )
- async def test_invalid_json_arguments_do_not_coalesce_with_valid_empty_object(
- invalid_arguments: dict[str, Any],
- raw_arguments: str,
- ):
- calls = 0
- def handler(request: EventRequest) -> dict[str, Any]:
- nonlocal calls
- calls += 1
- return {"event_id": request.id}
- executor = _executor([_definition("example.strict", handler)])
- batch = await executor.execute(
- [
- EventRequest(
- id="valid",
- name="example.strict",
- arguments={},
- source=EventSource.PROVIDER_RESOLVED,
- raw_arguments="{}",
- ),
- EventRequest(
- id="invalid",
- name="example.strict",
- arguments=invalid_arguments,
- source=EventSource.PROVIDER_RESOLVED,
- raw_arguments=raw_arguments,
- ),
- ],
- scope="strict-json-batch",
- )
- assert batch.coalesced_count == 0
- assert [result.status for result in batch.results] == [
- EventStatus.SUCCESS,
- EventStatus.INVALID_ARGUMENTS,
- ]
- assert calls == 1
- @pytest.mark.asyncio
- @pytest.mark.parametrize(
- ("invalid_arguments", "raw_arguments"),
- [
- ({"value": float("nan")}, '{"value":NaN}'),
- ({"value": object()}, '{"value":"object"}'),
- ],
- )
- async def test_invalid_json_arguments_do_not_replay_valid_success(
- invalid_arguments: dict[str, Any],
- raw_arguments: str,
- ):
- executor = _executor(
- [_definition("example.strict", lambda request: {"event_id": request.id})]
- )
- valid = EventRequest(
- id="same-id",
- name="example.strict",
- arguments={},
- source=EventSource.PROVIDER_RESOLVED,
- raw_arguments="{}",
- )
- invalid = EventRequest(
- id="same-id",
- name="example.strict",
- arguments=invalid_arguments,
- source=EventSource.PROVIDER_RESOLVED,
- raw_arguments=raw_arguments,
- )
- first = await executor.execute([valid], scope="strict-json-replay")
- second = await executor.execute([invalid], scope="strict-json-replay")
- assert first.results[0].status is EventStatus.SUCCESS
- assert second.replayed_count == 0
- assert second.results[0].status is EventStatus.INVALID_ARGUMENTS
- @pytest.mark.asyncio
- async def test_non_json_normalizer_output_has_distinct_canonical_key():
- calls = 0
- def normalize(arguments: dict[str, Any]) -> dict[str, Any]:
- if arguments.get("invalid"):
- return {"value": object()}
- return {}
- def handler(request: EventRequest) -> dict[str, Any]:
- nonlocal calls
- calls += 1
- return {"event_id": request.id}
- executor = _executor(
- [_definition("example.normalized", handler, normalizer=normalize)]
- )
- batch = await executor.execute(
- [
- EventRequest(
- id="valid",
- name="example.normalized",
- arguments={},
- source=EventSource.PROVIDER_RESOLVED,
- ),
- EventRequest(
- id="invalid",
- name="example.normalized",
- arguments={"invalid": True},
- source=EventSource.PROVIDER_RESOLVED,
- ),
- ],
- scope="normalizer-invalid-json",
- )
- assert batch.coalesced_count == 0
- assert [result.status for result in batch.results] == [
- EventStatus.SUCCESS,
- EventStatus.INVALID_ARGUMENTS,
- ]
- assert calls == 1
- @pytest.mark.asyncio
- async def test_concurrent_replay_waiters_share_execution_and_cancel_independently():
- calls = 0
- started = asyncio.Event()
- release = asyncio.Event()
- async def handler(request: EventRequest) -> dict[str, Any]:
- nonlocal calls
- calls += 1
- started.set()
- await release.wait()
- return {"event_id": request.id}
- executor = _executor([_definition("example.shared", handler)])
- request = _request("same-id", "example.shared")
- first = asyncio.create_task(executor.execute([request], scope="shared-scope"))
- second = asyncio.create_task(executor.execute([request], scope="shared-scope"))
- await asyncio.wait_for(started.wait(), timeout=0.2)
- first.cancel()
- with pytest.raises(asyncio.CancelledError):
- await first
- release.set()
- second_batch = await asyncio.wait_for(second, timeout=0.2)
- assert calls == 1
- assert second_batch.results[0].status is EventStatus.SUCCESS
- @pytest.mark.asyncio
- async def test_release_scope_removes_replay_entries_for_reuse():
- calls = 0
- def handler(request: EventRequest) -> dict[str, Any]:
- nonlocal calls
- calls += 1
- return {"call": calls}
- executor = _executor([_definition("example.cached", handler)])
- request = _request("same-id", "example.cached")
- await executor.execute([request], scope="released-scope")
- executor.release_scope("released-scope")
- batch = await executor.execute([request], scope="released-scope")
- assert calls == 2
- assert batch.replayed_count == 0
- @pytest.mark.asyncio
- async def test_release_scope_starts_new_generation_before_old_task_finishes():
- calls = 0
- started = [asyncio.Event(), asyncio.Event()]
- releases = [asyncio.Event(), asyncio.Event()]
- async def handler(request: EventRequest) -> dict[str, Any]:
- nonlocal calls
- call_index = calls
- calls += 1
- started[call_index].set()
- await releases[call_index].wait()
- return {"call": call_index + 1, "event_id": request.id}
- executor = _executor([_definition("example.generated", handler)])
- request = _request("same-id", "example.generated")
- first = asyncio.create_task(
- executor.execute([request], scope="generation-scope")
- )
- await asyncio.wait_for(started[0].wait(), timeout=0.2)
- executor.release_scope("generation-scope")
- second = asyncio.create_task(
- executor.execute([request], scope="generation-scope")
- )
- await asyncio.sleep(0.02)
- new_generation_started = started[1].is_set()
- releases[0].set()
- releases[1].set()
- first_batch, second_batch = await asyncio.gather(first, second)
- assert new_generation_started is True
- assert first_batch.results[0].payload["call"] == 1
- assert second_batch.results[0].payload["call"] == 2
- @pytest.mark.asyncio
- async def test_old_generation_completion_cannot_replace_new_inflight_or_cache():
- calls = 0
- started = [asyncio.Event(), asyncio.Event(), asyncio.Event()]
- releases = [asyncio.Event(), asyncio.Event(), asyncio.Event()]
- async def handler(request: EventRequest) -> dict[str, Any]:
- nonlocal calls
- call_index = calls
- calls += 1
- started[call_index].set()
- await releases[call_index].wait()
- return {"call": call_index + 1, "event_id": request.id}
- executor = _executor([_definition("example.generated", handler)])
- request = _request("same-id", "example.generated")
- old = asyncio.create_task(executor.execute([request], scope="reused-scope"))
- await asyncio.wait_for(started[0].wait(), timeout=0.2)
- executor.release_scope("reused-scope")
- current = asyncio.create_task(
- executor.execute([request], scope="reused-scope")
- )
- await asyncio.sleep(0.02)
- new_generation_started = started[1].is_set()
- if not new_generation_started:
- releases[0].set()
- await asyncio.gather(old, current)
- assert new_generation_started is True
- releases[0].set()
- old_batch = await asyncio.wait_for(old, timeout=0.2)
- waiter = asyncio.create_task(
- executor.execute([request], scope="reused-scope")
- )
- await asyncio.sleep(0.02)
- assert calls == 2
- releases[1].set()
- current_batch, waiter_batch = await asyncio.gather(current, waiter)
- replayed = await executor.execute([request], scope="reused-scope")
- assert old_batch.results[0].payload["call"] == 1
- assert current_batch.results[0].payload["call"] == 2
- assert waiter_batch.results[0].payload["call"] == 2
- assert replayed.replayed_count == 1
- assert replayed.results[0].payload["call"] == 2
- @pytest.mark.asyncio
- async def test_replay_cache_stays_bounded_across_250_scopes():
- executor = _executor(
- [_definition("example.cached", lambda request: {"id": request.id})],
- max_replay_entries=100,
- )
- request = _request("same-id", "example.cached")
- for index in range(250):
- await executor.execute([request], scope=f"scope-{index}")
- assert len(executor._replay_results) <= 100
- @pytest.mark.asyncio
- async def test_terminal_uses_independent_grace_after_sibling_batch_timeout():
- terminal_calls: list[str] = []
- async def blocked(request: EventRequest) -> dict[str, Any]:
- await asyncio.Event().wait()
- return {"event_id": request.id}
- def terminate(request: EventRequest) -> dict[str, Any]:
- terminal_calls.append(request.id)
- return {"terminated": True}
- executor = _executor(
- [
- _definition("example.blocked", blocked),
- _definition(
- "example.terminate",
- terminate,
- terminal=True,
- timeout_seconds=0.05,
- ),
- ],
- batch_timeout_seconds=0.01,
- terminal_grace_seconds=0.1,
- )
- batch = await executor.execute(
- [
- _request("sibling", "example.blocked"),
- _request("terminal", "example.terminate"),
- ],
- scope="terminal-grace",
- )
- assert [result.status for result in batch.results] == [
- EventStatus.TIMEOUT,
- EventStatus.SUCCESS,
- ]
- assert batch.deadline_exceeded is True
- assert terminal_calls == ["terminal"]
- @pytest.mark.asyncio
- async def test_terminal_grace_is_bounded_by_terminal_event_timeout():
- async def blocked_terminal(request: EventRequest) -> dict[str, Any]:
- await asyncio.Event().wait()
- return {"event_id": request.id}
- executor = _executor(
- [
- _definition(
- "example.terminate",
- blocked_terminal,
- terminal=True,
- timeout_seconds=0.01,
- )
- ],
- terminal_grace_seconds=0.2,
- )
- batch = await executor.execute(
- [_request("terminal", "example.terminate")],
- scope="terminal-timeout",
- )
- assert batch.results[0].status is EventStatus.TIMEOUT
- assert batch.results[0].error == "event timed out after 0.01 seconds"
- assert batch.deadline_exceeded is False
- @pytest.mark.asyncio
- async def test_replayed_terminal_grace_timeout_preserves_deadline_flag():
- async def blocked_terminal(request: EventRequest) -> dict[str, Any]:
- await asyncio.Event().wait()
- return {"event_id": request.id}
- executor = _executor(
- [_definition("example.terminate", blocked_terminal, terminal=True)],
- terminal_grace_seconds=0.01,
- )
- request = _request("terminal", "example.terminate")
- first = await executor.execute([request], scope="terminal-replay")
- replayed = await executor.execute([request], scope="terminal-replay")
- assert first.results[0].status is EventStatus.TIMEOUT
- assert first.results[0].error == "batch deadline exceeded"
- assert first.deadline_exceeded is False
- assert replayed.replayed_count == 1
- assert replayed.deadline_exceeded is False
- @pytest.mark.asyncio
- async def test_replayed_sibling_batch_timeout_preserves_deadline_flag():
- async def blocked(request: EventRequest) -> dict[str, Any]:
- await asyncio.Event().wait()
- return {"event_id": request.id}
- executor = _executor(
- [_definition("example.blocked", blocked)],
- batch_timeout_seconds=0.01,
- )
- request = _request("sibling", "example.blocked")
- first = await executor.execute([request], scope="sibling-replay")
- replayed = await executor.execute([request], scope="sibling-replay")
- assert first.results[0].status is EventStatus.TIMEOUT
- assert first.deadline_exceeded is True
- assert replayed.replayed_count == 1
- assert replayed.deadline_exceeded is True
- @pytest.mark.asyncio
- async def test_blocking_sync_handler_timeout_returns_before_thread_finishes():
- started = threading.Event()
- release = threading.Event()
- finished = threading.Event()
- def blocking_handler(request: EventRequest) -> dict[str, Any]:
- started.set()
- release.wait(timeout=1)
- finished.set()
- return {"event_id": request.id}
- executor = _executor(
- [
- _definition(
- "example.blocking",
- blocking_handler,
- timeout_seconds=0.02,
- )
- ],
- batch_timeout_seconds=0.2,
- )
- loop = asyncio.get_running_loop()
- started_at = loop.time()
- batch = await executor.execute(
- [_request("blocking", "example.blocking")],
- scope="blocking-handler",
- )
- elapsed = loop.time() - started_at
- release.set()
- await asyncio.to_thread(finished.wait, 0.2)
- assert started.is_set()
- assert batch.results[0].status is EventStatus.TIMEOUT
- assert elapsed < 0.1
- assert finished.is_set()
- @pytest.mark.asyncio
- async def test_timed_out_sync_handler_keeps_conflict_lock_until_thread_finishes():
- worker_started = threading.Event()
- worker_release = threading.Event()
- follower_started = asyncio.Event()
- def blocking_handler(request: EventRequest) -> dict[str, Any]:
- worker_started.set()
- worker_release.wait(timeout=1)
- return {"event_id": request.id}
- async def follower_handler(request: EventRequest) -> dict[str, Any]:
- follower_started.set()
- return {"event_id": request.id}
- executor = _executor(
- [
- _definition(
- "example.blocking",
- blocking_handler,
- conflict_keys=("shared",),
- timeout_seconds=0.02,
- ),
- _definition(
- "example.follower",
- follower_handler,
- conflict_keys=("shared",),
- timeout_seconds=0.2,
- ),
- ]
- )
- timed_out = await executor.execute(
- [_request("blocking", "example.blocking")],
- scope="sync-conflict-blocking",
- )
- follower = asyncio.create_task(
- executor.execute(
- [_request("follower", "example.follower")],
- scope="sync-conflict-follower",
- )
- )
- await asyncio.sleep(0.03)
- entered_before_worker_finished = follower_started.is_set()
- worker_release.set()
- follower_batch = await asyncio.wait_for(follower, timeout=0.3)
- assert worker_started.is_set()
- assert timed_out.results[0].status is EventStatus.TIMEOUT
- assert entered_before_worker_finished is False
- assert follower_batch.results[0].status is EventStatus.SUCCESS
- @pytest.mark.asyncio
- async def test_timed_out_sync_handler_keeps_parallel_slot_until_thread_finishes():
- worker_started = threading.Event()
- worker_release = threading.Event()
- follower_started = asyncio.Event()
- def blocking_handler(request: EventRequest) -> dict[str, Any]:
- worker_started.set()
- worker_release.wait(timeout=1)
- return {"event_id": request.id}
- async def follower_handler(request: EventRequest) -> dict[str, Any]:
- follower_started.set()
- return {"event_id": request.id}
- executor = _executor(
- [
- _definition(
- "example.blocking",
- blocking_handler,
- timeout_seconds=0.02,
- ),
- _definition(
- "example.follower",
- follower_handler,
- timeout_seconds=0.2,
- ),
- ],
- max_parallel_events=1,
- )
- timed_out = await executor.execute(
- [_request("blocking", "example.blocking")],
- scope="sync-slot-blocking",
- )
- follower = asyncio.create_task(
- executor.execute(
- [_request("follower", "example.follower")],
- scope="sync-slot-follower",
- )
- )
- await asyncio.sleep(0.03)
- entered_before_worker_finished = follower_started.is_set()
- worker_release.set()
- follower_batch = await asyncio.wait_for(follower, timeout=0.3)
- assert worker_started.is_set()
- assert timed_out.results[0].status is EventStatus.TIMEOUT
- assert entered_before_worker_finished is False
- assert follower_batch.results[0].status is EventStatus.SUCCESS
- @pytest.mark.asyncio
- async def test_timed_out_async_handler_releases_resources_after_cancellation():
- cancelled = asyncio.Event()
- follower_started = asyncio.Event()
- async def blocked_handler(request: EventRequest) -> dict[str, Any]:
- try:
- await asyncio.Event().wait()
- finally:
- cancelled.set()
- return {"event_id": request.id}
- async def follower_handler(request: EventRequest) -> dict[str, Any]:
- follower_started.set()
- return {"event_id": request.id}
- executor = _executor(
- [
- _definition(
- "example.blocked",
- blocked_handler,
- conflict_keys=("shared",),
- timeout_seconds=0.01,
- ),
- _definition(
- "example.follower",
- follower_handler,
- conflict_keys=("shared",),
- timeout_seconds=0.2,
- ),
- ],
- max_parallel_events=1,
- )
- timed_out = await executor.execute(
- [_request("blocked", "example.blocked")],
- scope="async-release-blocked",
- )
- follower = await executor.execute(
- [_request("follower", "example.follower")],
- scope="async-release-follower",
- )
- assert timed_out.results[0].status is EventStatus.TIMEOUT
- assert cancelled.is_set()
- assert follower_started.is_set()
- assert follower.results[0].status is EventStatus.SUCCESS
|