refactor: migrate vector sync from asyncio.Queue to anyio memory object streams

Replace asyncio.Queue with anyio.create_memory_object_stream() throughout
the vector sync system for better library consistency and improved shutdown
semantics.

## Changes Made

**scanner.py**:
- Changed parameter type from `asyncio.Queue` to `MemoryObjectSendStream[DocumentTask]`
- Replaced all `await document_queue.put()` calls with `await send_stream.send()`
- Wrapped scanner loop in `async with send_stream:` context manager for automatic cleanup
- Updated log messages: "Queued" → "Sent"
- Removed `import asyncio` (no longer needed)

**processor.py**:
- Changed parameter type from `asyncio.Queue` to `MemoryObjectReceiveStream[DocumentTask]`
- Replaced `asyncio.wait_for(document_queue.get(), timeout=1.0)` with `anyio.fail_after(1.0)` + `await receive_stream.receive()`
- Removed all `document_queue.task_done()` calls (not needed with streams)
- Added `anyio.EndOfStream` exception handling for graceful shutdown when scanner closes
- Removed `import asyncio` (no longer needed)

**app.py**:
- Removed `import asyncio` from top-level imports
- Added `from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream`
- Updated AppContext dataclass:
  - Replaced `document_queue: Optional[asyncio.Queue]` with:
    - `document_send_stream: Optional[MemoryObjectSendStream]`
    - `document_receive_stream: Optional[MemoryObjectReceiveStream]`
- Updated `app_lifespan_basic()`:
  - Replaced `asyncio.Queue(maxsize=...)` with `anyio.create_memory_object_stream(max_buffer_size=...)`
  - Pass `send_stream` to scanner_task
  - Pass `receive_stream.clone()` to each processor_task (enables multiple consumers)
  - Updated AppContext yield to include both streams
- Updated `starlette_lifespan()`:
  - Same changes as app_lifespan_basic for streamable-http transport
  - Removed `import asyncio as asyncio_module` (no longer needed)
  - Updated app.state storage to use send_stream and receive_stream

**semantic.py**:
- Updated `nc_get_vector_sync_status()` tool:
  - Access `document_receive_stream` instead of `document_queue` from lifespan context
  - Use `stream_stats.current_buffer_used` instead of `queue.qsize()` for pending count
  - More reliable metrics (qsize() was not guaranteed accurate)

## Benefits

1. **Library Consistency**: Pure anyio throughout codebase (was mixing asyncio.Queue with anyio.Event and anyio.create_task_group)
2. **Graceful Shutdown**: `async with send_stream:` automatically closes stream on exit, signaling EndOfStream to all processors
3. **Better Timeout Handling**: `anyio.fail_after()` is more idiomatic than `asyncio.wait_for()`
4. **Stream Cloning**: Easy to add multiple consumers via `receive_stream.clone()`
5. **Better Statistics**: `.statistics()` provides accurate buffer metrics (qsize() was unreliable)
6. **Type Safety**: Separate send/receive types prevent accidental misuse
7. **No task_done() tracking**: Streams handle completion automatically

## Testing

-  All 69 unit tests passing
-  All 5 smoke tests passing
-  No regressions in functionality
-  Graceful shutdown behavior improved

## References

