feat: add IngestTransport port for local/distributed ingest backends
Finish the hexagonal ports-&-adapters split started in #183. The producer side already had a TaskProducer port + adapters, but the consumer side was unabstracted and the INGEST_QUEUE selection leaked into a duplicated `if use_postgres:` branch across both app.py lifespan paths. Introduce an IngestTransport ABC (vector/queue/transport.py) that bundles the producer with running (or not running) the in-process consumer pool, built by a single build_transport() factory: - LocalTransport (INGEST_QUEUE=memory): in-process anyio stream drained by an N-worker pool that run_consumers starts. - DistributedTransport (INGEST_QUEUE=postgres): wraps ProcrastinateTaskProducer; run_consumers is a no-op because the consumer is the external `worker` role. Both lifespan paths now call build_transport + _wire_vector_sync_state (new helper that centralizes the app.state / module-singleton / browser-app writes) + transport.run_consumers + transport.aclose(), with no INGEST_QUEUE branching and no getattr drain probe. Adding a future backend (Redis/NATS/SQS) is one new adapter + one build_transport arm, with no app.py or scanner change. Preserves the single-tenant parallelism invariant (one shared multiplexed queue + N-worker pool, per-document not per-user dispatch) and documents it in ADR-028. The worker CLI is unchanged (it is the external consumer). Refs: Deck #196 (Deck #197 tracks the explicit parallelism regression test) Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Chris Coutinho
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
55ea8dd358
commit
e7bcdb1950
+131
-141
@@ -9,7 +9,7 @@ import traceback
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import cast
|
||||
from typing import Any, cast
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import anyio
|
||||
@@ -134,11 +134,11 @@ from nextcloud_mcp_server.vector.placeholder import sweep_orphan_placeholders
|
||||
from nextcloud_mcp_server.vector.processor import processor_task
|
||||
from nextcloud_mcp_server.vector.qdrant_client import get_qdrant_client
|
||||
from nextcloud_mcp_server.vector.queue import (
|
||||
MemoryTaskProducer,
|
||||
IngestTransport,
|
||||
TaskProducer,
|
||||
build_producer,
|
||||
build_transport,
|
||||
)
|
||||
from nextcloud_mcp_server.vector.scanner import DocumentTask, scanner_task
|
||||
from nextcloud_mcp_server.vector.scanner import scanner_task
|
||||
from nextcloud_mcp_server.vector.webhook_receiver import handle_nextcloud_webhook
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -347,6 +347,47 @@ class VectorSyncState:
|
||||
_vector_sync_state = VectorSyncState()
|
||||
|
||||
|
||||
def _wire_vector_sync_state(
|
||||
app: Starlette,
|
||||
transport: IngestTransport,
|
||||
shutdown_event: anyio.Event,
|
||||
scanner_wake_event: anyio.Event,
|
||||
) -> None:
|
||||
"""Publish the ingest transport + sync events to every state surface.
|
||||
|
||||
Both lifespan paths (single-user and multi-user) previously duplicated these
|
||||
writes three times each: ``app.state``, the module singleton
|
||||
``_vector_sync_state`` (read by FastMCP session lifespans), and the mounted
|
||||
``/app`` browser sub-app. ``document_send_stream``/``document_receive_stream``
|
||||
come from the transport — ``None`` in postgres mode (no in-process stream) —
|
||||
and ``task_producer`` is the transport's producer in both modes (Deck #183,
|
||||
ADR-028).
|
||||
"""
|
||||
send_stream = transport.send_stream
|
||||
receive_stream = transport.receive_stream
|
||||
task_producer = transport.producer
|
||||
|
||||
def _apply(state: Any) -> None:
|
||||
state.document_send_stream = send_stream
|
||||
state.document_receive_stream = receive_stream
|
||||
state.task_producer = task_producer
|
||||
state.shutdown_event = shutdown_event
|
||||
state.scanner_wake_event = scanner_wake_event
|
||||
|
||||
# app.state (Starlette) + the module singleton share the same attribute names.
|
||||
_apply(app.state)
|
||||
_apply(_vector_sync_state)
|
||||
logger.info("Vector sync state stored in module singleton")
|
||||
|
||||
# Also share with the mounted /app browser sub-app, if present.
|
||||
for route in app.routes:
|
||||
if isinstance(route, Mount) and route.path == "/app":
|
||||
browser_app = cast(Starlette, route.app)
|
||||
_apply(browser_app.state)
|
||||
logger.info("Vector sync state shared with browser_app for /app")
|
||||
break
|
||||
|
||||
|
||||
@dataclass
|
||||
class AppContext:
|
||||
"""Application context for BasicAuth mode."""
|
||||
@@ -1663,60 +1704,21 @@ def get_app(transport: str = "streamable-http", enabled_apps: list[str] | None =
|
||||
# Orphan-sweep before scanner starts — card #101.
|
||||
await _sweep_orphan_placeholders_if_enabled()
|
||||
|
||||
# Initialize shared state. INGEST_QUEUE selects the transport
|
||||
# (Deck #183): ``memory`` uses the in-process anyio stream + the
|
||||
# in-process processor pool (SQLite/dev); ``postgres`` defers jobs to
|
||||
# the per-tenant Postgres via procrastinate and runs no in-process
|
||||
# consumer (the separate ``worker`` role drains the queue).
|
||||
use_postgres = settings.ingest_queue == "postgres"
|
||||
# Initialize the ingest transport. INGEST_QUEUE selects the backend
|
||||
# (Deck #183, ADR-028): ``memory`` builds an in-process anyio stream
|
||||
# drained by an in-process pool (SQLite/dev); ``postgres`` defers jobs
|
||||
# to the per-tenant Postgres via procrastinate and runs no in-process
|
||||
# consumer (the separate ``worker`` role drains the queue). The
|
||||
# transport hides that choice — this path is now backend-agnostic.
|
||||
shutdown_event = anyio.Event()
|
||||
scanner_wake_event = anyio.Event()
|
||||
|
||||
send_stream = None
|
||||
receive_stream = None
|
||||
task_producer: TaskProducer
|
||||
if use_postgres:
|
||||
# Open the connector once (build_producer) and reuse it to create
|
||||
# procrastinate's tables before the scanner can defer — a single
|
||||
# open/close cycle, matching the worker command.
|
||||
producer = await build_producer(settings)
|
||||
await producer.ensure_schema()
|
||||
task_producer = producer
|
||||
logger.info("Ingest queue: postgres (procrastinate); worker drains it")
|
||||
else:
|
||||
send_stream, receive_stream = anyio.create_memory_object_stream[
|
||||
DocumentTask
|
||||
](max_buffer_size=settings.vector_sync_queue_max_size)
|
||||
task_producer = MemoryTaskProducer(send_stream)
|
||||
transport = await build_transport(settings)
|
||||
task_producer = transport.producer
|
||||
|
||||
# Store in app state for access from routes (ADR-007). In postgres
|
||||
# mode there is no in-memory stream, so document_send/receive_stream
|
||||
# stay None; task_producer is the procrastinate producer.
|
||||
app.state.document_send_stream = send_stream
|
||||
app.state.document_receive_stream = receive_stream
|
||||
app.state.task_producer = task_producer
|
||||
app.state.shutdown_event = shutdown_event
|
||||
app.state.scanner_wake_event = scanner_wake_event
|
||||
|
||||
# Also store in module singleton for FastMCP session lifespans
|
||||
_vector_sync_state.document_send_stream = send_stream
|
||||
_vector_sync_state.document_receive_stream = receive_stream
|
||||
_vector_sync_state.task_producer = task_producer
|
||||
_vector_sync_state.shutdown_event = shutdown_event
|
||||
_vector_sync_state.scanner_wake_event = scanner_wake_event
|
||||
logger.info("Vector sync state stored in module singleton")
|
||||
|
||||
# Also share with browser_app for /app route
|
||||
for route in app.routes:
|
||||
if isinstance(route, Mount) and route.path == "/app":
|
||||
browser_app = cast(Starlette, route.app)
|
||||
browser_app.state.document_send_stream = send_stream
|
||||
browser_app.state.document_receive_stream = receive_stream
|
||||
browser_app.state.task_producer = task_producer
|
||||
browser_app.state.shutdown_event = shutdown_event
|
||||
browser_app.state.scanner_wake_event = scanner_wake_event
|
||||
logger.info("Vector sync state shared with browser_app for /app")
|
||||
break
|
||||
# Publish to app.state (ADR-007), the module singleton (FastMCP
|
||||
# session lifespans), and the /app browser sub-app in one place.
|
||||
_wire_vector_sync_state(app, transport, shutdown_event, scanner_wake_event)
|
||||
|
||||
# Start background tasks using anyio TaskGroup
|
||||
async with anyio.create_task_group() as tg:
|
||||
@@ -1730,19 +1732,28 @@ def get_app(transport: str = "streamable-http", enabled_apps: list[str] | None =
|
||||
username,
|
||||
)
|
||||
|
||||
# The in-process processor pool runs only in memory mode; in
|
||||
# postgres mode the out-of-process worker is the consumer.
|
||||
if not use_postgres:
|
||||
assert receive_stream is not None
|
||||
for i in range(settings.vector_sync_processor_workers):
|
||||
await tg.start(
|
||||
processor_task,
|
||||
i,
|
||||
receive_stream.clone(),
|
||||
shutdown_event,
|
||||
client,
|
||||
username,
|
||||
)
|
||||
# Start the in-process consumer pool. ``run_consumers`` is a
|
||||
# no-op for the distributed (postgres) backend — its consumer is
|
||||
# the out-of-process ``worker`` role. The closure binds this
|
||||
# mode's shared client+username and forwards anyio's injected
|
||||
# ``task_status`` so ``tg.start`` observes each worker's
|
||||
# readiness. One shared receive stream + N workers ⇒ a single
|
||||
# multiplexed queue processed with N-way parallelism (ADR-028).
|
||||
async def spawn_worker(
|
||||
worker_id, receive_stream, *, task_status=anyio.TASK_STATUS_IGNORED
|
||||
):
|
||||
await processor_task(
|
||||
worker_id,
|
||||
receive_stream,
|
||||
shutdown_event,
|
||||
client,
|
||||
username,
|
||||
task_status=task_status,
|
||||
)
|
||||
|
||||
await transport.run_consumers(
|
||||
tg, spawn_worker, settings.vector_sync_processor_workers
|
||||
)
|
||||
|
||||
# Expose this long-lived task group to request-path code that
|
||||
# wants to spawn background work (e.g. ADR-019 verify-on-read
|
||||
@@ -1752,7 +1763,9 @@ def get_app(transport: str = "streamable-http", enabled_apps: list[str] | None =
|
||||
|
||||
logger.info(
|
||||
"Background sync tasks started: 1 scanner + %s processors (queue=%s)",
|
||||
0 if use_postgres else settings.vector_sync_processor_workers,
|
||||
0
|
||||
if settings.ingest_queue == "postgres"
|
||||
else settings.vector_sync_processor_workers,
|
||||
settings.ingest_queue,
|
||||
)
|
||||
|
||||
@@ -1766,10 +1779,11 @@ def get_app(transport: str = "streamable-http", enabled_apps: list[str] | None =
|
||||
shutdown_event.set()
|
||||
# Request path must not spawn into a cancelling group.
|
||||
_vector_sync_state.eviction_task_group = None
|
||||
# Close the procrastinate connector pool (postgres mode).
|
||||
_drain = getattr(task_producer, "drain", None)
|
||||
if use_postgres and _drain is not None:
|
||||
await _drain()
|
||||
# Tear down backend-owned resources (closes the
|
||||
# procrastinate connector pool in postgres mode; no-op
|
||||
# for the memory stream, which task-group cancellation
|
||||
# closes).
|
||||
await transport.aclose()
|
||||
await client.close()
|
||||
# TaskGroup automatically cancels all tasks on exit
|
||||
|
||||
@@ -1877,65 +1891,26 @@ def get_app(transport: str = "streamable-http", enabled_apps: list[str] | None =
|
||||
except Exception as e:
|
||||
logger.warning("App password cleanup failed (non-fatal): %s", e)
|
||||
|
||||
# Initialize shared state. INGEST_QUEUE selects the transport
|
||||
# (Deck #183): ``memory`` uses the in-process anyio stream + the
|
||||
# in-process processor pool; ``postgres`` defers jobs via
|
||||
# procrastinate and runs no in-process consumer (the separate
|
||||
# ``worker`` role drains the queue).
|
||||
use_postgres = settings.ingest_queue == "postgres"
|
||||
# Initialize the ingest transport. INGEST_QUEUE selects the
|
||||
# backend (Deck #183, ADR-028): ``memory`` builds an in-process
|
||||
# anyio stream drained by an in-process pool; ``postgres`` defers
|
||||
# jobs via procrastinate and runs no in-process consumer (the
|
||||
# separate ``worker`` role drains the queue). The transport hides
|
||||
# that choice — this path is now backend-agnostic.
|
||||
shutdown_event = anyio.Event()
|
||||
scanner_wake_event = anyio.Event()
|
||||
|
||||
# User state tracking for user manager
|
||||
user_states: dict = {}
|
||||
|
||||
send_stream = None
|
||||
receive_stream = None
|
||||
task_producer: TaskProducer
|
||||
if use_postgres:
|
||||
# Single open/close cycle: build_producer opens the connector
|
||||
# and ensure_schema reuses it to create procrastinate's tables
|
||||
# before any scanner defers (matches the worker command).
|
||||
producer = await build_producer(settings)
|
||||
await producer.ensure_schema()
|
||||
task_producer = producer
|
||||
logger.info(
|
||||
"Ingest queue: postgres (procrastinate); worker drains it"
|
||||
)
|
||||
else:
|
||||
send_stream, receive_stream = anyio.create_memory_object_stream[
|
||||
DocumentTask
|
||||
](max_buffer_size=settings.vector_sync_queue_max_size)
|
||||
task_producer = MemoryTaskProducer(send_stream)
|
||||
transport = await build_transport(settings)
|
||||
task_producer = transport.producer
|
||||
|
||||
# Store in app state for access from routes (ADR-007)
|
||||
app.state.document_send_stream = send_stream
|
||||
app.state.document_receive_stream = receive_stream
|
||||
app.state.task_producer = task_producer
|
||||
app.state.shutdown_event = shutdown_event
|
||||
app.state.scanner_wake_event = scanner_wake_event
|
||||
|
||||
# Also store in module singleton for FastMCP session lifespans
|
||||
_vector_sync_state.document_send_stream = send_stream
|
||||
_vector_sync_state.document_receive_stream = receive_stream
|
||||
_vector_sync_state.task_producer = task_producer
|
||||
_vector_sync_state.shutdown_event = shutdown_event
|
||||
_vector_sync_state.scanner_wake_event = scanner_wake_event
|
||||
logger.info("Vector sync state stored in module singleton")
|
||||
|
||||
# Also share with browser_app for /app route
|
||||
for route in app.routes:
|
||||
if isinstance(route, Mount) and route.path == "/app":
|
||||
browser_app = cast(Starlette, route.app)
|
||||
browser_app.state.document_send_stream = send_stream
|
||||
browser_app.state.document_receive_stream = receive_stream
|
||||
browser_app.state.task_producer = task_producer
|
||||
browser_app.state.shutdown_event = shutdown_event
|
||||
browser_app.state.scanner_wake_event = scanner_wake_event
|
||||
logger.info(
|
||||
"Vector sync state shared with browser_app for /app"
|
||||
)
|
||||
break
|
||||
# Publish to app.state (ADR-007), the module singleton (FastMCP
|
||||
# session lifespans), and the /app browser sub-app in one place.
|
||||
_wire_vector_sync_state(
|
||||
app, transport, shutdown_event, scanner_wake_event
|
||||
)
|
||||
|
||||
# Background sync authenticates as each provisioned user via
|
||||
# locally-stored Nextcloud app passwords (Login Flow v2 /
|
||||
@@ -1960,18 +1935,31 @@ def get_app(transport: str = "streamable-http", enabled_apps: list[str] | None =
|
||||
tg,
|
||||
)
|
||||
|
||||
# In-process processor pool runs only in memory mode; in
|
||||
# postgres mode the out-of-process worker consumes.
|
||||
if not use_postgres:
|
||||
assert receive_stream is not None
|
||||
for i in range(settings.vector_sync_processor_workers):
|
||||
await tg.start(
|
||||
oauth_processor_task,
|
||||
i,
|
||||
receive_stream.clone(),
|
||||
shutdown_event,
|
||||
nextcloud_host_for_sync,
|
||||
)
|
||||
# Start the in-process consumer pool. ``run_consumers`` is a
|
||||
# no-op for the distributed (postgres) backend — the
|
||||
# out-of-process ``worker`` role consumes there. The closure
|
||||
# binds this mode's nextcloud_host (per-document credential
|
||||
# resolution) and forwards anyio's injected ``task_status``.
|
||||
# One shared receive stream + N workers ⇒ a single
|
||||
# multiplexed queue draining every user's documents with
|
||||
# N-way parallelism, never one user at a time (ADR-028).
|
||||
async def spawn_worker(
|
||||
worker_id,
|
||||
receive_stream,
|
||||
*,
|
||||
task_status=anyio.TASK_STATUS_IGNORED,
|
||||
):
|
||||
await oauth_processor_task(
|
||||
worker_id,
|
||||
receive_stream,
|
||||
shutdown_event,
|
||||
nextcloud_host_for_sync,
|
||||
task_status=task_status,
|
||||
)
|
||||
|
||||
await transport.run_consumers(
|
||||
tg, spawn_worker, settings.vector_sync_processor_workers
|
||||
)
|
||||
|
||||
# Expose this long-lived task group to request-path code
|
||||
# that wants to spawn background work (e.g. ADR-019
|
||||
@@ -1981,7 +1969,9 @@ def get_app(transport: str = "streamable-http", enabled_apps: list[str] | None =
|
||||
|
||||
logger.info(
|
||||
"Background sync tasks started: 1 user manager + %s processors (queue=%s)",
|
||||
0 if use_postgres else settings.vector_sync_processor_workers,
|
||||
0
|
||||
if settings.ingest_queue == "postgres"
|
||||
else settings.vector_sync_processor_workers,
|
||||
settings.ingest_queue,
|
||||
)
|
||||
|
||||
@@ -1995,10 +1985,10 @@ def get_app(transport: str = "streamable-http", enabled_apps: list[str] | None =
|
||||
shutdown_event.set()
|
||||
# Request path must not spawn into a cancelling group.
|
||||
_vector_sync_state.eviction_task_group = None
|
||||
# Close the procrastinate connector pool (postgres).
|
||||
_drain = getattr(task_producer, "drain", None)
|
||||
if use_postgres and _drain is not None:
|
||||
await _drain()
|
||||
# Tear down backend-owned resources (closes the
|
||||
# procrastinate connector pool in postgres mode;
|
||||
# no-op for the memory stream).
|
||||
await transport.aclose()
|
||||
# Close token broker HTTP client
|
||||
if token_broker._http_client:
|
||||
await token_broker._http_client.aclose()
|
||||
|
||||
@@ -325,6 +325,11 @@ def worker(concurrency: int | None):
|
||||
get_procrastinate_app,
|
||||
)
|
||||
|
||||
# This is the consumer side of the distributed (postgres) ingest backend.
|
||||
# Unlike the in-process anyio pool, the worker talks to procrastinate's App
|
||||
# directly (run_worker_async), so it does NOT go through IngestTransport —
|
||||
# DistributedTransport.run_consumers is a deliberate no-op precisely because
|
||||
# this separate process is the consumer (see vector/queue/transport.py).
|
||||
workers = concurrency or settings.vector_sync_processor_workers
|
||||
app = get_procrastinate_app()
|
||||
|
||||
|
||||
@@ -3,5 +3,21 @@
|
||||
from .factory import build_producer
|
||||
from .memory import MemoryTaskProducer
|
||||
from .ports import TaskProducer
|
||||
from .transport import (
|
||||
DistributedTransport,
|
||||
IngestTransport,
|
||||
LocalTransport,
|
||||
SpawnWorker,
|
||||
build_transport,
|
||||
)
|
||||
|
||||
__all__ = ["MemoryTaskProducer", "TaskProducer", "build_producer"]
|
||||
__all__ = [
|
||||
"DistributedTransport",
|
||||
"IngestTransport",
|
||||
"LocalTransport",
|
||||
"MemoryTaskProducer",
|
||||
"SpawnWorker",
|
||||
"TaskProducer",
|
||||
"build_producer",
|
||||
"build_transport",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,202 @@
|
||||
"""Ingest-path transport (design §10, hexagonal; Deck #183 follow-up, ADR-028).
|
||||
|
||||
The :class:`TaskProducer` port (see ``ports.py``) is *one* side of ingest — the
|
||||
sink the scanner/webhook send a ``DocumentTask`` to. An :class:`IngestTransport`
|
||||
is the composition object that owns *both* sides of one backend:
|
||||
|
||||
- the ``producer`` to wire into ``app.state`` / hand to the scanner, and
|
||||
- how the **consumer** side runs for *this* process.
|
||||
|
||||
Two adapters, selected by ``INGEST_QUEUE`` via :func:`build_transport`:
|
||||
|
||||
- :class:`LocalTransport` (``memory`` — the SQLite/dev default): an in-process
|
||||
anyio ``MemoryObjectStream`` drained by a pool of in-process workers that
|
||||
:meth:`run_consumers` starts.
|
||||
- :class:`DistributedTransport` (``postgres``): wraps the
|
||||
:class:`ProcrastinateTaskProducer`; :meth:`run_consumers` is a no-op because
|
||||
the consumer is a *separate* process — the ``nextcloud-mcp-server worker``
|
||||
role drains the queue (see ``cli.py``).
|
||||
|
||||
There is deliberately no consumer *port* (mirroring ``ports.py``): in memory
|
||||
mode the in-process pool is the consumer, in postgres mode the external worker
|
||||
is. The transport just encapsulates "build the producer + run (or don't run)
|
||||
the in-process consumers" so the server lifespan has a single branch-free shape
|
||||
and a new backend (Redis/NATS/SQS) drops in as one more adapter + one
|
||||
:func:`build_transport` arm — no ``app.py`` or scanner change.
|
||||
|
||||
Why an ABC here but a ``Protocol`` for ``TaskProducer``: the producer port is a
|
||||
Protocol so anyio's third-party ``MemoryObjectSendStream`` satisfies it
|
||||
structurally; the transport has exactly two in-house adapters that share the
|
||||
``producer`` storage and the ``receive_stream``/``run_consumers``/``aclose``
|
||||
defaults, so a concrete ABC is simpler and checks more cleanly under ``ty``.
|
||||
|
||||
The single-tenant parallelism invariant lives here: :class:`LocalTransport`
|
||||
hands every worker a ``clone()`` of *one* shared receive stream, so a tenant's
|
||||
users are processed by an N-worker pool off a single multiplexed queue
|
||||
(per-document, not per-user, dispatch) — never one user fully then the next.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import abc
|
||||
import logging
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import anyio
|
||||
from anyio.abc import TaskGroup
|
||||
from anyio.streams.memory import (
|
||||
MemoryObjectReceiveStream,
|
||||
MemoryObjectSendStream,
|
||||
)
|
||||
|
||||
from .factory import build_producer
|
||||
from .memory import MemoryTaskProducer
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ...config import Settings
|
||||
from ..scanner import DocumentTask
|
||||
from .ports import TaskProducer
|
||||
from .procrastinate import ProcrastinateTaskProducer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# A worker-spawn callback supplied by the lifespan: given a worker index and a
|
||||
# *fresh* receive handle, start one in-process consumer in the lifespan's task
|
||||
# group. The lifespan owns this closure so the transport never learns about auth
|
||||
# modes (single shared nc_client+username vs per-document credential resolution
|
||||
# by host). anyio's ``TaskGroup.start`` injects a ``task_status`` keyword, which
|
||||
# the closure must forward to the underlying ``processor_task`` /
|
||||
# ``multi_user_processor_task`` (else ``start`` blocks forever); hence ``...``.
|
||||
SpawnWorker = Callable[..., Awaitable[None]]
|
||||
|
||||
|
||||
class IngestTransport(abc.ABC):
|
||||
"""Owns one ingest backend: the producer + running its in-process consumers.
|
||||
|
||||
Built once per server lifespan by :func:`build_transport`. The
|
||||
:class:`TaskProducer` port is unchanged; this is a higher-level composition
|
||||
object owned only by the lifespan (the worker CLI talks to procrastinate
|
||||
directly — see module docstring).
|
||||
"""
|
||||
|
||||
@property
|
||||
@abc.abstractmethod
|
||||
def producer(self) -> TaskProducer:
|
||||
"""The :class:`TaskProducer` to wire into ``app.state`` / the scanner."""
|
||||
|
||||
@property
|
||||
def send_stream(self) -> MemoryObjectSendStream[DocumentTask] | None:
|
||||
"""Memory backend's raw send end; ``None`` for distributed backends.
|
||||
|
||||
Exposed only so the lifespan can keep populating
|
||||
``_vector_sync_state.document_send_stream`` (which the integration
|
||||
conftest saves/closes as a singleton). Producers send via
|
||||
:attr:`producer`, not this.
|
||||
"""
|
||||
return None
|
||||
|
||||
@property
|
||||
def receive_stream(self) -> MemoryObjectReceiveStream[DocumentTask] | None:
|
||||
"""Memory backend's receive end (queue-depth surface); ``None`` for
|
||||
distributed backends, which have no in-process stream (``ingest_status``
|
||||
reads procrastinate job counts via the producer instead)."""
|
||||
return None
|
||||
|
||||
async def run_consumers(
|
||||
self, task_group: TaskGroup, spawn_worker: SpawnWorker, count: int
|
||||
) -> None:
|
||||
"""Start the in-process consumer pool in ``task_group``.
|
||||
|
||||
No-op by default: distributed backends are drained by the external
|
||||
``worker`` role, so there is nothing to start in the API process.
|
||||
"""
|
||||
return None
|
||||
|
||||
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.
|
||||
"""
|
||||
return None
|
||||
|
||||
|
||||
class LocalTransport(IngestTransport):
|
||||
"""In-process anyio memory stream + processor pool (``INGEST_QUEUE=memory``).
|
||||
|
||||
Builds the paired send/receive streams up front and owns the receive end;
|
||||
:meth:`run_consumers` hands each worker an independent ``clone()`` so every
|
||||
receiver observes end-of-stream when the scanner's send handles all close.
|
||||
"""
|
||||
|
||||
def __init__(self, max_buffer_size: float):
|
||||
send_stream, receive_stream = anyio.create_memory_object_stream["DocumentTask"](
|
||||
max_buffer_size=max_buffer_size
|
||||
)
|
||||
self._send_stream = send_stream
|
||||
self._receive_stream = receive_stream
|
||||
self._producer = MemoryTaskProducer(send_stream)
|
||||
|
||||
@property
|
||||
def producer(self) -> TaskProducer:
|
||||
return self._producer
|
||||
|
||||
@property
|
||||
def send_stream(self) -> MemoryObjectSendStream[DocumentTask]:
|
||||
return self._send_stream
|
||||
|
||||
@property
|
||||
def receive_stream(self) -> MemoryObjectReceiveStream[DocumentTask]:
|
||||
return self._receive_stream
|
||||
|
||||
async def run_consumers(
|
||||
self, task_group: TaskGroup, spawn_worker: SpawnWorker, count: int
|
||||
) -> None:
|
||||
# One shared receive stream, N workers each draining a clone → a single
|
||||
# multiplexed queue processed with N-way parallelism across all users in
|
||||
# the tenant (per-document dispatch). ``start`` (not ``start_soon``)
|
||||
# waits for each worker's ``task_status.started()`` readiness, matching
|
||||
# the prior inline lifespan behaviour.
|
||||
for i in range(count):
|
||||
await task_group.start(spawn_worker, i, self._receive_stream.clone())
|
||||
|
||||
|
||||
class DistributedTransport(IngestTransport):
|
||||
"""Postgres/procrastinate producer; consumers are the external worker role.
|
||||
|
||||
The producer's connector pool is opened by :func:`build_producer` and owned
|
||||
by the lifespan; :meth:`run_consumers` is the inherited no-op (the
|
||||
``nextcloud-mcp-server worker`` process drains the queue) and :meth:`aclose`
|
||||
closes the pool once on shutdown.
|
||||
"""
|
||||
|
||||
def __init__(self, producer: ProcrastinateTaskProducer):
|
||||
self._producer = producer
|
||||
|
||||
@property
|
||||
def producer(self) -> TaskProducer:
|
||||
return self._producer
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self._producer.drain()
|
||||
|
||||
|
||||
async def build_transport(settings: Settings) -> IngestTransport:
|
||||
"""Build the ingest transport for the configured ``INGEST_QUEUE`` backend.
|
||||
|
||||
- ``postgres`` → :class:`DistributedTransport`. Reuses :func:`build_producer`
|
||||
(which opens the connector pool) and applies procrastinate's schema once on
|
||||
that same open pool before any scanner can defer — a single open/close
|
||||
cycle, matching the ``worker`` command.
|
||||
- ``memory`` (SQLite/dev default) → :class:`LocalTransport`.
|
||||
"""
|
||||
if settings.ingest_queue == "postgres":
|
||||
producer = await build_producer(settings)
|
||||
await producer.ensure_schema()
|
||||
logger.info("Ingest queue: postgres (procrastinate); worker drains it")
|
||||
return DistributedTransport(producer)
|
||||
|
||||
return LocalTransport(max_buffer_size=settings.vector_sync_queue_max_size)
|
||||
Reference in New Issue
Block a user