commit 5747e9a183061842156d35e4ac4816e3e0b66d1f Author: Leo Vasanko Date: Fri Jan 2 22:16:16 2026 +0000 Initial commit of v2. diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..9a5db9e --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +.* +*.lock +*.egg-info +/dist +/build +__pycache__ +!.gitignore diff --git a/README.md b/README.md new file mode 100644 index 0000000..16f0e66 --- /dev/null +++ b/README.md @@ -0,0 +1,125 @@ +# The World's Fastest Random Generator v2 + +I was disappointed with the sad state of random number generators. Many languages don't ship anything useful and some are stuck with whatever the OS provides. Most existing implementations are slow, often maxing out at a few hundred megabytes per second, which becomes a real bottleneck in high-throughput systems. + +Secondly, [flaws have been found](https://numpy.org/doc/stable/reference/random/upgrading-pcg64.html) in popular non-cryptographic algorithms such as Mersenne Twister and PCG-style generators. These issues range from detectable structure to repeating sequences, making them unsuitable for serious or long-running workloads. + +The cryptographic alternative is simply better. Proper CSPRNGs avoid these pitfalls entirely and, on modern hardware, can now be *faster* than legacy non-cryptographic designs. There is little reason left to accept weaker guarantees for worse performance. + +This version uses **AEGIS**, a modern authenticated encryption primitive that leverages AES hardware acceleration. AEGIS provides extremely high throughput while retaining strong cryptographic properties, significantly outperforming our Legacy implementation and serving as an ideal foundation for a high-performance CSPRNG. + +## Quick start + +Install [UV](https://docs.astral.sh/uv/install/) and install the CLI tool with it: + +```sh +uv tool install randquik +``` + +Now you can create random data with it: +```sh +randquik -t8 --len 1TB > /dev/null +``` + +You can try how it performs on your machine in different scenarios with: +```sh +randquik --benchmark +``` + +Wipe an entire file without altering its size: +```sh +randquik -o sensitive.dat +``` + +Piping and redirection: +```sh +randquik | hexdump -C | head +``` + + + +## Features + +- Blazing-fast CSPRNG built on AEGIS single or multithreaded +- Deterministic seeding: same seed, same byte stream +- Seekable random stream and file output +- Flexible I/O: piping, files, and `mmap` +- Built-in benchmarking and dry-run modes +- Full screen console graphics for speed display + +Below are the most important features with example commands. + +### Performance + +You'll be looking at up to 100 GB/s raw generation speed, making this some orders of magnitude faster than your traditional random number generation. Single-threaded performance still is 10-20 times faster than other options that don't have threading. + +Your output, e.g. writing a file, will always be the bottle neck, not your random generator, but modern SSDs allow up to 10 GB/s write speeds already. + +### Size units + +Many options such as `--len`, `--seek` accept human-readable units. The table uses conventional, capitalized forms (e.g. `KiB`, `MB`), but you may write them in lower case and with or without the trailing `B` (for example `1m` or `5gi`). + +| SI unit | Binary unit | Meaning | Bytes factor | +|--------:|------------:|------------------------|-----------------------| +| 100 | — | 100 bytes | 1 | +| 1kB | 1KiB | Kilobyte / kibibyte | 1_000 / 1_024 | +| 1MB | 1MiB | Megabyte / mebibyte | 1_000_000 / 1_048_576 | +| 1GB | 1GiB | Gigabyte / gibibyte | 1_000_000_000 / 1_073_741_824 | +| 1TB | 1TiB | Terabyte / tebibyte | 1_000_000_000_000 / 1_099_511_627_776 | +| 1PB | 1PiB | Petabyte / pebibyte | 1_000_... / 1_125_... | +| — | 1sect | Sectors of output device | 512 (typical), 4096 (rarely) | + +Notes: +- Units are case-insensitive on input: `1GiB`, `1gib`, `1gi`, and `1giB` are all accepted. +- The trailing `B` is optional for parsing (`1g` and `1gb` are equivalent). +- Underscores are ignored, so `1_000Mi` works the same as `1000Mi`. +- `sect` uses the sector size of the output path when available, otherwise 512 bytes. + +### 2. Deterministic seeding + +You can provide an explicit seed string. The same seed, length, and options will always produce the same output bytes: +```sh +randquik -l 64MiB -s my-seed-string -o chunk.bin +``` + +This is useful for reproducible tests or simulations. + +### 4. Seekable stream and file + +It is possible to seek to any byte position in the stream without delay. + +- `--iseek`: seek the input random stream +- `--oseek`: seek in the output file +- `--seek`: set both input and output seek to the same position + +Example: resume as if 5 terabytes had already been written, and continue writing to `out.dat`: +```sh +randquik --seek 5T --len 1G -o out.dat +``` + +Bytes prior to seek position are kept as they were while the file is expanded to fit all the data starting at five terabytes mark (using sparse allocation so it doesn't actually consume 5 terabytes). + +Wipe all data on a disk (e.g. USB drive) but keep the partition table: +```sh +randquik -oseek 2048sect -o /dev/sde +``` + +### 6. Benchmark and dry-run modes + +Benchmark different modes and thread counts. Prints the options that perform the best on your system: +```sh +randquik --benchmark +``` + +To do a single run without actually writing bytes to disk or anywhere, use `--dry`: +```sh +randquik --len 50GiB -t8 --dry +``` + +## Legacy + +The original implementation is preserved in the `legacy` git branch. + +That version was once the fastest CSPRNG available, built around ChaCha20 with SIMD Assembly and C code written by me. While historically significant, it has been greatly surpassed by the current AEGIS-based design in performance. + +The legacy branch remains available for reference and benchmarking, but version 2 is the recommended implementation. diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..ad851d4 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,65 @@ +[project] +name = "randquik" +dynamic = ["version"] +description = "CLI tool for extremely fast random bytes." +readme = "README.md" +requires-python = ">=3.13" +authors = [ + { name = "Leo Vasanko" } +] +keywords = ["random", "CSPRNG", "AEGIS", "shred", "benchmark"] +classifiers = [ + "Development Status :: 4 - Beta", + "Environment :: Console", + "Intended Audience :: Developers", + "Intended Audience :: System Administrators", + "Topic :: Security :: Cryptography", + "Topic :: Utilities", +] +dependencies = [ + "aeg>=0.4.3", +] + +[project.scripts] +randquik = "randquik.cli:main" + +[project.urls] +Homepage = "https://git.zi.fi/LeoVasanko/randquik" +Repository = "https://github.com/LeoVasanko/randquik" + +[build-system] +requires = ["hatchling", "hatch-vcs"] +build-backend = "hatchling.build" + +[tool.hatch.version] +source = "vcs" + +[tool.hatch.version.raw-options] +local_scheme = "no-local-version" + +[tool.hatch.build.hooks.vcs] +version-file = "randquik/_version.py" + +[tool.hatch.build.targets.wheel] +packages = ["randquik"] + +[tool.ruff] +line-length = 100 +target-version = "py313" + +[tool.ruff.lint] +select = ["E", "F", "W", "I", "UP", "B", "C4", "SIM", "PTH", "RUF"] +ignore = ["E501"] + +[tool.ruff.lint.isort] +known-first-party = ["randquik"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +python_files = ["test_*.py"] + +[tool.mypy] +python_version = "3.13" +warn_return_any = true +warn_unused_configs = true +ignore_missing_imports = true diff --git a/randquik/__init__.py b/randquik/__init__.py new file mode 100644 index 0000000..d43a503 --- /dev/null +++ b/randquik/__init__.py @@ -0,0 +1,24 @@ +"""RandQuik - High-performance cryptographic random data generator. + +This package provides fast random data generation using AEGIS ciphers, +with support for multi-threaded generation and various I/O modes. +""" + +from randquik.crypto import derive_key, generate_random_seed +from randquik.progress import ProgressDisplay +from randquik.utils import format_size, format_time, parse_size + +try: + from randquik._version import __version__ +except ImportError: + __version__ = "0.0.0.dev0" + +__all__ = [ + "ProgressDisplay", + "__version__", + "derive_key", + "format_size", + "format_time", + "generate_random_seed", + "parse_size", +] diff --git a/randquik/__main__.py b/randquik/__main__.py new file mode 100644 index 0000000..e21b1b1 --- /dev/null +++ b/randquik/__main__.py @@ -0,0 +1,6 @@ +"""Entry point for running as `python -m randquik`.""" + +from randquik.cli import main + +if __name__ == "__main__": + main() diff --git a/randquik/_version.py b/randquik/_version.py new file mode 100644 index 0000000..12016ec --- /dev/null +++ b/randquik/_version.py @@ -0,0 +1,34 @@ +# file generated by setuptools-scm +# don't change, don't track in version control + +__all__ = [ + "__version__", + "__version_tuple__", + "version", + "version_tuple", + "__commit_id__", + "commit_id", +] + +TYPE_CHECKING = False +if TYPE_CHECKING: + from typing import Tuple + from typing import Union + + VERSION_TUPLE = Tuple[Union[int, str], ...] + COMMIT_ID = Union[str, None] +else: + VERSION_TUPLE = object + COMMIT_ID = object + +version: str +__version__: str +__version_tuple__: VERSION_TUPLE +version_tuple: VERSION_TUPLE +commit_id: COMMIT_ID +__commit_id__: COMMIT_ID + +__version__ = version = '0.1.dev0' +__version_tuple__ = version_tuple = (0, 1, 'dev0') + +__commit_id__ = commit_id = None diff --git a/randquik/benchmark.py b/randquik/benchmark.py new file mode 100644 index 0000000..c300c63 --- /dev/null +++ b/randquik/benchmark.py @@ -0,0 +1,152 @@ +"""Benchmark functions for measuring performance.""" + +import contextlib +import os +import pathlib +import re +import subprocess +import sys +import time + +from randquik.utils import sparse_range + +__all__ = ["bench_mode", "run_benchmark"] + + +def bench_mode( + tcounts: list[int], + io_mode: str, + length: str, + alg: str | None, + bench_file: pathlib.Path | None, +) -> list[tuple[int, float, list[str]]]: + max_repeats = 5 + max_time = 0.5 # Quit early if more than 500ms has passed + results = [] + + # Use a file in current folder for file modes + if "file" in io_mode: + iocmd = ["-o", str(bench_file)] + elif "dry" in io_mode: + iocmd = ["--dry"] + elif "null" in io_mode: + iocmd = ["-o", os.devnull] + else: + raise ValueError(f"Unknown io_mode: {io_mode}") + + if "mmap" in io_mode: + iocmd.append("--mmap") + + # Print iocmd at start of row + print(f"{' '.join(iocmd)[:20]:<20}", end="", flush=True) + + for workers in tcounts: + speeds = [] + worker_start = time.perf_counter() + for rep in range(max_repeats): + if rep > 0 and (time.perf_counter() - worker_start) > max_time: + break + cmd = [ + sys.executable, + "-m", + "randquik", + f"-l{length}", + f"-t{workers}", + *([f"-a{alg}"] if alg else []), + *iocmd, + ] + + try: + proc = subprocess.run( + cmd, + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + text=False, + ) + except KeyboardInterrupt: + print(" Interrupted\n\n", end="", flush=True) # 6 not 8 to account for ^C + print(f">>> {' '.join(cmd)}", file=sys.stderr) + sys.exit(1) + stderr = proc.stderr.decode(errors="ignore") + if proc.returncode != 0: + print(f"{'ERROR':>8}\n\n", end="", flush=True) + print(f">>> {' '.join(cmd)}\n{stderr}", file=sys.stderr) + sys.exit(1) + m2 = re.findall(r"([0-9]+\.[0-9]+)\s+GB/s", stderr) + if m2: + speeds.append(float(m2[-1])) + + if speeds: + sorted_speeds = sorted(speeds) + median = sorted_speeds[len(speeds) // 2] + print(f"{median:>8.2f}", end="", flush=True) + results.append((workers, median, iocmd)) + else: + print(f"{'---':>8}", end="", flush=True) + + print() # newline after row + + # Cleanup bench file + if bench_file: + with contextlib.suppress(OSError): + bench_file.unlink() + + return results + + +def run_benchmark(args): + """Run comprehensive benchmark across all I/O modes.""" + # Check if output file already exists + bench_file = pathlib.Path(args.output or "test.dat") + if bench_file.exists(): + if not args.output: + raise ValueError( + f"File test.dat already exists. Use -o {bench_file} to benchmark over it or choose another name." + ) + bench_file.unlink() + + length = args.len or "128MiB" + max_threads = args.threads if args.threads is not None else os.cpu_count() + + all_results = {} + tcounts = sparse_range(max_threads) + + # Print header row + print(f"{'randquik':<20}", end="") + for w in tcounts: + print(f"{'-t' + str(w):>8}", end="") + print() + print("-" * (20 + 8 * len(tcounts))) + + for io_mode in ["dry", "dry-mmap", "null", "file", "file-mmap"]: + results = bench_mode(tcounts, io_mode, length, alg=args.alg, bench_file=bench_file) + all_results[io_mode] = results + + print("-" * (20 + 8 * len(tcounts))) + + # Find best for file output + best_speed = 0.0 + best_iocmd = None + best_threads = 0 + for io_mode in ["file", "file-mmap"]: + for w, sp, iocmd in all_results.get(io_mode, []): + if sp > best_speed: + best_speed = sp + best_iocmd = iocmd + best_threads = w + + # Find fastest generation speed + gen_speed = 0.0 + for io_mode in all_results: + for _w, sp, _iocmd in all_results.get(io_mode, []): + if sp > gen_speed: + gen_speed = sp + + if best_iocmd: + threads = f" -t{best_threads}" if best_threads != 1 else "" + print( + f"\n>>> Fastest wrote {best_speed:.2f} GB/s, plain RNG {gen_speed:.0f} GB/s\n" + f"randquik {' '.join(best_iocmd)}{threads}\n" + ) + else: + print("\nNo file output results collected.", file=sys.stderr) diff --git a/randquik/cli.py b/randquik/cli.py new file mode 100644 index 0000000..9a6b1ec --- /dev/null +++ b/randquik/cli.py @@ -0,0 +1,305 @@ +"""Command-line interface for RandQuik.""" + +import argparse +import gc +import mmap +import os +import sys +import time + +import aeg +import tracerite + +from randquik.benchmark import run_benchmark +from randquik.crypto import derive_key, generate_random_seed +from randquik.io import open_fd +from randquik.progress import ProgressDisplay +from randquik.utils import ( + parse_size, + print_summary, +) +from randquik.workers import ( + BLOCK_SIZE, + FdProducer, + MmapProducer, +) + +tracerite.load() + +__all__ = ["main"] + +# Disable GC for performance +gc.disable() + +ciph: aeg.Cipher = None # type: ignore (set in main after parsing args) + +DEFAULT_ALG = "AEGIS-128X2" + + +def prepare_seed(args): + """Prepare seed and determine if it was generated.""" + generated_seed = args.seed is None + seed = generate_random_seed() if generated_seed else args.seed + return seed, generated_seed + + +def prepare_key(seed, keybytes): + """Derive key from seed.""" + key = derive_key(seed, keybytes) + return key + + +def parse_seeks(args): + """Parse seek values from arguments.""" + try: + iseek = parse_size(args.iseek, args.output) or 0 + oseek = parse_size(args.oseek, args.output) or 0 + if args.seek: + iseek = oseek = parse_size(args.seek, args.output) or 0 + except ValueError as e: + raise ValueError(f"Error parsing seek: {e}") from None + return iseek, oseek + + +def _mmap(output: str | None, oseek: int, length: int, *, dry=False) -> mmap.mmap: + with open_fd(output, length, oseek=oseek, dry=dry) as fd: + try: + return mmap.mmap(fd, length) + except (OSError, ValueError) as e: + if fd == -1: + raise ValueError("Cannot mmap all memory: specify --len SIZE") from e + if not output: + raise ValueError("Cannot mmap stdout: remove --mmap or use -o FILE") from e + raise ValueError(f"Cannot mmap {output}: {e}") from e + + +def singlethreaded(args, total_bytes, oseek, start_time, key, seed_for_display): + if args.verbose: + mode_desc = "infinite output" if total_bytes is None else "workers=0" + print( + f"Special mode: {mode_desc} — single-buffer, single-threaded generation", + file=sys.stderr, + ) + + with open_fd( + args.output, 0 if total_bytes is None else total_bytes, dry=args.dry, oseek=oseek + ) as fd: + written = 0 + buf = bytearray(BLOCK_SIZE) + view = memoryview(buf) + nonce = bytearray(ciph.NONCEBYTES) + + progress_state = {"written": 0} + progress = ProgressDisplay( + total_bytes, + start_time, + progress_state, + infinite=total_bytes is None, + seed=seed_for_display, + ) + if not args.quiet: + progress.start() + + try: + while total_bytes is None or written < total_bytes: + start = written + if total_bytes is None: + size = BLOCK_SIZE + else: + end = min(start + BLOCK_SIZE, total_bytes) + size = end - start + chunk = view[:size] + ciph.stream(key, nonce, size, into=chunk) + if not args.dry: + os.write(fd, chunk) + ciph.nonce_increment(nonce) + written += size + progress_state["written"] = written + finally: + progress.stop() + elapsed = time.perf_counter() - start_time + action = "generated" if args.dry else "wrote" + if not args.quiet and total_bytes is not None: + print_summary(written, elapsed, action, seed=seed_for_display) + + +def _main(): + """Internal main function that may raise exceptions.""" + parser = argparse.ArgumentParser(description="Generate random bytes using AEGIS ciphers") + parser.add_argument("-s", "--seed", help="Alphanumeric seed string", type=str) + parser.add_argument( + "-l", + "--len", + help="Length to generate (e.g. 1g, 100mi, 1000sect)", + type=str, + default=None, + ) + parser.add_argument("-o", "--output", help="Output file (default: stdout)", type=str) + parser.add_argument( + "-t", + "--threads", + help="Number of worker threads (benchmark: upper limit)", + type=int, + default=None, + ) + parser.add_argument( + "-a", + "--alg", + help=f"Cipher algorithm (default: {DEFAULT_ALG})", + type=str, + default=None, + ) + parser.add_argument( + "--mmap", + action="store_true", + help="Use file-backed mmap for output instead of writing via fd", + ) + parser.add_argument( + "--benchmark", + action="store_true", + help="Run benchmark (generates 1GB and reports speed)", + ) + parser.add_argument( + "--dry", + action="store_true", + help="Dry run: open output but skip writes (for benchmarking)", + ) + parser.add_argument( + "--seek", + type=str, + default=None, + help="Seek both input stream and output to position (e.g. 1g, 100mi)", + ) + parser.add_argument( + "--iseek", + type=str, + default=None, + help="Seek input random stream to position (overrides --seek)", + ) + parser.add_argument( + "--oseek", + type=str, + default=None, + help="Seek output file to position (overrides --seek)", + ) + parser.add_argument( + "-q", + "--quiet", + action="store_true", + help="Quiet mode: suppress all output except errors", + ) + parser.add_argument( + "-v", + "--verbose", + action="store_true", + help="Verbose mode: show I/O mode and timing statistics", + ) + + args = parser.parse_args() + + # Normalize "-" output to None (stdout) + if args.output == "-": + args.output = None + + global ciph + ciph = aeg.cipher(args.alg or DEFAULT_ALG) + + # Validate and process args + seed, generated_seed = prepare_seed(args) + key = prepare_key(seed, ciph.KEYBYTES) + iseek, oseek = parse_seeks(args) + total_bytes = parse_size(args.len) # None if not specified, 0 if -l0 + # Seed hint for generated seeds + seed_for_display = seed if generated_seed else None + + start_time = time.perf_counter() + + if args.benchmark: + if args.seed is not None: + raise ValueError("Cannot specify seed in benchmark mode") + if iseek or oseek: + raise ValueError("Cannot use seek options in benchmark mode") + run_benchmark(args) + return + + # Single-threaded mode (workers == 0) + if args.threads == 0: + return singlethreaded(args, total_bytes, oseek, start_time, key, seed_for_display) + + workers = args.threads if args.threads is not None else 1 + # File-backed mmap output + if args.mmap: + with ( + _mmap( + args.output, + oseek, + length=oseek + (total_bytes if total_bytes is not None else 0), + dry=args.dry, + ) as mm, + ): + producer = MmapProducer( + workers, + key, + ciph, + total_bytes, + mm, + use_madvise=True, + oseek=oseek, + iseek=iseek, + ) + progress_state = {"written": 0} + producer.progress_state = progress_state + producer.start() + + progress = ProgressDisplay( + total_bytes, start_time, progress_state, seed=seed_for_display + ) + if not args.quiet: + progress.start() + + try: + producer.join() + finally: + progress.stop() + producer.cleanup() + elapsed = time.perf_counter() - start_time + if not args.quiet: + print_summary(producer.written, elapsed, "wrote", seed=seed_for_display) + return + + # Standard file mode with ring buffers + with open_fd(args.output, total_bytes, dry=args.dry, oseek=oseek) as fd: + producer = FdProducer(workers, key, ciph, total_bytes, fd, dry=args.dry, iseek=iseek) + + if args.verbose: + dry_str = " (dry run)" if args.dry else "" + print( + f"I/O mode: {workers} workers, {producer.num_slots} buffers, sequential writes{dry_str}", + file=sys.stderr, + ) + + progress_state = {"written": 0} + progress = ProgressDisplay(total_bytes, start_time, progress_state, seed=seed_for_display) + if not args.quiet: + progress.start() + + try: + producer.start() + producer.run(progress_state) + finally: + producer.stop() + progress.stop() + elapsed = time.perf_counter() - start_time + if not args.quiet: + print_summary(producer.written, elapsed, "wrote", seed=seed_for_display) + + +def main(): + """Main entry point for the CLI with exception handling.""" + try: + _main() + except KeyboardInterrupt: + sys.exit(1) + except ValueError as e: + print(f"Error: {e}", file=sys.stderr) + sys.exit(1) diff --git a/randquik/crypto.py b/randquik/crypto.py new file mode 100644 index 0000000..8ae6c91 --- /dev/null +++ b/randquik/crypto.py @@ -0,0 +1,22 @@ +"""Cryptographic functions for key derivation and seed generation.""" + +import hashlib +import secrets +import string + +__all__ = [ + "derive_key", + "generate_random_seed", +] + + +def generate_random_seed() -> str: + """Generate a random alphanumeric seed string.""" + chars = string.ascii_letters + string.digits + return "".join(secrets.choice(chars) for _ in range(16)) + + +def derive_key(seed: str, key_bytes: int) -> bytes: + """Derive a key from a seed string using SHA-512.""" + assert 16 <= key_bytes <= 64, "Only 128-512 bits supported" + return hashlib.sha512(seed.encode()).digest()[:key_bytes] diff --git a/randquik/io.py b/randquik/io.py new file mode 100644 index 0000000..1b006c1 --- /dev/null +++ b/randquik/io.py @@ -0,0 +1,148 @@ +"""File I/O helpers for output handling and mmap operations.""" + +import contextlib +import ctypes +import mmap +import os +import pathlib +import sys +from collections.abc import Generator + +__all__ = [ + "HAS_MADVISE", + "madvise_buffer", + "madvise_file", + "open_fd", + "open_memoryview", +] + +# Platform-specific constants for madvise +MADV_DONTNEED = 4 # Linux/macOS +MADV_SEQUENTIAL = 2 # Hint for sequential access +MADV_WILLNEED = 3 # Pre-fault pages +MADV_RANDOM = getattr(mmap, "MADV_RANDOM", 1) # Not available on Windows + +# Try to get O_DIRECT (Linux only, not available on macOS) +O_DIRECT = getattr(os, "O_DIRECT", 0) + +# Load libc for madvise (not available on Windows) +HAS_MADVISE = False +_madvise = None +if sys.platform != "win32": + try: + if sys.platform == "darwin": + _libc = ctypes.CDLL("libc.dylib", use_errno=True) + else: + _libc = ctypes.CDLL("libc.so.6", use_errno=True) + _madvise = _libc.madvise + _madvise.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int] + _madvise.restype = ctypes.c_int + HAS_MADVISE = True + except (OSError, AttributeError): + pass + + +def _madvise_call(mm: mmap.mmap, advice: int, offset: int = 0, length: int = 0): + """Call madvise with specified advice.""" + if length == 0: + length = len(mm) + mm.madvise(advice, offset, length) + + +def madvise_buffer(mm: mmap.mmap, offset: int = 0, length: int = 0): + """Mark mmap region for random access (avoid caching).""" + _madvise_call(mm, MADV_RANDOM, offset, length) + + +def madvise_file(mm: mmap.mmap, offset: int = 0, length: int = 0): + """Mark mmap region for sequential access (pre-fault pages).""" + _madvise_call(mm, MADV_WILLNEED, offset, length) + + +def _open_output( + output_path: str, + total_bytes: int | None, + oseek: int = 0, +) -> int: + """Open output file descriptor, preallocate and apply platform hints.""" + if output_path: + flags = os.O_RDWR | os.O_CREAT + fd = os.open(str(pathlib.Path(output_path)), flags, 0o644) + else: + if sys.stdout.isatty(): + raise ValueError("Refusing to write binary data to terminal. Use -o to specify a file.") + fd = sys.stdout.fileno() + + required_size = oseek + (total_bytes if total_bytes is not None else 0) + current_size = os.fstat(fd).st_size + if required_size > current_size: + os.ftruncate(fd, required_size) + + # Seek to output position + if oseek > 0: + try: + os.lseek(fd, oseek, os.SEEK_SET) + except OSError as e: + raise ValueError( + f"Cannot oseek in {output_path or 'stdout'}. Use only --iseek or specify a seekable file." + ) from e + + # macOS: try to bypass unified buffer cache (F_NOCACHE) + if sys.platform == "darwin": + try: + import fcntl + + fcntl.fcntl(fd, fcntl.F_NOCACHE, 1) + except (OSError, AttributeError, ImportError): + pass + + return fd + + +@contextlib.contextmanager +def open_fd( + output_path: str | None, + total_bytes: int | None, + dry: bool = False, + oseek: int = 0, +) -> Generator[int]: + """Context manager for output file descriptor. + + Args: + output_path: Path to output file, or None for stdout + total_bytes: Total bytes to write + dry: If True, skip truncation/preallocation and tty check + oseek: Seek position for output + + Yields: + Integer file descriptor + """ + if dry: + yield -1 + return + if not output_path: + yield sys.stdout.fileno() + return + fd = _open_output(output_path, total_bytes, oseek) + try: + yield fd + finally: + with contextlib.suppress(Exception): + os.close(fd) + + +@contextlib.contextmanager +def open_memoryview(buf) -> Generator[memoryview]: + """Context manager for memoryview. + + Args: + buf: Buffer to create memoryview from + + Yields: + memoryview object + """ + view = memoryview(buf) + try: + yield view + finally: + view.release() diff --git a/randquik/progress.py b/randquik/progress.py new file mode 100644 index 0000000..979bd4e --- /dev/null +++ b/randquik/progress.py @@ -0,0 +1,606 @@ +"""Full-screen progress display with speed graph.""" + +import math +import os +import sys +import threading +import time + +from randquik.utils import format_time + +__all__ = ["ProgressDisplay"] + +# Unicode block characters for graph (8 levels per cell) +GRAPH_BLOCKS = " ▁▂▃▄▅▆▇█" + + +class ProgressDisplay: + """Full-screen progress display with speed graph, updated every 100ms. + + Only active when stderr is a tty. Reads progress from a shared state dict + with a single 'written' key. All display logic is encapsulated here. + + The graph fills from left to right as progress advances, doubling as both + a progress bar and a speed-over-time visualization. + """ + + def __init__( + self, + total_bytes: int | None, + start_time: float, + state: dict, + infinite: bool | None = None, + seed: str | None = None, + ): + self.total_bytes = total_bytes + self.start_time = start_time + self.state = state # Must have 'written' key + self.infinite = infinite if infinite is not None else total_bytes is None + self.seed = seed + self.active = sys.stderr.isatty() + self._stop = threading.Event() + self._thread = None + self._last_written = 0 + self._last_time = start_time + # Speed history for graph - fixed size, filled from left as progress advances + self._graph_width = 80 # Will be updated on first render + self._speed_history: list[float] = [] # Stores GB/s values, one per column + self._max_speed: float = 0.1 # Start with small value to avoid div by zero + # For infinite mode: track time of each speed sample + self._time_history: list[float] = [] + + def start(self): + if not self.active: + return + self._save_terminal_state() + self._thread = threading.Thread(target=self._run, daemon=True) + self._thread.start() + + def stop(self): + if not self.active or self._thread is None: + return + self._stop.set() + self._thread.join(timeout=0.5) + self._restore_terminal_state() + + def _save_terminal_state(self): + """Save terminal state and enter alternate screen buffer.""" + sys.stderr.write("\x1b[?1049h") # Enter alternate screen buffer + sys.stderr.write("\x1b[?25l") # Hide cursor + sys.stderr.flush() + + def _restore_terminal_state(self): + """Restore terminal state and exit alternate screen buffer.""" + sys.stderr.write("\x1b[?25h") # Show cursor + sys.stderr.write("\x1b[?1049l") # Exit alternate screen buffer + sys.stderr.flush() + + def _get_terminal_size(self) -> tuple[int, int]: + """Return (columns, rows).""" + try: + size = os.get_terminal_size(sys.stderr.fileno()) + return size.columns, size.lines + except (OSError, ValueError): + return 80, 24 + + def _render_graph_row( + self, + values: list[float], + max_val: float, + row: int, + total_rows: int, + width: int, + avg_speed: float = 0, + ) -> str: + """Render one row of the graph using Unicode blocks. + + Row 0 is top, total_rows-1 is bottom. Each cell can show 8 levels. + Values list may be shorter than width (unfilled area shown as dim bar at avg_speed). + """ + filled_chars = [] + unfilled_chars = [] + filled_cols = len(values) + # Calculate the row threshold for average speed + avg_normalized = (avg_speed / max_val) * total_rows * 8 if max_val > 0 else 0 + row_bottom = (total_rows - row - 1) * 8 + row_top = row_bottom + 8 + + # Build filled portion + for i in range(filled_cols): + v = values[i] + # Normalize value to 0..total_rows*8 range + normalized = (v / max_val) * total_rows * 8 if max_val > 0 else 0 + if normalized <= row_bottom: + filled_chars.append(" ") + elif normalized >= row_top: + filled_chars.append("█") + else: + level = int(normalized - row_bottom) + filled_chars.append(GRAPH_BLOCKS[min(level, 8)]) + + # Build unfilled portion (dim grey at avg_speed level) + unfilled_width = width - filled_cols + if unfilled_width > 0: + if avg_normalized <= row_bottom: + unfilled_char = " " + elif avg_normalized >= row_top: + unfilled_char = "█" + else: + level = int(avg_normalized - row_bottom) + unfilled_char = GRAPH_BLOCKS[min(level, 8)] + unfilled_chars = [unfilled_char] * unfilled_width + + # Combine with color codes only at transitions + filled_str = "".join(filled_chars) + unfilled_str = "".join(unfilled_chars) + if unfilled_str: + return f"{filled_str}\x1b[0m\x1b[38;5;234m{unfilled_str}\x1b[0m\x1b[33m" + return filled_str + + def _render_full_screen(self) -> str: + """Render the full screen display.""" + if self.infinite: + return self._render_infinite_screen() + return self._render_progress_screen() + + def _render_infinite_screen(self) -> str: + """Render screen for infinite mode (no known total).""" + cols, rows = self._get_terminal_size() + written = self.state.get("written", 0) + now = time.perf_counter() + elapsed = now - self.start_time + + # Calculate graph width (leave room for Y-axis labels) + graph_width = cols - 8 + self._graph_width = graph_width + + # Calculate speeds + overall_speed = written / elapsed if elapsed > 0 else 0 + dt = now - self._last_time + instant_speed = (written - self._last_written) / dt if dt > 0 else 0 + self._last_written = written + self._last_time = now + + # Update speed history + speed_gbs = instant_speed / 1_000_000_000 + self._speed_history.append(speed_gbs) + self._time_history.append(elapsed) + + # Update max speed + if speed_gbs > self._max_speed: + self._max_speed = speed_gbs + scale_max = self._nice_scale(self._max_speed) + + # Build output + lines = [] + lines.append("\x1b[2J\x1b[H") + + # Compact header with stats + spinner = "◐◓◑◒"[int(elapsed * 4) % 4] + written_gb = written / 1_000_000_000 + current_speed = instant_speed / 1_000_000_000 + seed_hint = f" \x1b[2m-s {self.seed}\x1b[0m" if self.seed else "" + header = ( + f" \x1b[1;36mRandQuik {spinner}\x1b[0m " + f"\x1b[2m│\x1b[0m {written_gb:6.2f}/∞ GB " + f"\x1b[2m@\x1b[0m {current_speed:5.2f} GB/s " + f"\x1b[2m│\x1b[0m {format_time(elapsed):>8}{seed_hint}" + ) + lines.append(header) + lines.append("") + + # Calculate graph dimensions (footer_lines includes GB/s label, time axis, and footer) + header_lines = len(lines) + footer_lines = 4 + graph_rows = max(3, rows - header_lines - footer_lines) + + # Downsample speed history for display - average samples within each column's time range + min_scale_time = 10.0 + scale_time = max(elapsed, min_scale_time) + col_width_time = scale_time / (graph_width - 1) if graph_width > 1 else scale_time + + display_values = [] + for col in range(graph_width): + col_time = col / (graph_width - 1) * scale_time if graph_width > 1 else 0 + if col_time > elapsed: + break + # Find all samples within this column's time range + col_start = col_time - col_width_time / 2 + col_end = col_time + col_width_time / 2 + samples = [ + self._speed_history[idx] + for idx, t in enumerate(self._time_history) + if col_start <= t <= col_end + ] + if samples: + display_values.append(sum(samples) / len(samples)) + elif self._speed_history: + # Fallback to closest if no samples in range + best_idx = 0 + best_diff = float("inf") + for idx, t in enumerate(self._time_history): + diff = abs(t - col_time) + if diff < best_diff: + best_diff = diff + best_idx = idx + display_values.append(self._speed_history[best_idx]) + + avg_speed_gbs = overall_speed / 1_000_000_000 + + # Add GB/s label above the graph + lines.append(" \x1b[36mGB/s\x1b[0m") + + # Compute nice Y-axis tick values + nice_ticks = set(self._nice_y_ticks(scale_max)) + labeled_values = set() # Track which tick values have been labeled + tolerance = scale_max / (graph_rows - 1) / 2 if graph_rows > 1 else 0.1 + + for row in range(graph_rows): + graph_line = self._render_graph_row( + display_values, + scale_max, + row, + graph_rows, + len(display_values), + avg_speed_gbs, + ) + graph_line = graph_line.ljust(graph_width) + + # Y-axis labels - only label nice tick values + row_value = ( + scale_max * (graph_rows - 1 - row) / (graph_rows - 1) if graph_rows > 1 else 0 + ) + label = " " + for tick in nice_ticks: + if abs(row_value - tick) < tolerance and tick not in labeled_values: + label = f"{self._format_label(tick):>4}" + labeled_values.add(tick) + break + lines.append(f" \x1b[36m{label}\x1b[0m \x1b[33m{graph_line}\x1b[0m") + + # Time axis + time_axis = self._build_infinite_time_axis(graph_width, scale_time) + lines.append(f" {''.join(time_axis)}") + + # Footer + lines.append(f"\x1b[2m{'[Ctrl+C to stop]':^{cols}}\x1b[0m") + + return "\n".join(lines) + + def _nice_scale(self, max_speed: float) -> float: + """Round up to next nice number for scale.""" + if max_speed <= 0.1: + return 0.1 + log_val = math.log10(max_speed) + power = math.floor(log_val) + mantissa = max_speed / (10**power) + nice_mantissa = math.ceil(mantissa) + if nice_mantissa > 9: + nice_mantissa = 1 + power += 1 + return nice_mantissa * (10**power) + + def _format_label(self, val: float) -> str: + """Format Y-axis label.""" + if val == 0: + return "0" + elif val >= 1: + return f"{val:.0f}" + else: + return f"{val:.1f}" + + def _nice_y_ticks(self, scale_max: float, max_ticks: int = 5) -> list[float]: + """Return nice Y-axis tick values from 0 to scale_max. + + Chooses a nice interval (1, 2, 5 × 10^N) that gives roughly max_ticks labels. + """ + if scale_max <= 0: + return [0] + + # Nice intervals: 0.1, 0.2, 0.5, 1, 2, 5, 10, 20, 50, ... + nice_bases = [1, 2, 5] + best_interval = scale_max + for exp in range(-1, 10): + for base in nice_bases: + interval = base * (10**exp) + num_ticks = scale_max / interval + if 2 <= num_ticks <= max_ticks: + best_interval = interval + break + else: + continue + break + + # Generate ticks from 0 to scale_max at best_interval + ticks = [] + val = 0.0 + while val <= scale_max + 1e-9: + ticks.append(val) + val += best_interval + return ticks + + def _build_infinite_time_axis(self, graph_width: int, scale_time: float) -> list[str]: + """Build time axis for infinite mode with nice interval labels. + + Shows from 0 to scale_time with nice interval markers. + """ + time_axis = [" "] * graph_width + + # Nice time intervals + nice_intervals = [ + 1, + 2, + 5, + 10, + 15, + 30, + 60, + 120, + 300, + 600, + 900, + 1800, + 3600, + 7200, + 18000, + 36000, + ] + + def format_time_short(secs): + """Format time for axis label.""" + if secs == 0: + return "0" + elif secs < 60: + return f"{int(secs)}s" + elif secs < 3600: + m = int(secs // 60) + return f"{m}m" + else: + h = int(secs // 3600) + return f"{h}h" + + # Find a nice interval that gives us ~4-8 labels + interval = nice_intervals[-1] + for ni in nice_intervals: + if scale_time / ni <= 8: + interval = ni + break + + # Place labels at nice intervals starting from 0 + t = 0 + while t <= scale_time: + col = int(t / scale_time * (graph_width - 1)) if scale_time > 0 else 0 + if 0 <= col < graph_width: + label = format_time_short(t) + label_start = max(0, col - len(label) // 2) + label_end = min(graph_width, label_start + len(label)) + if all(c == " " for c in time_axis[label_start:label_end]): + for i, ch in enumerate(label): + if label_start + i < graph_width: + time_axis[label_start + i] = ch + t += interval + + return time_axis + + def _render_progress_screen(self) -> str: + """Render the full screen display for finite progress.""" + cols, rows = self._get_terminal_size() + written = self.state.get("written", 0) + now = time.perf_counter() + elapsed = now - self.start_time + + # Calculate graph width (leave room for Y-axis labels) + graph_width = cols - 8 + self._graph_width = graph_width + + # Calculate speeds + overall_speed = written / elapsed if elapsed > 0 else 0 + dt = now - self._last_time + instant_speed = (written - self._last_written) / dt if dt > 0 else 0 + self._last_written = written + self._last_time = now + + # ETA and total estimated time + remaining = self.total_bytes - written + # ETA based on current instant speed for responsiveness + eta = remaining / instant_speed if instant_speed > 0 else -1 + # Graph X-axis scaling based on average speed for stability + avg_eta = remaining / overall_speed if overall_speed > 0 else -1 + estimated_total_time = elapsed + avg_eta if avg_eta > 0 else elapsed + + # Progress percentage (for display) + pct = min(100, written * 100 / self.total_bytes) if self.total_bytes > 0 else 0 + + # Time-based progress: columns represent time, not percentage + # Graph fills based on elapsed / estimated_total_time + # Use ceiling so column appears when we've started it, not when complete + if estimated_total_time > 0: + time_pct = min(100, elapsed * 100 / estimated_total_time) + else: + time_pct = 100 + target_cols = min(graph_width, int(graph_width * time_pct / 100) + 1) if time_pct > 0 else 0 + + # Update speed history - add new sample if we've advanced to a new column + speed_gbs = instant_speed / 1_000_000_000 + if len(self._speed_history) < target_cols: + # Fill in any skipped columns with the current speed + while len(self._speed_history) < target_cols: + self._speed_history.append(speed_gbs) + elif len(self._speed_history) > 0 and target_cols > 0: + # Update the current column with latest speed (smoothing) + self._speed_history[-1] = (self._speed_history[-1] + speed_gbs) / 2 + + # Update max speed - use "nice" scale values (1, 2, 3, ..., 9 × 10^N), minimum 0.1 GB/s + if speed_gbs > self._max_speed: + self._max_speed = speed_gbs + # Round up to next "nice" number: 0.1, 0.2, ..., 0.9, 1, 2, ..., 9, 10, 20, ... + if self._max_speed <= 0.1: + scale_max = 0.1 + else: + # Find the power of 10 just below max_speed + log_val = math.log10(self._max_speed) + power = math.floor(log_val) + # Get the leading digit and round up + mantissa = self._max_speed / (10**power) + nice_mantissa = math.ceil(mantissa) + if nice_mantissa > 9: + nice_mantissa = 1 + power += 1 + scale_max = nice_mantissa * (10**power) + + # Build output + lines = [] + lines.append("\x1b[2J\x1b[H") + + # Compact header with stats + spinner = "◐◓◑◒"[int(elapsed * 4) % 4] + current_speed = instant_speed / 1_000_000_000 + written_gb = written / 1_000_000_000 + total_gb = self.total_bytes / 1_000_000_000 + seed_hint = f" \x1b[2m-s {self.seed}\x1b[0m" if self.seed else "" + header = ( + f" \x1b[1;36mRandQuik {spinner}\x1b[0m " + f"\x1b[2m│\x1b[0m {written_gb:6.2f}\x1b[2m/\x1b[0m{total_gb:.2f} GB " + f"\x1b[2m@\x1b[0m {current_speed:5.2f} GB/s " + f"ETA {format_time(eta):>8}{seed_hint}" + ) + lines.append(header) + lines.append("") + + # Calculate graph dimensions (footer_lines includes GB/s label, time axis, and footer) + header_lines = len(lines) + footer_lines = 4 + graph_rows = max(3, rows - header_lines - footer_lines) + + # Render graph rows with Y-axis + avg_speed_gbs = overall_speed / 1_000_000_000 + + # Add GB/s label above the graph + lines.append(" \x1b[36mGB/s\x1b[0m") + + # Compute nice Y-axis tick values + nice_ticks = set(self._nice_y_ticks(scale_max)) + labeled_values = set() # Track which tick values have been labeled + tolerance = scale_max / (graph_rows - 1) / 2 if graph_rows > 1 else 0.1 + + for row in range(graph_rows): + graph_line = self._render_graph_row( + self._speed_history, + scale_max, + row, + graph_rows, + graph_width, + avg_speed_gbs, + ) + # Y-axis labels - only label nice tick values + row_value = ( + scale_max * (graph_rows - 1 - row) / (graph_rows - 1) if graph_rows > 1 else 0 + ) + label = " " + for tick in nice_ticks: + if abs(row_value - tick) < tolerance and tick not in labeled_values: + label = f"{self._format_label(tick):>4}" + labeled_values.add(tick) + break + lines.append(f" \x1b[36m{label}\x1b[0m \x1b[33m{graph_line}\x1b[0m") + + # Time labels on X-axis with nice intervals + time_axis = self._build_time_axis(graph_width, estimated_total_time) + lines.append(f" {''.join(time_axis)}") + + # Footer + lines.append(f"\x1b[2m{'[Ctrl+C to abort]':^{cols}}\x1b[0m") + + # Position percentage at top of bar at current progress point + # Find the height of the bar at current progress (last value in speed_history) + current_speed_gbs = self._speed_history[-1] if self._speed_history else 0 + # Normalize to find which row the top of the bar is at + # Bar fills from bottom up; row 0 is top, graph_rows-1 is bottom + if scale_max > 0 and current_speed_gbs > 0: + # How many rows from bottom does the bar fill? + bar_height_fraction = current_speed_gbs / scale_max + # The top of the bar is at this row (0 = top, graph_rows-1 = bottom) + bar_top_row = int(graph_rows * (1 - bar_height_fraction)) + bar_top_row = max(0, min(graph_rows - 1, bar_top_row)) + else: + bar_top_row = graph_rows - 1 # At bottom if no speed + + # Graph starts at row 4 (1-indexed: header=1, empty=2, GB/s=3, first graph row=4) + # Column offset is 6 (space + 4 char Y-label + space) + progress_col = int(pct / 100 * (graph_width - 1)) + 7 # +7 for Y-axis offset (1-indexed) + pct_label = f"{int(pct)}%" + # Center the label on the progress point + pct_col = max(7, progress_col - len(pct_label) // 2) + # Calculate screen row (1-indexed for ANSI) + pct_row = 4 + bar_top_row + pct_position = f"\x1b[{pct_row};{pct_col}H\x1b[1;37m{pct_label}\x1b[0m" + + return "\n".join(lines) + pct_position + + def _build_time_axis(self, graph_width: int, estimated_total_time: float) -> list[str]: + """Build time axis with nice interval labels.""" + + def nice_time_interval(total_secs): + """Return a nice interval for time axis labels.""" + nice_intervals = [ + 1, + 2, + 5, + 10, + 15, + 30, + 60, + 120, + 300, + 600, + 900, + 1800, + 3600, + 7200, + 18000, + 36000, + ] + for interval in nice_intervals: + if total_secs / interval <= 8: + return interval + return 36000 + + def format_time_short(secs): + """Format time for axis label.""" + if secs < 60: + return f"{int(secs)}s" + elif secs < 3600: + return f"{int(secs // 60)}m" + else: + return f"{int(secs // 3600)}h" + + time_axis = [" "] * graph_width + if estimated_total_time > 0: + interval = nice_time_interval(estimated_total_time) + t = 0 + while t <= estimated_total_time: + col = ( + int(t / estimated_total_time * (graph_width - 1)) + if estimated_total_time > 0 + else 0 + ) + if col < graph_width: + label = format_time_short(t) + start = max(0, col - len(label) // 2) + end = min(graph_width, start + len(label)) + if all(c == " " for c in time_axis[start:end]): + for i, ch in enumerate(label): + if start + i < graph_width: + time_axis[start + i] = ch + t += interval + else: + time_axis = list(f"{'0s':<{graph_width}}") + + return time_axis + + def _run(self): + """Background thread: update display every 100ms.""" + while not self._stop.wait(0.1): + screen = self._render_full_screen() + sys.stderr.write(screen) + sys.stderr.flush() diff --git a/randquik/utils.py b/randquik/utils.py new file mode 100644 index 0000000..16ba95c --- /dev/null +++ b/randquik/utils.py @@ -0,0 +1,237 @@ +"""Utility functions for formatting and parsing.""" + +import pathlib +import re +import sys +import time + +__all__ = [ + "format_size", + "format_time", + "get_output_size", + "get_sector_size", + "parse_size", + "print_summary", +] + + +def format_size(size: float) -> str: + """Format bytes as human-readable size.""" + for unit in ["B", "kB", "MB", "GB", "TB"]: + if abs(size) < 1000: + return f"{size:.0f}{unit}" + size /= 1000 + return f"{size:.0f}PB" + + +def format_time(seconds: float) -> str: + """Format seconds as human-readable time.""" + if seconds < 0: + return "--" + if seconds < 1: + return f"{seconds * 1000:.0f}ms" + if seconds < 90: + return f"{seconds:.0f}s" + elif seconds < 5400: # 90 minutes + m = int(seconds / 60) + return f"{m}m" + elif seconds < 172800: # 48 hours + h = int(seconds / 3600) + return f"{h}h" + else: + d = int(seconds / 86400) + return f"{d}d" + + +def print_summary(written: int, elapsed: float, action: str = "wrote", seed: str | None = None): + """Print a nice one-liner summary.""" + speed_gbs = (written / 1_000_000_000) / elapsed if elapsed > 0 else 0 + size_str = format_size(written) + time_str = format_time(elapsed) + seed_hint = f" [-s {seed}]" if seed else "" + print( + f"RandQuik {action} {size_str} in {time_str} ({speed_gbs:.2f} GB/s){seed_hint}", + file=sys.stderr, + ) + + +# Cache for sector size lookup (path -> size) +_sector_size_cache: dict[str, int] = {} + + +def get_sector_size(path: str | pathlib.Path) -> int: + """Get sector size for a block device, or 512 as fallback.""" + import os + import stat + + try: + st = pathlib.Path(path).stat() + if not stat.S_ISBLK(st.st_mode): + return 512 + except OSError: + return 512 + + try: + fd = os.open(str(path), os.O_RDONLY) + try: + import fcntl + import struct + + if sys.platform == "darwin": + # macOS: DKIOCGETBLOCKSIZE = 0x40046418 + DKIOCGETBLOCKSIZE = 0x40046418 + buf = fcntl.ioctl(fd, DKIOCGETBLOCKSIZE, b"\x00" * 4) + return struct.unpack("I", buf)[0] + else: + # Linux: BLKSSZGET = 0x1268 + BLKSSZGET = 0x1268 + buf = fcntl.ioctl(fd, BLKSSZGET, b"\x00" * 4) + return struct.unpack("i", buf)[0] + finally: + os.close(fd) + except (OSError, ImportError, Exception): + return 512 + + +def get_output_size(path: str | pathlib.Path | None) -> int | None: + """Get the size of output file or block device. + + Returns None for stdout or non-existent files. + Returns the size in bytes for existing files or block devices. + """ + import os + import stat + + if not path: + return None # stdout + + try: + st = pathlib.Path(path).stat() + except OSError: + return None # doesn't exist yet + + if stat.S_ISBLK(st.st_mode): + # Block device - get size via ioctl + try: + fd = os.open(str(path), os.O_RDONLY) + try: + import fcntl + import struct + + if sys.platform == "darwin": + # macOS: DKIOCGETBLOCKCOUNT and DKIOCGETBLOCKSIZE + DKIOCGETBLOCKCOUNT = 0x40086419 + DKIOCGETBLOCKSIZE = 0x40046418 + count_buf = fcntl.ioctl(fd, DKIOCGETBLOCKCOUNT, b"\x00" * 8) + size_buf = fcntl.ioctl(fd, DKIOCGETBLOCKSIZE, b"\x00" * 4) + block_count = struct.unpack("Q", count_buf)[0] + block_size = struct.unpack("I", size_buf)[0] + return block_count * block_size + else: + # Linux: BLKGETSIZE64 = 0x80081272 + BLKGETSIZE64 = 0x80081272 + buf = fcntl.ioctl(fd, BLKGETSIZE64, b"\x00" * 8) + return struct.unpack("Q", buf)[0] + finally: + os.close(fd) + except (OSError, ImportError, Exception): + return None + elif stat.S_ISREG(st.st_mode): + return st.st_size + else: + return None + + +def parse_size(length: str | None, output_path: str | pathlib.Path | None = None) -> int | None: + """Parse size string with SI/IEC prefixes. + + Supports: + - Plain numbers: 1000, 1_000_000 + - SI prefixes: k, m, g, t, p (powers of 1000) + - IEC prefixes: ki, mi, gi, ti, pi (powers of 1024) + - Optional 'b' suffix: kb, kib, mb, mib, etc. + - Special unit 'sect' = device sector size (detected, fallback 512) + - Case insensitive + + Examples: 1k, 1ki, 1kb, 1kib, 100m, 100mi, 1g, 1gi, 10sect + """ + if length is None: + return None + s = length.strip().lower().replace("_", "") + + # Handle sect unit + m = re.match(r"^(\d+)\s*sect?s?$", s) + if m: + if output_path: + path_str = str(output_path) + if path_str not in _sector_size_cache: + _sector_size_cache[path_str] = get_sector_size(output_path) + sector_size = _sector_size_cache.get(path_str, 512) + else: + sector_size = 512 + return int(m.group(1)) * sector_size + + # SI/IEC prefixes + si_prefixes = {"k": 1000, "m": 1000**2, "g": 1000**3, "t": 1000**4, "p": 1000**5} + iec_prefixes = { + "ki": 1024, + "mi": 1024**2, + "gi": 1024**3, + "ti": 1024**4, + "pi": 1024**5, + } + + # Try IEC first (ki, mi, etc.) - must check before SI + m = re.match(r"^(\d+(?:\.\d+)?)\s*(ki|mi|gi|ti|pi)b?$", s) + if m: + num, prefix = m.groups() + return int(float(num) * iec_prefixes[prefix]) + + # Try SI (k, m, g, etc.) + m = re.match(r"^(\d+(?:\.\d+)?)\s*([kmgtp])b?$", s) + if m: + num, prefix = m.groups() + return int(float(num) * si_prefixes[prefix]) + + # Plain number + m = re.match(r"^(\d+)$", s) + if m: + return int(m.group(1)) + + raise ValueError(f"Invalid size format: {length}") + + +def sparse_range(n: int, max_items: int = 9) -> list[int]: + """Generate a sparse range from 1 to N for benchmarking thread counts.""" + if n < 1: + return [1] + if n <= max_items - 1: + return list(range(n + 1)) + + keep = 3 # dense prefix: 1,2,3 + out = list(range(keep + 1)) + + remaining = max_items - keep + step = max(1, n // (remaining - 1)) + + for k in range(1, remaining): + v = k * step + if v > out[-1]: + out.append(v) + + if out[-1] != n: + out[-1] = n + + return out + + +def stopwatch(): + t = time.perf_counter() + + def _generator(): + nonlocal t + while True: + t, t0 = time.perf_counter(), t + yield t - t0 + + return _generator() diff --git a/randquik/workers.py b/randquik/workers.py new file mode 100644 index 0000000..84ec6e0 --- /dev/null +++ b/randquik/workers.py @@ -0,0 +1,298 @@ +"""Worker threads and ring buffer management for parallel generation.""" + +import logging +import os +import sys +import threading + +from randquik.io import madvise_file +from randquik.utils import stopwatch + +__all__ = [ + "BLOCK_SIZE", + "FdProducer", + "MmapProducer", +] + +BLOCK_SIZE = 1 << 20 + + +class FdProducer: + """Multi-threaded producer with ring buffer for sequential file output.""" + + def __init__( + self, + workers: int, + key: bytes, + ciph, + total_bytes: int | None, + fd: int, + dry: bool = False, + iseek: int = 0, + block_size: int = BLOCK_SIZE, + ): + self.workers = workers + self.key = key + self.ciph = ciph + self.total_bytes = total_bytes + self.fd = fd + self.dry = dry + self.iseek = iseek + self.block_size = block_size + + # iseek handling: which block to start at, and offset within first block + self.start_block = iseek // block_size + self.start_offset = iseek % block_size + + # Ring buffer state + self.num_slots = workers + 1 + self._buf = bytearray(self.num_slots * block_size) + self._lock = threading.Lock() + self.needdata = threading.Condition(self._lock) + self.needspace = threading.Condition(self._lock) + self._ready = [False] * self.num_slots + self._genpos = 0 + self._conpos = 0 + self._quit = threading.Event() + + self._threads: list[threading.Thread] = [] + self.written = 0 + self.wait_time = 0.0 + self.write_time = 0.0 + + def start(self): + """Start worker threads.""" + for _ in range(self.workers): + t = threading.Thread(target=self._worker, daemon=True) + self._threads.append(t) + t.start() + + def _worker_round(self, view) -> bool: + """Process one block. Returns True if work was done, False if should quit.""" + with self._lock: + while self._genpos - self._conpos >= self.num_slots: + if self._quit.is_set(): + return False + self.needspace.wait() + if self._quit.is_set(): + return False + block_num = self._genpos + self._genpos += 1 + slot = block_num % self.num_slots + buf = view[slot * self.block_size : (slot + 1) * self.block_size] + # Use start_block + block_num as actual nonce to handle iseek + actual_block = self.start_block + block_num + self.ciph.stream(self.key, actual_block.to_bytes(self.ciph.NONCEBYTES, "little"), into=buf) + with self._lock: + self._ready[slot] = True + self.needdata.notify_all() + return True + + def _worker(self): + # Thread-local memoryview of the shared buffer + view = memoryview(self._buf) + try: + while self._worker_round(view): + pass + except BaseException as e: + logging.exception("Worker thread exception: %s", e) + finally: + view.release() + self._quit.set() + with self._lock: + self.needdata.notify_all() + + def run(self, progress_state: dict | None = None): + """Consume blocks and write to fd. Call start() first.""" + view = memoryview(self._buf) + timer = stopwatch() + is_first_block = True + try: + while self.total_bytes is None or self.written < self.total_bytes: + with self._lock: + slot = self._conpos % self.num_slots + while not self._ready[slot]: + if self._quit.is_set(): + return + self.needdata.wait() + self._ready[slot] = False + + self.wait_time += next(timer) + buf = view[slot * self.block_size : (slot + 1) * self.block_size] + # Handle first block: skip start_offset bytes + if is_first_block and self.start_offset > 0: + buf_start = self.start_offset + is_first_block = False + else: + buf_start = 0 + + to_write = min( + self.block_size - buf_start, + self.total_bytes - self.written + if self.total_bytes is not None + else self.block_size - buf_start, + ) + if not self.dry: + os.write(self.fd, buf[buf_start : buf_start + to_write]) + self.write_time += next(timer) + self.written += to_write + if progress_state is not None: + progress_state["written"] = self.written + + with self._lock: + self._conpos += 1 + self.needspace.notify_all() + finally: + view.release() + + def stop(self): + """Signal workers to stop and wait for them.""" + with self._lock: + self._quit.set() + self.needspace.notify_all() + for t in self._threads: + t.join() + + +class MmapProducer: + """Multi-threaded producer that writes directly into mmap.""" + + def __init__( + self, + workers: int, + key: bytes, + ciph, + total_bytes: int | None, + mm_raw=None, + use_madvise: bool = False, + oseek: int = 0, + iseek: int = 0, + block_size: int = BLOCK_SIZE, + dry: bool = False, + ): + self.workers = workers + self.key = key + self.ciph = ciph + self.total_bytes = total_bytes + self.use_madvise = use_madvise + self.oseek = oseek + self.iseek = iseek + self.block_size = block_size + self.dry = dry + self.written = 0 + + # For dry runs, create an anonymous mmap + import mmap as mmap_module + + if dry: + length = oseek + (total_bytes if total_bytes is not None else 0) + self.mm_raw = mmap_module.mmap(-1, length) + self._owns_mmap = True + else: + self.mm_raw = mm_raw + self._owns_mmap = False + + # iseek handling: which block to start at, and offset within first block + self.start_block = iseek // block_size + self.start_offset = iseek % block_size + + if total_bytes is None: + self.num_blocks = sys.maxsize // block_size + else: + self.num_blocks = (total_bytes + block_size - 1) // block_size + # If first block is partial, we need one more block to cover total_bytes + if self.start_offset > 0 and total_bytes > 0: + self.num_blocks = (self.start_offset + total_bytes + block_size - 1) // block_size + self._lock = threading.Lock() + self._next_block = 0 + self._quit = threading.Event() + self._threads: list[threading.Thread] = [] + self.progress_state: dict | None = None + + def start(self): + """Start worker threads.""" + for _ in range(self.workers): + t = threading.Thread(target=self._worker, daemon=True) + self._threads.append(t) + t.start() + + def _worker(self): + # Create thread-local memoryview from mmap + view = memoryview(self.mm_raw) + # Temporary buffer for first/last partial blocks + tmp_buf = bytearray(self.block_size) + try: + while True: + with self._lock: + if self._next_block >= self.num_blocks: + return + block_num = self._next_block + self._next_block += 1 + + # Calculate actual nonce (accounting for iseek) + actual_block = self.start_block + block_num + nonce = actual_block.to_bytes(self.ciph.NONCEBYTES, "little") + + # Calculate byte range within output + # block_num=0 corresponds to output byte 0 + # But if start_offset > 0, first block is partial + if block_num == 0 and self.start_offset > 0: + # First partial block: generate full block, copy from start_offset + self.ciph.stream(self.key, nonce, self.block_size, into=tmp_buf) + if self.total_bytes is None: + copy_len = self.block_size - self.start_offset + else: + copy_len = min(self.block_size - self.start_offset, self.total_bytes) + out_start = self.oseek + out_end = out_start + copy_len + view[out_start:out_end] = tmp_buf[ + self.start_offset : self.start_offset + copy_len + ] + written = copy_len + else: + # Full block or last partial block + # Output position: account for first block being partial + if self.start_offset > 0: + out_start = ( + self.oseek + + (self.block_size - self.start_offset) + + (block_num - 1) * self.block_size + ) + else: + out_start = self.oseek + block_num * self.block_size + if self.total_bytes is None: + out_end = out_start + self.block_size + else: + out_end = min(out_start + self.block_size, self.oseek + self.total_bytes) + size = out_end - out_start + self.ciph.stream(self.key, nonce, size, into=view[out_start:out_end]) + written = size + + if self.use_madvise and self.mm_raw: + madvise_file(self.mm_raw, out_start, written) + with self._lock: + self.written += written + if self.progress_state is not None: + self.progress_state["written"] = self.written + except BaseException as e: + logging.exception("Worker thread exception: %s", e) + finally: + view.release() + self._quit.set() + self._quit.wait() + + def join(self): + """Wait for all worker threads to finish.""" + for t in self._threads: + t.join() + + def stop(self): + """Signal workers to stop.""" + self._quit.set() + + def cleanup(self): + """Release references to mmap resources.""" + if self._owns_mmap and self.mm_raw is not None: + self.mm_raw.close() + self.mm_raw = None + self._threads.clear()