Merge pull request #15 from rhel-lightspeed/chore/structlog-and-lock-cleanup

chore: adopt structlog and add asyncio.Lock to engine singleton

Author: major

CHANGELOG.md

@@ -24,9 +24,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
-- Deferred heavy `docs2db-api` imports (torch, transformers) in `engine.py` to reduce module load time
-- Migrated startup health check and engine shutdown to FastMCP lifespan hook for proper async lifecycle management
-- Simplified `__main__.py` — removed manual `asyncio.run()` calls for health check and cleanup
+- Deferred heavy `docs2db-api` imports (torch, transformers) in `engine.py` to reduce module load time (RSPEED-3047)
+- Migrated startup health check and engine shutdown to FastMCP lifespan hook for proper async lifecycle management (RSPEED-3047)
+- Simplified `__main__.py` — removed manual `asyncio.run()` calls for health check and cleanup (RSPEED-3047)
+- Adopted structlog throughout `engine.py`, `server.py`, `search_documents.py` — consistent keyword-arg logging, no f-strings
+- Added `asyncio.Lock` to `get_engine()` singleton to guard against concurrent initialization
+- `shutdown_engine()` now awaits `engine.close()` before clearing the singleton reference
 - Upgraded fastmcp from 2.x to 3.x (`>=3.3.1, <4`)
 - Removed standalone `mcp` dependency (pulled in transitively by fastmcp 3.x)
 - Updated smoke tests to match FastMCP 3.x decorator behavior (`@mcp.tool` now returns the original function)

src/docs2db_mcp/__main__.py

@@ -4,20 +4,20 @@
 import os
 import sys
 
+import structlog
+
 
 # Must read transport before importing modules that configure logging
 transport = os.environ.get("DOCS2DB_MCP_TRANSPORT", "sse")
 
 if transport == "sse":
-    # stdout not used by MCP protocol, logging can go there
     logging.basicConfig(
         level=logging.INFO,
         format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
         stream=sys.stderr,
     )
 else:
     # stdout is reserved for MCP protocol, must redirect all logging to stderr
