From 017095df3696faeffc15aff3a9b2cac6d050f631 Mon Sep 17 00:00:00 2001 From: Leo Vasanko Date: Fri, 2 Jan 2026 23:51:43 +0000 Subject: [PATCH] Miscellaneous fixed, cleanup and linting. --- README.md | 8 +- pyproject.toml | 23 ---- randquik/_version.py | 34 ----- randquik/cli.py | 14 ++- randquik/io.py | 3 +- randquik/workers.py | 291 ++++++++++++++++++++++++++++++++++++------- 6 files changed, 258 insertions(+), 115 deletions(-) delete mode 100644 randquik/_version.py diff --git a/README.md b/README.md index 47e0660..ac6f3ed 100644 --- a/README.md +++ b/README.md @@ -62,12 +62,6 @@ Many options such as `--len`, `--seek` accept human-readable units. The table us | 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. - ### Seeding for repeatable output You can provide an explicit seed string, always providing the same output, which can be useful e.g. for memory/disk testing where the data needs to be read back and verified. @@ -76,7 +70,7 @@ You can provide an explicit seed string, always providing the same output, which randquik -l 64MiB -s my-seed-string -o chunk.bin ``` -If no seed is provided, a secure random one will be generated and printed on console (unless hidden by `-q`). +If no seed is provided, a secure new random one is created and printed on console (unless hidden by `-q`). ### Seekable random stream and output file diff --git a/pyproject.toml b/pyproject.toml index ad851d4..878f7ea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,32 +34,9 @@ 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/_version.py b/randquik/_version.py deleted file mode 100644 index 12016ec..0000000 --- a/randquik/_version.py +++ /dev/null @@ -1,34 +0,0 @@ -# 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/cli.py b/randquik/cli.py index e2dbe6d..c6bd3b1 100644 --- a/randquik/cli.py +++ b/randquik/cli.py @@ -8,7 +8,6 @@ import sys import time import aeg -import tracerite from randquik.benchmark import run_benchmark from randquik.crypto import derive_key, generate_random_seed @@ -24,8 +23,6 @@ from randquik.workers import ( MmapProducer, ) -tracerite.load() - __all__ = ["main"] # Disable GC for performance @@ -261,15 +258,20 @@ def _main(): 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) + if args.verbose: + stats = producer.get_combined_stats() + print(stats.format_report(f"Workers (×{workers})"), file=sys.stderr) + producer.cleanup() 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) + producer = FdProducer( + workers, key, ciph, total_bytes, fd, dry=args.dry, iseek=iseek, profile=args.verbose + ) if args.verbose: dry_str = " (dry run)" if args.dry else "" @@ -292,6 +294,8 @@ def _main(): elapsed = time.perf_counter() - start_time if not args.quiet: print_summary(producer.written, elapsed, "wrote", seed=seed_for_display) + if args.verbose: + print(producer.format_stats_report(), file=sys.stderr) def main(): diff --git a/randquik/io.py b/randquik/io.py index 1b006c1..7674471 100644 --- a/randquik/io.py +++ b/randquik/io.py @@ -76,7 +76,8 @@ def _open_output( 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) + with contextlib.suppress(OSError): + os.ftruncate(fd, required_size) # Seek to output position if oseek > 0: diff --git a/randquik/workers.py b/randquik/workers.py index 84ec6e0..6afa63d 100644 --- a/randquik/workers.py +++ b/randquik/workers.py @@ -4,6 +4,7 @@ import logging import os import sys import threading +from dataclasses import dataclass from randquik.io import madvise_file from randquik.utils import stopwatch @@ -12,13 +13,87 @@ __all__ = [ "BLOCK_SIZE", "FdProducer", "MmapProducer", + "WorkerStats", ] + +@dataclass +class WorkerStats: + """Timing statistics for a worker thread.""" + + worker_id: int = -1 + # Lock timing breakdown + lock_acquire_time: float = 0.0 # Time to acquire the lock (contention) + lock_wait_space_time: float = 0.0 # Time waiting for has_space condition + lock_claim_time: float = 0.0 # Time inside lock claiming block number + lock_notify_time: float = 0.0 # Time inside lock marking ready + notify + # Work timing + crypto_time: float = 0.0 + madvise_time: float = 0.0 + # Counters + blocks_processed: int = 0 + bytes_generated: int = 0 + wait_cycles: int = 0 # How many times we had to wait for space + + def total_time(self) -> float: + """Total measured time.""" + return ( + self.lock_acquire_time + + self.lock_wait_space_time + + self.lock_claim_time + + self.lock_notify_time + + self.crypto_time + + self.madvise_time + ) + + def format_report(self, label: str | None = None) -> str: + """Format a human-readable report.""" + total = self.total_time() + if total == 0: + return f"Worker {self.worker_id}: no data" + + if label is None: + label = f"Worker {self.worker_id}" + + def pct(val: float) -> str: + return f"{100 * val / total:.1f}%" if total > 0 else "--" + + def ms(val: float) -> str: + return f"{val * 1000:.1f}ms" + + lock_total = ( + self.lock_acquire_time + + self.lock_wait_space_time + + self.lock_claim_time + + self.lock_notify_time + ) + + lines = [ + f"{label} ({self.blocks_processed} blocks, {self.bytes_generated / 1e6:.1f} MB):", + f" crypto: {ms(self.crypto_time):>10} ({pct(self.crypto_time)})", + f" lock total: {ms(lock_total):>10} ({pct(lock_total)})", + f" acquire: {ms(self.lock_acquire_time):>10} ({pct(self.lock_acquire_time)})", + f" wait space: {ms(self.lock_wait_space_time):>10} ({pct(self.lock_wait_space_time)}) [{self.wait_cycles} cycles]", + f" claim: {ms(self.lock_claim_time):>10} ({pct(self.lock_claim_time)})", + f" notify: {ms(self.lock_notify_time):>10} ({pct(self.lock_notify_time)})", + ] + if self.madvise_time > 0: + lines.append(f" madvise: {ms(self.madvise_time):>10} ({pct(self.madvise_time)})") + lines.append(f" total: {ms(total):>10}") + return "\n".join(lines) + + BLOCK_SIZE = 1 << 20 class FdProducer: - """Multi-threaded producer with ring buffer for sequential file output.""" + """Multi-threaded producer with ring buffer for sequential file output. + + Uses efficient synchronization: + - Single lock with two conditions (has_data, has_space) + - Workers wait on has_space, consumer waits on has_data + - Crypto runs outside the lock + """ def __init__( self, @@ -30,6 +105,7 @@ class FdProducer: dry: bool = False, iseek: int = 0, block_size: int = BLOCK_SIZE, + profile: bool = False, ): self.workers = workers self.key = key @@ -39,68 +115,141 @@ class FdProducer: self.dry = dry self.iseek = iseek self.block_size = block_size + self.profile = profile # 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 + # Ring buffer state - more slots reduce wait time + # With N workers, we want enough buffers so workers rarely wait + self.num_slots = workers * 4 self._buf = bytearray(self.num_slots * block_size) + + # Single lock with conditions (simpler, faster than semaphore + events) 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._has_data = threading.Condition(self._lock) + self._has_space = threading.Condition(self._lock) + + self._genpos = 0 # next block to generate + self._conpos = 0 # next block to consume + self._ready = [False] * self.num_slots # which slots have data + self._quit = False self._threads: list[threading.Thread] = [] self.written = 0 self.wait_time = 0.0 self.write_time = 0.0 + # Per-worker stats, collected after threads finish + self._worker_stats: list[WorkerStats] = [] + self._stats_lock = threading.Lock() + def start(self): """Start worker threads.""" - for _ in range(self.workers): - t = threading.Thread(target=self._worker, daemon=True) + for i in range(self.workers): + t = threading.Thread(target=self._worker, args=(i,), 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_fast(self, view): + """Fast worker loop without profiling.""" + while True: + with self._lock: + # Wait for a slot + while self._genpos - self._conpos >= self.num_slots: + if self._quit: + return + self._has_space.wait() + if self._quit: + return + block_num = self._genpos + self._genpos += 1 - def _worker(self): - # Thread-local memoryview of the shared buffer + # Generate outside lock + slot = block_num % self.num_slots + buf = view[slot * self.block_size : (slot + 1) * self.block_size] + actual_block = self.start_block + block_num + self.ciph.stream( + self.key, actual_block.to_bytes(self.ciph.NONCEBYTES, "little"), into=buf + ) + + # Mark ready + with self._lock: + self._ready[slot] = True + self._has_data.notify() + + def _worker_profile(self, worker_id: int, view, stats: WorkerStats, timer): + """Worker loop with detailed profiling.""" + stats.worker_id = worker_id + while True: + # Measure lock acquisition (contention) + next(timer) + self._lock.acquire() + stats.lock_acquire_time += next(timer) + + try: + # Measure time waiting for space + while self._genpos - self._conpos >= self.num_slots: + if self._quit: + self._lock.release() + return + next(timer) + self._has_space.wait() + stats.lock_wait_space_time += next(timer) + stats.wait_cycles += 1 + + if self._quit: + self._lock.release() + return + + # Measure claiming block number + next(timer) + block_num = self._genpos + self._genpos += 1 + stats.lock_claim_time += next(timer) + finally: + self._lock.release() + + slot = block_num % self.num_slots + buf = view[slot * self.block_size : (slot + 1) * self.block_size] + actual_block = self.start_block + block_num + + # Measure crypto + next(timer) + self.ciph.stream( + self.key, actual_block.to_bytes(self.ciph.NONCEBYTES, "little"), into=buf + ) + stats.crypto_time += next(timer) + stats.blocks_processed += 1 + stats.bytes_generated += self.block_size + + # Measure notify (lock acquire + mark ready + notify) + next(timer) + with self._lock: + self._ready[slot] = True + self._has_data.notify() + stats.lock_notify_time += next(timer) + + def _worker(self, worker_id: int): view = memoryview(self._buf) + stats = WorkerStats(worker_id=worker_id) if self.profile else None try: - while self._worker_round(view): - pass + if self.profile: + timer = stopwatch() + self._worker_profile(worker_id, view, stats, timer) + else: + self._worker_fast(view) except BaseException as e: logging.exception("Worker thread exception: %s", e) finally: view.release() - self._quit.set() + if stats: + with self._stats_lock: + self._worker_stats.append(stats) with self._lock: - self.needdata.notify_all() + self._quit = True + self._has_data.notify_all() def run(self, progress_state: dict | None = None): """Consume blocks and write to fd. Call start() first.""" @@ -112,10 +261,12 @@ class FdProducer: with self._lock: slot = self._conpos % self.num_slots while not self._ready[slot]: - if self._quit.is_set(): + if self._quit: return - self.needdata.wait() + self._has_data.wait() self._ready[slot] = False + self._conpos += 1 + self._has_space.notify() self.wait_time += next(timer) buf = view[slot * self.block_size : (slot + 1) * self.block_size] @@ -138,21 +289,42 @@ class FdProducer: 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() + self._quit = True + self._has_space.notify_all() for t in self._threads: t.join() + def get_worker_stats(self) -> list[WorkerStats]: + """Get stats for each worker. Call after stop().""" + with self._stats_lock: + # Sort by worker_id for consistent output + return sorted(self._worker_stats, key=lambda s: s.worker_id) + + def format_stats_report(self) -> str: + """Format a complete stats report for all workers.""" + lines = [] + stats_list = self.get_worker_stats() + + if not stats_list: + return "No worker stats available" + + # Per-worker stats + for stats in stats_list: + lines.append(stats.format_report()) + + # Consumer stats + lines.append("Consumer:") + lines.append(f" wait time: {self.wait_time * 1000:.1f}ms") + lines.append(f" write time: {self.write_time * 1000:.1f}ms") + + return "\n".join(lines) + class MmapProducer: """Multi-threaded producer that writes directly into mmap.""" @@ -209,6 +381,10 @@ class MmapProducer: self._threads: list[threading.Thread] = [] self.progress_state: dict | None = None + # Per-worker stats, collected after threads finish + self._worker_stats: list[WorkerStats] = [] + self._stats_lock = threading.Lock() + def start(self): """Start worker threads.""" for _ in range(self.workers): @@ -221,9 +397,13 @@ class MmapProducer: view = memoryview(self.mm_raw) # Temporary buffer for first/last partial blocks tmp_buf = bytearray(self.block_size) + stats = WorkerStats() + timer = stopwatch() try: while True: + next(timer) # reset with self._lock: + stats.lock_acquire_time += next(timer) if self._next_block >= self.num_blocks: return block_num = self._next_block @@ -236,6 +416,7 @@ class MmapProducer: # Calculate byte range within output # block_num=0 corresponds to output byte 0 # But if start_offset > 0, first block is partial + next(timer) # reset for crypto timing 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) @@ -268,16 +449,28 @@ class MmapProducer: self.ciph.stream(self.key, nonce, size, into=view[out_start:out_end]) written = size + stats.crypto_time += next(timer) + stats.blocks_processed += 1 + stats.bytes_generated += written + if self.use_madvise and self.mm_raw: + next(timer) # reset madvise_file(self.mm_raw, out_start, written) + stats.madvise_time += next(timer) + + next(timer) # reset with self._lock: self.written += written if self.progress_state is not None: self.progress_state["written"] = self.written + stats.notify_time += next(timer) except BaseException as e: logging.exception("Worker thread exception: %s", e) finally: view.release() + # Collect stats before signaling quit + with self._stats_lock: + self._worker_stats.append(stats) self._quit.set() self._quit.wait() @@ -296,3 +489,11 @@ class MmapProducer: self.mm_raw.close() self.mm_raw = None self._threads.clear() + + def get_combined_stats(self) -> WorkerStats: + """Get combined stats from all workers. Call after join().""" + combined = WorkerStats() + with self._stats_lock: + for s in self._worker_stats: + combined.merge(s) + return combined