diff --git a/src/basic_memory/config.py b/src/basic_memory/config.py index 23811399..cae95c9a 100644 --- a/src/basic_memory/config.py +++ b/src/basic_memory/config.py @@ -94,6 +94,12 @@ class BasicMemoryConfig(BaseSettings): gt=0, ) + streaming_checksum_threshold_mb: int = Field( + default=1, + description="File size threshold in MB for using streaming checksum computation. Files larger than this will be processed in chunks to reduce memory usage. Default: 1MB", + gt=0, + ) + kebab_filenames: bool = Field( default=False, description="Format for generated filenames. False preserves spaces and special chars, True converts them to hyphens for consistency with permalinks", diff --git a/src/basic_memory/file_utils.py b/src/basic_memory/file_utils.py index fa5c0678..e665f8fb 100644 --- a/src/basic_memory/file_utils.py +++ b/src/basic_memory/file_utils.py @@ -1,5 +1,6 @@ """Utilities for file operations.""" +import asyncio import hashlib from pathlib import Path import re @@ -12,6 +13,27 @@ from loguru import logger from basic_memory.utils import FilePath +def get_streaming_checksum_threshold() -> int: + """Get the streaming checksum threshold from config. + + Returns threshold in bytes. Defaults to 1MB if config is unavailable. + """ + try: + from basic_memory.config import ConfigManager + + config = ConfigManager().config + return config.streaming_checksum_threshold_mb * 1024 * 1024 + except Exception: + # Default to 1MB if config unavailable (e.g., during tests) + return 1024 * 1024 + + +# Default threshold for streaming vs in-memory checksum computation (1MB) +STREAMING_CHECKSUM_THRESHOLD = 1024 * 1024 # 1MB in bytes +# Chunk size for streaming reads (64KB - optimal for most file systems) +STREAMING_CHUNK_SIZE = 64 * 1024 # 64KB + + class FileError(Exception): """Base exception for file operations.""" @@ -52,6 +74,42 @@ async def compute_checksum(content: Union[str, bytes]) -> str: raise FileError(f"Failed to compute checksum: {e}") +async def compute_checksum_streaming(file_path: Path, chunk_size: int = STREAMING_CHUNK_SIZE) -> str: + """ + Compute SHA-256 checksum using streaming for large files. + + This function reads the file in chunks to avoid loading the entire file into memory, + which is critical for large binary files (PDFs, images, videos). + + Args: + file_path: Path to file to hash + chunk_size: Size of chunks to read (default 64KB) + + Returns: + SHA-256 hex digest + + Raises: + FileError: If checksum computation fails + """ + try: + hasher = hashlib.sha256() + + def read_chunks(): + """Read file in chunks synchronously (for thread pool execution).""" + with open(file_path, "rb") as f: + while chunk := f.read(chunk_size): + hasher.update(chunk) + + # Run blocking I/O in thread pool to avoid blocking event loop + loop = asyncio.get_event_loop() + await loop.run_in_executor(None, read_chunks) + + return hasher.hexdigest() + except Exception as e: # pragma: no cover + logger.error(f"Failed to compute streaming checksum for {file_path}: {e}") + raise FileError(f"Failed to compute streaming checksum: {e}") + + async def ensure_directory(path: FilePath) -> None: """ Ensure directory exists, creating if necessary. diff --git a/src/basic_memory/sync/sync_service.py b/src/basic_memory/sync/sync_service.py index 506fd437..e3f863ba 100644 --- a/src/basic_memory/sync/sync_service.py +++ b/src/basic_memory/sync/sync_service.py @@ -16,7 +16,11 @@ from sqlalchemy.exc import IntegrityError from basic_memory import db from basic_memory.config import BasicMemoryConfig, ConfigManager -from basic_memory.file_utils import has_frontmatter +from basic_memory.file_utils import ( + has_frontmatter, + compute_checksum_streaming, + get_streaming_checksum_threshold, +) from basic_memory.ignore_utils import load_bmignore_patterns, should_ignore_path from basic_memory.markdown import EntityParser, MarkdownProcessor from basic_memory.models import Entity, Project @@ -154,6 +158,8 @@ class SyncService: """Compute file checksum in thread pool to avoid blocking the event loop. Uses semaphore to limit concurrent file reads and prevent OOM on large projects. + For files larger than STREAMING_CHECKSUM_THRESHOLD (1MB), uses streaming + checksum computation to avoid loading entire file into memory. """ def _sync_compute_checksum(path_str: str) -> str: @@ -175,6 +181,27 @@ class SyncService: return hashlib.sha256(content_bytes).hexdigest() async with self._file_semaphore: + path_obj = self.file_service.base_path / path + + # Check file size to decide between streaming and in-memory checksum + try: + file_size = path_obj.stat().st_size + threshold = get_streaming_checksum_threshold() + + # For large files (>threshold), use streaming checksum to avoid memory issues + if file_size > threshold: + logger.debug( + f"Using streaming checksum for large file: path={path}, " + f"size_mb={file_size / (1024 * 1024):.2f}, " + f"threshold_mb={threshold / (1024 * 1024):.2f}" + ) + return await compute_checksum_streaming(path_obj) + + except Exception as e: + # If we can't get file size, fall back to original method + logger.warning(f"Failed to check file size for {path}: {e}, using non-streaming checksum") + + # For smaller files, use the original in-memory method loop = asyncio.get_event_loop() return await loop.run_in_executor(self._thread_pool, _sync_compute_checksum, path) diff --git a/tests/utils/test_file_utils.py b/tests/utils/test_file_utils.py index c7526a0d..07943751 100644 --- a/tests/utils/test_file_utils.py +++ b/tests/utils/test_file_utils.py @@ -11,6 +11,7 @@ from basic_memory.file_utils import ( FileWriteError, ParseError, compute_checksum, + compute_checksum_streaming, ensure_directory, has_frontmatter, parse_frontmatter, @@ -19,6 +20,7 @@ from basic_memory.file_utils import ( sanitize_for_folder, update_frontmatter, write_file_atomic, + STREAMING_CHECKSUM_THRESHOLD, ) @@ -327,3 +329,97 @@ def test_sanitize_for_filename_removes_invalid_characters(): ) def test_sanitize_for_folder_edge_cases(input_folder, expected): assert sanitize_for_folder(input_folder) == expected + + +@pytest.mark.asyncio +async def test_compute_checksum_streaming_small_file(tmp_path: Path): + """Test streaming checksum computation on a small file.""" + test_file = tmp_path / "small.txt" + content = "test content for streaming" + test_file.write_text(content, encoding="utf-8") + + # Compute checksum using streaming + streaming_checksum = await compute_checksum_streaming(test_file) + + # Compute checksum using regular method for comparison + regular_checksum = await compute_checksum(content) + + # Both methods should produce the same checksum + assert streaming_checksum == regular_checksum + assert len(streaming_checksum) == 64 # SHA-256 produces 64 char hex string + + +@pytest.mark.asyncio +async def test_compute_checksum_streaming_large_file(tmp_path: Path): + """Test streaming checksum computation on a large file (>1MB).""" + test_file = tmp_path / "large.bin" + + # Create a file larger than STREAMING_CHECKSUM_THRESHOLD (1MB) + # Use 2MB to ensure we're testing the streaming path + large_content = b"x" * (2 * 1024 * 1024) # 2MB of 'x' bytes + test_file.write_bytes(large_content) + + # Compute checksum using streaming + streaming_checksum = await compute_checksum_streaming(test_file) + + # Compute checksum using regular method for comparison + regular_checksum = await compute_checksum(large_content) + + # Both methods should produce the same checksum + assert streaming_checksum == regular_checksum + assert len(streaming_checksum) == 64 + + +@pytest.mark.asyncio +async def test_compute_checksum_streaming_binary_file(tmp_path: Path): + """Test streaming checksum on binary file content.""" + test_file = tmp_path / "binary.bin" + + # Create binary content with various byte values + binary_content = bytes(range(256)) * 1000 # ~256KB of binary data + test_file.write_bytes(binary_content) + + # Compute checksum using streaming + streaming_checksum = await compute_checksum_streaming(test_file) + + # Compute checksum using regular method for comparison + regular_checksum = await compute_checksum(binary_content) + + # Both methods should produce the same checksum + assert streaming_checksum == regular_checksum + + +@pytest.mark.asyncio +async def test_compute_checksum_streaming_empty_file(tmp_path: Path): + """Test streaming checksum on an empty file.""" + test_file = tmp_path / "empty.txt" + test_file.write_text("", encoding="utf-8") + + # Compute checksum using streaming + streaming_checksum = await compute_checksum_streaming(test_file) + + # Compute checksum using regular method for comparison + regular_checksum = await compute_checksum("") + + # Both methods should produce the same checksum + assert streaming_checksum == regular_checksum + + +@pytest.mark.asyncio +async def test_compute_checksum_streaming_custom_chunk_size(tmp_path: Path): + """Test streaming checksum with custom chunk size.""" + test_file = tmp_path / "custom.txt" + content = "test content " * 10000 # ~130KB + test_file.write_text(content, encoding="utf-8") + + # Test with various chunk sizes + chunk_sizes = [1024, 4096, 65536] # 1KB, 4KB, 64KB + + checksums = [] + for chunk_size in chunk_sizes: + checksum = await compute_checksum_streaming(test_file, chunk_size=chunk_size) + checksums.append(checksum) + + # All chunk sizes should produce the same checksum + assert len(set(checksums)) == 1 + assert checksums[0] == await compute_checksum(content)