-    # set CRITICAL level before importing docs2db-api to minimize startup logs
     os.environ.setdefault("DOCS2DB_LOG_LEVEL", "CRITICAL")
     logging.basicConfig(
         level=logging.INFO,
@@ -30,27 +30,28 @@
 from docs2db_mcp.server import mcp  # noqa: E402
 
 
-logger = logging.getLogger(__name__)
+logger = structlog.get_logger(__name__)
 
 
 def main() -> None:
     """Run the MCP server."""
-    logger.info("Starting docs2db MCP server on %s:%s", CONFIG.host, CONFIG.port)
-    logger.info("Transport: %s", CONFIG.transport)
-    logger.info("Database: %s:%s/%s", CONFIG.db_host, CONFIG.db_port, CONFIG.db_database)
     logger.info(
-        "RAG settings: threshold=%s, max_chunks=%s, reranking=%s",
-        CONFIG.rag_similarity_threshold,
-        CONFIG.rag_max_chunks,
-        CONFIG.rag_enable_reranking,
+        "Starting docs2db MCP server",
+        host=CONFIG.host,
+        port=CONFIG.port,
+        transport=CONFIG.transport,
+        database=f"{CONFIG.db_host}:{CONFIG.db_port}/{CONFIG.db_database}",
+        rag_threshold=CONFIG.rag_similarity_threshold,
+        rag_max_chunks=CONFIG.rag_max_chunks,
+        rag_reranking=CONFIG.rag_enable_reranking,
     )
 
     try:
         mcp.run(transport=CONFIG.transport, **CONFIG.transport_kwargs)
     except KeyboardInterrupt:
         logger.info("Received keyboard interrupt")
     except Exception as e:
-        logger.error("Server error: %s", e, exc_info=True)
+        logger.error("Server error", error=str(e), exc_info=True)
         sys.exit(1)
 
 

src/docs2db_mcp/engine.py

@@ -6,20 +6,23 @@
 
 from __future__ import annotations
 
-import logging
+import asyncio
 
 from typing import TYPE_CHECKING
 
+import structlog
+
 from docs2db_mcp.config import CONFIG
 
 
 if TYPE_CHECKING:
     from docs2db_api.rag.engine import UniversalRAGEngine
 
 
-logger = logging.getLogger(__name__)
+logger = structlog.get_logger(__name__)
 
 _engine: UniversalRAGEngine | None = None
+_engine_lock = asyncio.Lock()
 
 
 async def get_engine() -> UniversalRAGEngine:
@@ -37,41 +40,39 @@ async def get_engine() -> UniversalRAGEngine:
     """
     global _engine
 
-    if _engine is None:
-        from docs2db_api.rag.engine import RAGConfig
-        from docs2db_api.rag.engine import UniversalRAGEngine
-
-        logger.info("Initializing UniversalRAGEngine")
-
-        # Configure database connection (dict format)
-        db_config = {
-            "host": CONFIG.db_host,
-            "port": str(CONFIG.db_port),
-            "database": CONFIG.db_database,
-            "user": CONFIG.db_user,
-            "password": CONFIG.db_password,
-        }
-
-        # Configure RAG (no LLM config - refinement disabled)
-        rag_config = RAGConfig()
-        rag_config.similarity_threshold = CONFIG.rag_similarity_threshold
-        rag_config.max_chunks = CONFIG.rag_max_chunks
-        rag_config.enable_reranking = CONFIG.rag_enable_reranking
-        rag_config.enable_question_refinement = False  # Disabled for tool calling
-
-        # Create and initialize engine
-        _engine = UniversalRAGEngine(
-            config=rag_config,
-            db_config=db_config,
-        )
-
-        try:
-            await _engine.start()
-            logger.info("UniversalRAGEngine initialized successfully")
-        except Exception as e:
-            logger.error("Failed to initialize RAG engine: %s", e)
-            _engine = None
-            raise
+    async with _engine_lock:
+        if _engine is None:
+            from docs2db_api.rag.engine import RAGConfig
+            from docs2db_api.rag.engine import UniversalRAGEngine
+
+            logger.info("Initializing UniversalRAGEngine")
+
+            db_config = {
+                "host": CONFIG.db_host,
+                "port": str(CONFIG.db_port),
+                "database": CONFIG.db_database,
+                "user": CONFIG.db_user,
+                "password": CONFIG.db_password,
+            }
+
+            rag_config = RAGConfig()
+            rag_config.similarity_threshold = CONFIG.rag_similarity_threshold
+            rag_config.max_chunks = CONFIG.rag_max_chunks
+            rag_config.enable_reranking = CONFIG.rag_enable_reranking
+            rag_config.enable_question_refinement = False
+
+            _engine = UniversalRAGEngine(
+                config=rag_config,
+                db_config=db_config,
+            )
+
+            try:
+                await _engine.start()
+                logger.info("UniversalRAGEngine initialized successfully")
+            except Exception as e:
+                logger.error("Failed to initialize RAG engine", error=str(e))
+                _engine = None
+                raise
 
     return _engine
 
@@ -80,11 +81,13 @@ async def shutdown_engine() -> None:
     """Shutdown the RAG engine and cleanup resources."""
     global _engine
 
-    if _engine is not None:
-        logger.info("Shutting down UniversalRAGEngine")
-        # UniversalRAGEngine doesn't have a stop() method
-        # Just set to None to allow garbage collection
-        _engine = None
+    async with _engine_lock:
+        if _engine is not None:
+            logger.info("Shutting down UniversalRAGEngine")
+            try:
+                await _engine.close()
+            finally:
+                _engine = None
 
 
 async def health_check() -> None:
@@ -96,27 +99,23 @@ async def health_check() -> None:
     logger.info("Performing startup health check...")
 
     try:
-        # Initialize engine and connect to database
         engine = await get_engine()
         logger.info("Database connection established")
 
-        # Perform a simple test query to verify the system works
-        test_query = "test"
         result = await engine.search_documents(
-            query=test_query,
+            query="test",
             max_chunks=1,
             similarity_threshold=0.0,
             enable_reranking=False,
         )
 
-        # Verify we got a valid response
         if result is None:
             msg = "Health check query returned None"
             raise Exception(msg)
 
-        logger.info("Test query successful (returned %d documents)", len(result.documents))
+        logger.info("Test query successful", document_count=len(result.documents))
         logger.info("Health check passed - system is ready")
 
     except Exception as e:
-        logger.error("Health check failed: %s", e, exc_info=True)
+        logger.error("Health check failed", error=str(e), exc_info=True)
         raise Exception(f"Startup health check failed - cannot connect to database or perform queries: {e}") from e

src/docs2db_mcp/server.py

@@ -1,19 +1,20 @@
 """FastMCP server instance and initialization."""
 
-import logging
 import sys
 
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 
+import structlog
+
 from fastmcp import FastMCP
 
 from docs2db_mcp.config import CONFIG
 from docs2db_mcp.engine import health_check
 from docs2db_mcp.engine import shutdown_engine
 
 
-logger = logging.getLogger(__name__)
+logger = structlog.get_logger(__name__)
 
 
 @asynccontextmanager

src/docs2db_mcp/tools/search_documents.py

@@ -1,16 +1,16 @@
 """search_documents tool implementation."""
 
-import logging
-
 from typing import Any
 
+import structlog
+
 from mcp.types import ToolAnnotations
 
 from docs2db_mcp.engine import get_engine
 from docs2db_mcp.server import mcp
 
 
-logger = logging.getLogger(__name__)
+logger = structlog.get_logger(__name__)
 
 
 @mcp.tool(
@@ -46,7 +46,7 @@ async def search_documents(
         - query_used: The original query (refinement disabled)
         - num_results: Number of results returned
     """
-    logger.info(f"Searching for: {query!r} (max_chunks={max_chunks})")
+    logger.info("Searching", query_length=len(query), max_chunks=max_chunks)
 
     try:
         engine = await get_engine()
@@ -73,7 +73,7 @@ async def search_documents(
             for doc in result.documents
         ]
 
-        logger.info(f"Found {len(chunks)} results")
+        logger.info("Search complete", num_results=len(chunks))
 
         return {
             "chunks": chunks,
@@ -82,7 +82,7 @@ async def search_documents(
         }
 
     except Exception as e:
-        logger.error(f"Search failed: {e}", exc_info=True)
+        logger.error("Search failed", error=str(e), exc_info=True)
         return {
             "chunks": [],
             "query_used": query,