From 269a6584056862951ca0a92c2e72c66ae3ad4083 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 13 Nov 2025 18:12:11 -0800 Subject: [PATCH 01/18] start package --- memtest/.gitignore | 19 ++ memtest/Cargo.toml | 20 ++ memtest/Makefile | 29 +++ memtest/README.md | 201 ++++++++++++++++++ memtest/pyproject.toml | 37 ++++ memtest/python/memtest/__init__.py | 185 ++++++++++++++++ memtest/python/memtest/__main__.py | 91 ++++++++ memtest/python/tests/__init__.py | 1 + memtest/python/tests/test_basic.py | 137 ++++++++++++ memtest/python/tests/test_integration.py | 127 +++++++++++ memtest/src/allocator.rs | 256 +++++++++++++++++++++++ memtest/src/lib.rs | 38 ++++ memtest/src/stats.rs | 86 ++++++++ 13 files changed, 1227 insertions(+) create mode 100644 memtest/.gitignore create mode 100644 memtest/Cargo.toml create mode 100644 memtest/Makefile create mode 100644 memtest/README.md create mode 100644 memtest/pyproject.toml create mode 100644 memtest/python/memtest/__init__.py create mode 100644 memtest/python/memtest/__main__.py create mode 100644 memtest/python/tests/__init__.py create mode 100644 memtest/python/tests/test_basic.py create mode 100644 memtest/python/tests/test_integration.py create mode 100644 memtest/src/allocator.rs create mode 100644 memtest/src/lib.rs create mode 100644 memtest/src/stats.rs diff --git a/memtest/.gitignore b/memtest/.gitignore new file mode 100644 index 00000000000..171315214e2 --- /dev/null +++ b/memtest/.gitignore @@ -0,0 +1,19 @@ +# Rust +target/ +Cargo.lock + +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.pytest_cache/ +*.egg-info/ +dist/ +build/ + +# IDE +.vscode/ +.idea/ +*.swp +*.swo diff --git a/memtest/Cargo.toml b/memtest/Cargo.toml new file mode 100644 index 00000000000..ba04385b935 --- /dev/null +++ b/memtest/Cargo.toml @@ -0,0 +1,20 @@ +[workspace] + +[package] +name = "lance-memtest" +version = "0.1.0" +edition = "2021" +authors = ["Lance Developers"] +description = "Memory allocation testing utilities for Python" +license = "Apache-2.0" + +[lib] +name = "memtest" +crate-type = ["cdylib"] + +[dependencies] +libc = "0.2" + +[profile.release] +lto = true +codegen-units = 1 diff --git a/memtest/Makefile b/memtest/Makefile new file mode 100644 index 00000000000..3639ddbf8cd --- /dev/null +++ b/memtest/Makefile @@ -0,0 +1,29 @@ +.PHONY: build test lint format clean + +build: + cargo build + cp target/debug/libmemtest.so python/memtest/ + pip install -e . + +build-release: + cargo build --release + cp target/release/libmemtest.so python/memtest/ + pip install -e . + +test: + pytest python/tests/ -v + +lint: + cargo clippy -- -D warnings + ruff check python/ + +format: + cargo fmt + ruff format python/ + +clean: + cargo clean + rm -rf target/ + find . -type d -name __pycache__ -exec rm -rf {} + + find . -type f -name "*.pyc" -delete + find . -type f -name "*.so" -delete diff --git a/memtest/README.md b/memtest/README.md new file mode 100644 index 00000000000..86fa89f1dcb --- /dev/null +++ b/memtest/README.md @@ -0,0 +1,201 @@ +# lance-memtest + +Memory allocation testing utilities for Python test suites. This package provides tools to track memory allocations made by the Python interpreter and any Python libraries during test execution. + +## Features + +- **LD_PRELOAD-based interposition**: Intercepts all `malloc`, `free`, `calloc`, and `realloc` calls +- **Zero overhead when not tracking**: No performance impact unless explicitly enabled +- **Thread-safe statistics**: Uses atomic operations for accurate multi-threaded tracking +- **Python and CLI interfaces**: Use programmatically or from the command line +- **Comprehensive metrics**: Track allocations, deallocations, current usage, and peak memory + +## Installation + +### From source + +```bash +cd memtest +maturin develop +``` + +### For development + +```bash +cd memtest +make build +``` + +## Usage + +### Python API + +#### Basic tracking + +```python +import memtest + +# Reset statistics +memtest.reset_stats() + +# Your code here +data = [0] * 1000000 + +# Get statistics +stats = memtest.get_stats() +print(f"Allocated: {stats['total_bytes_allocated']} bytes") +print(f"Peak usage: {stats['peak_bytes']} bytes") +``` + +#### Context manager + +```python +import memtest + +with memtest.track() as get_stats: + # Allocate some memory + data = [0] * 1000000 + + # Get stats within the context + stats = get_stats() + print(f"Allocated: {stats['total_bytes_allocated']} bytes") +``` + +#### Pretty printing + +```python +import memtest + +# ... run some code ... + +memtest.print_stats() +``` + +Output: +``` +Memory Allocation Statistics: + Total allocations: 1,234 + Total deallocations: 1,100 + Total bytes allocated: 128.5 KB + Total bytes freed: 120.0 KB + Current memory usage: 8.5 KB + Peak memory usage: 15.2 KB +``` + +### Command Line Interface + +#### Run a command with tracking + +```bash +lance-memtest run python myscript.py +lance-memtest run pytest tests/ +``` + +#### Get the library path + +```bash +# Print path to the .so file +lance-memtest path + +# Use with LD_PRELOAD manually +export LD_PRELOAD=$(lance-memtest path) +python myscript.py +``` + +#### View current statistics + +```bash +lance-memtest stats +``` + +### Integration with pytest + +```python +import pytest +import memtest + +@pytest.fixture(autouse=True) +def track_memory(): + """Automatically track memory for all tests.""" + memtest.reset_stats() + yield + stats = memtest.get_stats() + + # Assert memory bounds + assert stats['peak_bytes'] < 100 * 1024 * 1024, "Test used more than 100MB" + +def test_my_function(): + result = my_function() + + # Check memory usage for this test + stats = memtest.get_stats() + print(f"Peak memory: {memtest.format_bytes(stats['peak_bytes'])}") +``` + +## Statistics + +The following metrics are tracked: + +- **`total_allocations`**: Total number of `malloc`/`calloc` calls +- **`total_deallocations`**: Total number of `free` calls +- **`total_bytes_allocated`**: Total bytes allocated across all calls +- **`total_bytes_deallocated`**: Total bytes freed across all calls +- **`current_bytes`**: Current memory usage (allocated - deallocated) +- **`peak_bytes`**: Peak memory usage observed + +## How It Works + +The package uses LD_PRELOAD to interpose the standard C library allocation functions (`malloc`, `free`, `calloc`, `realloc`). When these functions are called by Python or any C extension: + +1. The interposed function records the allocation size +2. Statistics are updated using atomic operations (thread-safe) +3. The original libc function is called to perform the actual allocation + +The Rust implementation ensures minimal overhead and uses a header-based approach to track allocation sizes. + +## Limitations + +- **Linux only**: LD_PRELOAD is a Linux-specific feature +- **Does not track Python object overhead**: Only tracks C-level allocations +- **Stack allocations not tracked**: Only heap allocations via malloc family +- **Reset affects all threads**: Statistics are global + +## Development + +### Build + +```bash +make build +``` + +### Run tests + +```bash +make test +``` + +### Format code + +```bash +make format +``` + +### Lint + +```bash +make lint +``` + +## Architecture + +The package consists of: + +1. **Rust interpose library** (`src/allocator.rs`): Interposes `malloc`/`free` family +2. **Statistics module** (`src/stats.rs`): Thread-safe atomic counters +3. **PyO3 bindings** (`src/lib.rs`): Exposes stats to Python +4. **Python wrapper** (`python/memtest/__init__.py`): High-level API +5. **CLI** (`python/memtest/__main__.py`): Command-line interface + +## License + +Apache-2.0 diff --git a/memtest/pyproject.toml b/memtest/pyproject.toml new file mode 100644 index 00000000000..396d7c442e0 --- /dev/null +++ b/memtest/pyproject.toml @@ -0,0 +1,37 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "lance-memtest" +version = "0.1.0" +description = "Memory allocation testing utilities for Python test suites" +readme = "README.md" +requires-python = ">=3.9" +license = { text = "Apache-2.0" } +authors = [ + { name = "Lance Developers" } +] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Rust", +] + +[project.scripts] +lance-memtest = "memtest.__main__:main" + +[tool.setuptools] +packages = ["memtest"] + +[tool.setuptools.package-dir] +memtest = "python/memtest" + +[tool.setuptools.package-data] +memtest = ["*.so", "*.dylib", "*.dll"] diff --git a/memtest/python/memtest/__init__.py b/memtest/python/memtest/__init__.py new file mode 100644 index 00000000000..0a4b58b4515 --- /dev/null +++ b/memtest/python/memtest/__init__.py @@ -0,0 +1,185 @@ +"""Memory allocation testing utilities for Python.""" + +import ctypes +from pathlib import Path +from typing import Dict, Optional +from contextlib import contextmanager + +__version__ = "0.1.0" + + +class _MemtestStats(ctypes.Structure): + """C struct matching MemtestStats in Rust.""" + + _fields_ = [ + ("total_allocations", ctypes.c_uint64), + ("total_deallocations", ctypes.c_uint64), + ("total_bytes_allocated", ctypes.c_uint64), + ("total_bytes_deallocated", ctypes.c_uint64), + ("current_bytes", ctypes.c_uint64), + ("peak_bytes", ctypes.c_uint64), + ] + + +def _load_library(): + """Load the memtest shared library.""" + # Find the library relative to this module + module_dir = Path(__file__).parent + + # Look for the library in common locations + possible_paths = [ + module_dir / "libmemtest.so", # Linux + module_dir / "libmemtest.dylib", # macOS + module_dir / "memtest.dll", # Windows + ] + + for lib_path in possible_paths: + if lib_path.exists(): + lib = ctypes.CDLL(str(lib_path)) + + # Define function signatures + lib.memtest_get_stats.argtypes = [ctypes.POINTER(_MemtestStats)] + lib.memtest_get_stats.restype = None + + lib.memtest_reset_stats.argtypes = [] + lib.memtest_reset_stats.restype = None + + return lib, lib_path + + raise RuntimeError("memtest library not found. Run 'make build' to build it.") + + +# Load library at module import +_lib, _lib_path = _load_library() + + +def get_library_path() -> Path: + """Get the path to the memtest shared library for use with LD_PRELOAD. + + Returns: + Path to the .so file that can be used with LD_PRELOAD + + Example: + >>> lib_path = get_library_path() + >>> os.environ['LD_PRELOAD'] = str(lib_path) + """ + return _lib_path + + +def get_stats() -> Dict[str, int]: + """Get current memory allocation statistics. + + Returns: + Dictionary containing: + - total_allocations: Total number of malloc/calloc calls + - total_deallocations: Total number of free calls + - total_bytes_allocated: Total bytes allocated + - total_bytes_deallocated: Total bytes freed + - current_bytes: Current memory usage (allocated - deallocated) + - peak_bytes: Peak memory usage observed + + Example: + >>> stats = get_stats() + >>> print(f"Current memory: {stats['current_bytes']} bytes") + >>> print(f"Peak memory: {stats['peak_bytes']} bytes") + """ + stats = _MemtestStats() + _lib.memtest_get_stats(ctypes.byref(stats)) + + return { + "total_allocations": stats.total_allocations, + "total_deallocations": stats.total_deallocations, + "total_bytes_allocated": stats.total_bytes_allocated, + "total_bytes_deallocated": stats.total_bytes_deallocated, + "current_bytes": stats.current_bytes, + "peak_bytes": stats.peak_bytes, + } + + +def reset_stats() -> None: + """Reset all allocation statistics to zero. + + This is useful for measuring allocations in a specific section of code. + + Example: + >>> reset_stats() + >>> # ... run code to measure ... + >>> stats = get_stats() + """ + _lib.memtest_reset_stats() + + +@contextmanager +def track(reset: bool = True): + """Context manager to track allocations within a code block. + + Args: + reset: Whether to reset statistics before entering the context + + Yields: + A function that returns current statistics + + Example: + >>> with track() as get: + ... data = [0] * 1000 + ... stats = get() + ... print(f"Allocated: {stats['total_bytes_allocated']} bytes") + """ + if reset: + reset_stats() + + yield get_stats + + +def format_bytes(num_bytes: int) -> str: + """Format byte count as human-readable string. + + Args: + num_bytes: Number of bytes + + Returns: + Formatted string (e.g., "1.5 MB") + """ + for unit in ["B", "KB", "MB", "GB", "TB"]: + if abs(num_bytes) < 1024.0: + return f"{num_bytes:.1f} {unit}" + num_bytes /= 1024.0 + return f"{num_bytes:.1f} PB" + + +def print_stats(stats: Optional[Dict[str, int]] = None) -> None: + """Print allocation statistics in a readable format. + + Args: + stats: Statistics dictionary. If None, fetches current stats. + + Example: + >>> print_stats() + Memory Allocation Statistics: + Total allocations: 1,234 + Total deallocations: 1,100 + Total bytes allocated: 128.5 KB + Total bytes freed: 120.0 KB + Current memory usage: 8.5 KB + Peak memory usage: 15.2 KB + """ + if stats is None: + stats = get_stats() + + print("Memory Allocation Statistics:") + print(f" Total allocations: {stats['total_allocations']:,}") + print(f" Total deallocations: {stats['total_deallocations']:,}") + print(f" Total bytes allocated: {format_bytes(stats['total_bytes_allocated'])}") + print(f" Total bytes freed: {format_bytes(stats['total_bytes_deallocated'])}") + print(f" Current memory usage: {format_bytes(stats['current_bytes'])}") + print(f" Peak memory usage: {format_bytes(stats['peak_bytes'])}") + + +__all__ = [ + "get_library_path", + "get_stats", + "reset_stats", + "track", + "format_bytes", + "print_stats", +] diff --git a/memtest/python/memtest/__main__.py b/memtest/python/memtest/__main__.py new file mode 100644 index 00000000000..97fa4a159a5 --- /dev/null +++ b/memtest/python/memtest/__main__.py @@ -0,0 +1,91 @@ +"""CLI for lance-memtest.""" + +import argparse +import os +import subprocess +import sys + +from memtest import get_library_path, print_stats + + +def cmd_path(args): + """Print the path to the memtest shared library.""" + lib_path = get_library_path() + print(lib_path) + return 0 + + +def cmd_run(args): + """Run a command with LD_PRELOAD set to track memory allocations.""" + lib_path = get_library_path() + + # Set up environment + env = os.environ.copy() + + # Prepend to LD_PRELOAD if it already exists + existing_preload = env.get("LD_PRELOAD", "") + if existing_preload: + env["LD_PRELOAD"] = f"{lib_path}:{existing_preload}" + else: + env["LD_PRELOAD"] = str(lib_path) + + # Run the command + try: + result = subprocess.run(args.command, env=env, shell=False) + return result.returncode + except FileNotFoundError: + print(f"Error: Command not found: {args.command[0]}", file=sys.stderr) + return 1 + except KeyboardInterrupt: + return 130 + + +def cmd_stats(args): + """Print current allocation statistics.""" + print_stats() + return 0 + + +def main(): + """Main CLI entry point.""" + parser = argparse.ArgumentParser( + prog="lance-memtest", + description="Memory allocation testing utilities for Python", + ) + + subparsers = parser.add_subparsers(dest="command", help="Command to run") + + # path command + path_parser = subparsers.add_parser( + "path", help="Print path to the memtest shared library" + ) + path_parser.set_defaults(func=cmd_path) + + # run command + run_parser = subparsers.add_parser( + "run", help="Run a command with memory tracking enabled" + ) + run_parser.add_argument("command", nargs="+", help="Command and arguments to run") + run_parser.set_defaults(func=cmd_run) + + # stats command + stats_parser = subparsers.add_parser( + "stats", help="Print current allocation statistics" + ) + stats_parser.set_defaults(func=cmd_stats) + + args = parser.parse_args() + + if not hasattr(args, "func"): + parser.print_help() + return 1 + + try: + return args.func(args) + except Exception as e: + print(f"Error: {e}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/memtest/python/tests/__init__.py b/memtest/python/tests/__init__.py new file mode 100644 index 00000000000..3263fffd5fe --- /dev/null +++ b/memtest/python/tests/__init__.py @@ -0,0 +1 @@ +"""Tests for lance-memtest.""" diff --git a/memtest/python/tests/test_basic.py b/memtest/python/tests/test_basic.py new file mode 100644 index 00000000000..9e83d5c32ad --- /dev/null +++ b/memtest/python/tests/test_basic.py @@ -0,0 +1,137 @@ +"""Basic tests for memtest functionality.""" + +import subprocess +import sys + +import memtest + + +def test_get_library_path(): + """Test that we can get the library path.""" + lib_path = memtest.get_library_path() + assert lib_path.exists() + assert lib_path.suffix == ".so" + + +def test_get_stats(): + """Test that we can get statistics.""" + stats = memtest.get_stats() + + assert isinstance(stats, dict) + assert "total_allocations" in stats + assert "total_deallocations" in stats + assert "total_bytes_allocated" in stats + assert "total_bytes_deallocated" in stats + assert "current_bytes" in stats + assert "peak_bytes" in stats + + # All values should be non-negative integers + for key, value in stats.items(): + assert isinstance(value, int) + assert value >= 0 + + +def test_reset_stats(): + """Test that we can reset statistics.""" + # Get initial stats + initial_stats = memtest.get_stats() + + # Reset + memtest.reset_stats() + + # All stats should be zero after reset + stats = memtest.get_stats() + assert stats["total_allocations"] == 0 + assert stats["total_deallocations"] == 0 + assert stats["total_bytes_allocated"] == 0 + assert stats["total_bytes_deallocated"] == 0 + assert stats["current_bytes"] == 0 + assert stats["peak_bytes"] == 0 + + +def test_track_context_manager(): + """Test the track context manager.""" + with memtest.track() as get_stats: + # Allocate some memory + data = [0] * 1000 + + # Get stats within the context + stats = get_stats() + + # We should see some allocations + assert stats["total_allocations"] > 0 + assert stats["total_bytes_allocated"] > 0 + + +def test_format_bytes(): + """Test byte formatting.""" + assert "B" in memtest.format_bytes(100) + assert "KB" in memtest.format_bytes(1024) + assert "MB" in memtest.format_bytes(1024 * 1024) + assert "GB" in memtest.format_bytes(1024 * 1024 * 1024) + + +def test_print_stats(): + """Test that print_stats doesn't crash.""" + # This should not raise an exception + memtest.print_stats() + + # Should also work with explicit stats + stats = memtest.get_stats() + memtest.print_stats(stats) + + +def test_allocation_tracking(): + """Test that allocations are actually tracked.""" + memtest.reset_stats() + + initial_stats = memtest.get_stats() + assert initial_stats["total_allocations"] == 0 + + # Allocate a large list + data = [0] * 10000 + + stats_after = memtest.get_stats() + + # We should see allocations (though the exact number depends on Python internals) + assert stats_after["total_allocations"] > 0 + assert stats_after["total_bytes_allocated"] > 0 + + # Peak should be at least as much as current + assert stats_after["peak_bytes"] >= stats_after["current_bytes"] + + +def test_cli_path(): + """Test the CLI path command.""" + result = subprocess.run( + [sys.executable, "-m", "memtest", "path"], + capture_output=True, + text=True, + ) + + assert result.returncode == 0 + assert ".so" in result.stdout + + +def test_cli_stats(): + """Test the CLI stats command.""" + result = subprocess.run( + [sys.executable, "-m", "memtest", "stats"], + capture_output=True, + text=True, + ) + + assert result.returncode == 0 + assert "Memory Allocation Statistics" in result.stdout + + +def test_cli_run(): + """Test the CLI run command.""" + result = subprocess.run( + [sys.executable, "-m", "memtest", "run", "python", "-c", "print('hello')"], + capture_output=True, + text=True, + ) + + assert result.returncode == 0 + assert "hello" in result.stdout diff --git a/memtest/python/tests/test_integration.py b/memtest/python/tests/test_integration.py new file mode 100644 index 00000000000..5fe0539e798 --- /dev/null +++ b/memtest/python/tests/test_integration.py @@ -0,0 +1,127 @@ +"""Integration tests for memtest with real allocations.""" + +import os +import subprocess +import sys +import tempfile + +import memtest + + +def test_preload_environment(): + """Test that LD_PRELOAD works correctly.""" + lib_path = memtest.get_library_path() + + # Create a small Python script that uses memtest + test_script = """ +import memtest + +memtest.reset_stats() + +# Allocate some data +data = [i for i in range(1000)] + +stats = memtest.get_stats() +print(f"Allocations: {stats['total_allocations']}") +print(f"Bytes: {stats['total_bytes_allocated']}") + +assert stats['total_allocations'] > 0, "Should see allocations" +assert stats['total_bytes_allocated'] > 0, "Should see bytes allocated" +""" + + with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f: + f.write(test_script) + script_path = f.name + + try: + env = os.environ.copy() + env["LD_PRELOAD"] = str(lib_path) + + result = subprocess.run( + [sys.executable, script_path], + env=env, + capture_output=True, + text=True, + ) + + assert result.returncode == 0, f"Script failed: {result.stderr}" + assert "Allocations:" in result.stdout + assert "Bytes:" in result.stdout + + finally: + os.unlink(script_path) + + +def test_repeated_allocations(): + """Test tracking repeated allocations and deallocations.""" + memtest.reset_stats() + + # Do several allocation/deallocation cycles + for i in range(10): + data = [0] * 1000 + del data + + stats = memtest.get_stats() + + # Should see multiple allocations + assert stats["total_allocations"] > 10 + assert stats["total_deallocations"] > 0 + assert stats["total_bytes_allocated"] > 0 + assert stats["total_bytes_deallocated"] > 0 + + +def test_peak_tracking(): + """Test that peak memory usage is tracked correctly.""" + memtest.reset_stats() + + # Allocate progressively larger arrays + arrays = [] + for size in [100, 1000, 10000]: + arrays.append([0] * size) + + stats = memtest.get_stats() + + # Peak should be higher than or equal to current + assert stats["peak_bytes"] >= stats["current_bytes"] + + # Free the arrays + arrays.clear() + + stats_after = memtest.get_stats() + + # Peak should remain the same (doesn't decrease) + assert stats_after["peak_bytes"] == stats["peak_bytes"] + + +def test_with_numpy(): + """Test tracking NumPy allocations if NumPy is available.""" + try: + import numpy as np + except ImportError: + pytest.skip("NumPy not available") + + memtest.reset_stats() + + # Create a large NumPy array + arr = np.zeros((1000, 1000), dtype=np.float64) + + stats = memtest.get_stats() + + # NumPy uses malloc internally, so we should see allocations + assert stats["total_allocations"] > 0 + assert stats["total_bytes_allocated"] > 0 + + +def test_context_manager_integration(): + """Test the context manager with real workload.""" + results = [] + + with memtest.track() as get_stats: + # Allocate in stages and track progress + for i in range(5): + data = [0] * 1000 + results.append(get_stats()) + + # Each measurement should show increasing allocations + for i in range(1, len(results)): + assert results[i]["total_allocations"] >= results[i - 1]["total_allocations"] diff --git a/memtest/src/allocator.rs b/memtest/src/allocator.rs new file mode 100644 index 00000000000..95e50c35f06 --- /dev/null +++ b/memtest/src/allocator.rs @@ -0,0 +1,256 @@ +use crate::stats::STATS; +use libc::{c_void, size_t}; +use std::cell::Cell; +use std::sync::Once; + +type MallocFn = unsafe extern "C" fn(size_t) -> *mut c_void; +type FreeFn = unsafe extern "C" fn(*mut c_void); +type CallocFn = unsafe extern "C" fn(size_t, size_t) -> *mut c_void; +type ReallocFn = unsafe extern "C" fn(*mut c_void, size_t) -> *mut c_void; + +static INIT: Once = Once::new(); +static mut REAL_MALLOC: Option = None; +static mut REAL_FREE: Option = None; +static mut REAL_CALLOC: Option = None; +static mut REAL_REALLOC: Option = None; + +const RTLD_NEXT: *mut c_void = -1isize as *mut c_void; + +thread_local! { + static IN_HOOK: Cell = const { Cell::new(false) }; +} + +extern "C" { + fn dlsym(handle: *mut c_void, symbol: *const libc::c_char) -> *mut c_void; +} + +/// Initialize the function pointers to the real allocation functions +unsafe fn init_real_functions() { + INIT.call_once(|| { + // Prevent recursion during initialization + IN_HOOK.with(|flag| flag.set(true)); + + REAL_MALLOC = Some(std::mem::transmute(dlsym( + RTLD_NEXT, + b"malloc\0".as_ptr() as *const libc::c_char, + ))); + REAL_FREE = Some(std::mem::transmute(dlsym( + RTLD_NEXT, + b"free\0".as_ptr() as *const libc::c_char, + ))); + REAL_CALLOC = Some(std::mem::transmute(dlsym( + RTLD_NEXT, + b"calloc\0".as_ptr() as *const libc::c_char, + ))); + REAL_REALLOC = Some(std::mem::transmute(dlsym( + RTLD_NEXT, + b"realloc\0".as_ptr() as *const libc::c_char, + ))); + + IN_HOOK.with(|flag| flag.set(false)); + }); +} + +/// Store allocation size in a header before the returned pointer +#[repr(C)] +struct AllocationHeader { + size: usize, +} + +const HEADER_SIZE: usize = std::mem::size_of::(); + +#[no_mangle] +pub unsafe extern "C" fn malloc(size: size_t) -> *mut c_void { + // Check if we're already in a hook to prevent recursion + let in_hook = IN_HOOK.with(|flag| { + if flag.get() { + true + } else { + flag.set(true); + false + } + }); + + if in_hook { + // We're in recursion, just call the real malloc + init_real_functions(); + if let Some(real_malloc) = REAL_MALLOC { + return real_malloc(size); + } + return std::ptr::null_mut(); + } + + init_real_functions(); + + let result = if let Some(real_malloc) = REAL_MALLOC { + let total_size = size + HEADER_SIZE; + let ptr = real_malloc(total_size); + + if !ptr.is_null() { + // Store size in header + let header = ptr as *mut AllocationHeader; + (*header).size = size; + + STATS.record_allocation(size); + + // Return pointer after header + ptr.add(HEADER_SIZE) + } else { + ptr + } + } else { + std::ptr::null_mut() + }; + + IN_HOOK.with(|flag| flag.set(false)); + result +} + +#[no_mangle] +pub unsafe extern "C" fn free(ptr: *mut c_void) { + if ptr.is_null() { + return; + } + + // Check if we're already in a hook to prevent recursion + let in_hook = IN_HOOK.with(|flag| { + if flag.get() { + true + } else { + flag.set(true); + false + } + }); + + init_real_functions(); + + if let Some(real_free) = REAL_FREE { + if in_hook { + // We're in recursion, just call the real free + real_free(ptr); + return; + } + + // Get the actual allocation pointer (before header) + let actual_ptr = (ptr as *mut u8).sub(HEADER_SIZE); + let header = actual_ptr as *mut AllocationHeader; + let size = (*header).size; + + STATS.record_deallocation(size); + + real_free(actual_ptr as *mut c_void); + + IN_HOOK.with(|flag| flag.set(false)); + } +} + +#[no_mangle] +pub unsafe extern "C" fn calloc(nmemb: size_t, size: size_t) -> *mut c_void { + let in_hook = IN_HOOK.with(|flag| { + if flag.get() { + true + } else { + flag.set(true); + false + } + }); + + if in_hook { + init_real_functions(); + if let Some(real_calloc) = REAL_CALLOC { + return real_calloc(nmemb, size); + } + return std::ptr::null_mut(); + } + + init_real_functions(); + + let result = if let Some(real_calloc) = REAL_CALLOC { + let total_size = nmemb * size; + let allocation_size = total_size + HEADER_SIZE; + + let ptr = real_calloc(allocation_size, 1); + + if !ptr.is_null() { + let header = ptr as *mut AllocationHeader; + (*header).size = total_size; + + STATS.record_allocation(total_size); + + ptr.add(HEADER_SIZE) + } else { + ptr + } + } else { + std::ptr::null_mut() + }; + + IN_HOOK.with(|flag| flag.set(false)); + result +} + +#[no_mangle] +pub unsafe extern "C" fn realloc(ptr: *mut c_void, size: size_t) -> *mut c_void { + let in_hook = IN_HOOK.with(|flag| { + if flag.get() { + true + } else { + flag.set(true); + false + } + }); + + if in_hook { + init_real_functions(); + if let Some(real_realloc) = REAL_REALLOC { + return real_realloc(ptr, size); + } + return std::ptr::null_mut(); + } + + init_real_functions(); + + let result = if let Some(real_realloc) = REAL_REALLOC { + if ptr.is_null() { + // realloc(NULL, size) is equivalent to malloc(size) + // Note: This will set the flag again, but that's handled + IN_HOOK.with(|flag| flag.set(false)); + return malloc(size); + } + + if size == 0 { + // realloc(ptr, 0) is equivalent to free(ptr) + IN_HOOK.with(|flag| flag.set(false)); + free(ptr); + return std::ptr::null_mut(); + } + + // Get old size from header + let actual_ptr = (ptr as *mut u8).sub(HEADER_SIZE); + let old_header = actual_ptr as *mut AllocationHeader; + let old_size = (*old_header).size; + + // Reallocate with new size + let total_size = size + HEADER_SIZE; + let new_ptr = real_realloc(actual_ptr as *mut c_void, total_size); + + if !new_ptr.is_null() { + // Update header with new size + let new_header = new_ptr as *mut AllocationHeader; + (*new_header).size = size; + + // Record the change in allocation + STATS.record_deallocation(old_size); + STATS.record_allocation(size); + + new_ptr.add(HEADER_SIZE) + } else { + new_ptr + } + } else { + std::ptr::null_mut() + }; + + IN_HOOK.with(|flag| flag.set(false)); + result +} diff --git a/memtest/src/lib.rs b/memtest/src/lib.rs new file mode 100644 index 00000000000..df712aa381d --- /dev/null +++ b/memtest/src/lib.rs @@ -0,0 +1,38 @@ +mod allocator; +mod stats; + +use stats::STATS; + +/// C-compatible statistics struct +#[repr(C)] +pub struct MemtestStats { + pub total_allocations: u64, + pub total_deallocations: u64, + pub total_bytes_allocated: u64, + pub total_bytes_deallocated: u64, + pub current_bytes: u64, + pub peak_bytes: u64, +} + +/// Get all statistics in a single call +#[no_mangle] +pub extern "C" fn memtest_get_stats(stats: *mut MemtestStats) { + if stats.is_null() { + return; + } + + unsafe { + (*stats).total_allocations = STATS.total_allocations.load(std::sync::atomic::Ordering::Relaxed); + (*stats).total_deallocations = STATS.total_deallocations.load(std::sync::atomic::Ordering::Relaxed); + (*stats).total_bytes_allocated = STATS.total_bytes_allocated.load(std::sync::atomic::Ordering::Relaxed); + (*stats).total_bytes_deallocated = STATS.total_bytes_deallocated.load(std::sync::atomic::Ordering::Relaxed); + (*stats).current_bytes = STATS.current_bytes.load(std::sync::atomic::Ordering::Relaxed); + (*stats).peak_bytes = STATS.peak_bytes.load(std::sync::atomic::Ordering::Relaxed); + } +} + +/// Reset all statistics to zero +#[no_mangle] +pub extern "C" fn memtest_reset_stats() { + STATS.reset(); +} diff --git a/memtest/src/stats.rs b/memtest/src/stats.rs new file mode 100644 index 00000000000..08e3c3d129a --- /dev/null +++ b/memtest/src/stats.rs @@ -0,0 +1,86 @@ +use std::sync::atomic::{AtomicU64, Ordering}; + +/// Global allocation statistics tracked using atomic operations for thread safety +pub struct AllocationStats { + pub total_allocations: AtomicU64, + pub total_deallocations: AtomicU64, + pub total_bytes_allocated: AtomicU64, + pub total_bytes_deallocated: AtomicU64, + pub current_bytes: AtomicU64, + pub peak_bytes: AtomicU64, +} + +impl AllocationStats { + pub const fn new() -> Self { + Self { + total_allocations: AtomicU64::new(0), + total_deallocations: AtomicU64::new(0), + total_bytes_allocated: AtomicU64::new(0), + total_bytes_deallocated: AtomicU64::new(0), + current_bytes: AtomicU64::new(0), + peak_bytes: AtomicU64::new(0), + } + } + + pub fn record_allocation(&self, size: usize) { + self.total_allocations.fetch_add(1, Ordering::Relaxed); + self.total_bytes_allocated + .fetch_add(size as u64, Ordering::Relaxed); + + let current = self.current_bytes.fetch_add(size as u64, Ordering::Relaxed) + size as u64; + + // Update peak if necessary + let mut peak = self.peak_bytes.load(Ordering::Relaxed); + while current > peak { + match self.peak_bytes.compare_exchange_weak( + peak, + current, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => break, + Err(p) => peak = p, + } + } + } + + pub fn record_deallocation(&self, size: usize) { + self.total_deallocations.fetch_add(1, Ordering::Relaxed); + self.total_bytes_deallocated + .fetch_add(size as u64, Ordering::Relaxed); + self.current_bytes.fetch_sub(size as u64, Ordering::Relaxed); + } + + pub fn reset(&self) { + self.total_allocations.store(0, Ordering::Relaxed); + self.total_deallocations.store(0, Ordering::Relaxed); + self.total_bytes_allocated.store(0, Ordering::Relaxed); + self.total_bytes_deallocated.store(0, Ordering::Relaxed); + self.current_bytes.store(0, Ordering::Relaxed); + self.peak_bytes.store(0, Ordering::Relaxed); + } + + pub fn get_snapshot(&self) -> StatsSnapshot { + StatsSnapshot { + total_allocations: self.total_allocations.load(Ordering::Relaxed), + total_deallocations: self.total_deallocations.load(Ordering::Relaxed), + total_bytes_allocated: self.total_bytes_allocated.load(Ordering::Relaxed), + total_bytes_deallocated: self.total_bytes_deallocated.load(Ordering::Relaxed), + current_bytes: self.current_bytes.load(Ordering::Relaxed), + peak_bytes: self.peak_bytes.load(Ordering::Relaxed), + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct StatsSnapshot { + pub total_allocations: u64, + pub total_deallocations: u64, + pub total_bytes_allocated: u64, + pub total_bytes_deallocated: u64, + pub current_bytes: u64, + pub peak_bytes: u64, +} + +/// Global statistics instance +pub static STATS: AllocationStats = AllocationStats::new(); From e4069fcd5d5963a1b67fb40ecb45718a22d3612d Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 13 Nov 2025 19:33:03 -0800 Subject: [PATCH 02/18] get tests passing --- memtest/Cargo.toml | 3 + memtest/python/tests/test_basic.py | 18 +- memtest/python/tests/test_integration.py | 7 +- memtest/src/allocator.rs | 296 ++++++++++------------- memtest/src/lib.rs | 29 ++- memtest/src/stats.rs | 11 +- 6 files changed, 163 insertions(+), 201 deletions(-) diff --git a/memtest/Cargo.toml b/memtest/Cargo.toml index ba04385b935..72405566ed6 100644 --- a/memtest/Cargo.toml +++ b/memtest/Cargo.toml @@ -8,6 +8,9 @@ authors = ["Lance Developers"] description = "Memory allocation testing utilities for Python" license = "Apache-2.0" +[lints.clippy] +arithmetic_side_effects = "deny" + [lib] name = "memtest" crate-type = ["cdylib"] diff --git a/memtest/python/tests/test_basic.py b/memtest/python/tests/test_basic.py index 9e83d5c32ad..625840b79b7 100644 --- a/memtest/python/tests/test_basic.py +++ b/memtest/python/tests/test_basic.py @@ -34,7 +34,7 @@ def test_get_stats(): def test_reset_stats(): """Test that we can reset statistics.""" # Get initial stats - initial_stats = memtest.get_stats() + _ = memtest.get_stats() # Reset memtest.reset_stats() @@ -53,7 +53,7 @@ def test_track_context_manager(): """Test the track context manager.""" with memtest.track() as get_stats: # Allocate some memory - data = [0] * 1000 + _ = [0] * 1000 # Get stats within the context stats = get_stats() @@ -89,7 +89,7 @@ def test_allocation_tracking(): assert initial_stats["total_allocations"] == 0 # Allocate a large list - data = [0] * 10000 + _ = [0] * 10000 stats_after = memtest.get_stats() @@ -123,15 +123,3 @@ def test_cli_stats(): assert result.returncode == 0 assert "Memory Allocation Statistics" in result.stdout - - -def test_cli_run(): - """Test the CLI run command.""" - result = subprocess.run( - [sys.executable, "-m", "memtest", "run", "python", "-c", "print('hello')"], - capture_output=True, - text=True, - ) - - assert result.returncode == 0 - assert "hello" in result.stdout diff --git a/memtest/python/tests/test_integration.py b/memtest/python/tests/test_integration.py index 5fe0539e798..f0ab8ac9f90 100644 --- a/memtest/python/tests/test_integration.py +++ b/memtest/python/tests/test_integration.py @@ -4,6 +4,7 @@ import subprocess import sys import tempfile +import pytest import memtest @@ -64,7 +65,7 @@ def test_repeated_allocations(): stats = memtest.get_stats() # Should see multiple allocations - assert stats["total_allocations"] > 10 + assert stats["total_allocations"] >= 10 assert stats["total_deallocations"] > 0 assert stats["total_bytes_allocated"] > 0 assert stats["total_bytes_deallocated"] > 0 @@ -103,7 +104,7 @@ def test_with_numpy(): memtest.reset_stats() # Create a large NumPy array - arr = np.zeros((1000, 1000), dtype=np.float64) + _ = np.zeros((1000, 1000), dtype=np.float64) stats = memtest.get_stats() @@ -119,7 +120,7 @@ def test_context_manager_integration(): with memtest.track() as get_stats: # Allocate in stages and track progress for i in range(5): - data = [0] * 1000 + _ = [0] * 1000 results.append(get_stats()) # Each measurement should show increasing allocations diff --git a/memtest/src/allocator.rs b/memtest/src/allocator.rs index 95e50c35f06..4b38cf22898 100644 --- a/memtest/src/allocator.rs +++ b/memtest/src/allocator.rs @@ -1,54 +1,56 @@ use crate::stats::STATS; use libc::{c_void, size_t}; -use std::cell::Cell; -use std::sync::Once; +use std::sync::atomic::{AtomicBool, AtomicPtr, Ordering}; type MallocFn = unsafe extern "C" fn(size_t) -> *mut c_void; type FreeFn = unsafe extern "C" fn(*mut c_void); type CallocFn = unsafe extern "C" fn(size_t, size_t) -> *mut c_void; type ReallocFn = unsafe extern "C" fn(*mut c_void, size_t) -> *mut c_void; -static INIT: Once = Once::new(); -static mut REAL_MALLOC: Option = None; -static mut REAL_FREE: Option = None; -static mut REAL_CALLOC: Option = None; -static mut REAL_REALLOC: Option = None; +static REAL_MALLOC: AtomicPtr = AtomicPtr::new(std::ptr::null_mut()); +static REAL_FREE: AtomicPtr = AtomicPtr::new(std::ptr::null_mut()); +static REAL_CALLOC: AtomicPtr = AtomicPtr::new(std::ptr::null_mut()); +static REAL_REALLOC: AtomicPtr = AtomicPtr::new(std::ptr::null_mut()); +static INITIALIZING: AtomicBool = AtomicBool::new(false); +static INITIALIZED: AtomicBool = AtomicBool::new(false); const RTLD_NEXT: *mut c_void = -1isize as *mut c_void; -thread_local! { - static IN_HOOK: Cell = const { Cell::new(false) }; -} - extern "C" { fn dlsym(handle: *mut c_void, symbol: *const libc::c_char) -> *mut c_void; } /// Initialize the function pointers to the real allocation functions unsafe fn init_real_functions() { - INIT.call_once(|| { - // Prevent recursion during initialization - IN_HOOK.with(|flag| flag.set(true)); - - REAL_MALLOC = Some(std::mem::transmute(dlsym( - RTLD_NEXT, - b"malloc\0".as_ptr() as *const libc::c_char, - ))); - REAL_FREE = Some(std::mem::transmute(dlsym( - RTLD_NEXT, - b"free\0".as_ptr() as *const libc::c_char, - ))); - REAL_CALLOC = Some(std::mem::transmute(dlsym( - RTLD_NEXT, - b"calloc\0".as_ptr() as *const libc::c_char, - ))); - REAL_REALLOC = Some(std::mem::transmute(dlsym( - RTLD_NEXT, - b"realloc\0".as_ptr() as *const libc::c_char, - ))); - - IN_HOOK.with(|flag| flag.set(false)); - }); + // If already initialized, return + if INITIALIZED.load(Ordering::Acquire) { + return; + } + + // Try to set initializing flag + if INITIALIZING + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + // Someone else is initializing, spin wait + while !INITIALIZED.load(Ordering::Acquire) { + std::hint::spin_loop(); + } + return; + } + + // We're the one initializing + let malloc_ptr = dlsym(RTLD_NEXT, b"malloc\0".as_ptr() as *const libc::c_char); + let free_ptr = dlsym(RTLD_NEXT, b"free\0".as_ptr() as *const libc::c_char); + let calloc_ptr = dlsym(RTLD_NEXT, b"calloc\0".as_ptr() as *const libc::c_char); + let realloc_ptr = dlsym(RTLD_NEXT, b"realloc\0".as_ptr() as *const libc::c_char); + + REAL_MALLOC.store(malloc_ptr, Ordering::Release); + REAL_FREE.store(free_ptr, Ordering::Release); + REAL_CALLOC.store(calloc_ptr, Ordering::Release); + REAL_REALLOC.store(realloc_ptr, Ordering::Release); + + INITIALIZED.store(true, Ordering::Release); } /// Store allocation size in a header before the returned pointer @@ -61,49 +63,37 @@ const HEADER_SIZE: usize = std::mem::size_of::(); #[no_mangle] pub unsafe extern "C" fn malloc(size: size_t) -> *mut c_void { - // Check if we're already in a hook to prevent recursion - let in_hook = IN_HOOK.with(|flag| { - if flag.get() { - true - } else { - flag.set(true); - false - } - }); - - if in_hook { - // We're in recursion, just call the real malloc - init_real_functions(); - if let Some(real_malloc) = REAL_MALLOC { - return real_malloc(size); - } + // If we're currently initializing, forward directly to avoid recursion + if INITIALIZING.load(Ordering::Acquire) && !INITIALIZED.load(Ordering::Acquire) { + // During initialization, dlsym might call malloc + // We can't use RTLD_NEXT here, so we'll just use a simple bump allocator + // or return null and hope dlsym handles it return std::ptr::null_mut(); } init_real_functions(); - let result = if let Some(real_malloc) = REAL_MALLOC { - let total_size = size + HEADER_SIZE; - let ptr = real_malloc(total_size); + let malloc_ptr = REAL_MALLOC.load(Ordering::Acquire); + if malloc_ptr.is_null() { + return std::ptr::null_mut(); + } - if !ptr.is_null() { - // Store size in header - let header = ptr as *mut AllocationHeader; - (*header).size = size; + let real_malloc: MallocFn = std::mem::transmute(malloc_ptr); + let total_size = size.saturating_add(HEADER_SIZE); + let ptr = real_malloc(total_size); - STATS.record_allocation(size); + if !ptr.is_null() { + // Store size in header + let header = ptr as *mut AllocationHeader; + (*header).size = size; - // Return pointer after header - ptr.add(HEADER_SIZE) - } else { - ptr - } - } else { - std::ptr::null_mut() - }; + STATS.record_allocation(size); - IN_HOOK.with(|flag| flag.set(false)); - result + // Return pointer after header + ptr.add(HEADER_SIZE) + } else { + ptr + } } #[no_mangle] @@ -112,145 +102,107 @@ pub unsafe extern "C" fn free(ptr: *mut c_void) { return; } - // Check if we're already in a hook to prevent recursion - let in_hook = IN_HOOK.with(|flag| { - if flag.get() { - true - } else { - flag.set(true); - false - } - }); + // If called during initialization, do nothing (malloc returned null anyway) + if INITIALIZING.load(Ordering::Acquire) && !INITIALIZED.load(Ordering::Acquire) { + return; + } init_real_functions(); - if let Some(real_free) = REAL_FREE { - if in_hook { - // We're in recursion, just call the real free - real_free(ptr); - return; - } + let free_ptr = REAL_FREE.load(Ordering::Acquire); + if free_ptr.is_null() { + return; + } - // Get the actual allocation pointer (before header) - let actual_ptr = (ptr as *mut u8).sub(HEADER_SIZE); - let header = actual_ptr as *mut AllocationHeader; - let size = (*header).size; + let real_free: FreeFn = std::mem::transmute(free_ptr); - STATS.record_deallocation(size); + // Get the actual allocation pointer (before header) + let actual_ptr = (ptr as *mut u8).sub(HEADER_SIZE); + let header = actual_ptr as *mut AllocationHeader; + let size = (*header).size; - real_free(actual_ptr as *mut c_void); + STATS.record_deallocation(size); - IN_HOOK.with(|flag| flag.set(false)); - } + real_free(actual_ptr as *mut c_void); } #[no_mangle] pub unsafe extern "C" fn calloc(nmemb: size_t, size: size_t) -> *mut c_void { - let in_hook = IN_HOOK.with(|flag| { - if flag.get() { - true - } else { - flag.set(true); - false - } - }); - - if in_hook { - init_real_functions(); - if let Some(real_calloc) = REAL_CALLOC { - return real_calloc(nmemb, size); - } + if INITIALIZING.load(Ordering::Acquire) && !INITIALIZED.load(Ordering::Acquire) { return std::ptr::null_mut(); } init_real_functions(); - let result = if let Some(real_calloc) = REAL_CALLOC { - let total_size = nmemb * size; - let allocation_size = total_size + HEADER_SIZE; + let calloc_ptr = REAL_CALLOC.load(Ordering::Acquire); + if calloc_ptr.is_null() { + return std::ptr::null_mut(); + } - let ptr = real_calloc(allocation_size, 1); + let real_calloc: CallocFn = std::mem::transmute(calloc_ptr); + let total_size = nmemb.saturating_mul(size); + let allocation_size = total_size.saturating_add(HEADER_SIZE); - if !ptr.is_null() { - let header = ptr as *mut AllocationHeader; - (*header).size = total_size; + let ptr = real_calloc(allocation_size, 1); - STATS.record_allocation(total_size); + if !ptr.is_null() { + let header = ptr as *mut AllocationHeader; + (*header).size = total_size; - ptr.add(HEADER_SIZE) - } else { - ptr - } - } else { - std::ptr::null_mut() - }; + STATS.record_allocation(total_size); - IN_HOOK.with(|flag| flag.set(false)); - result + ptr.add(HEADER_SIZE) + } else { + ptr + } } #[no_mangle] pub unsafe extern "C" fn realloc(ptr: *mut c_void, size: size_t) -> *mut c_void { - let in_hook = IN_HOOK.with(|flag| { - if flag.get() { - true - } else { - flag.set(true); - false - } - }); - - if in_hook { - init_real_functions(); - if let Some(real_realloc) = REAL_REALLOC { - return real_realloc(ptr, size); - } + if INITIALIZING.load(Ordering::Acquire) && !INITIALIZED.load(Ordering::Acquire) { return std::ptr::null_mut(); } init_real_functions(); - let result = if let Some(real_realloc) = REAL_REALLOC { - if ptr.is_null() { - // realloc(NULL, size) is equivalent to malloc(size) - // Note: This will set the flag again, but that's handled - IN_HOOK.with(|flag| flag.set(false)); - return malloc(size); - } + let realloc_ptr = REAL_REALLOC.load(Ordering::Acquire); + if realloc_ptr.is_null() { + return std::ptr::null_mut(); + } - if size == 0 { - // realloc(ptr, 0) is equivalent to free(ptr) - IN_HOOK.with(|flag| flag.set(false)); - free(ptr); - return std::ptr::null_mut(); - } + let real_realloc: ReallocFn = std::mem::transmute(realloc_ptr); - // Get old size from header - let actual_ptr = (ptr as *mut u8).sub(HEADER_SIZE); - let old_header = actual_ptr as *mut AllocationHeader; - let old_size = (*old_header).size; + if ptr.is_null() { + // realloc(NULL, size) is equivalent to malloc(size) + return malloc(size); + } - // Reallocate with new size - let total_size = size + HEADER_SIZE; - let new_ptr = real_realloc(actual_ptr as *mut c_void, total_size); + if size == 0 { + // realloc(ptr, 0) is equivalent to free(ptr) + free(ptr); + return std::ptr::null_mut(); + } - if !new_ptr.is_null() { - // Update header with new size - let new_header = new_ptr as *mut AllocationHeader; - (*new_header).size = size; + // Get old size from header + let actual_ptr = (ptr as *mut u8).sub(HEADER_SIZE); + let old_header = actual_ptr as *mut AllocationHeader; + let old_size = (*old_header).size; - // Record the change in allocation - STATS.record_deallocation(old_size); - STATS.record_allocation(size); + // Reallocate with new size + let total_size = size.saturating_add(HEADER_SIZE); + let new_ptr = real_realloc(actual_ptr as *mut c_void, total_size); - new_ptr.add(HEADER_SIZE) - } else { - new_ptr - } - } else { - std::ptr::null_mut() - }; + if !new_ptr.is_null() { + // Update header with new size + let new_header = new_ptr as *mut AllocationHeader; + (*new_header).size = size; - IN_HOOK.with(|flag| flag.set(false)); - result + // Record the change in allocation + STATS.record_deallocation(old_size); + STATS.record_allocation(size); + + new_ptr.add(HEADER_SIZE) + } else { + new_ptr + } } diff --git a/memtest/src/lib.rs b/memtest/src/lib.rs index df712aa381d..4c869864552 100644 --- a/memtest/src/lib.rs +++ b/memtest/src/lib.rs @@ -15,20 +15,31 @@ pub struct MemtestStats { } /// Get all statistics in a single call +/// +/// # Safety +/// The `stats` pointer must be valid and properly aligned #[no_mangle] -pub extern "C" fn memtest_get_stats(stats: *mut MemtestStats) { +pub unsafe extern "C" fn memtest_get_stats(stats: *mut MemtestStats) { if stats.is_null() { return; } - unsafe { - (*stats).total_allocations = STATS.total_allocations.load(std::sync::atomic::Ordering::Relaxed); - (*stats).total_deallocations = STATS.total_deallocations.load(std::sync::atomic::Ordering::Relaxed); - (*stats).total_bytes_allocated = STATS.total_bytes_allocated.load(std::sync::atomic::Ordering::Relaxed); - (*stats).total_bytes_deallocated = STATS.total_bytes_deallocated.load(std::sync::atomic::Ordering::Relaxed); - (*stats).current_bytes = STATS.current_bytes.load(std::sync::atomic::Ordering::Relaxed); - (*stats).peak_bytes = STATS.peak_bytes.load(std::sync::atomic::Ordering::Relaxed); - } + (*stats).total_allocations = STATS + .total_allocations + .load(std::sync::atomic::Ordering::Relaxed); + (*stats).total_deallocations = STATS + .total_deallocations + .load(std::sync::atomic::Ordering::Relaxed); + (*stats).total_bytes_allocated = STATS + .total_bytes_allocated + .load(std::sync::atomic::Ordering::Relaxed); + (*stats).total_bytes_deallocated = STATS + .total_bytes_deallocated + .load(std::sync::atomic::Ordering::Relaxed); + (*stats).current_bytes = STATS + .current_bytes + .load(std::sync::atomic::Ordering::Relaxed); + (*stats).peak_bytes = STATS.peak_bytes.load(std::sync::atomic::Ordering::Relaxed); } /// Reset all statistics to zero diff --git a/memtest/src/stats.rs b/memtest/src/stats.rs index 08e3c3d129a..b482e8f22a1 100644 --- a/memtest/src/stats.rs +++ b/memtest/src/stats.rs @@ -27,7 +27,8 @@ impl AllocationStats { self.total_bytes_allocated .fetch_add(size as u64, Ordering::Relaxed); - let current = self.current_bytes.fetch_add(size as u64, Ordering::Relaxed) + size as u64; + let prev = self.current_bytes.fetch_add(size as u64, Ordering::Relaxed); + let current = prev.saturating_add(size as u64); // Update peak if necessary let mut peak = self.peak_bytes.load(Ordering::Relaxed); @@ -48,7 +49,13 @@ impl AllocationStats { self.total_deallocations.fetch_add(1, Ordering::Relaxed); self.total_bytes_deallocated .fetch_add(size as u64, Ordering::Relaxed); - self.current_bytes.fetch_sub(size as u64, Ordering::Relaxed); + + // Use fetch_update to perform saturating subtraction atomically + self.current_bytes + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { + Some(current.saturating_sub(size as u64)) + }) + .ok(); } pub fn reset(&self) { From b33e3213334de4cf44f21ce220205fff0fbee47d Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 13 Nov 2025 19:41:14 -0800 Subject: [PATCH 03/18] cleanup --- memtest/src/allocator.rs | 8 ++++---- memtest/src/stats.rs | 36 +----------------------------------- 2 files changed, 5 insertions(+), 39 deletions(-) diff --git a/memtest/src/allocator.rs b/memtest/src/allocator.rs index 4b38cf22898..3d404679890 100644 --- a/memtest/src/allocator.rs +++ b/memtest/src/allocator.rs @@ -40,10 +40,10 @@ unsafe fn init_real_functions() { } // We're the one initializing - let malloc_ptr = dlsym(RTLD_NEXT, b"malloc\0".as_ptr() as *const libc::c_char); - let free_ptr = dlsym(RTLD_NEXT, b"free\0".as_ptr() as *const libc::c_char); - let calloc_ptr = dlsym(RTLD_NEXT, b"calloc\0".as_ptr() as *const libc::c_char); - let realloc_ptr = dlsym(RTLD_NEXT, b"realloc\0".as_ptr() as *const libc::c_char); + let malloc_ptr = dlsym(RTLD_NEXT, c"malloc".as_ptr() as *const libc::c_char); + let free_ptr = dlsym(RTLD_NEXT, c"free".as_ptr() as *const libc::c_char); + let calloc_ptr = dlsym(RTLD_NEXT, c"calloc".as_ptr() as *const libc::c_char); + let realloc_ptr = dlsym(RTLD_NEXT, c"realloc".as_ptr() as *const libc::c_char); REAL_MALLOC.store(malloc_ptr, Ordering::Release); REAL_FREE.store(free_ptr, Ordering::Release); diff --git a/memtest/src/stats.rs b/memtest/src/stats.rs index b482e8f22a1..76c0253e843 100644 --- a/memtest/src/stats.rs +++ b/memtest/src/stats.rs @@ -29,20 +29,7 @@ impl AllocationStats { let prev = self.current_bytes.fetch_add(size as u64, Ordering::Relaxed); let current = prev.saturating_add(size as u64); - - // Update peak if necessary - let mut peak = self.peak_bytes.load(Ordering::Relaxed); - while current > peak { - match self.peak_bytes.compare_exchange_weak( - peak, - current, - Ordering::Relaxed, - Ordering::Relaxed, - ) { - Ok(_) => break, - Err(p) => peak = p, - } - } + self.peak_bytes.fetch_max(current, Ordering::Relaxed); } pub fn record_deallocation(&self, size: usize) { @@ -66,27 +53,6 @@ impl AllocationStats { self.current_bytes.store(0, Ordering::Relaxed); self.peak_bytes.store(0, Ordering::Relaxed); } - - pub fn get_snapshot(&self) -> StatsSnapshot { - StatsSnapshot { - total_allocations: self.total_allocations.load(Ordering::Relaxed), - total_deallocations: self.total_deallocations.load(Ordering::Relaxed), - total_bytes_allocated: self.total_bytes_allocated.load(Ordering::Relaxed), - total_bytes_deallocated: self.total_bytes_deallocated.load(Ordering::Relaxed), - current_bytes: self.current_bytes.load(Ordering::Relaxed), - peak_bytes: self.peak_bytes.load(Ordering::Relaxed), - } - } -} - -#[derive(Debug, Clone, Copy)] -pub struct StatsSnapshot { - pub total_allocations: u64, - pub total_deallocations: u64, - pub total_bytes_allocated: u64, - pub total_bytes_deallocated: u64, - pub current_bytes: u64, - pub peak_bytes: u64, } /// Global statistics instance From 9ebdb37213edfadfe3b2057196159ea8eb224a20 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 13 Nov 2025 19:51:17 -0800 Subject: [PATCH 04/18] simplify --- memtest/README.md | 196 +++-------------------------- memtest/python/memtest/__main__.py | 82 +----------- 2 files changed, 18 insertions(+), 260 deletions(-) diff --git a/memtest/README.md b/memtest/README.md index 86fa89f1dcb..4d0766d0732 100644 --- a/memtest/README.md +++ b/memtest/README.md @@ -2,200 +2,34 @@ Memory allocation testing utilities for Python test suites. This package provides tools to track memory allocations made by the Python interpreter and any Python libraries during test execution. -## Features - -- **LD_PRELOAD-based interposition**: Intercepts all `malloc`, `free`, `calloc`, and `realloc` calls -- **Zero overhead when not tracking**: No performance impact unless explicitly enabled -- **Thread-safe statistics**: Uses atomic operations for accurate multi-threaded tracking -- **Python and CLI interfaces**: Use programmatically or from the command line -- **Comprehensive metrics**: Track allocations, deallocations, current usage, and peak memory - -## Installation - -### From source - -```bash -cd memtest -maturin develop -``` - -### For development - -```bash -cd memtest -make build -``` - ## Usage -### Python API +Install with: -#### Basic tracking - -```python -import memtest - -# Reset statistics -memtest.reset_stats() - -# Your code here -data = [0] * 1000000 - -# Get statistics -stats = memtest.get_stats() -print(f"Allocated: {stats['total_bytes_allocated']} bytes") -print(f"Peak usage: {stats['peak_bytes']} bytes") +```shell +make build-release ``` -#### Context manager - -```python -import memtest - -with memtest.track() as get_stats: - # Allocate some memory - data = [0] * 1000000 +To activate the memory tracking, you need to set the `LD_PRELOAD` environment variable: - # Get stats within the context - stats = get_stats() - print(f"Allocated: {stats['total_bytes_allocated']} bytes") +```shell +export LD_PRELOAD=$(lance-memtest) ``` -#### Pretty printing +Then you can write Python code that tracks memory allocations: ```python import memtest -# ... run some code ... - -memtest.print_stats() -``` - -Output: -``` -Memory Allocation Statistics: - Total allocations: 1,234 - Total deallocations: 1,100 - Total bytes allocated: 128.5 KB - Total bytes freed: 120.0 KB - Current memory usage: 8.5 KB - Peak memory usage: 15.2 KB -``` - -### Command Line Interface - -#### Run a command with tracking +def test_memory(): + with memtest.track() as get_stats: + # Your code that allocates memory + data = [0] * 1000000 -```bash -lance-memtest run python myscript.py -lance-memtest run pytest tests/ + stats = get_stats() + assert stats['peak_bytes'] < 10**7 # Assert peak memory usage ``` -#### Get the library path - -```bash -# Print path to the .so file -lance-memtest path - -# Use with LD_PRELOAD manually -export LD_PRELOAD=$(lance-memtest path) -python myscript.py -``` - -#### View current statistics - -```bash -lance-memtest stats -``` - -### Integration with pytest - -```python -import pytest -import memtest - -@pytest.fixture(autouse=True) -def track_memory(): - """Automatically track memory for all tests.""" - memtest.reset_stats() - yield - stats = memtest.get_stats() - - # Assert memory bounds - assert stats['peak_bytes'] < 100 * 1024 * 1024, "Test used more than 100MB" - -def test_my_function(): - result = my_function() - - # Check memory usage for this test - stats = memtest.get_stats() - print(f"Peak memory: {memtest.format_bytes(stats['peak_bytes'])}") -``` - -## Statistics - -The following metrics are tracked: - -- **`total_allocations`**: Total number of `malloc`/`calloc` calls -- **`total_deallocations`**: Total number of `free` calls -- **`total_bytes_allocated`**: Total bytes allocated across all calls -- **`total_bytes_deallocated`**: Total bytes freed across all calls -- **`current_bytes`**: Current memory usage (allocated - deallocated) -- **`peak_bytes`**: Peak memory usage observed - -## How It Works - -The package uses LD_PRELOAD to interpose the standard C library allocation functions (`malloc`, `free`, `calloc`, `realloc`). When these functions are called by Python or any C extension: - -1. The interposed function records the allocation size -2. Statistics are updated using atomic operations (thread-safe) -3. The original libc function is called to perform the actual allocation - -The Rust implementation ensures minimal overhead and uses a header-based approach to track allocation sizes. - -## Limitations - -- **Linux only**: LD_PRELOAD is a Linux-specific feature -- **Does not track Python object overhead**: Only tracks C-level allocations -- **Stack allocations not tracked**: Only heap allocations via malloc family -- **Reset affects all threads**: Statistics are global - -## Development - -### Build - -```bash -make build -``` - -### Run tests - -```bash -make test -``` - -### Format code - -```bash -make format -``` - -### Lint - -```bash -make lint -``` - -## Architecture - -The package consists of: - -1. **Rust interpose library** (`src/allocator.rs`): Interposes `malloc`/`free` family -2. **Statistics module** (`src/stats.rs`): Thread-safe atomic counters -3. **PyO3 bindings** (`src/lib.rs`): Exposes stats to Python -4. **Python wrapper** (`python/memtest/__init__.py`): High-level API -5. **CLI** (`python/memtest/__main__.py`): Command-line interface - -## License +## How this works -Apache-2.0 +The library uses dynamic linking to intercept memory allocation calls (like `malloc`, `free`, etc.) made by the Python interpreter and its extensions. It keeps track of the total number of allocations, deallocations, and the peak memory usage during the execution of your code. diff --git a/memtest/python/memtest/__main__.py b/memtest/python/memtest/__main__.py index 97fa4a159a5..262b845da25 100644 --- a/memtest/python/memtest/__main__.py +++ b/memtest/python/memtest/__main__.py @@ -1,91 +1,15 @@ """CLI for lance-memtest.""" -import argparse -import os -import subprocess import sys +from memtest import get_library_path -from memtest import get_library_path, print_stats - -def cmd_path(args): - """Print the path to the memtest shared library.""" +def main(): + """Main CLI entry point - print path to shared library.""" lib_path = get_library_path() print(lib_path) return 0 -def cmd_run(args): - """Run a command with LD_PRELOAD set to track memory allocations.""" - lib_path = get_library_path() - - # Set up environment - env = os.environ.copy() - - # Prepend to LD_PRELOAD if it already exists - existing_preload = env.get("LD_PRELOAD", "") - if existing_preload: - env["LD_PRELOAD"] = f"{lib_path}:{existing_preload}" - else: - env["LD_PRELOAD"] = str(lib_path) - - # Run the command - try: - result = subprocess.run(args.command, env=env, shell=False) - return result.returncode - except FileNotFoundError: - print(f"Error: Command not found: {args.command[0]}", file=sys.stderr) - return 1 - except KeyboardInterrupt: - return 130 - - -def cmd_stats(args): - """Print current allocation statistics.""" - print_stats() - return 0 - - -def main(): - """Main CLI entry point.""" - parser = argparse.ArgumentParser( - prog="lance-memtest", - description="Memory allocation testing utilities for Python", - ) - - subparsers = parser.add_subparsers(dest="command", help="Command to run") - - # path command - path_parser = subparsers.add_parser( - "path", help="Print path to the memtest shared library" - ) - path_parser.set_defaults(func=cmd_path) - - # run command - run_parser = subparsers.add_parser( - "run", help="Run a command with memory tracking enabled" - ) - run_parser.add_argument("command", nargs="+", help="Command and arguments to run") - run_parser.set_defaults(func=cmd_run) - - # stats command - stats_parser = subparsers.add_parser( - "stats", help="Print current allocation statistics" - ) - stats_parser.set_defaults(func=cmd_stats) - - args = parser.parse_args() - - if not hasattr(args, "func"): - parser.print_help() - return 1 - - try: - return args.func(args) - except Exception as e: - print(f"Error: {e}", file=sys.stderr) - return 1 - - if __name__ == "__main__": sys.exit(main()) From 82590740d5e22facff6a769d0b68d999bfbcb93e Mon Sep 17 00:00:00 2001 From: Will Jones Date: Fri, 14 Nov 2025 11:17:00 -0800 Subject: [PATCH 05/18] wip --- memtest/src/allocator.rs | 254 +++++++++++++++------------------------ python/pyproject.toml | 1 + 2 files changed, 99 insertions(+), 156 deletions(-) diff --git a/memtest/src/allocator.rs b/memtest/src/allocator.rs index 3d404679890..9d8f5ee548f 100644 --- a/memtest/src/allocator.rs +++ b/memtest/src/allocator.rs @@ -17,192 +17,134 @@ static INITIALIZED: AtomicBool = AtomicBool::new(false); const RTLD_NEXT: *mut c_void = -1isize as *mut c_void; extern "C" { - fn dlsym(handle: *mut c_void, symbol: *const libc::c_char) -> *mut c_void; + #[link_name = "__libc_malloc"] + fn libc_malloc( size: size_t ) -> *mut c_void; + #[link_name = "__libc_calloc"] + fn libc_calloc( count: size_t, element_size: size_t ) -> *mut c_void; + #[link_name = "__libc_realloc"] + fn libc_realloc( ptr: *mut c_void, size: size_t ) -> *mut c_void; + #[link_name = "__libc_free"] + fn libc_free( ptr: *mut c_void ); + #[link_name = "__libc_memalign"] + fn libc_memalign( alignment: size_t, size: size_t ) -> *mut c_void; } -/// Initialize the function pointers to the real allocation functions -unsafe fn init_real_functions() { - // If already initialized, return - if INITIALIZED.load(Ordering::Acquire) { - return; - } - - // Try to set initializing flag - if INITIALIZING - .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) - .is_err() - { - // Someone else is initializing, spin wait - while !INITIALIZED.load(Ordering::Acquire) { - std::hint::spin_loop(); - } - return; - } +// Implementations of standard allocation functions +// To track the size on free, we store the size at the start of the allocated block +// and return a pointer offset by the size of u64 (8 bytes). - // We're the one initializing - let malloc_ptr = dlsym(RTLD_NEXT, c"malloc".as_ptr() as *const libc::c_char); - let free_ptr = dlsym(RTLD_NEXT, c"free".as_ptr() as *const libc::c_char); - let calloc_ptr = dlsym(RTLD_NEXT, c"calloc".as_ptr() as *const libc::c_char); - let realloc_ptr = dlsym(RTLD_NEXT, c"realloc".as_ptr() as *const libc::c_char); - - REAL_MALLOC.store(malloc_ptr, Ordering::Release); - REAL_FREE.store(free_ptr, Ordering::Release); - REAL_CALLOC.store(calloc_ptr, Ordering::Release); - REAL_REALLOC.store(realloc_ptr, Ordering::Release); - - INITIALIZED.store(true, Ordering::Release); +fn extract(virtual_ptr: *mut c_void) -> (usize, *mut c_void) { + let actual_ptr = (virtual_ptr as *mut u8).sub(8) as *mut u8; + let size_ptr = actual_ptr as *mut u64; + let size = unsafe { *size_ptr } as usize; + (size, actual_ptr as *mut c_void) } -/// Store allocation size in a header before the returned pointer -#[repr(C)] -struct AllocationHeader { - size: usize, +/// Take a allocated pointer and size, store the size, and return the adjusted pointer +fn to_virtual(actual_ptr: *mut c_void, size: usize) -> *mut c_void { + if actual_ptr.is_null() { + return std::ptr::null_mut(); + } + let ptr = actual_ptr as *mut u8; + unsafe { + *(ptr as *mut u64) = size as u64; + } + ptr.add(8) as *mut c_void } -const HEADER_SIZE: usize = std::mem::size_of::(); - #[no_mangle] pub unsafe extern "C" fn malloc(size: size_t) -> *mut c_void { - // If we're currently initializing, forward directly to avoid recursion - if INITIALIZING.load(Ordering::Acquire) && !INITIALIZED.load(Ordering::Acquire) { - // During initialization, dlsym might call malloc - // We can't use RTLD_NEXT here, so we'll just use a simple bump allocator - // or return null and hope dlsym handles it - return std::ptr::null_mut(); - } - - init_real_functions(); + STATS.record_allocation(size); + to_virtual(libc_malloc(size + 8), size) +} - let malloc_ptr = REAL_MALLOC.load(Ordering::Acquire); - if malloc_ptr.is_null() { +#[no_mangle] +pub unsafe extern "C" fn calloc(size: size_t, element_size: size_t) -> *mut c_void { + let Some(total_size) = size.checked_mul(element_size) else { return std::ptr::null_mut(); - } - - let real_malloc: MallocFn = std::mem::transmute(malloc_ptr); - let total_size = size.saturating_add(HEADER_SIZE); - let ptr = real_malloc(total_size); - - if !ptr.is_null() { - // Store size in header - let header = ptr as *mut AllocationHeader; - (*header).size = size; - - STATS.record_allocation(size); - - // Return pointer after header - ptr.add(HEADER_SIZE) - } else { - ptr - } + }; + STATS.record_allocation(total_size); + to_virtual(libc_calloc(total_size + 8, 1), total_size) } #[no_mangle] pub unsafe extern "C" fn free(ptr: *mut c_void) { - if ptr.is_null() { + let actual_ptr = if ptr.is_null() { return; - } - - // If called during initialization, do nothing (malloc returned null anyway) - if INITIALIZING.load(Ordering::Acquire) && !INITIALIZED.load(Ordering::Acquire) { - return; - } - - init_real_functions(); - - let free_ptr = REAL_FREE.load(Ordering::Acquire); - if free_ptr.is_null() { - return; - } - - let real_free: FreeFn = std::mem::transmute(free_ptr); - - // Get the actual allocation pointer (before header) - let actual_ptr = (ptr as *mut u8).sub(HEADER_SIZE); - let header = actual_ptr as *mut AllocationHeader; - let size = (*header).size; - + } else { + (ptr as *mut u8).sub(8) as *mut c_void + }; + let (size, ) + let size_ptr = (actual_ptr as *mut u8) as *mut u64; + let size = *size_ptr as size_t; STATS.record_deallocation(size); - - real_free(actual_ptr as *mut c_void); + libc_free(actual_ptr as *mut c_void); } #[no_mangle] -pub unsafe extern "C" fn calloc(nmemb: size_t, size: size_t) -> *mut c_void { - if INITIALIZING.load(Ordering::Acquire) && !INITIALIZED.load(Ordering::Acquire) { - return std::ptr::null_mut(); - } - - init_real_functions(); - - let calloc_ptr = REAL_CALLOC.load(Ordering::Acquire); - if calloc_ptr.is_null() { - return std::ptr::null_mut(); - } - - let real_calloc: CallocFn = std::mem::transmute(calloc_ptr); - let total_size = nmemb.saturating_mul(size); - let allocation_size = total_size.saturating_add(HEADER_SIZE); - - let ptr = real_calloc(allocation_size, 1); - - if !ptr.is_null() { - let header = ptr as *mut AllocationHeader; - (*header).size = total_size; - - STATS.record_allocation(total_size); - - ptr.add(HEADER_SIZE) - } else { +pub unsafe extern "C" fn realloc(ptr: *mut c_void, size: size_t) -> *mut c_void { + let actual_ptr = if ptr.is_null() { ptr + } else { + (ptr as *mut u8).sub(8) as *mut c_void + }; + let old_size = if !ptr.is_null() { + let size_ptr = (actual_ptr as *mut u8) as *mut u64; + *size_ptr as size_t + } else { + 0 + }; + STATS.record_deallocation(old_size); + STATS.record_allocation(size); + let new_ptr = libc_realloc(actual_ptr, size + 8); + if new_ptr.is_null() { + return std::ptr::null_mut(); } + new_ptr.add(8) as *mut c_void } #[no_mangle] -pub unsafe extern "C" fn realloc(ptr: *mut c_void, size: size_t) -> *mut c_void { - if INITIALIZING.load(Ordering::Acquire) && !INITIALIZED.load(Ordering::Acquire) { - return std::ptr::null_mut(); - } - - init_real_functions(); - - let realloc_ptr = REAL_REALLOC.load(Ordering::Acquire); - if realloc_ptr.is_null() { - return std::ptr::null_mut(); - } - - let real_realloc: ReallocFn = std::mem::transmute(realloc_ptr); +pub unsafe extern "C" fn memalign(alignment: size_t, size: size_t) -> *mut c_void { + STATS.record_allocation(size); + libc_memalign(alignment, size) +} +#[no_mangle] +pub unsafe extern "C" fn posix_memalign( + memptr: *mut *mut c_void, + alignment: size_t, + size: size_t, +) -> i32 { + let ptr = libc_memalign(alignment, size); if ptr.is_null() { - // realloc(NULL, size) is equivalent to malloc(size) - return malloc(size); + return libc::ENOMEM; } + STATS.record_allocation(size); + *memptr = ptr; + 0 +} - if size == 0 { - // realloc(ptr, 0) is equivalent to free(ptr) - free(ptr); - return std::ptr::null_mut(); - } - - // Get old size from header - let actual_ptr = (ptr as *mut u8).sub(HEADER_SIZE); - let old_header = actual_ptr as *mut AllocationHeader; - let old_size = (*old_header).size; - - // Reallocate with new size - let total_size = size.saturating_add(HEADER_SIZE); - let new_ptr = real_realloc(actual_ptr as *mut c_void, total_size); - - if !new_ptr.is_null() { - // Update header with new size - let new_header = new_ptr as *mut AllocationHeader; - (*new_header).size = size; +#[no_mangle] +pub unsafe extern "C" fn aligned_alloc(alignment: size_t, size: size_t) -> *mut c_void { + // Do we need to adjust this for alignment? + let effective_size = size + 8; + STATS.record_allocation(size); + let size = libc_memalign(alignment, size); +} - // Record the change in allocation - STATS.record_deallocation(old_size); - STATS.record_allocation(size); +#[no_mangle] +pub unsafe extern "C" fn valloc(size: size_t) -> *mut c_void { + STATS.record_allocation(size); + libc_memalign(libc::sysconf(libc::_SC_PAGESIZE) as size_t, size) +} - new_ptr.add(HEADER_SIZE) - } else { - new_ptr +#[no_mangle] +pub unsafe extern "C" fn reallocarray( old_ptr: *mut c_void, count: size_t, element_size: size_t ) -> *mut c_void { + let size = count.checked_mul(element_size); + if size.is_none() { + return std::ptr::null_mut(); } + let size = size.unwrap(); + STATS.record_allocation(size); + libc_realloc( old_ptr, size ) } diff --git a/python/pyproject.toml b/python/pyproject.toml index bffb76c33d7..2218bd0e50f 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -61,6 +61,7 @@ tests = [ "tensorflow; sys_platform == 'linux'", "tqdm", "datafusion>=50.1", + # TODO: Make memtest a dependency. ] dev = ["ruff==0.4.1", "pyright"] benchmarks = ["pytest-benchmark"] From 03b419b8e4e937c12952d81279a8428095b04e5d Mon Sep 17 00:00:00 2001 From: Will Jones Date: Fri, 14 Nov 2025 13:22:10 -0800 Subject: [PATCH 06/18] try this --- memtest/Cargo.toml | 2 +- memtest/src/allocator.rs | 211 +++++++++----- memtest/tests/integration_test.rs | 447 ++++++++++++++++++++++++++++++ 3 files changed, 585 insertions(+), 75 deletions(-) create mode 100644 memtest/tests/integration_test.rs diff --git a/memtest/Cargo.toml b/memtest/Cargo.toml index 72405566ed6..ef4cd5736ab 100644 --- a/memtest/Cargo.toml +++ b/memtest/Cargo.toml @@ -13,7 +13,7 @@ arithmetic_side_effects = "deny" [lib] name = "memtest" -crate-type = ["cdylib"] +crate-type = ["cdylib", "rlib"] [dependencies] libc = "0.2" diff --git a/memtest/src/allocator.rs b/memtest/src/allocator.rs index 9d8f5ee548f..a13f6f327da 100644 --- a/memtest/src/allocator.rs +++ b/memtest/src/allocator.rs @@ -1,61 +1,108 @@ use crate::stats::STATS; use libc::{c_void, size_t}; -use std::sync::atomic::{AtomicBool, AtomicPtr, Ordering}; - -type MallocFn = unsafe extern "C" fn(size_t) -> *mut c_void; -type FreeFn = unsafe extern "C" fn(*mut c_void); -type CallocFn = unsafe extern "C" fn(size_t, size_t) -> *mut c_void; -type ReallocFn = unsafe extern "C" fn(*mut c_void, size_t) -> *mut c_void; - -static REAL_MALLOC: AtomicPtr = AtomicPtr::new(std::ptr::null_mut()); -static REAL_FREE: AtomicPtr = AtomicPtr::new(std::ptr::null_mut()); -static REAL_CALLOC: AtomicPtr = AtomicPtr::new(std::ptr::null_mut()); -static REAL_REALLOC: AtomicPtr = AtomicPtr::new(std::ptr::null_mut()); -static INITIALIZING: AtomicBool = AtomicBool::new(false); -static INITIALIZED: AtomicBool = AtomicBool::new(false); - -const RTLD_NEXT: *mut c_void = -1isize as *mut c_void; extern "C" { #[link_name = "__libc_malloc"] - fn libc_malloc( size: size_t ) -> *mut c_void; + fn libc_malloc(size: size_t) -> *mut c_void; #[link_name = "__libc_calloc"] - fn libc_calloc( count: size_t, element_size: size_t ) -> *mut c_void; + fn libc_calloc(count: size_t, element_size: size_t) -> *mut c_void; #[link_name = "__libc_realloc"] - fn libc_realloc( ptr: *mut c_void, size: size_t ) -> *mut c_void; + fn libc_realloc(ptr: *mut c_void, size: size_t) -> *mut c_void; #[link_name = "__libc_free"] - fn libc_free( ptr: *mut c_void ); + fn libc_free(ptr: *mut c_void); #[link_name = "__libc_memalign"] - fn libc_memalign( alignment: size_t, size: size_t ) -> *mut c_void; + fn libc_memalign(alignment: size_t, size: size_t) -> *mut c_void; +} + +// Magic number to identify our allocations +const MAGIC: u64 = 0xDEADBEEF_CAFEBABE; + +/// Header stored before each tracked allocation +#[repr(C)] +struct AllocationHeader { + magic: u64, + size: u64, + alignment: u64, + /// For aligned allocations, stores the actual pointer returned by libc_memalign + /// For unaligned allocations, this is unused (but present for consistent size) + actual_ptr: u64, } -// Implementations of standard allocation functions -// To track the size on free, we store the size at the start of the allocated block -// and return a pointer offset by the size of u64 (8 bytes). +const HEADER_SIZE: usize = std::mem::size_of::(); -fn extract(virtual_ptr: *mut c_void) -> (usize, *mut c_void) { - let actual_ptr = (virtual_ptr as *mut u8).sub(8) as *mut u8; - let size_ptr = actual_ptr as *mut u64; - let size = unsafe { *size_ptr } as usize; - (size, actual_ptr as *mut c_void) +/// Check if a pointer was allocated by us +unsafe fn is_ours(virtual_ptr: *mut c_void) -> bool { + if virtual_ptr.is_null() { + return false; + } + let header_ptr = (virtual_ptr as *mut u8).sub(HEADER_SIZE) as *const AllocationHeader; + (*header_ptr).magic == MAGIC } -/// Take a allocated pointer and size, store the size, and return the adjusted pointer -fn to_virtual(actual_ptr: *mut c_void, size: usize) -> *mut c_void { +/// Extract size, alignment, and actual pointer from a virtual pointer +unsafe fn extract(virtual_ptr: *mut c_void) -> (usize, usize, *mut c_void) { + let header_ptr = (virtual_ptr as *mut u8).sub(HEADER_SIZE) as *const AllocationHeader; + let header = &*header_ptr; + + let size = header.size as usize; + let alignment = header.alignment as usize; + + let actual_ptr = if alignment > 0 { + // For aligned allocations, the actual pointer is stored in the header + header.actual_ptr as *mut c_void + } else { + // For unaligned allocations, the actual pointer is the header itself + header_ptr as *mut c_void + }; + + (size, alignment, actual_ptr) +} + +/// Take an allocated pointer and size, store header, and return the adjusted pointer +unsafe fn to_virtual(actual_ptr: *mut c_void, size: usize, alignment: usize) -> *mut c_void { if actual_ptr.is_null() { return std::ptr::null_mut(); } - let ptr = actual_ptr as *mut u8; - unsafe { - *(ptr as *mut u64) = size as u64; + + if alignment > 0 { + // For aligned allocations: + // 1. Find the first aligned position after we have room for the header + // 2. Store the header just before that position + // 3. Store the actual_ptr in the header so we can free it later + + let actual_addr = actual_ptr as usize; + // Find the first address >= actual_addr + HEADER_SIZE that is aligned + let min_virtual_addr = actual_addr.saturating_add(HEADER_SIZE); + let virtual_addr = (min_virtual_addr.saturating_add(alignment).saturating_sub(1)) + & !(alignment.saturating_sub(1)); + + // Write header just before the aligned virtual address + let header_ptr = (virtual_addr.saturating_sub(HEADER_SIZE)) as *mut AllocationHeader; + *header_ptr = AllocationHeader { + magic: MAGIC, + size: size as u64, + alignment: alignment as u64, + actual_ptr: actual_addr as u64, + }; + + virtual_addr as *mut c_void + } else { + // Unaligned allocation - header is at the start + let header_ptr = actual_ptr as *mut AllocationHeader; + *header_ptr = AllocationHeader { + magic: MAGIC, + size: size as u64, + alignment: 0, + actual_ptr: 0, // Unused for unaligned allocations + }; + (actual_ptr as *mut u8).add(HEADER_SIZE) as *mut c_void } - ptr.add(8) as *mut c_void } #[no_mangle] pub unsafe extern "C" fn malloc(size: size_t) -> *mut c_void { STATS.record_allocation(size); - to_virtual(libc_malloc(size + 8), size) + to_virtual(libc_malloc(size.saturating_add(HEADER_SIZE)), size, 0) } #[no_mangle] @@ -64,49 +111,60 @@ pub unsafe extern "C" fn calloc(size: size_t, element_size: size_t) -> *mut c_vo return std::ptr::null_mut(); }; STATS.record_allocation(total_size); - to_virtual(libc_calloc(total_size + 8, 1), total_size) + to_virtual( + libc_calloc(total_size.saturating_add(HEADER_SIZE), 1), + total_size, + 0, + ) } #[no_mangle] pub unsafe extern "C" fn free(ptr: *mut c_void) { - let actual_ptr = if ptr.is_null() { + if ptr.is_null() { return; + } + + if is_ours(ptr) { + // It's ours - extract size and track + let (size, _alignment, actual_ptr) = extract(ptr); + STATS.record_deallocation(size); + libc_free(actual_ptr); } else { - (ptr as *mut u8).sub(8) as *mut c_void - }; - let (size, ) - let size_ptr = (actual_ptr as *mut u8) as *mut u64; - let size = *size_ptr as size_t; - STATS.record_deallocation(size); - libc_free(actual_ptr as *mut c_void); + // Not ours - just free it without tracking + libc_free(ptr); + } } #[no_mangle] pub unsafe extern "C" fn realloc(ptr: *mut c_void, size: size_t) -> *mut c_void { - let actual_ptr = if ptr.is_null() { - ptr - } else { - (ptr as *mut u8).sub(8) as *mut c_void - }; - let old_size = if !ptr.is_null() { - let size_ptr = (actual_ptr as *mut u8) as *mut u64; - *size_ptr as size_t + let (old_size, actual_ptr) = if ptr.is_null() || !is_ours(ptr) { + // Either null or not ours - don't track + if ptr.is_null() { + (0, std::ptr::null_mut()) + } else { + // Not ours - just realloc without tracking + return libc_realloc(ptr, size); + } } else { - 0 + let (s, _align, a) = extract(ptr); + (s, a) }; + STATS.record_deallocation(old_size); STATS.record_allocation(size); - let new_ptr = libc_realloc(actual_ptr, size + 8); - if new_ptr.is_null() { - return std::ptr::null_mut(); - } - new_ptr.add(8) as *mut c_void + + let new_ptr = libc_realloc(actual_ptr, size.saturating_add(HEADER_SIZE)); + to_virtual(new_ptr, size, 0) } #[no_mangle] pub unsafe extern "C" fn memalign(alignment: size_t, size: size_t) -> *mut c_void { STATS.record_allocation(size); - libc_memalign(alignment, size) + // Allocate extra space for header + padding to maintain alignment + // We need: header (24 bytes) + actual_ptr (8 bytes) + padding to reach alignment + let extra = alignment.saturating_add(HEADER_SIZE).saturating_add(8); + let actual_ptr = libc_memalign(alignment, size.saturating_add(extra)); + to_virtual(actual_ptr, size, alignment) } #[no_mangle] @@ -115,36 +173,41 @@ pub unsafe extern "C" fn posix_memalign( alignment: size_t, size: size_t, ) -> i32 { - let ptr = libc_memalign(alignment, size); - if ptr.is_null() { + STATS.record_allocation(size); + let extra = alignment.saturating_add(HEADER_SIZE).saturating_add(8); + let actual_ptr = libc_memalign(alignment, size.saturating_add(extra)); + if actual_ptr.is_null() { return libc::ENOMEM; } - STATS.record_allocation(size); - *memptr = ptr; + *memptr = to_virtual(actual_ptr, size, alignment); 0 } #[no_mangle] pub unsafe extern "C" fn aligned_alloc(alignment: size_t, size: size_t) -> *mut c_void { - // Do we need to adjust this for alignment? - let effective_size = size + 8; STATS.record_allocation(size); - let size = libc_memalign(alignment, size); + let extra = alignment.saturating_add(HEADER_SIZE).saturating_add(8); + let actual_ptr = libc_memalign(alignment, size.saturating_add(extra)); + to_virtual(actual_ptr, size, alignment) } #[no_mangle] pub unsafe extern "C" fn valloc(size: size_t) -> *mut c_void { STATS.record_allocation(size); - libc_memalign(libc::sysconf(libc::_SC_PAGESIZE) as size_t, size) + let page_size = libc::sysconf(libc::_SC_PAGESIZE) as size_t; + let extra = page_size.saturating_add(HEADER_SIZE).saturating_add(8); + let actual_ptr = libc_memalign(page_size, size.saturating_add(extra)); + to_virtual(actual_ptr, size, page_size) } #[no_mangle] -pub unsafe extern "C" fn reallocarray( old_ptr: *mut c_void, count: size_t, element_size: size_t ) -> *mut c_void { - let size = count.checked_mul(element_size); - if size.is_none() { +pub unsafe extern "C" fn reallocarray( + old_ptr: *mut c_void, + count: size_t, + element_size: size_t, +) -> *mut c_void { + let Some(size) = count.checked_mul(element_size) else { return std::ptr::null_mut(); - } - let size = size.unwrap(); - STATS.record_allocation(size); - libc_realloc( old_ptr, size ) + }; + realloc(old_ptr, size) } diff --git a/memtest/tests/integration_test.rs b/memtest/tests/integration_test.rs new file mode 100644 index 00000000000..b83b50cd3d9 --- /dev/null +++ b/memtest/tests/integration_test.rs @@ -0,0 +1,447 @@ +use libc::{c_void, size_t}; +use std::ptr; + +// Import from the library we're testing +use memtest::{memtest_get_stats, memtest_reset_stats, MemtestStats}; + +extern "C" { + fn malloc(size: size_t) -> *mut c_void; + fn calloc(count: size_t, element_size: size_t) -> *mut c_void; + fn realloc(ptr: *mut c_void, size: size_t) -> *mut c_void; + fn free(ptr: *mut c_void); + fn memalign(alignment: size_t, size: size_t) -> *mut c_void; + fn posix_memalign(memptr: *mut *mut c_void, alignment: size_t, size: size_t) -> i32; + fn aligned_alloc(alignment: size_t, size: size_t) -> *mut c_void; +} + +fn get_stats() -> MemtestStats { + let mut stats = MemtestStats { + total_allocations: 0, + total_deallocations: 0, + total_bytes_allocated: 0, + total_bytes_deallocated: 0, + current_bytes: 0, + peak_bytes: 0, + }; + unsafe { + memtest_get_stats(&mut stats as *mut MemtestStats); + } + stats +} + +fn reset_stats() { + memtest_reset_stats(); +} + +#[test] +fn test_malloc_free() { + unsafe { + reset_stats(); + let stats_after_reset = get_stats(); + + let size = 1024; + let ptr = malloc(size); + assert!(!ptr.is_null()); + + let stats_after_alloc = get_stats(); + // Check delta from reset + assert_eq!( + stats_after_alloc + .total_allocations + .saturating_sub(stats_after_reset.total_allocations), + 1 + ); + assert_eq!( + stats_after_alloc + .total_bytes_allocated + .saturating_sub(stats_after_reset.total_bytes_allocated), + size as u64 + ); + + free(ptr); + + let stats_after_free = get_stats(); + assert_eq!( + stats_after_free + .total_deallocations + .saturating_sub(stats_after_reset.total_deallocations), + 1 + ); + assert_eq!( + stats_after_free + .total_bytes_deallocated + .saturating_sub(stats_after_reset.total_bytes_deallocated), + size as u64 + ); + } +} + +#[test] +fn test_calloc_free() { + unsafe { + reset_stats(); + let stats_baseline = get_stats(); + + let count = 10; + let element_size = 100; + let total_size = count * element_size; + + let ptr = calloc(count, element_size); + assert!(!ptr.is_null()); + + // Verify memory is zeroed + let slice = std::slice::from_raw_parts(ptr as *const u8, total_size); + assert!(slice.iter().all(|&b| b == 0)); + + let stats = get_stats(); + assert_eq!( + stats + .total_allocations + .saturating_sub(stats_baseline.total_allocations), + 1 + ); + assert_eq!( + stats + .total_bytes_allocated + .saturating_sub(stats_baseline.total_bytes_allocated), + total_size as u64 + ); + + free(ptr); + + let stats = get_stats(); + assert_eq!( + stats + .total_deallocations + .saturating_sub(stats_baseline.total_deallocations), + 1 + ); + } +} + +#[test] +fn test_realloc() { + reset_stats(); + + unsafe { + // Start with malloc + let ptr1 = malloc(100); + assert!(!ptr1.is_null()); + + let stats = get_stats(); + assert_eq!(stats.total_allocations, 1); + assert_eq!(stats.total_bytes_allocated, 100); + + // Grow the allocation + let ptr2 = realloc(ptr1, 200); + assert!(!ptr2.is_null()); + + let stats = get_stats(); + assert_eq!(stats.total_allocations, 2); // realloc counts as new allocation + assert_eq!(stats.total_deallocations, 1); // old allocation freed + assert_eq!(stats.total_bytes_allocated, 300); // 100 + 200 + assert_eq!(stats.total_bytes_deallocated, 100); + assert_eq!(stats.current_bytes, 200); + + // Shrink the allocation + let ptr3 = realloc(ptr2, 50); + assert!(!ptr3.is_null()); + + let stats = get_stats(); + assert_eq!(stats.total_allocations, 3); + assert_eq!(stats.total_deallocations, 2); + assert_eq!(stats.current_bytes, 50); + + free(ptr3); + + let stats = get_stats(); + assert_eq!(stats.current_bytes, 0); + } +} + +#[test] +fn test_realloc_null_is_malloc() { + reset_stats(); + + unsafe { + // realloc(NULL, size) should behave like malloc + let ptr = realloc(ptr::null_mut(), 100); + assert!(!ptr.is_null()); + + let stats = get_stats(); + assert_eq!(stats.total_allocations, 1); + assert_eq!(stats.total_bytes_allocated, 100); + + free(ptr); + } +} + +#[test] +fn test_peak_tracking() { + unsafe { + reset_stats(); + let stats_baseline = get_stats(); + + let ptr1 = malloc(1000); + let ptr2 = malloc(500); + let ptr3 = malloc(2000); + + let stats = get_stats(); + let current_bytes = stats + .current_bytes + .saturating_sub(stats_baseline.current_bytes); + let peak_bytes = stats.peak_bytes.saturating_sub(stats_baseline.peak_bytes); + assert_eq!(current_bytes, 3500); + assert_eq!(peak_bytes, 3500); + + free(ptr3); + + let stats = get_stats(); + let current_bytes = stats + .current_bytes + .saturating_sub(stats_baseline.current_bytes); + let peak_bytes = stats.peak_bytes.saturating_sub(stats_baseline.peak_bytes); + assert_eq!(current_bytes, 1500); + assert_eq!(peak_bytes, 3500); // Peak should remain + + let ptr4 = malloc(1000); + + let stats = get_stats(); + let current_bytes = stats + .current_bytes + .saturating_sub(stats_baseline.current_bytes); + let peak_bytes = stats.peak_bytes.saturating_sub(stats_baseline.peak_bytes); + assert_eq!(current_bytes, 2500); + assert_eq!(peak_bytes, 3500); // Still the peak + + free(ptr1); + free(ptr2); + free(ptr4); + } +} + +#[test] +fn test_memalign() { + unsafe { + reset_stats(); + let stats_baseline = get_stats(); + + let alignment = 128; + let size = 1024; + + let ptr = memalign(alignment, size); + assert!(!ptr.is_null()); + + // Verify alignment + assert_eq!(ptr as usize % alignment, 0); + + let stats = get_stats(); + assert_eq!( + stats + .total_allocations + .saturating_sub(stats_baseline.total_allocations), + 1 + ); + assert_eq!( + stats + .total_bytes_allocated + .saturating_sub(stats_baseline.total_bytes_allocated), + size as u64 + ); + + free(ptr); + + let stats = get_stats(); + assert_eq!( + stats + .total_deallocations + .saturating_sub(stats_baseline.total_deallocations), + 1 + ); + } +} + +#[test] +fn test_posix_memalign() { + unsafe { + reset_stats(); + let stats_baseline = get_stats(); + + let alignment = 256; + let size = 2048; + let mut ptr: *mut c_void = ptr::null_mut(); + + let ret = posix_memalign(&mut ptr as *mut *mut c_void, alignment, size); + assert_eq!(ret, 0); + assert!(!ptr.is_null()); + + // Verify alignment + assert_eq!(ptr as usize % alignment, 0); + + let stats = get_stats(); + assert_eq!( + stats + .total_allocations + .saturating_sub(stats_baseline.total_allocations), + 1 + ); + assert_eq!( + stats + .total_bytes_allocated + .saturating_sub(stats_baseline.total_bytes_allocated), + size as u64 + ); + + free(ptr); + } +} + +#[test] +fn test_aligned_alloc() { + reset_stats(); + + unsafe { + let alignment = 64; + let size = 512; + + let ptr = aligned_alloc(alignment, size); + assert!(!ptr.is_null()); + + // Verify alignment + assert_eq!(ptr as usize % alignment, 0); + + let stats = get_stats(); + assert_eq!(stats.total_allocations, 1); + assert_eq!(stats.total_bytes_allocated, size as u64); + + free(ptr); + } +} + +#[test] +fn test_large_alignment() { + reset_stats(); + + unsafe { + // Test with page-sized alignment (4096 bytes) + let alignment = 4096; + let size = 8192; + + let ptr = memalign(alignment, size); + assert!(!ptr.is_null()); + assert_eq!(ptr as usize % alignment, 0); + + // Write to the memory to ensure it's actually usable + let slice = std::slice::from_raw_parts_mut(ptr as *mut u8, size); + slice[0] = 42; + slice[size - 1] = 43; + assert_eq!(slice[0], 42); + assert_eq!(slice[size - 1], 43); + + let stats = get_stats(); + assert_eq!(stats.total_allocations, 1); + assert_eq!(stats.total_bytes_allocated, size as u64); + + free(ptr); + + let stats = get_stats(); + assert_eq!(stats.current_bytes, 0); + } +} + +#[test] +fn test_mixed_aligned_unaligned() { + reset_stats(); + + unsafe { + let ptr1 = malloc(1000); // Unaligned + let ptr2 = memalign(128, 2000); // Aligned + let ptr3 = malloc(500); // Unaligned + let ptr4 = aligned_alloc(64, 1500); // Aligned + + let stats = get_stats(); + assert_eq!(stats.total_allocations, 4); + assert_eq!(stats.total_bytes_allocated, 5000); + assert_eq!(stats.current_bytes, 5000); + + // Verify alignments + assert_eq!(ptr2 as usize % 128, 0); + assert_eq!(ptr4 as usize % 64, 0); + + free(ptr1); + free(ptr2); + free(ptr3); + free(ptr4); + + let stats = get_stats(); + assert_eq!(stats.total_deallocations, 4); + assert_eq!(stats.current_bytes, 0); + } +} + +#[test] +fn test_free_null() { + reset_stats(); + + unsafe { + // Freeing null should not crash or affect stats + free(ptr::null_mut()); + + let stats = get_stats(); + assert_eq!(stats.total_deallocations, 0); + } +} + +#[test] +fn test_reset_stats() { + unsafe { + let ptr1 = malloc(1000); + let ptr2 = malloc(2000); + + let stats = get_stats(); + assert!(stats.total_allocations > 0); + assert!(stats.total_bytes_allocated > 0); + + reset_stats(); + + let stats = get_stats(); + assert_eq!(stats.total_allocations, 0); + assert_eq!(stats.total_deallocations, 0); + assert_eq!(stats.total_bytes_allocated, 0); + assert_eq!(stats.total_bytes_deallocated, 0); + assert_eq!(stats.current_bytes, 0); + assert_eq!(stats.peak_bytes, 0); + + // Clean up (stats won't count these since we reset) + free(ptr1); + free(ptr2); + } +} + +#[test] +fn test_alignment_with_write() { + reset_stats(); + + unsafe { + // Test that aligned allocations are actually writable + let alignment = 256; + let size = 1024; + + let ptr = memalign(alignment, size); + assert!(!ptr.is_null()); + assert_eq!(ptr as usize % alignment, 0); + + // Write pattern to memory + let slice = std::slice::from_raw_parts_mut(ptr as *mut u8, size); + for (i, byte) in slice.iter_mut().enumerate() { + *byte = (i % 256) as u8; + } + + // Verify pattern + for (i, byte) in slice.iter().enumerate() { + assert_eq!(*byte, (i % 256) as u8); + } + + free(ptr); + } +} From 9d6b55e1e14ce17ebe69ac45f7a93baf6aa153e5 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Fri, 14 Nov 2025 13:34:24 -0800 Subject: [PATCH 07/18] python test --- python/python/tests/test_memory.py | 32 ++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) create mode 100644 python/python/tests/test_memory.py diff --git a/python/python/tests/test_memory.py b/python/python/tests/test_memory.py new file mode 100644 index 00000000000..f674514e123 --- /dev/null +++ b/python/python/tests/test_memory.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +from pathlib import Path + +import lance +import memtest +import pyarrow as pa + + +def test_insert_memory(tmp_path: Path): + def batch_generator(): + # 5MB batches -> 100MB total + for _ in range(20): + yield pa.RecordBatch.from_arrays( + [pa.array([b"x" * 1024 * 1024] * 5)], names=["data"] + ) + + reader = pa.RecordBatchReader.from_batches( + schema=pa.schema([("data", pa.binary())]), + batches=batch_generator(), + ) + + with memtest.track() as get_stats: + ds = lance.write_dataset( + reader, + tmp_path / "test.lance", + ) + stats = get_stats() + + assert stats["peak_bytes"] >= 5 * 1024 * 1024 + assert stats["peak_bytes"] < 30 * 1024 * 1024 From fd9c04cdc9d70f33bc65288ac4c84b26909dcb68 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 11 Dec 2025 10:38:20 -0800 Subject: [PATCH 08/18] fix missing API --- memtest/src/allocator.rs | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/memtest/src/allocator.rs b/memtest/src/allocator.rs index a13f6f327da..d3c9fca75cc 100644 --- a/memtest/src/allocator.rs +++ b/memtest/src/allocator.rs @@ -211,3 +211,19 @@ pub unsafe extern "C" fn reallocarray( }; realloc(old_ptr, size) } + +#[no_mangle] +pub unsafe extern "C" fn malloc_usable_size(ptr: *mut c_void) -> size_t { + if ptr.is_null() { + return 0; + } + + if is_ours(ptr) { + let (size, _, _) = extract(ptr); + size + } else { + // Not our allocation - return 0 as we don't know the size + // (there's no __libc_malloc_usable_size to call) + 0 + } +} From 131beaf634baed3155532082902acb89e7cafa0d Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 11 Dec 2025 12:30:17 -0800 Subject: [PATCH 09/18] wip --- memtest/pyproject.toml | 3 + memtest/python/memtest/__init__.py | 21 ++ memtest/python/memtest/pytest_plugin.py | 203 ++++++++++++++++++ memtest/python/tests/test_benchmark_plugin.py | 44 ++++ python/python/ci_benchmarks/README.md | 83 +++++++ .../ci_benchmarks/benchmarks/test_search.py | 32 +++ python/python/ci_benchmarks/datagen/basic.py | 2 +- .../python/ci_benchmarks/datagen/lineitems.py | 2 +- 8 files changed, 388 insertions(+), 2 deletions(-) create mode 100644 memtest/python/memtest/pytest_plugin.py create mode 100644 memtest/python/tests/test_benchmark_plugin.py create mode 100644 python/python/ci_benchmarks/README.md diff --git a/memtest/pyproject.toml b/memtest/pyproject.toml index 396d7c442e0..4cc765b8b0f 100644 --- a/memtest/pyproject.toml +++ b/memtest/pyproject.toml @@ -35,3 +35,6 @@ memtest = "python/memtest" [tool.setuptools.package-data] memtest = ["*.so", "*.dylib", "*.dll"] + +[project.entry-points.pytest11] +memtest = "memtest.pytest_plugin" diff --git a/memtest/python/memtest/__init__.py b/memtest/python/memtest/__init__.py index 0a4b58b4515..8e38bcadec2 100644 --- a/memtest/python/memtest/__init__.py +++ b/memtest/python/memtest/__init__.py @@ -175,6 +175,26 @@ def print_stats(stats: Optional[Dict[str, int]] = None) -> None: print(f" Peak memory usage: {format_bytes(stats['peak_bytes'])}") +def is_preloaded() -> bool: + """Check if libmemtest.so is preloaded and actively tracking allocations. + + Returns: + True if the library is preloaded via LD_PRELOAD, False otherwise. + + Example: + >>> if is_preloaded(): + ... stats = get_stats() + ... print(f"Tracking {stats['total_allocations']} allocations") + """ + try: + stats = get_stats() + # If we can get stats and there's been any activity, we're preloaded + # Even with no activity, if the library loads we're preloaded + return True + except Exception: + return False + + __all__ = [ "get_library_path", "get_stats", @@ -182,4 +202,5 @@ def print_stats(stats: Optional[Dict[str, int]] = None) -> None: "track", "format_bytes", "print_stats", + "is_preloaded", ] diff --git a/memtest/python/memtest/pytest_plugin.py b/memtest/python/memtest/pytest_plugin.py new file mode 100644 index 00000000000..5ab70330430 --- /dev/null +++ b/memtest/python/memtest/pytest_plugin.py @@ -0,0 +1,203 @@ +"""Pytest plugin for memory tracking during benchmarks. + +This plugin provides a `memory_benchmark` fixture that wraps pytest-benchmark +to track memory allocations during the actual benchmark execution. + +The plugin auto-detects if libmemtest.so is preloaded. If not, the +`memory_benchmark` fixture simply passes through to the regular `benchmark` +fixture. + +Usage: + def test_something(memory_benchmark): + memory_benchmark(my_function, arg1, arg2) + +Output: + - Terminal summary with memory stats per benchmark + - BMF JSON file for bencher.dev upload (--memory-json option) +""" + +import json +from functools import wraps +from typing import Any, Callable, Dict, Optional + +import pytest + +from . import format_bytes, get_stats, is_preloaded, reset_stats + +# Global storage for memory results across all tests +_memory_results: Dict[str, Dict[str, int]] = {} + + +def pytest_addoption(parser: pytest.Parser) -> None: + """Add command-line options for memory tracking.""" + group = parser.getgroup("memory", "memory tracking options") + group.addoption( + "--memory-json", + action="store", + default=None, + metavar="PATH", + help="Output path for memory stats JSON in Bencher Metric Format (BMF)", + ) + + +class MemoryTrackingBenchmark: + """Wrapper around pytest-benchmark that tracks memory during execution.""" + + def __init__(self, benchmark: Any, test_name: str): + self._benchmark = benchmark + self._test_name = test_name + self._peak_memory = 0 + self._total_allocations = 0 + + def _wrap_function(self, func: Callable) -> Callable: + """Wrap a function to track memory around each invocation.""" + + @wraps(func) + def wrapper(*args, **kwargs): + reset_stats() + result = func(*args, **kwargs) + stats = get_stats() + # Track max peak across iterations + self._peak_memory = max(self._peak_memory, stats["peak_bytes"]) + self._total_allocations += stats["total_allocations"] + return result + + return wrapper + + def __call__(self, func: Callable, *args, **kwargs) -> Any: + """Run benchmark with memory tracking.""" + wrapped = self._wrap_function(func) + return self._benchmark(wrapped, *args, **kwargs) + + def pedantic( + self, + func: Callable, + args: tuple = (), + kwargs: Optional[Dict] = None, + setup: Optional[Callable] = None, + teardown: Optional[Callable] = None, + rounds: int = 1, + warmup_rounds: int = 0, + iterations: int = 1, + ) -> Any: + """Run pedantic benchmark with memory tracking.""" + kwargs = kwargs or {} + wrapped = self._wrap_function(func) + return self._benchmark.pedantic( + wrapped, + args=args, + kwargs=kwargs, + setup=setup, + teardown=teardown, + rounds=rounds, + warmup_rounds=warmup_rounds, + iterations=iterations, + ) + + @property + def group(self): + return self._benchmark.group + + @group.setter + def group(self, value): + self._benchmark.group = value + + @property + def name(self): + return self._benchmark.name + + @property + def extra_info(self): + return self._benchmark.extra_info + + @extra_info.setter + def extra_info(self, value): + self._benchmark.extra_info = value + + def get_memory_stats(self) -> Dict[str, int]: + """Get the collected memory statistics.""" + return { + "peak_bytes": self._peak_memory, + "total_allocations": self._total_allocations, + } + + +@pytest.fixture +def memory_benchmark(benchmark, request): + """Fixture that wraps benchmark to track memory during execution. + + If libmemtest.so is not preloaded, this fixture simply returns the + regular benchmark fixture unchanged. + + Usage: + def test_something(memory_benchmark): + memory_benchmark(my_function, arg1, arg2) + + def test_pedantic(memory_benchmark): + memory_benchmark.pedantic(my_function, rounds=5, iterations=10) + """ + if not is_preloaded(): + # Not preloaded - just return regular benchmark + yield benchmark + return + + test_name = request.node.name + tracker = MemoryTrackingBenchmark(benchmark, test_name) + + yield tracker + + # Store results after test completes + stats = tracker.get_memory_stats() + if stats["peak_bytes"] > 0 or stats["total_allocations"] > 0: + _memory_results[test_name] = stats + + +def pytest_terminal_summary(terminalreporter, exitstatus: int, config) -> None: + """Print memory statistics summary at the end of the test run.""" + if not _memory_results: + return + + terminalreporter.write_sep("=", "Memory Statistics") + + # Calculate column widths + name_width = max(len(name) for name in _memory_results.keys()) + name_width = max(name_width, len("Test")) + + # Header + terminalreporter.write_line( + f"{'Test':<{name_width}} {'Peak Memory':>12} {'Allocations':>12}" + ) + terminalreporter.write_line("-" * (name_width + 28)) + + # Results sorted by peak memory (descending) + sorted_results = sorted( + _memory_results.items(), key=lambda x: x[1]["peak_bytes"], reverse=True + ) + + for test_name, stats in sorted_results: + peak = format_bytes(stats["peak_bytes"]) + allocs = f"{stats['total_allocations']:,}" + terminalreporter.write_line(f"{test_name:<{name_width}} {peak:>12} {allocs:>12}") + + terminalreporter.write_line("") + + +def pytest_sessionfinish(session, exitstatus: int) -> None: + """Write memory results to JSON file if --memory-json was specified.""" + if not _memory_results: + return + + output_path = session.config.getoption("--memory-json") + if not output_path: + return + + # Convert to Bencher Metric Format (BMF) + bmf_output = {} + for test_name, stats in _memory_results.items(): + bmf_output[test_name] = { + "peak_memory_bytes": {"value": stats["peak_bytes"]}, + "total_allocations": {"value": stats["total_allocations"]}, + } + + with open(output_path, "w") as f: + json.dump(bmf_output, f, indent=2) diff --git a/memtest/python/tests/test_benchmark_plugin.py b/memtest/python/tests/test_benchmark_plugin.py new file mode 100644 index 00000000000..f4d1dd3337d --- /dev/null +++ b/memtest/python/tests/test_benchmark_plugin.py @@ -0,0 +1,44 @@ +"""Test the pytest-benchmark memory tracking plugin.""" + +import pytest + + +def allocate_memory(size_mb: int) -> list: + """Allocate approximately size_mb of memory.""" + # Each int in Python takes about 28 bytes, but in a list it's stored as a pointer + # A list of zeros: each element is ~8 bytes for the pointer + shared int object + # For a rough approximation, 1MB ~= 125000 elements + return [0] * (size_mb * 125000) + + +def test_basic_memory_tracking(memory_benchmark): + """Test that memory is tracked during benchmark execution.""" + + def workload(): + data = allocate_memory(10) # ~10MB + return len(data) + + result = memory_benchmark(workload) + assert result == 10 * 125000 + + +def test_pedantic_mode(memory_benchmark): + """Test memory tracking with pedantic mode.""" + + def workload(): + data = allocate_memory(5) # ~5MB + return sum(data) + + result = memory_benchmark.pedantic(workload, rounds=3, iterations=1) + assert result == 0 + + +def test_with_arguments(memory_benchmark): + """Test memory tracking with function arguments.""" + + def workload(multiplier: int, base_size: int = 1): + data = allocate_memory(base_size * multiplier) + return len(data) + + result = memory_benchmark(workload, 2, base_size=3) # 6MB + assert result == 6 * 125000 diff --git a/python/python/ci_benchmarks/README.md b/python/python/ci_benchmarks/README.md new file mode 100644 index 00000000000..a8742949a4e --- /dev/null +++ b/python/python/ci_benchmarks/README.md @@ -0,0 +1,83 @@ +# CI Benchmarks + +This directory contains benchmarks that run in CI and report results to [bencher.dev](https://bencher.dev). + +## Structure + +``` +ci_benchmarks/ +├── benchmarks/ # Benchmark tests +│ ├── test_scan.py +│ ├── test_search.py +│ └── test_random_access.py +├── datagen/ # Dataset generation scripts +│ ├── gen_all.py # Generate all datasets +│ ├── basic.py # 10M row dataset +│ └── lineitems.py # TPC-H lineitem dataset +└── datasets.py # Dataset URI resolver (local vs GCS) +``` + +## Running Benchmarks Locally + +### 1. Generate test datasets + +```bash +python python/ci_benchmarks/datagen/gen_all.py +``` + +This creates datasets in `~/lance-benchmarks-ci-datasets/`. + +### 2. Run benchmarks + +```bash +pytest python/ci_benchmarks/ --benchmark-only +``` + +To save results as JSON: + +```bash +pytest python/ci_benchmarks/ --benchmark-json results.json +``` + +## Running with Memory Tracking (Linux-only) + +To track memory allocations during benchmarks, use the `lance-memtest` library with `LD_PRELOAD`. + +### 1. Install lance-memtest + +```bash +pip install lance-memtest +``` + +### 2. Run with memory tracking + +```bash +LD_PRELOAD=$(lance-memtest) pytest python/ci_benchmarks/ \ + --benchmark-json timing_results.json \ + --memory-json memory_results.json +``` + +This produces: +- `timing_results.json` - Standard pytest-benchmark timing results +- `memory_results.json` - Memory stats in Bencher Metric Format (BMF) + +### 3. Using memory_benchmark fixture + +For benchmarks that need memory tracking, use the `memory_benchmark` fixture instead of `benchmark`: + +```python +def test_full_scan(memory_benchmark, dataset): + memory_benchmark(dataset.to_table) +``` + +When `LD_PRELOAD` is not set, `memory_benchmark` passes through to the regular `benchmark` fixture. + +## Uploading to Bencher + +```bash +# Upload timing results +bencher run --adapter python_pytest --file timing_results.json + +# Upload memory results +bencher run --adapter json --file memory_results.json +``` diff --git a/python/python/ci_benchmarks/benchmarks/test_search.py b/python/python/ci_benchmarks/benchmarks/test_search.py index 7f0eb2f84b3..d257c835d8c 100644 --- a/python/python/ci_benchmarks/benchmarks/test_search.py +++ b/python/python/ci_benchmarks/benchmarks/test_search.py @@ -244,3 +244,35 @@ def clear_timer(): benchmark.pedantic( bench, warmup_rounds=1, rounds=1, iterations=1, setup=clear_timer ) + +@pytest.mark.io_memory_benchmark() +@pytest.mark.parametrize("filt", BASIC_BTREE_FILTERS, ids=BASIC_BTREE_FILTER_LABELS) +@pytest.mark.parametrize("payload", ["small_strings", "integers"]) +def test_basic_btree_search(benchmark, filt: str | None, payload: str): + dataset_uri = get_dataset_uri("basic") + ds = lance.dataset(dataset_uri) + + columns = [] + if payload is not None: + columns = [payload] + + def bench(dataset): + dataset.to_table( + columns=[payload], + filter=filt, + with_row_id=True, + batch_size=32 * 1024, + ) + + benchmark(bench, ds) + +def benchmark(bench, ds, warmup: bool=True): + if warmup: + bench(ds) + ds.io_stats_incremental() + with memtest.track() as get_stats: + bench(ds) + memory_stats = get_stats() + io_stats = ds.io_stats_incremental() + # TODO: Record the I/O and memory stats + diff --git a/python/python/ci_benchmarks/datagen/basic.py b/python/python/ci_benchmarks/datagen/basic.py index c14d7dcb47a..fa8a8845894 100644 --- a/python/python/ci_benchmarks/datagen/basic.py +++ b/python/python/ci_benchmarks/datagen/basic.py @@ -12,7 +12,7 @@ from ci_benchmarks.datasets import get_dataset_uri -NUM_ROWS = 10_000_000 +NUM_ROWS = 100_000 NUM_BATCHES = 100 ROWS_PER_BATCH = NUM_ROWS // NUM_BATCHES diff --git a/python/python/ci_benchmarks/datagen/lineitems.py b/python/python/ci_benchmarks/datagen/lineitems.py index b91c1c3b422..ddd11049172 100644 --- a/python/python/ci_benchmarks/datagen/lineitems.py +++ b/python/python/ci_benchmarks/datagen/lineitems.py @@ -9,7 +9,7 @@ from ci_benchmarks.datasets import get_dataset_uri -NUM_ROWS = 59986052 +NUM_ROWS = 100_000 def _gen_data(scale_factor: int): From 57410be9685a6c30171874edea53fe86a81757b7 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 11 Dec 2025 12:51:43 -0800 Subject: [PATCH 10/18] better benchmark --- memtest/pyproject.toml | 3 - memtest/python/memtest/pytest_plugin.py | 203 ------------- memtest/python/tests/test_benchmark_plugin.py | 44 --- python/python/ci_benchmarks/README.md | 68 +++-- python/python/ci_benchmarks/benchmark.py | 278 ++++++++++++++++++ .../ci_benchmarks/benchmarks/test_search.py | 19 +- python/python/ci_benchmarks/conftest.py | 5 + python/python/tests/test_memory.py | 2 +- 8 files changed, 333 insertions(+), 289 deletions(-) delete mode 100644 memtest/python/memtest/pytest_plugin.py delete mode 100644 memtest/python/tests/test_benchmark_plugin.py create mode 100644 python/python/ci_benchmarks/benchmark.py create mode 100644 python/python/ci_benchmarks/conftest.py diff --git a/memtest/pyproject.toml b/memtest/pyproject.toml index 4cc765b8b0f..396d7c442e0 100644 --- a/memtest/pyproject.toml +++ b/memtest/pyproject.toml @@ -35,6 +35,3 @@ memtest = "python/memtest" [tool.setuptools.package-data] memtest = ["*.so", "*.dylib", "*.dll"] - -[project.entry-points.pytest11] -memtest = "memtest.pytest_plugin" diff --git a/memtest/python/memtest/pytest_plugin.py b/memtest/python/memtest/pytest_plugin.py deleted file mode 100644 index 5ab70330430..00000000000 --- a/memtest/python/memtest/pytest_plugin.py +++ /dev/null @@ -1,203 +0,0 @@ -"""Pytest plugin for memory tracking during benchmarks. - -This plugin provides a `memory_benchmark` fixture that wraps pytest-benchmark -to track memory allocations during the actual benchmark execution. - -The plugin auto-detects if libmemtest.so is preloaded. If not, the -`memory_benchmark` fixture simply passes through to the regular `benchmark` -fixture. - -Usage: - def test_something(memory_benchmark): - memory_benchmark(my_function, arg1, arg2) - -Output: - - Terminal summary with memory stats per benchmark - - BMF JSON file for bencher.dev upload (--memory-json option) -""" - -import json -from functools import wraps -from typing import Any, Callable, Dict, Optional - -import pytest - -from . import format_bytes, get_stats, is_preloaded, reset_stats - -# Global storage for memory results across all tests -_memory_results: Dict[str, Dict[str, int]] = {} - - -def pytest_addoption(parser: pytest.Parser) -> None: - """Add command-line options for memory tracking.""" - group = parser.getgroup("memory", "memory tracking options") - group.addoption( - "--memory-json", - action="store", - default=None, - metavar="PATH", - help="Output path for memory stats JSON in Bencher Metric Format (BMF)", - ) - - -class MemoryTrackingBenchmark: - """Wrapper around pytest-benchmark that tracks memory during execution.""" - - def __init__(self, benchmark: Any, test_name: str): - self._benchmark = benchmark - self._test_name = test_name - self._peak_memory = 0 - self._total_allocations = 0 - - def _wrap_function(self, func: Callable) -> Callable: - """Wrap a function to track memory around each invocation.""" - - @wraps(func) - def wrapper(*args, **kwargs): - reset_stats() - result = func(*args, **kwargs) - stats = get_stats() - # Track max peak across iterations - self._peak_memory = max(self._peak_memory, stats["peak_bytes"]) - self._total_allocations += stats["total_allocations"] - return result - - return wrapper - - def __call__(self, func: Callable, *args, **kwargs) -> Any: - """Run benchmark with memory tracking.""" - wrapped = self._wrap_function(func) - return self._benchmark(wrapped, *args, **kwargs) - - def pedantic( - self, - func: Callable, - args: tuple = (), - kwargs: Optional[Dict] = None, - setup: Optional[Callable] = None, - teardown: Optional[Callable] = None, - rounds: int = 1, - warmup_rounds: int = 0, - iterations: int = 1, - ) -> Any: - """Run pedantic benchmark with memory tracking.""" - kwargs = kwargs or {} - wrapped = self._wrap_function(func) - return self._benchmark.pedantic( - wrapped, - args=args, - kwargs=kwargs, - setup=setup, - teardown=teardown, - rounds=rounds, - warmup_rounds=warmup_rounds, - iterations=iterations, - ) - - @property - def group(self): - return self._benchmark.group - - @group.setter - def group(self, value): - self._benchmark.group = value - - @property - def name(self): - return self._benchmark.name - - @property - def extra_info(self): - return self._benchmark.extra_info - - @extra_info.setter - def extra_info(self, value): - self._benchmark.extra_info = value - - def get_memory_stats(self) -> Dict[str, int]: - """Get the collected memory statistics.""" - return { - "peak_bytes": self._peak_memory, - "total_allocations": self._total_allocations, - } - - -@pytest.fixture -def memory_benchmark(benchmark, request): - """Fixture that wraps benchmark to track memory during execution. - - If libmemtest.so is not preloaded, this fixture simply returns the - regular benchmark fixture unchanged. - - Usage: - def test_something(memory_benchmark): - memory_benchmark(my_function, arg1, arg2) - - def test_pedantic(memory_benchmark): - memory_benchmark.pedantic(my_function, rounds=5, iterations=10) - """ - if not is_preloaded(): - # Not preloaded - just return regular benchmark - yield benchmark - return - - test_name = request.node.name - tracker = MemoryTrackingBenchmark(benchmark, test_name) - - yield tracker - - # Store results after test completes - stats = tracker.get_memory_stats() - if stats["peak_bytes"] > 0 or stats["total_allocations"] > 0: - _memory_results[test_name] = stats - - -def pytest_terminal_summary(terminalreporter, exitstatus: int, config) -> None: - """Print memory statistics summary at the end of the test run.""" - if not _memory_results: - return - - terminalreporter.write_sep("=", "Memory Statistics") - - # Calculate column widths - name_width = max(len(name) for name in _memory_results.keys()) - name_width = max(name_width, len("Test")) - - # Header - terminalreporter.write_line( - f"{'Test':<{name_width}} {'Peak Memory':>12} {'Allocations':>12}" - ) - terminalreporter.write_line("-" * (name_width + 28)) - - # Results sorted by peak memory (descending) - sorted_results = sorted( - _memory_results.items(), key=lambda x: x[1]["peak_bytes"], reverse=True - ) - - for test_name, stats in sorted_results: - peak = format_bytes(stats["peak_bytes"]) - allocs = f"{stats['total_allocations']:,}" - terminalreporter.write_line(f"{test_name:<{name_width}} {peak:>12} {allocs:>12}") - - terminalreporter.write_line("") - - -def pytest_sessionfinish(session, exitstatus: int) -> None: - """Write memory results to JSON file if --memory-json was specified.""" - if not _memory_results: - return - - output_path = session.config.getoption("--memory-json") - if not output_path: - return - - # Convert to Bencher Metric Format (BMF) - bmf_output = {} - for test_name, stats in _memory_results.items(): - bmf_output[test_name] = { - "peak_memory_bytes": {"value": stats["peak_bytes"]}, - "total_allocations": {"value": stats["total_allocations"]}, - } - - with open(output_path, "w") as f: - json.dump(bmf_output, f, indent=2) diff --git a/memtest/python/tests/test_benchmark_plugin.py b/memtest/python/tests/test_benchmark_plugin.py deleted file mode 100644 index f4d1dd3337d..00000000000 --- a/memtest/python/tests/test_benchmark_plugin.py +++ /dev/null @@ -1,44 +0,0 @@ -"""Test the pytest-benchmark memory tracking plugin.""" - -import pytest - - -def allocate_memory(size_mb: int) -> list: - """Allocate approximately size_mb of memory.""" - # Each int in Python takes about 28 bytes, but in a list it's stored as a pointer - # A list of zeros: each element is ~8 bytes for the pointer + shared int object - # For a rough approximation, 1MB ~= 125000 elements - return [0] * (size_mb * 125000) - - -def test_basic_memory_tracking(memory_benchmark): - """Test that memory is tracked during benchmark execution.""" - - def workload(): - data = allocate_memory(10) # ~10MB - return len(data) - - result = memory_benchmark(workload) - assert result == 10 * 125000 - - -def test_pedantic_mode(memory_benchmark): - """Test memory tracking with pedantic mode.""" - - def workload(): - data = allocate_memory(5) # ~5MB - return sum(data) - - result = memory_benchmark.pedantic(workload, rounds=3, iterations=1) - assert result == 0 - - -def test_with_arguments(memory_benchmark): - """Test memory tracking with function arguments.""" - - def workload(multiplier: int, base_size: int = 1): - data = allocate_memory(base_size * multiplier) - return len(data) - - result = memory_benchmark(workload, 2, base_size=3) # 6MB - assert result == 6 * 125000 diff --git a/python/python/ci_benchmarks/README.md b/python/python/ci_benchmarks/README.md index a8742949a4e..13efa77ff2c 100644 --- a/python/python/ci_benchmarks/README.md +++ b/python/python/ci_benchmarks/README.md @@ -14,6 +14,8 @@ ci_benchmarks/ │ ├── gen_all.py # Generate all datasets │ ├── basic.py # 10M row dataset │ └── lineitems.py # TPC-H lineitem dataset +├── benchmark.py # IO/memory benchmark infrastructure +├── conftest.py # Pytest configuration └── datasets.py # Dataset URI resolver (local vs GCS) ``` @@ -27,57 +29,75 @@ python python/ci_benchmarks/datagen/gen_all.py This creates datasets in `~/lance-benchmarks-ci-datasets/`. -### 2. Run benchmarks +### 2. Run pytest-benchmark tests ```bash pytest python/ci_benchmarks/ --benchmark-only ``` -To save results as JSON: +To save timing results as JSON: ```bash pytest python/ci_benchmarks/ --benchmark-json results.json ``` -## Running with Memory Tracking (Linux-only) +## IO/Memory Benchmarks -To track memory allocations during benchmarks, use the `lance-memtest` library with `LD_PRELOAD`. +The `io_memory_benchmark` marker provides benchmarks that track both IO statistics +and memory allocations during the benchmark execution (not setup/teardown). -### 1. Install lance-memtest +### Writing IO/Memory Benchmarks -```bash -pip install lance-memtest +```python +@pytest.mark.io_memory_benchmark() +def test_full_scan(io_mem_benchmark): + dataset_uri = get_dataset_uri("basic") + ds = lance.dataset(dataset_uri) + + def bench(dataset): + dataset.to_table() + + io_mem_benchmark(bench, ds) ``` -### 2. Run with memory tracking +The `io_mem_benchmark` fixture: +- Runs an optional warmup iteration (not measured) +- Tracks IO stats via `dataset.io_stats_incremental()` +- Optionally tracks memory via `lance-memtest` if preloaded + +### Running IO/Memory Benchmarks +Without memory tracking: ```bash -LD_PRELOAD=$(lance-memtest) pytest python/ci_benchmarks/ \ - --benchmark-json timing_results.json \ - --memory-json memory_results.json +pytest python/ci_benchmarks/benchmarks/test_search.py::test_io_mem_basic_btree_search -v ``` -This produces: -- `timing_results.json` - Standard pytest-benchmark timing results -- `memory_results.json` - Memory stats in Bencher Metric Format (BMF) - -### 3. Using memory_benchmark fixture +With memory tracking (Linux only): +```bash +LD_PRELOAD=$(lance-memtest) pytest python/ci_benchmarks/benchmarks/test_search.py::test_io_mem_basic_btree_search -v +``` -For benchmarks that need memory tracking, use the `memory_benchmark` fixture instead of `benchmark`: +### Output -```python -def test_full_scan(memory_benchmark, dataset): - memory_benchmark(dataset.to_table) +Terminal output shows a summary table: +``` +======================== IO/Memory Benchmark Statistics ======================== +Test Peak Mem Allocs Read IOPS Read Bytes +--------------------------------------------------------------------------------------- +test_io_mem_basic_btree_search[...] 3.6 MB 135,387 2 1.8 MB ``` -When `LD_PRELOAD` is not set, `memory_benchmark` passes through to the regular `benchmark` fixture. +To save results as JSON (Bencher Metric Format): +```bash +pytest ... --benchmark-stats-json stats.json +``` ## Uploading to Bencher ```bash -# Upload timing results +# Upload timing results (from pytest-benchmark) bencher run --adapter python_pytest --file timing_results.json -# Upload memory results -bencher run --adapter json --file memory_results.json +# Upload IO/memory stats +bencher run --adapter json --file stats.json ``` diff --git a/python/python/ci_benchmarks/benchmark.py b/python/python/ci_benchmarks/benchmark.py new file mode 100644 index 00000000000..e735b8e382f --- /dev/null +++ b/python/python/ci_benchmarks/benchmark.py @@ -0,0 +1,278 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +""" +Custom benchmark infrastructure for tracking IO and memory stats. + +This module provides an `io_memory_benchmark` marker and fixture that tracks: +- Peak memory usage +- Total allocations +- Read IOPS and bytes +- Write IOPS and bytes + +Usage: + @pytest.mark.io_memory_benchmark() + def test_something(benchmark): + def workload(dataset): + dataset.to_table() + benchmark(workload, dataset) +""" + +import json +from dataclasses import dataclass +from typing import Any, Callable, List + +import pytest + +# Try to import memtest, but don't fail if not available +try: + import memtest + + MEMTEST_AVAILABLE = memtest.is_preloaded() +except ImportError: + MEMTEST_AVAILABLE = False + + +@dataclass +class BenchmarkStats: + """Statistics collected during a benchmark run.""" + + # Memory stats (only populated if memtest is preloaded) + peak_bytes: int = 0 + total_allocations: int = 0 + + # IO stats + read_iops: int = 0 + read_bytes: int = 0 + write_iops: int = 0 + write_bytes: int = 0 + + +@dataclass +class BenchmarkResult: + """Result of a single benchmark test.""" + + name: str + stats: BenchmarkStats + + +# Global storage for benchmark results +_benchmark_results: List[BenchmarkResult] = [] + + +def _format_bytes(num_bytes: int) -> str: + """Format byte count as human-readable string.""" + for unit in ["B", "KB", "MB", "GB", "TB"]: + if abs(num_bytes) < 1024.0: + return f"{num_bytes:.1f} {unit}" + num_bytes /= 1024.0 + return f"{num_bytes:.1f} PB" + + +class IOMemoryBenchmark: + """Benchmark fixture that tracks IO and memory during execution.""" + + def __init__(self, test_name: str): + self._test_name = test_name + self._stats = BenchmarkStats() + + def __call__( + self, + func: Callable, + dataset: Any, + warmup: bool = True, + ) -> Any: + """ + Run a benchmark function with IO and memory tracking. + + Parameters + ---------- + func : Callable + The function to benchmark. Should accept a dataset as first argument. + dataset : lance.LanceDataset + The dataset to pass to the function. + warmup : bool, default True + Whether to run a warmup iteration before measuring. + + Returns + ------- + Any + The return value of the benchmark function. + """ + # Warmup run (not measured) + if warmup: + func(dataset) + + # Reset IO stats before the measured run + dataset.io_stats_incremental() + + # Run with memory tracking if available + if MEMTEST_AVAILABLE: + memtest.reset_stats() + result = func(dataset) + mem_stats = memtest.get_stats() + self._stats.peak_bytes = mem_stats["peak_bytes"] + self._stats.total_allocations = mem_stats["total_allocations"] + else: + result = func(dataset) + + # Capture IO stats + io_stats = dataset.io_stats_incremental() + self._stats.read_iops = io_stats.read_iops + self._stats.read_bytes = io_stats.read_bytes + self._stats.write_iops = io_stats.write_iops + self._stats.write_bytes = io_stats.written_bytes + + return result + + def get_stats(self) -> BenchmarkStats: + """Get the collected statistics.""" + return self._stats + + +@pytest.fixture +def io_mem_benchmark(request): + """ + Fixture that provides IO and memory benchmarking. + + Only active for tests marked with @pytest.mark.io_memory_benchmark(). + For other tests, returns a no-op benchmark that just calls the function. + + Usage: + @pytest.mark.io_memory_benchmark() + def test_something(io_mem_benchmark): + def workload(dataset): + dataset.to_table() + io_mem_benchmark(workload, dataset) + """ + marker = request.node.get_closest_marker("io_memory_benchmark") + + if marker is None: + # Not an io_memory_benchmark test, return a simple passthrough + class PassthroughBenchmark: + def __call__(self, func, dataset, warmup=True): + return func(dataset) + + yield PassthroughBenchmark() + return + + test_name = request.node.name + tracker = IOMemoryBenchmark(test_name) + + yield tracker + + # Store results after test completes + stats = tracker.get_stats() + _benchmark_results.append(BenchmarkResult(name=test_name, stats=stats)) + + +def pytest_configure(config): + """Register the io_memory_benchmark marker.""" + config.addinivalue_line( + "markers", + "io_memory_benchmark(): Mark test as an IO/memory benchmark", + ) + + +def pytest_addoption(parser): + """Add command-line options for benchmark output.""" + group = parser.getgroup("io_memory_benchmark", "IO/memory benchmark options") + group.addoption( + "--benchmark-stats-json", + action="store", + default=None, + metavar="PATH", + help="Output path for benchmark stats JSON in Bencher Metric Format (BMF)", + ) + + +def pytest_terminal_summary(terminalreporter, exitstatus, config): + """Print benchmark statistics summary at the end of the test run.""" + if not _benchmark_results: + return + + terminalreporter.write_sep("=", "IO/Memory Benchmark Statistics") + + # Calculate column widths + name_width = max(len(r.name) for r in _benchmark_results) + name_width = max(name_width, len("Test")) + + # Header + if MEMTEST_AVAILABLE: + terminalreporter.write_line( + f"{'Test':<{name_width}} {'Peak Mem':>10} {'Allocs':>10} " + f"{'Read IOPS':>10} {'Read Bytes':>12} " + f"{'Write IOPS':>10} {'Write Bytes':>12}" + ) + terminalreporter.write_line("-" * (name_width + 72)) + else: + terminalreporter.write_line( + f"{'Test':<{name_width}} " + f"{'Read IOPS':>10} {'Read Bytes':>12} " + f"{'Write IOPS':>10} {'Write Bytes':>12}" + ) + terminalreporter.write_line("-" * (name_width + 50)) + + # Results sorted by read bytes (descending) + sorted_results = sorted( + _benchmark_results, key=lambda r: r.stats.read_bytes, reverse=True + ) + + for result in sorted_results: + s = result.stats + if MEMTEST_AVAILABLE: + terminalreporter.write_line( + f"{result.name:<{name_width}} " + f"{_format_bytes(s.peak_bytes):>10} " + f"{s.total_allocations:>10,} " + f"{s.read_iops:>10,} " + f"{_format_bytes(s.read_bytes):>12} " + f"{s.write_iops:>10,} " + f"{_format_bytes(s.write_bytes):>12}" + ) + else: + terminalreporter.write_line( + f"{result.name:<{name_width}} " + f"{s.read_iops:>10,} " + f"{_format_bytes(s.read_bytes):>12} " + f"{s.write_iops:>10,} " + f"{_format_bytes(s.write_bytes):>12}" + ) + + if not MEMTEST_AVAILABLE: + terminalreporter.write_line("") + terminalreporter.write_line( + "Note: Memory tracking not available. " + "Run with LD_PRELOAD=$(lance-memtest) to enable." + ) + + terminalreporter.write_line("") + + +def pytest_sessionfinish(session, exitstatus): + """Write benchmark results to JSON file if --benchmark-stats-json was specified.""" + if not _benchmark_results: + return + + output_path = session.config.getoption("--benchmark-stats-json") + if not output_path: + return + + # Convert to Bencher Metric Format (BMF) + bmf_output = {} + for result in _benchmark_results: + s = result.stats + bmf_output[result.name] = { + "read_iops": {"value": s.read_iops}, + "read_bytes": {"value": s.read_bytes}, + "write_iops": {"value": s.write_iops}, + "write_bytes": {"value": s.write_bytes}, + } + if MEMTEST_AVAILABLE: + bmf_output[result.name]["peak_memory_bytes"] = {"value": s.peak_bytes} + bmf_output[result.name]["total_allocations"] = { + "value": s.total_allocations + } + + with open(output_path, "w") as f: + json.dump(bmf_output, f, indent=2) diff --git a/python/python/ci_benchmarks/benchmarks/test_search.py b/python/python/ci_benchmarks/benchmarks/test_search.py index d257c835d8c..40cd9f05645 100644 --- a/python/python/ci_benchmarks/benchmarks/test_search.py +++ b/python/python/ci_benchmarks/benchmarks/test_search.py @@ -245,10 +245,12 @@ def clear_timer(): bench, warmup_rounds=1, rounds=1, iterations=1, setup=clear_timer ) + @pytest.mark.io_memory_benchmark() @pytest.mark.parametrize("filt", BASIC_BTREE_FILTERS, ids=BASIC_BTREE_FILTER_LABELS) @pytest.mark.parametrize("payload", ["small_strings", "integers"]) -def test_basic_btree_search(benchmark, filt: str | None, payload: str): +def test_io_mem_basic_btree_search(io_mem_benchmark, filt: str | None, payload: str): + """Benchmark btree search with IO and memory tracking.""" dataset_uri = get_dataset_uri("basic") ds = lance.dataset(dataset_uri) @@ -258,21 +260,10 @@ def test_basic_btree_search(benchmark, filt: str | None, payload: str): def bench(dataset): dataset.to_table( - columns=[payload], + columns=columns, filter=filt, with_row_id=True, batch_size=32 * 1024, ) - - benchmark(bench, ds) - -def benchmark(bench, ds, warmup: bool=True): - if warmup: - bench(ds) - ds.io_stats_incremental() - with memtest.track() as get_stats: - bench(ds) - memory_stats = get_stats() - io_stats = ds.io_stats_incremental() - # TODO: Record the I/O and memory stats + io_mem_benchmark(bench, ds) diff --git a/python/python/ci_benchmarks/conftest.py b/python/python/ci_benchmarks/conftest.py new file mode 100644 index 00000000000..7ea42b773bb --- /dev/null +++ b/python/python/ci_benchmarks/conftest.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +# Import the benchmark plugin to register hooks and fixtures +pytest_plugins = ["ci_benchmarks.benchmark"] diff --git a/python/python/tests/test_memory.py b/python/python/tests/test_memory.py index f674514e123..e25ab465c10 100644 --- a/python/python/tests/test_memory.py +++ b/python/python/tests/test_memory.py @@ -22,7 +22,7 @@ def batch_generator(): ) with memtest.track() as get_stats: - ds = lance.write_dataset( + lance.write_dataset( reader, tmp_path / "test.lance", ) From c4bdea3f6b2fd65318dc90bec58849735f057a12 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 11 Dec 2025 13:14:29 -0800 Subject: [PATCH 11/18] upload --- .github/workflows/ci-benchmarks.yml | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ci-benchmarks.yml b/.github/workflows/ci-benchmarks.yml index 3277c7e7f71..8156e0478e1 100644 --- a/.github/workflows/ci-benchmarks.yml +++ b/.github/workflows/ci-benchmarks.yml @@ -46,6 +46,12 @@ jobs: source venv/bin/activate pip install maturin duckdb requests pytest pytest-benchmark maturin develop --locked --release + - name: Build memtest + run: | + source venv/bin/activate + cd ../memtest + cargo build --release + pip install -e . - name: Generate datasets run: | python -m venv venv @@ -53,8 +59,20 @@ jobs: python python/ci_benchmarks/datagen/gen_all.py - name: Run benchmarks run: | - python -m venv venv source venv/bin/activate bencher run --project weston-lancedb --token ${{ secrets.LANCE_BENCHER_TOKEN }} --adapter python_pytest \ --branch main --testbed google-genoa --err --file results.json "python -mpytest --benchmark-json \ results.json python/ci_benchmarks" + - name: Run IO/memory benchmarks + run: | + source venv/bin/activate + LIB_PATH=$(lance-memtest) + LD_PRELOAD=$LIB_PATH pytest python/ci_benchmarks \ + -k "io_mem_" \ + --benchmark-stats-json io_mem_stats.json + - name: Upload IO/memory stats to bencher + run: | + source venv/bin/activate + bencher run --project weston-lancedb --token ${{ secrets.LANCE_BENCHER_TOKEN }} \ + --adapter json --branch main --testbed google-genoa \ + --err --file io_mem_stats.json From 12906e5fbb6ed3d398d80b8e9672e74df6759bf9 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 11 Dec 2025 14:34:19 -0800 Subject: [PATCH 12/18] cleanup --- memtest/python/memtest/__init__.py | 11 ++++------- python/python/ci_benchmarks/README.md | 4 ++++ 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/memtest/python/memtest/__init__.py b/memtest/python/memtest/__init__.py index 8e38bcadec2..45e7f4b42d1 100644 --- a/memtest/python/memtest/__init__.py +++ b/memtest/python/memtest/__init__.py @@ -186,13 +186,10 @@ def is_preloaded() -> bool: ... stats = get_stats() ... print(f"Tracking {stats['total_allocations']} allocations") """ - try: - stats = get_stats() - # If we can get stats and there's been any activity, we're preloaded - # Even with no activity, if the library loads we're preloaded - return True - except Exception: - return False + import os + + ld_preload = os.environ.get("LD_PRELOAD", "") + return "libmemtest" in ld_preload __all__ = [ diff --git a/python/python/ci_benchmarks/README.md b/python/python/ci_benchmarks/README.md index 13efa77ff2c..b2ddb5e72c2 100644 --- a/python/python/ci_benchmarks/README.md +++ b/python/python/ci_benchmarks/README.md @@ -92,6 +92,10 @@ To save results as JSON (Bencher Metric Format): pytest ... --benchmark-stats-json stats.json ``` +## Investigating memory use for a particular benchmark + + + ## Uploading to Bencher ```bash From 787f9c0c530600e2a8d0ea8df3359fd6f44f8386 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 11 Dec 2025 15:08:41 -0800 Subject: [PATCH 13/18] this is fixed --- .github/workflows/ci-benchmarks.yml | 3 +- memtest/python/memtest/__init__.py | 59 +++++++++++++++++++++++++++-- memtest/python/memtest/__main__.py | 32 +++++++++++++--- 3 files changed, 84 insertions(+), 10 deletions(-) diff --git a/.github/workflows/ci-benchmarks.yml b/.github/workflows/ci-benchmarks.yml index 8156e0478e1..7a6e8573c93 100644 --- a/.github/workflows/ci-benchmarks.yml +++ b/.github/workflows/ci-benchmarks.yml @@ -49,12 +49,11 @@ jobs: - name: Build memtest run: | source venv/bin/activate - cd ../memtest + cd memtest cargo build --release pip install -e . - name: Generate datasets run: | - python -m venv venv source venv/bin/activate python python/ci_benchmarks/datagen/gen_all.py - name: Run benchmarks diff --git a/memtest/python/memtest/__init__.py b/memtest/python/memtest/__init__.py index 45e7f4b42d1..200bb319203 100644 --- a/memtest/python/memtest/__init__.py +++ b/memtest/python/memtest/__init__.py @@ -1,12 +1,24 @@ """Memory allocation testing utilities for Python.""" import ctypes +import platform +import warnings from pathlib import Path from typing import Dict, Optional from contextlib import contextmanager __version__ = "0.1.0" +# Platform support check +_SUPPORTED_PLATFORM = platform.system() == "Linux" +if not _SUPPORTED_PLATFORM: + warnings.warn( + f"lance-memtest only supports Linux (current platform: {platform.system()}). " + "Memory statistics will not be available.", + RuntimeWarning, + stacklevel=2, + ) + class _MemtestStats(ctypes.Structure): """C struct matching MemtestStats in Rust.""" @@ -23,6 +35,9 @@ class _MemtestStats(ctypes.Structure): def _load_library(): """Load the memtest shared library.""" + if not _SUPPORTED_PLATFORM: + return None, None + # Find the library relative to this module module_dir = Path(__file__).parent @@ -53,15 +68,29 @@ def _load_library(): _lib, _lib_path = _load_library() -def get_library_path() -> Path: +def _empty_stats() -> Dict[str, int]: + """Return empty stats for unsupported platforms.""" + return { + "total_allocations": 0, + "total_deallocations": 0, + "total_bytes_allocated": 0, + "total_bytes_deallocated": 0, + "current_bytes": 0, + "peak_bytes": 0, + } + + +def get_library_path() -> Optional[Path]: """Get the path to the memtest shared library for use with LD_PRELOAD. Returns: - Path to the .so file that can be used with LD_PRELOAD + Path to the .so file that can be used with LD_PRELOAD, or None on + unsupported platforms. Example: >>> lib_path = get_library_path() - >>> os.environ['LD_PRELOAD'] = str(lib_path) + >>> if lib_path: + ... os.environ['LD_PRELOAD'] = str(lib_path) """ return _lib_path @@ -78,11 +107,16 @@ def get_stats() -> Dict[str, int]: - current_bytes: Current memory usage (allocated - deallocated) - peak_bytes: Peak memory usage observed + On unsupported platforms, all values will be 0. + Example: >>> stats = get_stats() >>> print(f"Current memory: {stats['current_bytes']} bytes") >>> print(f"Peak memory: {stats['peak_bytes']} bytes") """ + if _lib is None: + return _empty_stats() + stats = _MemtestStats() _lib.memtest_get_stats(ctypes.byref(stats)) @@ -100,12 +134,15 @@ def reset_stats() -> None: """Reset all allocation statistics to zero. This is useful for measuring allocations in a specific section of code. + On unsupported platforms, this is a no-op. Example: >>> reset_stats() >>> # ... run code to measure ... >>> stats = get_stats() """ + if _lib is None: + return _lib.memtest_reset_stats() @@ -192,6 +229,21 @@ def is_preloaded() -> bool: return "libmemtest" in ld_preload +def is_supported() -> bool: + """Check if memory tracking is supported on this platform. + + Returns: + True if on Linux (the only supported platform), False otherwise. + + Example: + >>> if is_supported(): + ... with track() as get: + ... # ... do work ... + ... stats = get() + """ + return _SUPPORTED_PLATFORM + + __all__ = [ "get_library_path", "get_stats", @@ -200,4 +252,5 @@ def is_preloaded() -> bool: "format_bytes", "print_stats", "is_preloaded", + "is_supported", ] diff --git a/memtest/python/memtest/__main__.py b/memtest/python/memtest/__main__.py index 262b845da25..9bdf49a1351 100644 --- a/memtest/python/memtest/__main__.py +++ b/memtest/python/memtest/__main__.py @@ -1,14 +1,36 @@ """CLI for lance-memtest.""" import sys -from memtest import get_library_path +from memtest import get_library_path, print_stats, is_supported def main(): - """Main CLI entry point - print path to shared library.""" - lib_path = get_library_path() - print(lib_path) - return 0 + """Main CLI entry point.""" + args = sys.argv[1:] + + if not args or args[0] == "path": + lib_path = get_library_path() + if lib_path is None: + print( + "lance-memtest is not supported on this platform", + file=sys.stderr, + ) + return 1 + print(lib_path) + return 0 + elif args[0] == "stats": + if not is_supported(): + print( + "lance-memtest is not supported on this platform", + file=sys.stderr, + ) + return 1 + print_stats() + return 0 + else: + print(f"Unknown command: {args[0]}", file=sys.stderr) + print("Usage: lance-memtest [path|stats]", file=sys.stderr) + return 1 if __name__ == "__main__": From 260891cd053c5ac76fc848ea1d8f13a67da80467 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 11 Dec 2025 15:21:18 -0800 Subject: [PATCH 14/18] cleanup --- memtest/Makefile | 2 +- python/pyproject.toml | 1 - .../ci_benchmarks/benchmarks/test_search.py | 47 ------------------- python/python/ci_benchmarks/datagen/basic.py | 2 +- .../python/ci_benchmarks/datagen/lineitems.py | 2 +- python/python/tests/test_memory.py | 8 +++- 6 files changed, 10 insertions(+), 52 deletions(-) diff --git a/memtest/Makefile b/memtest/Makefile index 3639ddbf8cd..071eb6f271d 100644 --- a/memtest/Makefile +++ b/memtest/Makefile @@ -11,7 +11,7 @@ build-release: pip install -e . test: - pytest python/tests/ -v + LD_PRELOAD=./python/memtest/libmemtest.so pytest python/tests/ -v lint: cargo clippy -- -D warnings diff --git a/python/pyproject.toml b/python/pyproject.toml index 2218bd0e50f..bffb76c33d7 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -61,7 +61,6 @@ tests = [ "tensorflow; sys_platform == 'linux'", "tqdm", "datafusion>=50.1", - # TODO: Make memtest a dependency. ] dev = ["ruff==0.4.1", "pyright"] benchmarks = ["pytest-benchmark"] diff --git a/python/python/ci_benchmarks/benchmarks/test_search.py b/python/python/ci_benchmarks/benchmarks/test_search.py index 40cd9f05645..cf23de5324d 100644 --- a/python/python/ci_benchmarks/benchmarks/test_search.py +++ b/python/python/ci_benchmarks/benchmarks/test_search.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The Lance Authors -import re import lance import pytest @@ -201,56 +200,10 @@ def test_basic_bitmap_search( do_basic_search(benchmark, filt, payload, use_cache) -IOPS = 0.0 - - -def set_iops(iops: float): - global IOPS - IOPS = iops - - -def iops_timer(): - return IOPS - - -@pytest.mark.benchmark(warmup=False, timer=iops_timer) -@pytest.mark.parametrize("filt", BASIC_BTREE_FILTERS, ids=BASIC_BTREE_FILTER_LABELS) -@pytest.mark.parametrize("payload", ["small_strings", "integers"]) -def test_iops_basic_btree_search(benchmark, filt: str | None, payload: str): - dataset_uri = get_dataset_uri("basic") - ds = lance.dataset(dataset_uri) - - columns = [] - if payload is not None: - columns = [payload] - - def bench(): - plan = ds.scanner( - columns=columns, - filter=filt, - with_row_id=True, - batch_size=32 * 1024, - ).analyze_plan() - iops = re.search(r"iops=(\d+)", plan) - if iops is not None: - set_iops(float(iops.group(1))) - else: - set_iops(0.0) - - def clear_timer(): - set_iops(0.0) - - # We still do a warmup since caching may reduce IOPS and not just latency - benchmark.pedantic( - bench, warmup_rounds=1, rounds=1, iterations=1, setup=clear_timer - ) - - @pytest.mark.io_memory_benchmark() @pytest.mark.parametrize("filt", BASIC_BTREE_FILTERS, ids=BASIC_BTREE_FILTER_LABELS) @pytest.mark.parametrize("payload", ["small_strings", "integers"]) def test_io_mem_basic_btree_search(io_mem_benchmark, filt: str | None, payload: str): - """Benchmark btree search with IO and memory tracking.""" dataset_uri = get_dataset_uri("basic") ds = lance.dataset(dataset_uri) diff --git a/python/python/ci_benchmarks/datagen/basic.py b/python/python/ci_benchmarks/datagen/basic.py index fa8a8845894..c14d7dcb47a 100644 --- a/python/python/ci_benchmarks/datagen/basic.py +++ b/python/python/ci_benchmarks/datagen/basic.py @@ -12,7 +12,7 @@ from ci_benchmarks.datasets import get_dataset_uri -NUM_ROWS = 100_000 +NUM_ROWS = 10_000_000 NUM_BATCHES = 100 ROWS_PER_BATCH = NUM_ROWS // NUM_BATCHES diff --git a/python/python/ci_benchmarks/datagen/lineitems.py b/python/python/ci_benchmarks/datagen/lineitems.py index ddd11049172..b79e37eda1c 100644 --- a/python/python/ci_benchmarks/datagen/lineitems.py +++ b/python/python/ci_benchmarks/datagen/lineitems.py @@ -9,7 +9,7 @@ from ci_benchmarks.datasets import get_dataset_uri -NUM_ROWS = 100_000 +NUM_ROWS = 59_986_052 def _gen_data(scale_factor: int): diff --git a/python/python/tests/test_memory.py b/python/python/tests/test_memory.py index e25ab465c10..3a59240f93d 100644 --- a/python/python/tests/test_memory.py +++ b/python/python/tests/test_memory.py @@ -4,11 +4,17 @@ from pathlib import Path import lance -import memtest import pyarrow as pa +import pytest + +pytest.importorskip( + "memtest", reason="memtest is not available. Please install from ../memtest" +) def test_insert_memory(tmp_path: Path): + import memtest + def batch_generator(): # 5MB batches -> 100MB total for _ in range(20): From ee9e68d2424e23ef757f85df76df8e9da09f37b6 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 11 Dec 2025 15:51:04 -0800 Subject: [PATCH 15/18] fix installation --- .github/workflows/ci-benchmarks.yml | 3 +-- .github/workflows/python.yml | 4 ++++ python/python/ci_benchmarks/README.md | 17 +++++++++++++++++ 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci-benchmarks.yml b/.github/workflows/ci-benchmarks.yml index 7a6e8573c93..ab1da5a121b 100644 --- a/.github/workflows/ci-benchmarks.yml +++ b/.github/workflows/ci-benchmarks.yml @@ -50,8 +50,7 @@ jobs: run: | source venv/bin/activate cd memtest - cargo build --release - pip install -e . + make build-release - name: Generate datasets run: | source venv/bin/activate diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml index 0abe4b80f85..cc8146260cf 100644 --- a/.github/workflows/python.yml +++ b/.github/workflows/python.yml @@ -114,6 +114,10 @@ jobs: with: workspaces: python prefix-key: ${{ env.CACHE_PREFIX }} + - name: Install memtest + run: | + cd ../memtest + make build-release - uses: ./.github/workflows/build_linux_wheel with: args: "--profile ci" diff --git a/python/python/ci_benchmarks/README.md b/python/python/ci_benchmarks/README.md index b2ddb5e72c2..dc493cdab1e 100644 --- a/python/python/ci_benchmarks/README.md +++ b/python/python/ci_benchmarks/README.md @@ -94,7 +94,24 @@ pytest ... --benchmark-stats-json stats.json ## Investigating memory use for a particular benchmark +To investigate memory use for a particular benchmark, you can use the `bytehound` library. +After installing it, you can run a benchmark with memory profiling enabled: +```shell +LD_PRELOAD=/usr/local/lib/libbytehound.so \ + pytest 'python/ci_benchmarks/benchmarks/test_search.py::test_io_mem_basic_btree_search[small_strings-equal]' -v +``` + +Then use the `bytehound` server to visualize the memory profiling data: + +```shell +bytehound server memory-profiling_*.dat +``` + +You can use time filters on the allocations view to see memory allocations at a specific point in time, +which can help you filter out allocations from setup. Once you have filters in place, you can use +the Flamegraph view (available from the menu in the upper right corner) to get a flamegraph of the +memory allocations in that time range. ## Uploading to Bencher From a0bee28afea6b58ad4021f6c5ee25e02ed0a2ab6 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 11 Dec 2025 16:02:09 -0800 Subject: [PATCH 16/18] fix tests --- .github/workflows/python.yml | 6 ++---- .github/workflows/run_tests/action.yml | 10 ++++++++++ 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml index cc8146260cf..6755d38087c 100644 --- a/.github/workflows/python.yml +++ b/.github/workflows/python.yml @@ -114,14 +114,12 @@ jobs: with: workspaces: python prefix-key: ${{ env.CACHE_PREFIX }} - - name: Install memtest - run: | - cd ../memtest - make build-release - uses: ./.github/workflows/build_linux_wheel with: args: "--profile ci" - uses: ./.github/workflows/run_tests + with: + memtest: true - name: Upload wheels as artifacts if: ${{ matrix.python-minor-version == '13' }} uses: actions/upload-artifact@v4 diff --git a/.github/workflows/run_tests/action.yml b/.github/workflows/run_tests/action.yml index 14c4b3d6f46..24ffa7f9a61 100644 --- a/.github/workflows/run_tests/action.yml +++ b/.github/workflows/run_tests/action.yml @@ -9,6 +9,10 @@ inputs: required: false description: "Skip pytorch tests" default: "false" + memtest: + required: false + description: "Run memtest" + default: "false" runs: using: "composite" steps: @@ -24,6 +28,12 @@ runs: run: | # Install cpu only pytorch pip install torch --index-url https://download.pytorch.org/whl/cpu + - name: Install memtest + working-directory: memtest + if: inputs.memtest == 'true' + run: | + make build-release + echo "LD_PRELOAD=$(lance-memtest)" >> $GITHUB_ENV - name: Run python tests shell: bash working-directory: python From 1ed1dcb0e40913de0935ebb848dc216e063244c6 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 11 Dec 2025 16:10:44 -0800 Subject: [PATCH 17/18] set shell --- .github/workflows/run_tests/action.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/run_tests/action.yml b/.github/workflows/run_tests/action.yml index 24ffa7f9a61..935a5a7e099 100644 --- a/.github/workflows/run_tests/action.yml +++ b/.github/workflows/run_tests/action.yml @@ -31,6 +31,7 @@ runs: - name: Install memtest working-directory: memtest if: inputs.memtest == 'true' + shell: bash run: | make build-release echo "LD_PRELOAD=$(lance-memtest)" >> $GITHUB_ENV From fe513cf36e244d08b26a334065d0ffc23b3d1339 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Mon, 15 Dec 2025 12:45:54 -0800 Subject: [PATCH 18/18] pr feedback --- memtest/python/memtest/__init__.py | 25 +++++++++---------------- memtest/python/memtest/__main__.py | 13 ++----------- python/python/ci_benchmarks/README.md | 10 ---------- python/python/tests/test_memory.py | 4 +--- 4 files changed, 12 insertions(+), 40 deletions(-) diff --git a/memtest/python/memtest/__init__.py b/memtest/python/memtest/__init__.py index 200bb319203..b5ffcac6537 100644 --- a/memtest/python/memtest/__init__.py +++ b/memtest/python/memtest/__init__.py @@ -41,25 +41,18 @@ def _load_library(): # Find the library relative to this module module_dir = Path(__file__).parent - # Look for the library in common locations - possible_paths = [ - module_dir / "libmemtest.so", # Linux - module_dir / "libmemtest.dylib", # macOS - module_dir / "memtest.dll", # Windows - ] - - for lib_path in possible_paths: - if lib_path.exists(): - lib = ctypes.CDLL(str(lib_path)) + lib_path = module_dir / "libmemtest.so" + if lib_path.exists(): + lib = ctypes.CDLL(str(lib_path)) - # Define function signatures - lib.memtest_get_stats.argtypes = [ctypes.POINTER(_MemtestStats)] - lib.memtest_get_stats.restype = None + # Define function signatures + lib.memtest_get_stats.argtypes = [ctypes.POINTER(_MemtestStats)] + lib.memtest_get_stats.restype = None - lib.memtest_reset_stats.argtypes = [] - lib.memtest_reset_stats.restype = None + lib.memtest_reset_stats.argtypes = [] + lib.memtest_reset_stats.restype = None - return lib, lib_path + return lib, lib_path raise RuntimeError("memtest library not found. Run 'make build' to build it.") diff --git a/memtest/python/memtest/__main__.py b/memtest/python/memtest/__main__.py index 9bdf49a1351..4e684e2c01f 100644 --- a/memtest/python/memtest/__main__.py +++ b/memtest/python/memtest/__main__.py @@ -1,7 +1,7 @@ """CLI for lance-memtest.""" import sys -from memtest import get_library_path, print_stats, is_supported +from memtest import get_library_path def main(): @@ -18,18 +18,9 @@ def main(): return 1 print(lib_path) return 0 - elif args[0] == "stats": - if not is_supported(): - print( - "lance-memtest is not supported on this platform", - file=sys.stderr, - ) - return 1 - print_stats() - return 0 else: print(f"Unknown command: {args[0]}", file=sys.stderr) - print("Usage: lance-memtest [path|stats]", file=sys.stderr) + print("Usage: lance-memtest [path]", file=sys.stderr) return 1 diff --git a/python/python/ci_benchmarks/README.md b/python/python/ci_benchmarks/README.md index dc493cdab1e..0245d29166f 100644 --- a/python/python/ci_benchmarks/README.md +++ b/python/python/ci_benchmarks/README.md @@ -112,13 +112,3 @@ You can use time filters on the allocations view to see memory allocations at a which can help you filter out allocations from setup. Once you have filters in place, you can use the Flamegraph view (available from the menu in the upper right corner) to get a flamegraph of the memory allocations in that time range. - -## Uploading to Bencher - -```bash -# Upload timing results (from pytest-benchmark) -bencher run --adapter python_pytest --file timing_results.json - -# Upload IO/memory stats -bencher run --adapter json --file stats.json -``` diff --git a/python/python/tests/test_memory.py b/python/python/tests/test_memory.py index 3a59240f93d..39485c13f35 100644 --- a/python/python/tests/test_memory.py +++ b/python/python/tests/test_memory.py @@ -7,14 +7,12 @@ import pyarrow as pa import pytest -pytest.importorskip( +memtest = pytest.importorskip( "memtest", reason="memtest is not available. Please install from ../memtest" ) def test_insert_memory(tmp_path: Path): - import memtest - def batch_generator(): # 5MB batches -> 100MB total for _ in range(20):