refactor: address PR #851 review round 3 (ingest transport)
- LocalTransport.run_consumers increments active_consumer_count per worker (instead of once after the loop) so the count is accurate if a later tg.start() raises mid-pool. - Add LocalTransport.aclose() to explicitly close its owned send/receive stream ends (belt-and-suspenders against unclosed-resource warnings; anyio aclose is idempotent, and by shutdown the scanner is already winding down). Reworded the base IngestTransport.aclose() docstring to point at the overrides. - Inline ingest_transport.producer at the scanner/user_manager call sites, dropping the single-use task_producer alias in both lifespan paths. - Annotate DistributedTransport._producer explicitly as ProcrastinateTaskProducer so the drain() coupling is visible and ty catches drift. - Add a unit test for LocalTransport.aclose() (closes the owned streams, idempotent). Refs: Deck #196 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
parent
2179e9ddb0
commit
c0c52c1b34
@@ -1720,7 +1720,6 @@ def get_app(transport: str = "streamable-http", enabled_apps: list[str] | None =
|
||||
# Named ingest_transport (not transport) to avoid shadowing the
|
||||
# get_app(transport=...) HTTP-transport parameter.
|
||||
ingest_transport = await build_transport(settings)
|
||||
task_producer = ingest_transport.producer
|
||||
|
||||
# Publish to app.state (ADR-007), the module singleton (FastMCP
|
||||
# session lifespans), and the /app browser sub-app in one place.
|
||||
@@ -1730,10 +1729,10 @@ def get_app(transport: str = "streamable-http", enabled_apps: list[str] | None =
|
||||
|
||||
# Start background tasks using anyio TaskGroup
|
||||
async with anyio.create_task_group() as tg:
|
||||
# Start scanner task (publishes to task_producer)
|
||||
# Start scanner task (publishes to the transport's producer)
|
||||
await tg.start(
|
||||
scanner_task,
|
||||
task_producer,
|
||||
ingest_transport.producer,
|
||||
shutdown_event,
|
||||
scanner_wake_event,
|
||||
client,
|
||||
@@ -1912,7 +1911,6 @@ def get_app(transport: str = "streamable-http", enabled_apps: list[str] | None =
|
||||
# Named ingest_transport (not transport) to avoid shadowing the
|
||||
# get_app(transport=...) HTTP-transport parameter.
|
||||
ingest_transport = await build_transport(settings)
|
||||
task_producer = ingest_transport.producer
|
||||
|
||||
# Publish to app.state (ADR-007), the module singleton (FastMCP
|
||||
# session lifespans), and the /app browser sub-app in one place.
|
||||
@@ -1930,11 +1928,11 @@ def get_app(transport: str = "streamable-http", enabled_apps: list[str] | None =
|
||||
# management API revoke endpoint (via app.state.oauth_context).
|
||||
async with anyio.create_task_group() as tg:
|
||||
# Start user manager task (supervises per-user scanners).
|
||||
# Each per-user scanner clones task_producer; for the bus
|
||||
# Each per-user scanner clones the producer; for the bus
|
||||
# producer clone() returns the shared connection.
|
||||
await tg.start(
|
||||
user_manager_task,
|
||||
task_producer,
|
||||
ingest_transport.producer,
|
||||
shutdown_event,
|
||||
scanner_wake_event,
|
||||
token_storage,
|
||||
|
||||
@@ -128,9 +128,9 @@ class IngestTransport(abc.ABC):
|
||||
async def aclose(self) -> None:
|
||||
"""Tear down backend-owned resources once on lifespan shutdown.
|
||||
|
||||
No-op by default (the memory stream is closed by task-group cancellation,
|
||||
exactly as before); :class:`DistributedTransport` closes its connector
|
||||
pool here.
|
||||
No-op by default; subclasses that own resources (a connector pool, stream
|
||||
handles) override this to release them — see
|
||||
:meth:`DistributedTransport.aclose` and :meth:`LocalTransport.aclose`.
|
||||
"""
|
||||
return None
|
||||
|
||||
@@ -178,7 +178,22 @@ class LocalTransport(IngestTransport):
|
||||
# the prior inline lifespan behaviour.
|
||||
for i in range(count):
|
||||
await task_group.start(spawn_worker, i, self._receive_stream.clone())
|
||||
self._active_consumer_count = count
|
||||
# Increment per-worker (not once after the loop) so the count is
|
||||
# accurate even if a later start() raises — a crash log then reflects
|
||||
# how many workers were actually live.
|
||||
self._active_consumer_count += 1
|
||||
|
||||
async def aclose(self) -> None:
|
||||
# Belt-and-suspenders cleanup of the two stream ends this transport owns,
|
||||
# so they don't linger until GC (which can emit unclosed-resource
|
||||
# warnings under the test runner / alternative runtimes). anyio's aclose
|
||||
# is idempotent, so the scanner's own ``async with`` on the send side
|
||||
# (single-user) closing it first is harmless; worker receive *clones* are
|
||||
# independent handles, closed by task-group cancellation. By shutdown the
|
||||
# ``shutdown_event`` is already set, so the scanner is winding down rather
|
||||
# than issuing fresh sends.
|
||||
await self._send_stream.aclose()
|
||||
await self._receive_stream.aclose()
|
||||
|
||||
|
||||
class DistributedTransport(IngestTransport):
|
||||
@@ -197,7 +212,10 @@ class DistributedTransport(IngestTransport):
|
||||
"""
|
||||
|
||||
def __init__(self, producer: ProcrastinateTaskProducer):
|
||||
self._producer = producer
|
||||
# Explicit (not inferred): aclose() calls drain(), which lives on the
|
||||
# concrete ProcrastinateTaskProducer, not the TaskProducer protocol —
|
||||
# the annotation keeps that coupling visible and lets ty catch drift.
|
||||
self._producer: ProcrastinateTaskProducer = producer
|
||||
|
||||
@property
|
||||
def producer(self) -> TaskProducer:
|
||||
|
||||
@@ -23,6 +23,7 @@ from nextcloud_mcp_server.vector.queue import (
|
||||
SpawnWorker,
|
||||
build_transport,
|
||||
)
|
||||
from nextcloud_mcp_server.vector.scanner import DocumentTask
|
||||
|
||||
pytestmark = pytest.mark.unit
|
||||
|
||||
@@ -94,6 +95,25 @@ class TestLocalTransport:
|
||||
assert all(s is not None for s in received_streams)
|
||||
assert len({id(s) for s in received_streams}) == 3
|
||||
|
||||
async def test_aclose_closes_owned_streams_idempotently(self):
|
||||
transport = LocalTransport(max_buffer_size=5)
|
||||
await transport.aclose()
|
||||
|
||||
# The send end is closed → the producer raises rather than silently
|
||||
# dropping (the producer wraps the same stream aclose() closed).
|
||||
with pytest.raises(anyio.ClosedResourceError):
|
||||
await transport.producer.send(
|
||||
DocumentTask(
|
||||
user_id="u",
|
||||
doc_id="1",
|
||||
doc_type="note",
|
||||
operation="index",
|
||||
modified_at=0,
|
||||
)
|
||||
)
|
||||
# Idempotent: closing again is a no-op, not an error.
|
||||
await transport.aclose()
|
||||
|
||||
|
||||
class TestDistributedTransport:
|
||||
async def test_run_consumers_is_noop(self):
|
||||
|
||||
Reference in New Issue
Block a user