feat: add memory optimizations to prevent OOM issues

Author: MickaelCa

.pre-commit-config.yaml

@@ -122,6 +122,7 @@ repos:
             loguru>=0.7.0,
             pathspec>=0.12.1,
             prometheus-client,
+            psutil>=5.9.0,
             pydantic,
             pytest-asyncio,
             pytest-mock,
@@ -150,6 +151,7 @@ repos:
             loguru>=0.7.0,
             pathspec>=0.12.1,
             prometheus-client,
+            psutil>=5.9.0,
             pydantic,
             pytest-asyncio,
             pytest-mock,

pyproject.toml

@@ -48,6 +48,7 @@ server = [
     "boto3>=1.28.0",  # AWS SDK for S3 support
     "fastapi[standard]>=0.109.1",  # Minimum safe release (https://osv.dev/vulnerability/PYSEC-2024-38)
     "prometheus-client",
+    "psutil>=5.9.0",  # Memory monitoring for optimization
     "sentry-sdk[fastapi]",
     "slowapi",
     "uvicorn>=0.11.7",  # Minimum safe release (https://osv.dev/vulnerability/PYSEC-2020-150)

requirements.txt

@@ -5,6 +5,7 @@ httpx
 loguru>=0.7.0
 pathspec>=0.12.1
 prometheus-client
+psutil>=5.9.0  # Memory monitoring for optimization
 pydantic
 python-dotenv
 sentry-sdk[fastapi]

src/gitingest/config.py

@@ -9,6 +9,10 @@
 MAX_TOTAL_SIZE_BYTES = 500 * 1024 * 1024  # Maximum size of output file (500 MB)
 DEFAULT_TIMEOUT = 60  # seconds
 
+# Memory optimization settings
+BATCH_SIZE = 100  # Process files in batches to reduce memory usage
+MEMORY_CHECK_INTERVAL = 50  # Check memory usage every N files
+
 OUTPUT_FILE_NAME = "digest.txt"
 
 TMP_BASE_PATH = Path(tempfile.gettempdir()) / "gitingest"

src/gitingest/ingestion.py

@@ -5,11 +5,12 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-from gitingest.config import MAX_DIRECTORY_DEPTH, MAX_FILES, MAX_TOTAL_SIZE_BYTES
+from gitingest.config import MAX_DIRECTORY_DEPTH, MAX_FILES, MAX_TOTAL_SIZE_BYTES, MEMORY_CHECK_INTERVAL
 from gitingest.output_formatter import format_node
 from gitingest.schemas import FileSystemNode, FileSystemNodeType, FileSystemStats
 from gitingest.utils.ingestion_utils import _should_exclude, _should_include
 from gitingest.utils.logging_config import get_logger
+from gitingest.utils.memory_utils import check_memory_pressure, force_garbage_collection, log_memory_stats
 
 if TYPE_CHECKING:
     from gitingest.schemas import IngestionQuery
@@ -51,6 +52,9 @@ def ingest_query(query: IngestionQuery) -> tuple[str, str, str]:
         },
     )
 
+    # Log initial memory usage
+    log_memory_stats("at ingestion start")
+
     subpath = Path(query.subpath.strip("/")).as_posix()
     path = query.local_path / subpath
 
@@ -117,6 +121,9 @@ def ingest_query(query: IngestionQuery) -> tuple[str, str, str]:
         },
     )
 
+    # Log final memory usage
+    log_memory_stats("at ingestion completion")
+
     return format_node(root_node, query=query)
 
 
