diff --git a/src/firefly_weave/connectors/dispatcher.py b/src/firefly_weave/connectors/dispatcher.py index 160e4af8..93b7cff0 100644 --- a/src/firefly_weave/connectors/dispatcher.py +++ b/src/firefly_weave/connectors/dispatcher.py @@ -34,6 +34,7 @@ from firefly_weave.contracts.access import Scope from firefly_weave.contracts.definitions import ContractModel from firefly_weave.contracts.workers import InstanceRequest +from firefly_weave.operations.execution import ReservedSlots from firefly_weave.sdk.worker import Worker from firefly_weave.workers.repository import WorkerRepository @@ -118,6 +119,9 @@ def __init__( image_digest, shutdown_seconds, ) + # One slot per execution the workers can run at once, plus a spare for a pure call that + # outlived its cancelled execution. Connector calls never hold a request work slot. + self.reservation = ReservedSlots(sum(entry.capacity for entry in entries) + 1) self.workers: list[Worker] = [] self.tasks: list[asyncio.Task[None]] = [] self.closing = False @@ -172,7 +176,7 @@ async def open(self) -> None: ) async def execute(lease: Any, scope: Scope = entry.scope, principal: UUID = entry.principal_id) -> Any: - return await self.service.execute(scope, principal, lease) + return await self.service.execute(scope, principal, lease, reservation=self.reservation) transport = ServiceTransport(self.service, entry.scope, entry.principal_id, instance.id) self.workers.append(Worker(transport, {ref: execute for ref in entry.task_types}, entry.capacity)) diff --git a/src/firefly_weave/connectors/execution.py b/src/firefly_weave/connectors/execution.py index 4a134763..10af5402 100644 --- a/src/firefly_weave/connectors/execution.py +++ b/src/firefly_weave/connectors/execution.py @@ -38,7 +38,7 @@ TaskLease, ) from firefly_weave.email.service import EmailService -from firefly_weave.operations.execution import request_execution +from firefly_weave.operations.execution import ReservedSlots, request_execution from firefly_weave.persistence.uow import UnitOfWork from firefly_weave.runtime.repository import RuntimeRepository from firefly_weave.workers.leases import TaskService @@ -67,9 +67,13 @@ async def operation(self, name: str, scope: Scope, principal_id: UUID, *args: An actor = await self.access.load_principal(principal_id, tx=tx) return await getattr(self.tasks, name)(tx, *args, actor=actor, scope=scope, context=AuditContext()) - async def execute(self, scope: Scope, principal_id: UUID, lease: TaskLease) -> JsonValue: + async def execute( + self, scope: Scope, principal_id: UUID, lease: TaskLease, *, reservation: ReservedSlots | None = None + ) -> JsonValue: self.registry.require_operational() - async with request_execution(): + # A request lease covers the whole connector call. In-process native workers bring their + # own reservation instead, whose slots only the call's pure work takes. + async with request_execution() if reservation is None else reservation.execution(): return await self._execute(scope, principal_id, lease) async def _execute(self, scope: Scope, principal_id: UUID, lease: TaskLease) -> JsonValue: diff --git a/src/firefly_weave/operations/execution.py b/src/firefly_weave/operations/execution.py index 2cf8fe87..15d29d0c 100644 --- a/src/firefly_weave/operations/execution.py +++ b/src/firefly_weave/operations/execution.py @@ -70,6 +70,7 @@ def release(self) -> None: _request_lease: ContextVar[_Lease | None] = ContextVar("weave_request_execution", default=None) _debug_lease: ContextVar[_Lease | None] = ContextVar("weave_debug_execution", default=None) +_reserved_slots: ContextVar[threading.BoundedSemaphore | None] = ContextVar("weave_reserved_execution", default=None) @asynccontextmanager @@ -95,6 +96,31 @@ async def guarded(*args: P.args, **kwargs: P.kwargs) -> R: return guarded +class ReservedSlots: + """Pure execution reserved for one lifespan owner, apart from request admission. + + In-flight requests can hold every work and control slot. Background work that borrowed them + was refused and read the refusal as a failed recovery quantum, a blocked receipt, or a failed + consumer turn. Inside `execution()`, each pure call takes one slot from this reservation + instead, and any request lease the task inherited is set aside. A call still running after its + caller was cancelled keeps its slot until it ends, so the size bounds the owner's pure threads; + a spare slot keeps one such call from refusing the owner's next one. + """ + + def __init__(self, slots: int) -> None: + self.size = slots + self.slots = threading.BoundedSemaphore(slots) + + @asynccontextmanager + async def execution(self) -> AsyncIterator[None]: + request, reserved = _request_lease.set(None), _reserved_slots.set(self.slots) + try: + yield + finally: + _reserved_slots.reset(reserved) + _request_lease.reset(request) + + @asynccontextmanager async def request_execution(*, control: bool = False) -> AsyncIterator[None]: lease = _Lease(control) @@ -125,7 +151,7 @@ async def inventory_execution() -> AsyncIterator[None]: async def execute_pure[Result](call: Callable[[], Result], *, control: bool = False) -> Result: inherited = _request_lease.get() - lease = inherited or _Lease(control) + lease = inherited or _Lease(control, slots=_reserved_slots.get()) try: lease.enter() debug = _debug_lease.get() diff --git a/src/firefly_weave/providers/loop.py b/src/firefly_weave/providers/loop.py index 8de45cb0..0a8e8f2c 100644 --- a/src/firefly_weave/providers/loop.py +++ b/src/firefly_weave/providers/loop.py @@ -25,11 +25,16 @@ from firefly_weave.access.scheduler import next_scope from firefly_weave.email.source import EmailSourceService +from firefly_weave.operations.execution import ReservedSlots from firefly_weave.persistence.migrations import check_schema from firefly_weave.persistence.uow import UnitOfWork from firefly_weave.providers.dispatcher import ProviderDispatcher from firefly_weave.settings import Settings +# Requests can hold every work and control slot. A provider cycle runs one pure call at a time; +# the spare slot absorbs a call that outlived its timed-out receipt, as the work pool did. +PROVIDER_EXECUTION = ReservedSlots(2) + class ProviderLoop: def __init__( @@ -75,18 +80,19 @@ async def open(self) -> None: async def cycle(self) -> None: assert self.engine is not None - async with asyncio.timeout(5): - async with self.engine.begin() as connection: - await connection.execute(text("SET LOCAL statement_timeout='4000ms'")) - await connection.execute(text("SET LOCAL lock_timeout='1000ms'")) - tenant = await connection.scalar(text("SELECT weave_provider_tenants(1)")) - if tenant is None: - return - authority = await next_scope(self.uow, tenant, provider=True) - if authority is not None: - await self.dispatcher.scan(authority, 10) - if self.email_sources is not None: - await self.email_sources.scan(authority, 1) + async with PROVIDER_EXECUTION.execution(): + async with asyncio.timeout(5): + async with self.engine.begin() as connection: + await connection.execute(text("SET LOCAL statement_timeout='4000ms'")) + await connection.execute(text("SET LOCAL lock_timeout='1000ms'")) + tenant = await connection.scalar(text("SELECT weave_provider_tenants(1)")) + if tenant is None: + return + authority = await next_scope(self.uow, tenant, provider=True) + if authority is not None: + await self.dispatcher.scan(authority, 10) + if self.email_sources is not None: + await self.email_sources.scan(authority, 1) async def poll(self) -> None: while True: diff --git a/src/firefly_weave/runtime/scheduler.py b/src/firefly_weave/runtime/scheduler.py index 8808339a..26b8eeb7 100644 --- a/src/firefly_weave/runtime/scheduler.py +++ b/src/firefly_weave/runtime/scheduler.py @@ -26,12 +26,25 @@ from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine from firefly_weave.access.scheduler import next_scope, tenant_page +from firefly_weave.definitions.models import CatalogError +from firefly_weave.operations.execution import ReservedSlots from firefly_weave.persistence.migrations import check_schema from firefly_weave.persistence.uow import UnitOfWork from firefly_weave.runtime.recovery import RecoveryService from firefly_weave.settings import Settings from firefly_weave.triggers.scheduler import Scheduler +# Requests can hold every work and control slot. A recovery cycle runs one pure call at a time; +# the spare slot absorbs a call that outlived its cancelled quantum, as the work pool did. +RECOVERY_EXECUTION = ReservedSlots(2) + + +def failure_cause(error: Exception) -> str: + """Name a failed scan by error class and catalog code, never by message or values.""" + if isinstance(error, CatalogError): + return f"{type(error).__name__} {error.code}" + return type(error).__name__ + class RecoveryLoop: def __init__( @@ -86,39 +99,46 @@ async def open(self) -> None: async def cycle(self) -> None: assert self.engine is not None - budget = min(self.settings.database_timeout_seconds, 5.0) - async with asyncio.timeout(budget): - tenants = await tenant_page(self.engine) - for tenant in tenants: - try: - async with asyncio.timeout(budget): - authority = await next_scope(self.uow, tenant) - if authority is None: + async with RECOVERY_EXECUTION.execution(): + budget = min(self.settings.database_timeout_seconds, 5.0) + async with asyncio.timeout(budget): + tenants = await tenant_page(self.engine) + for tenant in tenants: + try: + async with asyncio.timeout(budget): + authority = await next_scope(self.uow, tenant) + if authority is None: + continue + # Separate transactions/time budgets reserve recovery even when + # starts are backpressured. A tenant gets one environment quantum. + async with asyncio.timeout(budget): + if self.operational(): + await self.recovery._scan_scheduled(authority, 10) + else: + await self.recovery._scan_terminal_scheduled(authority, 10) + except Exception as error: + logging.getLogger(__name__).error( + "Tenant recovery scan failed (%s); traversal will continue", failure_cause(error) + ) continue - # Separate transactions/time budgets reserve recovery even when - # starts are backpressured. A tenant gets one environment quantum. - async with asyncio.timeout(budget): - if self.operational(): - await self.recovery._scan_scheduled(authority, 10) - else: - await self.recovery._scan_terminal_scheduled(authority, 10) - except Exception: - logging.getLogger(__name__).error("Tenant recovery scan failed; traversal will continue") - continue - if not self.operational(): - continue - try: - async with asyncio.timeout(budget): - await self.schedules._scan_scheduled(authority, 10) - except Exception: - logging.getLogger(__name__).error("Tenant schedule scan failed; traversal will continue") + if not self.operational(): + continue + try: + async with asyncio.timeout(budget): + await self.schedules._scan_scheduled(authority, 10) + except Exception as error: + logging.getLogger(__name__).error( + "Tenant schedule scan failed (%s); traversal will continue", failure_cause(error) + ) async def poll(self) -> None: while True: try: await self.cycle() - except Exception: - logging.getLogger(__name__).error("Recovery catalog scan failed; retrying on next poll") + except Exception as error: + logging.getLogger(__name__).error( + "Recovery catalog scan failed (%s); retrying on next poll", failure_cause(error) + ) await asyncio.sleep(self.settings.scheduler_poll_seconds) async def close(self) -> None: diff --git a/src/firefly_weave/triggers/kafka_loop.py b/src/firefly_weave/triggers/kafka_loop.py index 0de257cb..4e65cfc1 100644 --- a/src/firefly_weave/triggers/kafka_loop.py +++ b/src/firefly_weave/triggers/kafka_loop.py @@ -27,11 +27,16 @@ from firefly_weave.access.scheduler import next_scope from firefly_weave.connectors.broker import BrokerConnectionConfig from firefly_weave.connectors.kafka_transport import BrokerClients, BrokerOwner +from firefly_weave.operations.execution import ReservedSlots from firefly_weave.persistence.migrations import check_schema from firefly_weave.persistence.uow import UnitOfWork from firefly_weave.settings import Settings from firefly_weave.triggers.kafka import KafkaSource, KafkaTrigger +# Requests can hold every work and control slot. At most `max_clients - 1` consumer turns run +# one pure call each, and the broker policy caps `max_clients` at 8, which leaves a spare slot. +TURN_EXECUTION = ReservedSlots(8) + class KafkaLoop: def __init__(self, settings: Settings, uow: UnitOfWork, triggers: KafkaTrigger, clients: BrokerClients) -> None: @@ -101,7 +106,7 @@ async def turn(self, source: KafkaSource) -> None: current = asyncio.current_task() assert current is not None try: - async with asyncio.timeout(60): + async with asyncio.timeout(60), TURN_EXECUTION.execution(): revision, secrets = await self.triggers.credentials(source.authority) import json diff --git a/tests/unit/connectors/test_local_build.py b/tests/unit/connectors/test_local_build.py index c8552d7b..eebd98d5 100644 --- a/tests/unit/connectors/test_local_build.py +++ b/tests/unit/connectors/test_local_build.py @@ -268,6 +268,28 @@ async def test_local_development_executor_uses_local_attestation_not_the_image(a } +async def test_native_handlers_execute_through_the_dispatcher_reservation(admission): + digest = "sha256:" + "f" * 64 + service = _Service(_release(digest)) + executions = [] + + async def execute(scope, principal_id, lease, *, reservation=None): + executions.append(reservation) + return {"ok": True} + + service.execute = execute # type: ignore[attr-defined] + native = NativeDispatcher(service, (_executor(),), digest, 1) # type: ignore[arg-type] + await native.open() + try: + [worker] = native.workers + [handler] = set(worker.handlers.values()) # type: ignore[attr-defined] + assert await handler(object()) == {"ok": True} + finally: + await native.close() + # Connector calls take the native reservation, never a request work slot. + assert executions == [native.reservation] + + async def test_image_executor_keeps_the_packaged_image_check(admission): digest = "sha256:" + "b" * 64 native = NativeDispatcher(_Service(_release(digest)), (_executor("image"),), digest, 1) diff --git a/tests/unit/operations/test_background_admission.py b/tests/unit/operations/test_background_admission.py new file mode 100644 index 00000000..33b2868a --- /dev/null +++ b/tests/unit/operations/test_background_admission.py @@ -0,0 +1,292 @@ +# Copyright 2026 Firefly Software Foundation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Author: Firefly Software Foundation +# SPDX-License-Identifier: Apache-2.0 + +"""Lifespan owners run pure work from their own reservations while requests hold every slot.""" + +import asyncio +import threading +from contextlib import asynccontextmanager +from types import SimpleNamespace +from uuid import uuid4 + +import pytest + +from firefly_weave.contracts.access import Scope +from firefly_weave.definitions.models import CatalogError +from firefly_weave.operations.execution import CONTROL_SLOTS, WORK_SLOTS, execute_pure, request_execution + + +@asynccontextmanager +async def saturated_requests(): + """Every work and control slot owned by an in-flight request in its own task.""" + finish = asyncio.Event() + admitted = [] + + async def request(control, ready): + async with request_execution(control=control): + ready.set() + await finish.wait() + + tasks = [] + for control, count in ((False, WORK_SLOTS), (True, CONTROL_SLOTS)): + for _ in range(count): + ready = asyncio.Event() + admitted.append(ready) + tasks.append(asyncio.create_task(request(control, ready))) + try: + async with asyncio.timeout(2): + for ready in admitted: + await ready.wait() + with pytest.raises(CatalogError, match="capacity"): + await execute_pure(lambda: "borrowed") + yield + finally: + finish.set() + await asyncio.gather(*tasks) + + +class Blocked: + """Pure calls that stay running until released, like a call that outlived its caller.""" + + def __init__(self): + self.entered, self.release = threading.Semaphore(0), threading.Event() + + def __call__(self): + self.entered.release() + return self.release.wait(5) + + async def started(self, count=1): + for _ in range(count): + assert await asyncio.to_thread(self.entered.acquire, True, 2) + + +async def admitted_eventually(reservation): + # A released thread gives its slot back just after its call returns. + async with asyncio.timeout(2): + while True: + try: + async with reservation.execution(): + return await execute_pure(lambda: "admitted") + except CatalogError: + await asyncio.sleep(0.01) + + +async def test_reserved_slots_admit_background_work_while_requests_hold_every_slot(): + from firefly_weave.operations.execution import ReservedSlots, request_execution + + reservation = ReservedSlots(1) + async with saturated_requests(), reservation.execution(): + assert await execute_pure(lambda: "work") == "work" + assert await execute_pure(lambda: "control", control=True) == "control" + # Request admission keeps its own ceilings while background work runs. + with pytest.raises(CatalogError, match="capacity"): + async with request_execution(): + raise AssertionError("request admitted beyond the work slots") + + +async def test_reserved_slots_bound_concurrent_pure_calls_to_their_size(): + from firefly_weave.operations.execution import ReservedSlots + + reservation, blocked = ReservedSlots(2), Blocked() + async with saturated_requests(), reservation.execution(): + # Unlike one request lease, concurrent calls in one unit each take their own slot. + calls = [asyncio.create_task(execute_pure(blocked)) for _ in range(2)] + try: + await blocked.started(2) + with pytest.raises(CatalogError, match="capacity"): + await execute_pure(lambda: "beyond the reservation") + finally: + blocked.release.set() + assert await asyncio.gather(*calls) == [True, True] + assert await admitted_eventually(reservation) == "admitted" + + +async def test_a_call_that_outlived_its_cancelled_unit_keeps_its_slot(): + from firefly_weave.operations.execution import ReservedSlots + + reservation, blocked = ReservedSlots(2), Blocked() + + async def unit(): + async with reservation.execution(): + await execute_pure(blocked) + + units = [] + async with saturated_requests(): + try: + units.append(asyncio.create_task(unit())) + await blocked.started() + units[0].cancel() + # The spare slot keeps one orphaned call from refusing the owner's next unit. + async with reservation.execution(): + assert await execute_pure(lambda: "next unit") == "next unit" + units.append(asyncio.create_task(unit())) + await blocked.started() + units[1].cancel() + async with reservation.execution(): + with pytest.raises(CatalogError, match="capacity"): + await execute_pure(lambda: "a third thread") + finally: + blocked.release.set() + await asyncio.gather(*units, return_exceptions=True) + assert await admitted_eventually(reservation) == "admitted" + + +async def test_reserved_slots_replace_a_closed_inherited_request_lease(): + from firefly_weave.operations.execution import ReservedSlots, request_execution + + reservation = ReservedSlots(1) + opened, ended = asyncio.Event(), asyncio.Event() + + async def background(): + opened.set() + await ended.wait() + async with reservation.execution(): + return await execute_pure(lambda: "ran") + + # Like a loop task opened while a request was being served. + async with request_execution(control=True): + task = asyncio.create_task(background()) + await opened.wait() + ended.set() + assert await task == "ran" + + +async def test_provider_cycle_dispatches_while_requests_hold_every_slot(monkeypatch): + from firefly_weave.providers.loop import ProviderLoop + + tenant, ran = uuid4(), [] + + class Connection: + async def execute(self, statement): + return None + + async def scalar(self, statement): + return tenant + + class Engine: + @asynccontextmanager + async def begin(self): + yield Connection() + + async def next_scope(uow, identifier, *, provider): + assert identifier == tenant and provider + return SimpleNamespace(scope=Scope(tenant_id=tenant)) + + async def provider_scan(authority, limit): + # Run starts reduce on the work pool; signal settlement can use the control pool. + ran.append(await execute_pure(lambda: "provider")) + ran.append(await execute_pure(lambda: "settled", control=True)) + + async def email_scan(authority, limit): + ran.append(await execute_pure(lambda: "email")) + + monkeypatch.setattr("firefly_weave.providers.loop.next_scope", next_scope) + loop = ProviderLoop( + SimpleNamespace(), # type: ignore[arg-type] + None, # type: ignore[arg-type] + SimpleNamespace(scan=provider_scan), # type: ignore[arg-type] + SimpleNamespace(scan=email_scan), # type: ignore[arg-type] + ) + loop.engine = Engine() # type: ignore[assignment] + + async with saturated_requests(): + await loop.cycle() + + assert ran == ["provider", "settled", "email"] + + +async def test_concurrent_broker_turns_run_pure_work_while_requests_hold_every_slot(caplog): + from firefly_weave.triggers.kafka_loop import KafkaLoop + + # Both pure calls must be running at the same time to pass the barrier. + together = threading.Barrier(2, timeout=2) + ran, released = [], [] + + async def credentials(authority): + ran.append(await execute_pure(lambda: together.wait() >= 0)) + raise RuntimeError("end the turn") + + async def release(authority): + released.append(authority) + + loop = KafkaLoop( + SimpleNamespace(), # type: ignore[arg-type] + None, # type: ignore[arg-type] + SimpleNamespace(credentials=credentials, release=release), # type: ignore[arg-type] + None, # type: ignore[arg-type] + ) + sources = [SimpleNamespace(authority=object()) for _ in range(2)] + + async with saturated_requests(): + await asyncio.gather(*(loop.turn(source) for source in sources)) # type: ignore[arg-type] + + assert ran == [True, True] + # Each turn still releases its source authority, in whichever order the turns end. + assert len(released) == 2 and set(released) == {source.authority for source in sources} + + +def test_broker_turn_reservation_covers_every_turn_the_policy_allows_plus_a_spare(): + from annotated_types import Le + + from firefly_weave.connectors.broker import BrokerPolicy + from firefly_weave.triggers.kafka_loop import TURN_EXECUTION + + # The loop runs at most `max_clients - 1` turns; raising the policy cap must resize the reservation. + [ceiling] = [item.le for item in BrokerPolicy.model_fields["max_clients"].metadata if isinstance(item, Le)] + assert TURN_EXECUTION.size == (ceiling - 1) + 1 + + +async def test_native_executions_take_their_reservation_instead_of_request_work_slots(monkeypatch): + # Two workers with capacity 2 and 1: three executions at once, plus a spare slot. + from firefly_weave.connectors.dispatcher import ExecutorConfig, NativeDispatcher + from firefly_weave.connectors.execution import ConnectorExecutionService + + scope = Scope(tenant_id=uuid4(), project_id=uuid4(), environment_id=uuid4()) + entries = tuple( + ExecutorConfig(scope=scope, principal_id=uuid4(), release_id=uuid4(), task_types=["http"], capacity=capacity) + for capacity in (2, 1) + ) + entered, finish = threading.Semaphore(0), threading.Event() + + async def connector_call(self, scope, principal_id, lease): + # The pure part of an execution, held open like a slow connector call. + return await execute_pure(lambda: entered.release() or finish.wait(5)) + + monkeypatch.setattr(ConnectorExecutionService, "_execute", connector_call) + service = ConnectorExecutionService( + None, # type: ignore[arg-type] + None, # type: ignore[arg-type] + None, # type: ignore[arg-type] + SimpleNamespace(require_operational=lambda: None), # type: ignore[arg-type] + None, # type: ignore[arg-type] + ) + native = NativeDispatcher(service, entries, "sha256:" + "a" * 64, 1) + + async with saturated_requests(): + executions = [ + asyncio.create_task(service.execute(scope, entry.principal_id, None, reservation=native.reservation)) # type: ignore[arg-type] + for entry in (entries[0], entries[0], entries[1], entries[1]) + ] + try: + for _ in executions: + assert await asyncio.to_thread(entered.acquire, True, 2) + # The configured native capacity plus the spare is the ceiling. + with pytest.raises(CatalogError, match="capacity"): + await service.execute(scope, entries[1].principal_id, None, reservation=native.reservation) # type: ignore[arg-type] + finally: + finish.set() + results = await asyncio.gather(*executions, return_exceptions=True) + assert results == [True, True, True, True] diff --git a/tests/unit/runtime/test_recovery_admission.py b/tests/unit/runtime/test_recovery_admission.py new file mode 100644 index 00000000..84090d40 --- /dev/null +++ b/tests/unit/runtime/test_recovery_admission.py @@ -0,0 +1,317 @@ +# Copyright 2026 Firefly Software Foundation. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Author: Firefly Software Foundation +# SPDX-License-Identifier: Apache-2.0 + +"""Recovery owns its pure execution slot apart from request admission and names its failures.""" + +import asyncio +import logging +import threading +from contextlib import asynccontextmanager +from datetime import UTC, datetime +from types import SimpleNamespace +from uuid import uuid4 + +from firefly_weave.access.scheduler import _CATALOG_AUTHORITY, _SchedulerScope +from firefly_weave.contracts.access import Scope +from firefly_weave.contracts.runtime import RecoveryReport +from firefly_weave.definitions.models import CatalogError +from firefly_weave.operations.execution import CONTROL_SLOTS, WORK_SLOTS, execute_pure, request_execution +from firefly_weave.persistence.uow import Transaction +from firefly_weave.runtime.kernel import transition_async +from firefly_weave.runtime.models import RuntimeEvent +from firefly_weave.runtime.recovery import RecoveryService +from firefly_weave.runtime.scheduler import RECOVERY_EXECUTION, RecoveryLoop + +SCHEDULER_LOGGER = "firefly_weave.runtime.scheduler" + + +class _Session: + """Just enough session for the catalog authority check and the candidate marker.""" + + def __init__(self, tenant): + self.tenant, self.info = tenant, {} + + def get_transaction(self): + return SimpleNamespace(is_active=True) + + async def execute(self, statement, params=None): + return None + + async def scalar(self, statement, params=None): + return str(self.tenant) + + +class _View: + def __init__(self, state): + self.state = state + + def model_copy(self, *, update): + return _View(update["state"]) + + +class _Harness: + """Real loop, recovery quanta and kernel admission; storage and the pure kernel are fakes.""" + + def __init__(self, monkeypatch, *, expired=1, deadlines=0, recovery_errors=(), schedule_errors=()): + self.tenant = uuid4() + self.scope = Scope(tenant_id=self.tenant, project_id=uuid4(), environment_id=uuid4()) + self.authority = _SchedulerScope(scope=self.scope, issuer=_CATALOG_AUTHORITY) + self.expired = [uuid4() for _ in range(expired)] + self.deadlines = deadlines + self.transitions, self.persisted, self.schedule_scans = [], [], [] + self.recovery_failures = [] + self.recovery_errors = list(recovery_errors) + self.schedule_errors = list(schedule_errors) + harness = self + + class Repository: + def __init__(self, tx, outbox): + pass + + async def expired_ready(self, limit): + return harness.expired[:limit] + + async def now(self): + return datetime.now(UTC) + + async def persist(self, view, event, result, digest, **timing): + harness.persisted.append(event.type) + + async def task_status(self, identifier, status): + harness.expired.remove(identifier) + + def transition(state, event, artifact): + harness.transitions.append(event.type) + return SimpleNamespace(state=state) + + monkeypatch.setattr("firefly_weave.runtime.recovery.RuntimeRepository", Repository) + monkeypatch.setattr("firefly_weave.runtime.recovery.view_of", lambda row: _View(row["state"])) + monkeypatch.setattr("firefly_weave.runtime.recovery.import_artifact", lambda artifact: None) + # `transition_async` resolves the pure reducer at call time, inside the admitted thread. + monkeypatch.setattr("firefly_weave.runtime.kernel.transition", transition) + + async def tenant_page(engine): + return [self.tenant] + + async def next_scope(uow, tenant): + return self.authority + + monkeypatch.setattr("firefly_weave.runtime.scheduler.tenant_page", tenant_page) + monkeypatch.setattr("firefly_weave.runtime.scheduler.next_scope", next_scope) + + @asynccontextmanager + async def transaction(scope, principal): + yield Transaction(session=_Session(scope.tenant_id), scope=scope) + + async def lock_task(tx, identifier, *, skip_locked): + row = {"id": uuid4(), "artifact": {}, "state": SimpleNamespace(status="waiting", accepted_sequence=3)} + return row, {"status": "ready", "node_id": "approval"} + + async def observe_unavailable(tx, row): + return False + + async def deadlines(tx, limit, *, terminal_only=False): + if self.recovery_errors: + raise self.recovery_errors.pop(0) + if not self.deadlines: + return RecoveryReport() + # Like `waits.settle`: a timed-out deadline reduces on the control pool. + self.deadlines -= 1 + state = SimpleNamespace(status="waiting", accepted_sequence=3) + event = RuntimeEvent(id=uuid4(), type="timed_out", timestamp=datetime.now(UTC), sequence=4) + await transition_async(state, event, None) # type: ignore[arg-type] + return RecoveryReport(timed_out=1) + + async def recovery_candidates(tx, limit): + return [] + + runtime = SimpleNamespace( + definitions=SimpleNamespace( + transaction=transaction, registry=SimpleNamespace(require_operational=lambda: None), outbox=None + ), + lock_task=lock_task, + observe_unavailable=observe_unavailable, + ) + self.recovery = RecoveryService( + runtime, + SimpleNamespace(recovery_candidates=recovery_candidates), + SimpleNamespace(scan=deadlines), + None, + ) + scan = self.recovery._scan_scheduled + + async def observed_scan(authority, limit): + try: + return await scan(authority, limit) + except Exception as error: + self.recovery_failures.append(error) + raise + + monkeypatch.setattr(self.recovery, "_scan_scheduled", observed_scan) + + async def schedule_scan(authority, limit): + self.schedule_scans.append(authority.scope) + if self.schedule_errors: + raise self.schedule_errors.pop(0) + + self.loop = RecoveryLoop( + SimpleNamespace(database_timeout_seconds=5.0, scheduler_poll_seconds=0.01), + None, + self.recovery, + SimpleNamespace(_scan_scheduled=schedule_scan), + ) + self.loop.engine = object() + + async def cycle_while_requests_hold_every_slot(self, *, control=False): + admitted = [asyncio.Event() for _ in range(CONTROL_SLOTS if control else WORK_SLOTS)] + finish = asyncio.Event() + + async def in_flight_request(ready): + # Each request owns one slot for its whole response, like `BodyBoundary._serve`. + async with request_execution(control=control): + ready.set() + await finish.wait() + + # Separate tasks, so the loop does not inherit a request lease from this context. + requests = [asyncio.create_task(in_flight_request(ready)) for ready in admitted] + try: + async with asyncio.timeout(2): + for ready in admitted: + await ready.wait() + await self.loop.cycle() + finally: + finish.set() + await asyncio.gather(*requests) + + +async def test_recovery_cycle_progresses_while_in_flight_requests_hold_every_work_slot(monkeypatch): + harness = _Harness(monkeypatch) + + await harness.cycle_while_requests_hold_every_slot() + + assert harness.recovery_failures == [] + assert harness.transitions == ["incident_opened"] and harness.persisted == ["incident_opened"] + assert harness.schedule_scans == [harness.scope] + + +async def test_recovery_deadlines_progress_while_in_flight_requests_hold_every_control_slot(monkeypatch): + harness = _Harness(monkeypatch, expired=0, deadlines=1) + + await harness.cycle_while_requests_hold_every_slot(control=True) + + assert harness.recovery_failures == [] + assert harness.transitions == ["timed_out"] + assert harness.schedule_scans == [harness.scope] + + +async def orphaned_recovery_calls(count): + """Recovery pure calls that outlived their cancelled quanta and still run in their threads.""" + entered, release = threading.Semaphore(0), threading.Event() + + async def quantum(): + async with RECOVERY_EXECUTION.execution(): + await execute_pure(lambda: entered.release() or release.wait(5)) + + quanta = [asyncio.create_task(quantum()) for _ in range(count)] + for _ in quanta: + assert await asyncio.to_thread(entered.acquire, True, 2) + for task in quanta: + task.cancel() + return release, quanta + + +async def test_recovery_tolerates_one_call_that_outlived_its_cancelled_quantum(monkeypatch): + harness = _Harness(monkeypatch) + release, quanta = await orphaned_recovery_calls(1) + try: + await harness.cycle_while_requests_hold_every_slot() + finally: + release.set() + await asyncio.gather(*quanta, return_exceptions=True) + + assert harness.recovery_failures == [] + assert harness.transitions == ["incident_opened"] + + +async def test_recovery_quanta_are_refused_while_orphaned_calls_hold_the_whole_reservation(monkeypatch, caplog): + caplog.set_level(logging.ERROR, logger=SCHEDULER_LOGGER) + harness = _Harness(monkeypatch) + release, quanta = await orphaned_recovery_calls(RECOVERY_EXECUTION.size) + try: + await harness.loop.cycle() + finally: + release.set() + await asyncio.gather(*quanta, return_exceptions=True) + + # The bound holds: no third recovery thread starts, and the refusal is named. + [failure] = harness.recovery_failures + assert (failure.status, failure.code) == (429, "WV-OPERATION-CAPACITY") + assert [record.getMessage() for record in caplog.records] == [ + "Tenant recovery scan failed (CatalogError WV-OPERATION-CAPACITY); traversal will continue" + ] + assert harness.transitions == [] and harness.schedule_scans == [] + async with asyncio.timeout(2): + while not harness.transitions: + # The released threads give their slots back just after their calls return. + await harness.loop.cycle() + await asyncio.sleep(0.01) + assert harness.transitions == ["incident_opened"] + + +async def test_poll_logs_a_failed_catalog_scan_by_class_and_code(monkeypatch, caplog): + caplog.set_level(logging.ERROR, logger=SCHEDULER_LOGGER) + harness = _Harness(monkeypatch) + + async def tenant_page(engine): + raise CatalogError(429, "WV-OPERATION-CAPACITY", "Catalog sentinel unavailable") + + monkeypatch.setattr("firefly_weave.runtime.scheduler.tenant_page", tenant_page) + poll = asyncio.create_task(harness.loop.poll()) + try: + async with asyncio.timeout(2): + while not caplog.records: + await asyncio.sleep(0.01) + finally: + poll.cancel() + await asyncio.gather(poll, return_exceptions=True) + + assert caplog.records[0].getMessage() == ( + "Recovery catalog scan failed (CatalogError WV-OPERATION-CAPACITY); retrying on next poll" + ) + assert "sentinel" not in caplog.text + + +async def test_tenant_scan_failures_log_error_class_and_code_without_values(monkeypatch, caplog): + caplog.set_level(logging.ERROR, logger=SCHEDULER_LOGGER) + harness = _Harness( + monkeypatch, + recovery_errors=[CatalogError(429, "WV-OPERATION-CAPACITY", "Operation capacity sentinel")], + schedule_errors=[ + CatalogError(409, "WV-OPERATION-UNAVAILABLE", "Schedule nightly-sentinel unavailable"), + RuntimeError("tenant-sentinel"), + ], + ) + + for _ in range(3): + await harness.loop.cycle() + + messages = [record.getMessage() for record in caplog.records if record.name == SCHEDULER_LOGGER] + assert messages == [ + "Tenant recovery scan failed (CatalogError WV-OPERATION-CAPACITY); traversal will continue", + "Tenant schedule scan failed (CatalogError WV-OPERATION-UNAVAILABLE); traversal will continue", + "Tenant schedule scan failed (RuntimeError); traversal will continue", + ] + assert "sentinel" not in caplog.text