- https://anyio.readthedocs.io/en/stable/why.html#queue-fix
- https://anyio.readthedocs.io/en/stable/streams.html#memory-object-streams

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Chris Coutinho
2025-11-09 06:43:44 +01:00
co-authored by Claude
parent 4b026e9aa0
commit 72232f937a
4 changed files with 83 additions and 78 deletions
+20 -16
View File
@@ -1,4 +1,3 @@
import asyncio
import logging import logging
import os import os
from collections.abc import AsyncIterator from collections.abc import AsyncIterator
@@ -13,6 +12,7 @@ import anyio
import click import click
import httpx import httpx
import uvicorn import uvicorn
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from mcp.server.auth.settings import AuthSettings from mcp.server.auth.settings import AuthSettings
from mcp.server.fastmcp import Context, FastMCP from mcp.server.fastmcp import Context, FastMCP
from pydantic import AnyHttpUrl from pydantic import AnyHttpUrl
@@ -211,7 +211,8 @@ class AppContext:
"""Application context for BasicAuth mode.""" """Application context for BasicAuth mode."""
client: NextcloudClient client: NextcloudClient
document_queue: Optional[asyncio.Queue] = None document_send_stream: Optional[MemoryObjectSendStream] = None
document_receive_stream: Optional[MemoryObjectReceiveStream] = None
shutdown_event: Optional[anyio.Event] = None shutdown_event: Optional[anyio.Event] = None
scanner_wake_event: Optional[anyio.Event] = None scanner_wake_event: Optional[anyio.Event] = None
@@ -404,7 +405,9 @@ async def app_lifespan_basic(server: FastMCP) -> AsyncIterator[AppContext]:
) )
# Initialize shared state # Initialize shared state
document_queue = asyncio.Queue(maxsize=settings.vector_sync_queue_max_size) send_stream, receive_stream = anyio.create_memory_object_stream(
max_buffer_size=settings.vector_sync_queue_max_size
)
shutdown_event = anyio.Event() shutdown_event = anyio.Event()
scanner_wake_event = anyio.Event() scanner_wake_event = anyio.Event()
@@ -413,19 +416,19 @@ async def app_lifespan_basic(server: FastMCP) -> AsyncIterator[AppContext]:
# Start scanner task # Start scanner task
tg.start_soon( tg.start_soon(
scanner_task, scanner_task,
document_queue, send_stream,
shutdown_event, shutdown_event,
scanner_wake_event, scanner_wake_event,
client, client,
username, username,
) )
# Start processor pool # Start processor pool (each gets a cloned receive stream)
for i in range(settings.vector_sync_processor_workers): for i in range(settings.vector_sync_processor_workers):
tg.start_soon( tg.start_soon(
processor_task, processor_task,
i, i,
document_queue, receive_stream.clone(),
shutdown_event, shutdown_event,
client, client,
username, username,
@@ -439,7 +442,8 @@ async def app_lifespan_basic(server: FastMCP) -> AsyncIterator[AppContext]:
try: try:
yield AppContext( yield AppContext(
client=client, client=client,
document_queue=document_queue, document_send_stream=send_stream,
document_receive_stream=receive_stream,
shutdown_event=shutdown_event, shutdown_event=shutdown_event,
scanner_wake_event=scanner_wake_event, scanner_wake_event=scanner_wake_event,
) )
@@ -1009,8 +1013,6 @@ def get_app(transport: str = "sse", enabled_apps: list[str] | None = None):
# Start background vector sync tasks for BasicAuth mode (ADR-007) # Start background vector sync tasks for BasicAuth mode (ADR-007)
# For streamable-http transport, FastMCP lifespan isn't automatically triggered # For streamable-http transport, FastMCP lifespan isn't automatically triggered
# so we manually start background tasks here if vector sync is enabled # so we manually start background tasks here if vector sync is enabled
import asyncio as asyncio_module
import anyio as anyio_module import anyio as anyio_module
settings = get_settings() settings = get_settings()
@@ -1029,21 +1031,23 @@ def get_app(transport: str = "sse", enabled_apps: list[str] | None = None):
client = NextcloudClient.from_env() client = NextcloudClient.from_env()
# Initialize shared state # Initialize shared state
document_queue = asyncio_module.Queue( send_stream, receive_stream = anyio_module.create_memory_object_stream(
maxsize=settings.vector_sync_queue_max_size max_buffer_size=settings.vector_sync_queue_max_size
) )
shutdown_event = anyio_module.Event() shutdown_event = anyio_module.Event()
scanner_wake_event = anyio_module.Event() scanner_wake_event = anyio_module.Event()
# Store in app state for access from routes (ADR-007) # Store in app state for access from routes (ADR-007)
app.state.document_queue = document_queue app.state.document_send_stream = send_stream
app.state.document_receive_stream = receive_stream
app.state.shutdown_event = shutdown_event app.state.shutdown_event = shutdown_event
app.state.scanner_wake_event = scanner_wake_event app.state.scanner_wake_event = scanner_wake_event
# Also share with browser_app for /user/page route # Also share with browser_app for /user/page route
for route in app.routes: for route in app.routes:
if isinstance(route, Mount) and route.path == "/user": if isinstance(route, Mount) and route.path == "/user":
route.app.state.document_queue = document_queue route.app.state.document_send_stream = send_stream
route.app.state.document_receive_stream = receive_stream
route.app.state.shutdown_event = shutdown_event route.app.state.shutdown_event = shutdown_event
route.app.state.scanner_wake_event = scanner_wake_event route.app.state.scanner_wake_event = scanner_wake_event
logger.info( logger.info(
@@ -1056,19 +1060,19 @@ def get_app(transport: str = "sse", enabled_apps: list[str] | None = None):
# Start scanner task # Start scanner task
tg.start_soon( tg.start_soon(
scanner_task, scanner_task,
document_queue, send_stream,
shutdown_event, shutdown_event,
scanner_wake_event, scanner_wake_event,
client, client,
username, username,
) )
# Start processor pool # Start processor pool (each gets a cloned receive stream)
for i in range(settings.vector_sync_processor_workers): for i in range(settings.vector_sync_processor_workers):
tg.start_soon( tg.start_soon(
processor_task, processor_task,
i, i,
document_queue, receive_stream.clone(),
shutdown_event, shutdown_event,
client, client,
username, username,
+11 -6
View File
@@ -381,12 +381,16 @@ def configure_semantic_tools(mcp: FastMCP):
) )
try: try:
# Get document queue from lifespan context # Get document receive stream from lifespan context
lifespan_ctx = ctx.request_context.lifespan_context lifespan_ctx = ctx.request_context.lifespan_context
document_queue = getattr(lifespan_ctx, "document_queue", None) document_receive_stream = getattr(
lifespan_ctx, "document_receive_stream", None
)
if document_queue is None: if document_receive_stream is None:
logger.debug("document_queue not available in lifespan context") logger.debug(
"document_receive_stream not available in lifespan context"
)
return VectorSyncStatusResponse( return VectorSyncStatusResponse(
indexed_count=0, indexed_count=0,
pending_count=0, pending_count=0,
@@ -394,8 +398,9 @@ def configure_semantic_tools(mcp: FastMCP):
enabled=True, enabled=True,
) )
# Get pending count from queue # Get pending count from stream statistics
pending_count = document_queue.qsize() stream_stats = document_receive_stream.statistics()
pending_count = stream_stats.current_buffer_used
# Get Qdrant client and query indexed count # Get Qdrant client and query indexed count
indexed_count = 0 indexed_count = 0
+15 -20
View File
@@ -1,14 +1,14 @@
"""Processor task for vector database synchronization. """Processor task for vector database synchronization.
Processes documents from queue: fetches content, generates embeddings, stores in Qdrant. Processes documents from stream: fetches content, generates embeddings, stores in Qdrant.
""" """
import asyncio
import logging import logging
import time import time
import uuid import uuid
import anyio import anyio
from anyio.streams.memory import MemoryObjectReceiveStream
from httpx import HTTPStatusError from httpx import HTTPStatusError
from qdrant_client.models import FieldCondition, Filter, MatchValue, PointStruct from qdrant_client.models import FieldCondition, Filter, MatchValue, PointStruct
@@ -24,27 +24,26 @@ logger = logging.getLogger(__name__)
async def processor_task( async def processor_task(
worker_id: int, worker_id: int,
document_queue: asyncio.Queue, receive_stream: MemoryObjectReceiveStream[DocumentTask],
shutdown_event: anyio.Event, shutdown_event: anyio.Event,
nc_client: NextcloudClient, nc_client: NextcloudClient,
user_id: str, user_id: str,
): ):
""" """
Process documents from queue concurrently. Process documents from stream concurrently.
Each processor task runs in a loop: Each processor task runs in a loop:
1. Pull document from queue (with timeout) 1. Receive document from stream (with timeout)
2. Fetch content from Nextcloud 2. Fetch content from Nextcloud
3. Tokenize and chunk text 3. Tokenize and chunk text
4. Generate embeddings (I/O bound - external API) 4. Generate embeddings (I/O bound - external API)
5. Upload vectors to Qdrant 5. Upload vectors to Qdrant
6. Mark task complete
Multiple processors run concurrently for I/O parallelism. Multiple processors run concurrently for I/O parallelism.
Args: Args:
worker_id: Worker identifier for logging worker_id: Worker identifier for logging
document_queue: Queue to pull documents from receive_stream: Stream to receive documents from
shutdown_event: Event signaling shutdown shutdown_event: Event signaling shutdown
nc_client: Authenticated Nextcloud client nc_client: Authenticated Nextcloud client
user_id: User being processed user_id: User being processed
@@ -54,32 +53,28 @@ async def processor_task(
while not shutdown_event.is_set(): while not shutdown_event.is_set():
try: try:
# Get document with timeout (allows checking shutdown) # Get document with timeout (allows checking shutdown)
doc_task = await asyncio.wait_for( with anyio.fail_after(1.0):
document_queue.get(), doc_task = await receive_stream.receive()
timeout=1.0,
)
# Process document # Process document
await process_document(doc_task, nc_client) await process_document(doc_task, nc_client)
# Mark complete except TimeoutError:
document_queue.task_done()
except asyncio.TimeoutError:
# No documents available, continue # No documents available, continue
continue continue
except anyio.EndOfStream:
# Scanner finished and closed stream, exit gracefully
logger.info(f"Processor {worker_id}: Scanner finished, exiting")
break
except Exception as e: except Exception as e:
logger.error( logger.error(
f"Processor {worker_id} error processing " f"Processor {worker_id} error processing "
f"{doc_task.doc_type}_{doc_task.doc_id}: {e}", f"{doc_task.doc_type}_{doc_task.doc_id}: {e}",
exc_info=True, exc_info=True,
) )
# Mark task done even on error to prevent queue blocking # Continue to next document (no task_done() needed with streams)
try:
document_queue.task_done()
except ValueError:
pass
logger.info(f"Processor {worker_id} stopped") logger.info(f"Processor {worker_id} stopped")
+20 -19
View File
@@ -3,12 +3,12 @@
Periodically scans enabled users' content and queues changed documents for processing. Periodically scans enabled users' content and queues changed documents for processing.
""" """
import asyncio
import logging import logging
import time import time
from dataclasses import dataclass from dataclasses import dataclass
import anyio import anyio
from anyio.streams.memory import MemoryObjectSendStream
from qdrant_client.models import FieldCondition, Filter, MatchValue from qdrant_client.models import FieldCondition, Filter, MatchValue
from nextcloud_mcp_server.client import NextcloudClient from nextcloud_mcp_server.client import NextcloudClient
@@ -35,7 +35,7 @@ _potentially_deleted: dict[tuple[str, str], float] = {}
async def scanner_task( async def scanner_task(
document_queue: asyncio.Queue, send_stream: MemoryObjectSendStream[DocumentTask],
shutdown_event: anyio.Event, shutdown_event: anyio.Event,
wake_event: anyio.Event, wake_event: anyio.Event,
nc_client: NextcloudClient, nc_client: NextcloudClient,
@@ -47,7 +47,7 @@ async def scanner_task(
For BasicAuth mode, scans a single user with credentials available at runtime. For BasicAuth mode, scans a single user with credentials available at runtime.
Args: Args:
document_queue: Queue to enqueue changed documents send_stream: Stream to send changed documents to processors
shutdown_event: Event signaling shutdown shutdown_event: Event signaling shutdown
wake_event: Event to trigger immediate scan wake_event: Event to trigger immediate scan
nc_client: Authenticated Nextcloud client nc_client: Authenticated Nextcloud client
@@ -56,12 +56,13 @@ async def scanner_task(
logger.info(f"Scanner task started for user: {user_id}") logger.info(f"Scanner task started for user: {user_id}")
settings = get_settings() settings = get_settings()
async with send_stream:
while not shutdown_event.is_set(): while not shutdown_event.is_set():
try: try:
# Scan user documents # Scan user documents
await scan_user_documents( await scan_user_documents(
user_id=user_id, user_id=user_id,
document_queue=document_queue, send_stream=send_stream,
nc_client=nc_client, nc_client=nc_client,
) )
@@ -77,23 +78,23 @@ async def scanner_task(
# Shutdown, exit loop # Shutdown, exit loop
break break
logger.info("Scanner task stopped") logger.info("Scanner task stopped - stream closed")
async def scan_user_documents( async def scan_user_documents(
user_id: str, user_id: str,
document_queue: asyncio.Queue, send_stream: MemoryObjectSendStream[DocumentTask],
nc_client: NextcloudClient, nc_client: NextcloudClient,
initial_sync: bool = False, initial_sync: bool = False,
): ):
""" """
Scan a single user's documents and queue changes. Scan a single user's documents and send changes to processor stream.
Args: Args:
user_id: User to scan user_id: User to scan
document_queue: Queue to enqueue changed documents send_stream: Stream to send changed documents to processors
nc_client: Authenticated Nextcloud client nc_client: Authenticated Nextcloud client
initial_sync: If True, queue all documents (first-time sync) initial_sync: If True, send all documents (first-time sync)
""" """
logger.info(f"Scanning documents for user: {user_id}") logger.info(f"Scanning documents for user: {user_id}")
@@ -102,9 +103,9 @@ async def scan_user_documents(
logger.debug(f"Found {len(notes)} notes for {user_id}") logger.debug(f"Found {len(notes)} notes for {user_id}")
if initial_sync: if initial_sync:
# Queue everything on first sync # Send everything on first sync
for note in notes: for note in notes:
await document_queue.put( await send_stream.send(
DocumentTask( DocumentTask(
user_id=user_id, user_id=user_id,
doc_id=str(note["id"]), doc_id=str(note["id"]),
@@ -113,7 +114,7 @@ async def scan_user_documents(
modified_at=note["modified"], modified_at=note["modified"],
) )
) )
logger.info(f"Queued {len(notes)} documents for initial sync: {user_id}") logger.info(f"Sent {len(notes)} documents for initial sync: {user_id}")
return return
# Get indexed state from Qdrant # Get indexed state from Qdrant
@@ -154,9 +155,9 @@ async def scan_user_documents(
) )
del _potentially_deleted[doc_key] del _potentially_deleted[doc_key]
# Queue if never indexed or modified since last index # Send if never indexed or modified since last index
if indexed_at is None or note["modified"] > indexed_at: if indexed_at is None or note["modified"] > indexed_at:
await document_queue.put( await send_stream.send(
DocumentTask( DocumentTask(
user_id=user_id, user_id=user_id,
doc_id=doc_id, doc_id=doc_id,
@@ -183,12 +184,12 @@ async def scan_user_documents(
time_missing = current_time - first_missing_time time_missing = current_time - first_missing_time
if time_missing >= grace_period: if time_missing >= grace_period:
# Grace period elapsed, queue for deletion # Grace period elapsed, send for deletion
logger.info( logger.info(
f"Document {doc_id} missing for {time_missing:.1f}s " f"Document {doc_id} missing for {time_missing:.1f}s "
f"(>{grace_period:.1f}s grace period), queueing deletion" f"(>{grace_period:.1f}s grace period), sending deletion"
) )
await document_queue.put( await send_stream.send(
DocumentTask( DocumentTask(
user_id=user_id, user_id=user_id,
doc_id=doc_id, doc_id=doc_id,
@@ -198,7 +199,7 @@ async def scan_user_documents(
) )
) )
queued += 1 queued += 1
# Remove from tracking after queueing deletion # Remove from tracking after sending deletion
del _potentially_deleted[doc_key] del _potentially_deleted[doc_key]
else: else:
logger.debug( logger.debug(
@@ -213,6 +214,6 @@ async def scan_user_documents(
_potentially_deleted[doc_key] = current_time _potentially_deleted[doc_key] = current_time
if queued > 0: if queued > 0:
logger.info(f"Queued {queued} documents for incremental sync: {user_id}") logger.info(f"Sent {queued} documents for incremental sync: {user_id}")
else: else:
logger.debug(f"No changes detected for {user_id}") logger.debug(f"No changes detected for {user_id}")