@@ -258,6 +265,15 @@ def _process_file(path: Path, parent_node: FileSystemNode, stats: FileSystemStat
     stats.total_files += 1
     stats.total_size += file_size
 
+    # Check memory usage periodically and force GC if needed
+    if stats.total_files % MEMORY_CHECK_INTERVAL == 0 and check_memory_pressure():
+        logger.warning(
+            "Memory pressure detected, forcing garbage collection",
+            extra={"files_processed": stats.total_files},
+        )
+        force_garbage_collection()
+        log_memory_stats(f"after processing {stats.total_files} files")
+
     child = FileSystemNode(
         name=path.name,
         type=FileSystemNodeType.FILE,

src/gitingest/output_formatter.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import ssl
+from io import StringIO
 from typing import TYPE_CHECKING
 
 import requests.exceptions
@@ -122,8 +123,32 @@ def _gather_file_contents(node: FileSystemNode) -> str:
     if node.type != FileSystemNodeType.DIRECTORY:
         return node.content_string
 
-    # Recursively gather contents of all files under the current directory
-    return "\n".join(_gather_file_contents(child) for child in node.children)
+    # Use StringIO for memory-efficient string concatenation
+    content_buffer = StringIO()
+    try:
+        _gather_file_contents_recursive(node, content_buffer)
+        return content_buffer.getvalue()
+    finally:
+        content_buffer.close()
+
+
+def _gather_file_contents_recursive(node: FileSystemNode, buffer: StringIO) -> None:
+    """Recursively gather file contents into a StringIO buffer to reduce memory usage.
+
+    Parameters
+    ----------
+    node : FileSystemNode
+        The current directory or file node being processed.
+    buffer : StringIO
+        Buffer to write content to.
+
+    """
+    if node.type != FileSystemNodeType.DIRECTORY:
+        buffer.write(node.content_string)
+        return
+
+    for child in node.children:
+        _gather_file_contents_recursive(child, buffer)
 
 
 def _create_tree_structure(

src/gitingest/schemas/filesystem.py

@@ -5,6 +5,7 @@
 import os
 from dataclasses import dataclass, field
 from enum import Enum, auto
+from io import StringIO
 from typing import TYPE_CHECKING
 
 from gitingest.utils.compat_func import readlink
@@ -155,7 +156,34 @@ def content(self) -> str:  # pylint: disable=too-many-return-statements
             return "Error: Unable to decode file with available encodings"
 
         try:
-            with self.path.open(encoding=good_enc) as fp:
-                return fp.read()
+            return self._read_file_content_streaming(good_enc)
         except (OSError, UnicodeDecodeError) as exc:
             return f"Error reading file with {good_enc!r}: {exc}"
+
+    def _read_file_content_streaming(self, encoding: str, chunk_size: int = 8192) -> str:
+        """Read file content using streaming to reduce memory usage.
+
+        Parameters
+        ----------
+        encoding : str
+            The encoding to use for reading the file.
+        chunk_size : int
+            Size of chunks to read at a time (default: 8192 bytes).
+
+        Returns
+        -------
+        str
+            The file content.
+
+        """
+        content_buffer = StringIO()
+        try:
+            with self.path.open(encoding=encoding) as fp:
+                while True:
+                    chunk = fp.read(chunk_size)
+                    if not chunk:
+                        break
+                    content_buffer.write(chunk)
+            return content_buffer.getvalue()
+        finally:
+            content_buffer.close()

src/gitingest/utils/memory_utils.py

@@ -0,0 +1,83 @@
+"""Memory utility functions for monitoring and optimization."""
+
+from __future__ import annotations
+
+import gc
+from typing import Any
+
+import psutil
+
+from gitingest.utils.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+def get_memory_usage() -> dict[str, Any]:
+    """Get current memory usage statistics.
+
+    Returns
+    -------
+    dict[str, Any]
+        Dictionary containing memory usage statistics in MB.
+
+    """
+    try:
+        process = psutil.Process()
+        memory_info = process.memory_info()
+
+        return {
+            "rss_mb": memory_info.rss / (1024 * 1024),  # Resident Set Size
+            "vms_mb": memory_info.vms / (1024 * 1024),  # Virtual Memory Size
+            "percent": process.memory_percent(),
+        }
+    except Exception as exc:
+        logger.warning("Failed to get memory usage", extra={"error": str(exc)})
+        return {"rss_mb": 0, "vms_mb": 0, "percent": 0}
+
+
+def force_garbage_collection() -> None:
+    """Force garbage collection to free up memory."""
+    try:
+        collected = gc.collect()
+        logger.debug("Forced garbage collection", extra={"objects_collected": collected})
+    except Exception as exc:
+        logger.warning("Failed to force garbage collection", extra={"error": str(exc)})
+
+
+def check_memory_pressure(threshold_mb: float = 3000) -> bool:
+    """Check if memory usage is above threshold.
+
+    Parameters
+    ----------
+    threshold_mb : float
+        Memory threshold in MB (default: 3000 MB = 3 GB).
+
+    Returns
+    -------
+    bool
+        True if memory usage is above threshold.
+
+    """
+    memory_stats = get_memory_usage()
+    return memory_stats["rss_mb"] > threshold_mb
+
+
+def log_memory_stats(context: str = "") -> None:
+    """Log current memory statistics.
+
+    Parameters
+    ----------
+    context : str
+        Context information for the log message.
+
+    """
+    memory_stats = get_memory_usage()
+    logger.info(
+        "Memory usage %s",
+        context,
+        extra={
+            "memory_rss_mb": round(memory_stats["rss_mb"], 2),
+            "memory_vms_mb": round(memory_stats["vms_mb"], 2),
+            "memory_percent": round(memory_stats["percent"], 2),
+        },
+    )

src/server/query_processor.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import shutil
+from io import StringIO
 from pathlib import Path
 from typing import TYPE_CHECKING, cast
 
@@ -302,7 +303,19 @@ async def process_query(
 
     try:
         summary, tree, content = ingest_query(query)
-        digest_content = tree + "\n" + content
+
+        # Clean up repository immediately after ingestion to free memory
+        _cleanup_repository(clone_config)
+
+        # Use StringIO for memory-efficient string concatenation
+        digest_buffer = StringIO()
+        try:
+            digest_buffer.write(tree)
+            digest_buffer.write("\n")
+            digest_buffer.write(content)
+            digest_content = digest_buffer.getvalue()
+        finally:
+            digest_buffer.close()
         _store_digest_content(query, clone_config, digest_content, summary, tree, content)
     except Exception as exc:
         _print_error(query.url, exc, max_file_size, pattern_type, pattern)
@@ -326,8 +339,7 @@ async def process_query(
 
     digest_url = _generate_digest_url(query)
 
-    # Clean up the repository after successful processing
-    _cleanup_repository(clone_config)
+    # Repository was already cleaned up after ingestion to free memory earlier
 
     return IngestSuccessResponse(
         repo_url=input_text,