Add a retention_days option to schedule database cleanup.

- Fixes the MDB_MAP_FULL error.

Author: luuquangvu

app/main.py

@@ -1,3 +1,4 @@
+import asyncio
 from contextlib import asynccontextmanager
 
 from fastapi import FastAPI
@@ -6,21 +7,64 @@
 from .server.chat import router as chat_router
 from .server.health import router as health_router
 from .server.middleware import add_cors_middleware, add_exception_handler
-from .services.pool import GeminiClientPool
+from .services import GeminiClientPool, LMDBConversationStore
+
+RETENTION_CLEANUP_INTERVAL_SECONDS = 6 * 60 * 60  # 6 hours
+
+
+async def _run_retention_cleanup(stop_event: asyncio.Event) -> None:
+    """
+    Periodically enforce LMDB retention policy until the stop_event is set.
+    """
+    store = LMDBConversationStore()
+    if store.retention_days <= 0:
+        logger.info("LMDB retention cleanup disabled; skipping scheduler.")
+        return
+
+    logger.info(
+        f"Starting LMDB retention cleanup task (retention={store.retention_days} day(s), interval={RETENTION_CLEANUP_INTERVAL_SECONDS} seconds)."
+    )
+
+    while not stop_event.is_set():
+        try:
+            store.cleanup_expired()
+        except Exception:
+            logger.exception("LMDB retention cleanup task failed.")
+
+        try:
+            await asyncio.wait_for(
+                stop_event.wait(),
+                timeout=RETENTION_CLEANUP_INTERVAL_SECONDS,
+            )
+        except asyncio.TimeoutError:
+            continue
+
+    logger.info("LMDB retention cleanup task stopped.")
 
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
+    cleanup_stop_event = asyncio.Event()
+    cleanup_task: asyncio.Task | None = None
     try:
         pool = GeminiClientPool()
         await pool.init()
+
+        cleanup_task = asyncio.create_task(_run_retention_cleanup(cleanup_stop_event))
+
+        logger.info(f"Gemini clients initialized: {[c.id for c in pool.clients]}.")
+        logger.info("Gemini API Server ready to serve requests.")
+        yield
     except Exception as e:
         logger.exception(f"Failed to initialize Gemini clients: {e}")
         raise
-
-    logger.success(f"Gemini clients initialized: {[c.id for c in pool.clients]}.")
-    logger.success("Gemini API Server ready to serve requests.")
-    yield
+    finally:
+        cleanup_stop_event.set()
+        if cleanup_task:
+            try:
+                await cleanup_task
+            except asyncio.CancelledError:
+                logger.debug("LMDB retention cleanup task cancelled during shutdown.")
 
 
 def create_app() -> FastAPI:

app/services/lmdb.py

@@ -1,7 +1,7 @@
 import hashlib
 import re
 from contextlib import contextmanager
-from datetime import datetime
+from datetime import datetime, timedelta
 from pathlib import Path
 from typing import Any, Dict, List, Optional
 
@@ -39,7 +39,12 @@ class LMDBConversationStore(metaclass=Singleton):
 
     HASH_LOOKUP_PREFIX = "hash:"
 
-    def __init__(self, db_path: Optional[str] = None, max_db_size: Optional[int] = None):
+    def __init__(
+        self,
+        db_path: Optional[str] = None,
+        max_db_size: Optional[int] = None,
+        retention_days: Optional[int] = None,
+    ):
         """
         Initialize LMDB store.
 
@@ -52,9 +57,12 @@ def __init__(self, db_path: Optional[str] = None, max_db_size: Optional[int] = N
             db_path = g_config.storage.path
         if max_db_size is None:
             max_db_size = g_config.storage.max_size
+        if retention_days is None:
+            retention_days = g_config.storage.retention_days
 
         self.db_path: Path = Path(db_path)
         self.max_db_size: int = max_db_size
+        self.retention_days: int = max(0, int(retention_days))
         self._env: lmdb.Environment | None = None
 
         self._ensure_db_path()
@@ -310,6 +318,78 @@ def keys(self, prefix: str = "", limit: Optional[int] = None) -> List[str]:
 
         return keys
 
+    def cleanup_expired(self, retention_days: Optional[int] = None) -> int:
+        """
+        Delete conversations older than the given retention period.
+
+        Args:
+            retention_days: Optional override for retention period in days.
+
+        Returns:
+            Number of conversations removed.
+        """
+        retention_value = (
+            self.retention_days if retention_days is None else max(0, int(retention_days))
+        )
+        if retention_value <= 0:
+            logger.debug("Retention cleanup skipped because retention is disabled.")
+            return 0
+
+        cutoff = datetime.now() - timedelta(days=retention_value)
+        expired_entries: list[tuple[str, ConversationInStore]] = []
+
+        try:
+            with self._get_transaction(write=False) as txn:
+                cursor = txn.cursor()
+
+                for key_bytes, value_bytes in cursor:
+                    key_str = key_bytes.decode("utf-8")
+                    if key_str.startswith(self.HASH_LOOKUP_PREFIX):
+                        continue
+
+                    try:
+                        storage_data = orjson.loads(value_bytes)  # type: ignore[arg-type]
+                        conv = ConversationInStore.model_validate(storage_data)
+                    except Exception as exc:
+                        logger.warning(f"Failed to decode record for key {key_str}: {exc}")
+                        continue
+
+                    timestamp = conv.created_at or conv.updated_at
+                    if not timestamp:
+                        continue
+
+                    if timestamp < cutoff:
+                        expired_entries.append((key_str, conv))
+        except Exception as exc:
+            logger.error(f"Failed to scan LMDB for retention cleanup: {exc}")
+            raise
+
+        if not expired_entries:
+            return 0
+
+        removed = 0
+        try:
+            with self._get_transaction(write=True) as txn:
+                for key_str, conv in expired_entries:
+                    key_bytes = key_str.encode("utf-8")
+                    if not txn.delete(key_bytes):
+                        continue
+
+                    message_hash = _hash_conversation(conv.client_id, conv.model, conv.messages)
+                    if message_hash and key_str != message_hash:
+                        txn.delete(f"{self.HASH_LOOKUP_PREFIX}{message_hash}".encode("utf-8"))
+                    removed += 1
+        except Exception as exc:
+            logger.error(f"Failed to delete expired conversations: {exc}")
+            raise
+
+        if removed:
+            logger.info(
+                f"LMDB retention cleanup removed {removed} conversation(s) older than {cutoff.isoformat()}."
+            )
+
+        return removed
+
     def stats(self) -> Dict[str, Any]:
         """
         Get database statistics.

app/utils/config.py

@@ -84,10 +84,15 @@ class StorageConfig(BaseModel):
         description="Path to the storage directory where data will be saved",
     )
     max_size: int = Field(
-        default=1024**2 * 128,  # 128 MB
+        default=1024**2 * 256,  # 256 MB
         ge=1,
         description="Maximum size of the storage in bytes",
     )
+    retention_days: int = Field(
+        default=14,
+        ge=0,
+        description="Number of days to retain conversations before automatic cleanup (0 disables cleanup)",
+    )
 
 
 class LoggingConfig(BaseModel):

config/config.yaml

@@ -29,7 +29,8 @@ gemini:
 
 storage:
   path: "data/lmdb"        # Database storage path
-  max_size: 134217728      # Maximum database size (128 MB)
+  max_size: 268435456      # Maximum database size (256 MB)
+  retention_days: 14       # Number of days to retain conversations before cleanup
 
 logging:
   level: "INFO"           # Log level: DEBUG, INFO, WARNING, ERROR