Initial commit of v2.

This commit is contained in:
2026-01-02 22:16:16 +00:00
commit 5747e9a183
13 changed files with 2029 additions and 0 deletions
+7
View File
@@ -0,0 +1,7 @@
.*
*.lock
*.egg-info
/dist
/build
__pycache__
!.gitignore
+125
View File
@@ -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.
+65
View File
@@ -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
+24
View File
@@ -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",
]
+6
View File
@@ -0,0 +1,6 @@
"""Entry point for running as `python -m randquik`."""
from randquik.cli import main
if __name__ == "__main__":
main()
+34
View File
@@ -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
+152
View File
@@ -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)
+305
View File
@@ -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)
+22
View File
@@ -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]
+148
View File
@@ -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()
+606
View File
@@ -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()
+237
View File
@@ -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()
+298
View File
@@ -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()