From 392dd99faf8df389e90e3b04152d6e288ba68b7c Mon Sep 17 00:00:00 2001 From: Andres Contreras Date: Thu, 8 Oct 2026 09:51:52 -0700 Subject: [PATCH 1/3] Name the error class and code when a tenant scan fails The recovery and schedule traversal logged only "Tenant recovery scan failed" and "Tenant schedule scan failed", so a capacity refusal could not be told apart from any other error. Both lines now name the error class and, for a CatalogError, its code, never the message or a value. A unit test drives the real RecoveryLoop cycle through the recovery quanta into transition_async while two in-flight requests hold every work slot. The cycle fails with WV-OPERATION-CAPACITY and skips that tenant's schedule scan; once the requests finish, it reaches the kernel. The test pins today's admission design, in which lifespan loops borrow the request work pool. --- src/firefly_weave/runtime/scheduler.py | 20 +- tests/unit/runtime/test_recovery_admission.py | 229 ++++++++++++++++++ 2 files changed, 245 insertions(+), 4 deletions(-) create mode 100644 tests/unit/runtime/test_recovery_admission.py diff --git a/src/firefly_weave/runtime/scheduler.py b/src/firefly_weave/runtime/scheduler.py index bdda7ec5..6da227eb 100644 --- a/src/firefly_weave/runtime/scheduler.py +++ b/src/firefly_weave/runtime/scheduler.py @@ -25,6 +25,7 @@ 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.persistence.migrations import check_schema from firefly_weave.persistence.uow import UnitOfWork from firefly_weave.runtime.recovery import RecoveryService @@ -32,6 +33,13 @@ from firefly_weave.triggers.scheduler import Scheduler +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__( self, @@ -99,16 +107,20 @@ async def cycle(self) -> None: 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") + except Exception as error: + logging.getLogger(__name__).error( + "Tenant recovery scan failed (%s); traversal will continue", failure_cause(error) + ) 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") + 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: diff --git a/tests/unit/runtime/test_recovery_admission.py b/tests/unit/runtime/test_recovery_admission.py new file mode 100644 index 00000000..ca8f242b --- /dev/null +++ b/tests/unit/runtime/test_recovery_admission.py @@ -0,0 +1,229 @@ +# 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's pure kernel work borrows request admission slots and names its failures.""" + +import asyncio +import logging +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 WORK_SLOTS, request_execution +from firefly_weave.persistence.uow import Transaction +from firefly_weave.runtime.recovery import RecoveryService +from firefly_weave.runtime.scheduler import 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, 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()] + self.transitions, self.persisted, self.schedule_scans = [], [], [] + self.recovery_failures = [] + 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): + return RecoveryReport() + + 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), + None, + self.recovery, + SimpleNamespace(_scan_scheduled=schedule_scan), + ) + self.loop.engine = object() + + async def cycle_while_requests_hold_every_work_slot(self): + admitted = [asyncio.Event() for _ in range(WORK_SLOTS)] + finish = asyncio.Event() + + async def in_flight_request(ready): + # Each mutation request owns one work slot for its whole response, like `BodyBoundary._serve`. + async with request_execution(): + 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_fails_when_in_flight_requests_hold_every_work_slot(monkeypatch): + # This pins today's admission design: lifespan loops borrow the request work pool. + # A reserved background lease should invert the first half of this test. + harness = _Harness(monkeypatch) + + await harness.cycle_while_requests_hold_every_work_slot() + + [failure] = harness.recovery_failures + assert isinstance(failure, CatalogError) + assert (failure.status, failure.code) == (429, "WV-OPERATION-CAPACITY") + assert str(failure) == "Pure execution capacity unavailable" + assert harness.transitions == [] and harness.persisted == [] + # The failed recovery quantum also skips this tenant's schedule scan for the cycle. + assert harness.schedule_scans == [] + + # Once the requests finish, the same cycle reaches the kernel and the schedule scan. + await harness.loop.cycle() + + assert len(harness.recovery_failures) == 1 + assert harness.transitions == ["incident_opened"] and harness.persisted == ["incident_opened"] + assert harness.schedule_scans == [harness.scope] + + +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, + schedule_errors=[ + CatalogError(409, "WV-OPERATION-UNAVAILABLE", "Schedule nightly-sentinel unavailable"), + RuntimeError("tenant-sentinel"), + ], + ) + + await harness.cycle_while_requests_hold_every_work_slot() + await harness.loop.cycle() + 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 From 340583cbf795fe193172363afd56b7df7fbcc544 Mon Sep 17 00:00:00 2001 From: Andres Contreras Date: Thu, 8 Oct 2026 10:29:41 -0700 Subject: [PATCH 2/3] Reserve pure execution for lifespan loops and native workers In-flight requests hold a work or control slot for their whole response, and lifespan loops borrowed those same no-queue pools. Two Studio mutations, or four reads, refused recovery quanta, deadline settlements, schedule starts, provider and email dispatch, and Kafka consumer turns. Native connector executions held a request work slot for the whole connector call, so two of them took every mutation slot from Studio and a third failed at entry. ReservedSlots gives each owner its own pool. Inside its execution() context, every pure call takes one slot from that pool instead of the request pools, and an inherited request lease is set aside. A call that outlived its cancelled caller keeps its slot until it ends, so the size bounds the owner's pure threads, and a spare slot absorbs one such call as the shared work pool did. Recovery and provider cycles get 2 slots, Kafka turns 8 (at most max_clients - 1 turns, max_clients <= 8), and the native dispatcher its configured capacity plus one. Native claims, heartbeats and settlements keep request admission and their replay policy. Request admission itself is unchanged. The recovery catalog failure log now names the error class and code as the tenant logs do. --- src/firefly_weave/connectors/dispatcher.py | 6 +- src/firefly_weave/connectors/execution.py | 10 +- src/firefly_weave/operations/execution.py | 28 ++- src/firefly_weave/providers/loop.py | 30 +-- src/firefly_weave/runtime/scheduler.py | 70 ++++--- src/firefly_weave/triggers/kafka_loop.py | 7 +- tests/unit/connectors/test_local_build.py | 22 +++ .../operations/test_background_admission.py | 186 ++++++++++++++++++ tests/unit/operations/test_execution.py | 136 +++++++++++++ tests/unit/runtime/test_recovery_admission.py | 144 +++++++++++--- 10 files changed, 562 insertions(+), 77 deletions(-) create mode 100644 tests/unit/operations/test_background_admission.py diff --git a/src/firefly_weave/connectors/dispatcher.py b/src/firefly_weave/connectors/dispatcher.py index d6317463..1e1a978c 100644 --- a/src/firefly_weave/connectors/dispatcher.py +++ b/src/firefly_weave/connectors/dispatcher.py @@ -33,6 +33,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 @@ -117,6 +118,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 @@ -171,7 +175,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 e25bd56c..58a195dc 100644 --- a/src/firefly_weave/operations/execution.py +++ b/src/firefly_weave/operations/execution.py @@ -68,6 +68,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 @@ -93,6 +94,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) @@ -106,7 +132,7 @@ async def request_execution(*, control: bool = False) -> 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 395acac0..652710d2 100644 --- a/src/firefly_weave/providers/loop.py +++ b/src/firefly_weave/providers/loop.py @@ -24,11 +24,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__( @@ -72,18 +77,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 6da227eb..6e5cc71b 100644 --- a/src/firefly_weave/runtime/scheduler.py +++ b/src/firefly_weave/runtime/scheduler.py @@ -26,12 +26,17 @@ 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.""" @@ -91,43 +96,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 as error: - logging.getLogger(__name__).error( - "Tenant recovery scan failed (%s); traversal will continue", failure_cause(error) - ) - 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) - ) + 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 1770bb18..ac91bc0d 100644 --- a/src/firefly_weave/triggers/kafka_loop.py +++ b/src/firefly_weave/triggers/kafka_loop.py @@ -26,11 +26,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: @@ -98,7 +103,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..f26a3db7 --- /dev/null +++ b/tests/unit/operations/test_background_admission.py @@ -0,0 +1,186 @@ +# 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) + + +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/operations/test_execution.py b/tests/unit/operations/test_execution.py index 2042389c..13508ecb 100644 --- a/tests/unit/operations/test_execution.py +++ b/tests/unit/operations/test_execution.py @@ -18,6 +18,7 @@ import asyncio import threading +from contextlib import asynccontextmanager import pytest @@ -133,3 +134,138 @@ async def run(): assert await asyncio.to_thread(exited.wait, 2) async with debug_execution(): assert await execute_pure(lambda: "debug-recovered") == "debug-recovered" + + +@asynccontextmanager +async def saturated_requests(): + """Every work and control slot owned by an in-flight request in its own task.""" + from firefly_weave.operations.execution import CONTROL_SLOTS, WORK_SLOTS, request_execution + + finish, admitted, tasks = asyncio.Event(), [], [] + + async def request(control, ready): + async with request_execution(control=control): + ready.set() + await finish.wait() + + try: + for control, count in ((False, WORK_SLOTS), (True, CONTROL_SLOTS)): + for _ in range(count): + admitted.append(asyncio.Event()) + tasks.append(asyncio.create_task(request(control, admitted[-1]))) + async with asyncio.timeout(2): + for ready in admitted: + await ready.wait() + yield + finally: + finish.set() + await asyncio.gather(*tasks, return_exceptions=True) + + +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(): + with pytest.raises(CatalogError, match="capacity"): + await execute_pure(lambda: "borrowed") + async with 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" diff --git a/tests/unit/runtime/test_recovery_admission.py b/tests/unit/runtime/test_recovery_admission.py index ca8f242b..84090d40 100644 --- a/tests/unit/runtime/test_recovery_admission.py +++ b/tests/unit/runtime/test_recovery_admission.py @@ -14,10 +14,11 @@ # Author: Firefly Software Foundation # SPDX-License-Identifier: Apache-2.0 -"""Recovery's pure kernel work borrows request admission slots and names its failures.""" +"""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 @@ -27,10 +28,12 @@ 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 WORK_SLOTS, request_execution +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 RecoveryLoop +from firefly_weave.runtime.scheduler import RECOVERY_EXECUTION, RecoveryLoop SCHEDULER_LOGGER = "firefly_weave.runtime.scheduler" @@ -62,13 +65,15 @@ def model_copy(self, *, update): class _Harness: """Real loop, recovery quanta and kernel admission; storage and the pure kernel are fakes.""" - def __init__(self, monkeypatch, schedule_errors=()): + 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()] + 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 @@ -119,7 +124,16 @@ async def observe_unavailable(tx, row): return False async def deadlines(tx, limit, *, terminal_only=False): - return RecoveryReport() + 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 [] @@ -154,20 +168,20 @@ async def schedule_scan(authority, limit): raise self.schedule_errors.pop(0) self.loop = RecoveryLoop( - SimpleNamespace(database_timeout_seconds=5.0), + 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_work_slot(self): - admitted = [asyncio.Event() for _ in range(WORK_SLOTS)] + 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 mutation request owns one work slot for its whole response, like `BodyBoundary._serve`. - async with request_execution(): + # Each request owns one slot for its whole response, like `BodyBoundary._serve`. + async with request_execution(control=control): ready.set() await finish.wait() @@ -183,42 +197,116 @@ async def in_flight_request(ready): await asyncio.gather(*requests) -async def test_recovery_cycle_fails_when_in_flight_requests_hold_every_work_slot(monkeypatch): - # This pins today's admission design: lifespan loops borrow the request work pool. - # A reserved background lease should invert the first half of this test. +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_work_slot() + 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 isinstance(failure, CatalogError) assert (failure.status, failure.code) == (429, "WV-OPERATION-CAPACITY") - assert str(failure) == "Pure execution capacity unavailable" - assert harness.transitions == [] and harness.persisted == [] - # The failed recovery quantum also skips this tenant's schedule scan for the cycle. - assert harness.schedule_scans == [] + 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"] - # Once the requests finish, the same cycle reaches the kernel and the schedule scan. - await harness.loop.cycle() - assert len(harness.recovery_failures) == 1 - assert harness.transitions == ["incident_opened"] and harness.persisted == ["incident_opened"] - assert harness.schedule_scans == [harness.scope] +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"), ], ) - await harness.cycle_while_requests_hold_every_work_slot() - await harness.loop.cycle() - await harness.loop.cycle() + for _ in range(3): + await harness.loop.cycle() messages = [record.getMessage() for record in caplog.records if record.name == SCHEDULER_LOGGER] assert messages == [ From 9925b0928b6d1df9092ed24987ac9ee3f5d2bada Mon Sep 17 00:00:00 2001 From: Andres Contreras Date: Thu, 8 Oct 2026 10:32:02 -0700 Subject: [PATCH 3/3] Move the reservation tests next to the loop admission tests The ReservedSlots tests now share the request saturation helper in test_background_admission.py instead of duplicating it at the end of test_execution.py, which the compatibility rescan change also extends. --- .../operations/test_background_admission.py | 106 ++++++++++++++ tests/unit/operations/test_execution.py | 136 ------------------ 2 files changed, 106 insertions(+), 136 deletions(-) diff --git a/tests/unit/operations/test_background_admission.py b/tests/unit/operations/test_background_admission.py index f26a3db7..33b2868a 100644 --- a/tests/unit/operations/test_background_admission.py +++ b/tests/unit/operations/test_background_admission.py @@ -58,6 +58,112 @@ async def request(control, ready): 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 diff --git a/tests/unit/operations/test_execution.py b/tests/unit/operations/test_execution.py index 13508ecb..2042389c 100644 --- a/tests/unit/operations/test_execution.py +++ b/tests/unit/operations/test_execution.py @@ -18,7 +18,6 @@ import asyncio import threading -from contextlib import asynccontextmanager import pytest @@ -134,138 +133,3 @@ async def run(): assert await asyncio.to_thread(exited.wait, 2) async with debug_execution(): assert await execute_pure(lambda: "debug-recovered") == "debug-recovered" - - -@asynccontextmanager -async def saturated_requests(): - """Every work and control slot owned by an in-flight request in its own task.""" - from firefly_weave.operations.execution import CONTROL_SLOTS, WORK_SLOTS, request_execution - - finish, admitted, tasks = asyncio.Event(), [], [] - - async def request(control, ready): - async with request_execution(control=control): - ready.set() - await finish.wait() - - try: - for control, count in ((False, WORK_SLOTS), (True, CONTROL_SLOTS)): - for _ in range(count): - admitted.append(asyncio.Event()) - tasks.append(asyncio.create_task(request(control, admitted[-1]))) - async with asyncio.timeout(2): - for ready in admitted: - await ready.wait() - yield - finally: - finish.set() - await asyncio.gather(*tasks, return_exceptions=True) - - -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(): - with pytest.raises(CatalogError, match="capacity"): - await execute_pure(lambda: "borrowed") - async with 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"