diff --git a/randquik/__init__.py b/randquik/__init__.py index d43a503..82e22c0 100644 --- a/randquik/__init__.py +++ b/randquik/__init__.py @@ -6,7 +6,8 @@ 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 +from randquik.stats import format_size, format_time +from randquik.utils import parse_size try: from randquik._version import __version__ diff --git a/randquik/cli.py b/randquik/cli.py index b596c4a..b248227 100644 --- a/randquik/cli.py +++ b/randquik/cli.py @@ -2,32 +2,20 @@ import argparse import gc -import os import sys -import time import aeg 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, -) +from randquik.utils import parse_size +from randquik.workers import run __all__ = ["main"] # Disable GC for performance gc.disable() -ciph: aeg.Cipher = None # type: ignore (set in main after parsing args) - DEFAULT_ALG = "AEGIS-128X2" @@ -56,56 +44,6 @@ def parse_seeks(args): return iseek, oseek -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") @@ -169,8 +107,9 @@ def _main(): parser.add_argument( "-v", "--verbose", - action="store_true", - help="Verbose mode: show I/O mode and timing statistics", + action="count", + default=0, + help="Verbose mode: -v for I/O mode, -vv for worker statistics", ) args = parser.parse_args() @@ -179,7 +118,6 @@ def _main(): if args.output == "-": args.output = None - global ciph ciph = aeg.cipher(args.alg or DEFAULT_ALG) # Validate and process args @@ -187,10 +125,8 @@ def _main(): 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() + # Always track the seed for commands, but only show repeat for generated seeds + seed_for_display = seed if args.benchmark: if args.seed is not None: @@ -200,41 +136,64 @@ def _main(): 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) + # Build continue command for interruption + action = "generated" if args.dry else "wrote" + continue_cmd = None + repeat_cmd = None + if args.output and seed_for_display: + # Will be updated with actual written bytes after run + continue_cmd = f"randquik -s {seed_for_display} --seek {{seek}} -o {args.output}" + if args.len: + continue_cmd += f" -l {args.len}" + # Build repeat command for randomly generated seeds (so user can reproduce) + if generated_seed and not args.quiet: + repeat_cmd = f"randquik -s {seed_for_display}" + if args.len: + repeat_cmd += f" -l {args.len}" + if args.output: + repeat_cmd += f" -o {args.output}" + # Run generation workers = args.threads if args.threads is not None else 1 + result = run( + output=args.output, + total_bytes=total_bytes, + iseek=iseek, + oseek=oseek, + key=key, + ciph=ciph, + workers=workers, + dry=args.dry, + quiet=args.quiet, + seed_for_display=seed_for_display, + action=action, + ) - # 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, profile=args.verbose - ) + # Set repeat command for generated seeds + result.repeat_cmd = repeat_cmd - 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, + # Update continue command with actual written bytes + if result.interrupted and result.written > 0 and args.output: + new_iseek = iseek + result.written + new_oseek = oseek + result.written + if new_iseek == new_oseek: + result.continue_cmd = ( + f"randquik -s {seed_for_display} --seek {new_iseek} -o {args.output}" ) + else: + result.continue_cmd = f"randquik -s {seed_for_display} --iseek {new_iseek} --oseek {new_oseek} -o {args.output}" + if args.len: + result.continue_cmd += f" -l {args.len}" - progress_state = {"written": 0} - progress = ProgressDisplay(total_bytes, start_time, progress_state, seed=seed_for_display) - if not args.quiet: - progress.start() + # Print summary + show_summary = not args.quiet or args.verbose >= 1 or result.interrupted + if show_summary and (total_bytes is not None or result.interrupted): + result.print_summary(verbose=args.verbose) + if args.verbose >= 2: + result.print_detailed_stats() - 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) - if args.verbose: - print(producer.format_stats_report(), file=sys.stderr) + if result.interrupted: + sys.exit(1) def main(): diff --git a/randquik/io.py b/randquik/io.py index 9c31993..f20fbed 100644 --- a/randquik/io.py +++ b/randquik/io.py @@ -19,7 +19,7 @@ def _open_output( oseek: int = 0, ) -> tuple[int, bool]: """Open output file descriptor, preallocate and apply platform hints. - + Returns: Tuple of (file descriptor, whether we created the file) """ diff --git a/randquik/progress.py b/randquik/progress.py index 979bd4e..9a5543d 100644 --- a/randquik/progress.py +++ b/randquik/progress.py @@ -6,7 +6,7 @@ import sys import threading import time -from randquik.utils import format_time +from randquik.stats import format_time __all__ = ["ProgressDisplay"] diff --git a/randquik/stats.py b/randquik/stats.py new file mode 100644 index 0000000..a5bf48a --- /dev/null +++ b/randquik/stats.py @@ -0,0 +1,260 @@ +"""Statistics collection and formatting for workers.""" + +import re +import sys +import time +from dataclasses import dataclass, field + +__all__ = [ + "ConsumerStats", + "RunResult", + "SingleThreadedStats", + "WorkerStats", + "format_size", + "format_time", + "format_worker_stats_report", + "stopwatch", +] + + +def stopwatch(): + """Generator that yields elapsed time since last yield.""" + t = time.perf_counter() + while True: + now = time.perf_counter() + yield now - t + t = now + + +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" + + +@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 + # 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 + ) + + +@dataclass +class ConsumerStats: + """Timing statistics for the consumer thread.""" + + wait_time: float = 0.0 + write_time: float = 0.0 + + def total_time(self) -> float: + return self.wait_time + self.write_time + + +@dataclass +class SingleThreadedStats: + """Timing statistics for single-threaded mode.""" + + crypto_time: float = 0.0 + write_time: float = 0.0 + + def total_time(self) -> float: + return self.crypto_time + self.write_time + + +@dataclass +class RunResult: + """Result from running a worker.""" + + written: int + elapsed: float + interrupted: bool + action: str = "wrote" + # Raw stats objects (optional, for verbose output) + consumer_stats: ConsumerStats | None = None + singlethreaded_stats: SingleThreadedStats | None = None + worker_stats: list[WorkerStats] = field(default_factory=list) + # For continue command on interrupt + continue_cmd: str | None = None + # For repeat command (when seed was randomly generated) + repeat_cmd: str | None = None + + def _format_io_stats(self) -> str | None: + """Format I/O timing stats for the summary line.""" + if self.singlethreaded_stats is not None: + st = self.singlethreaded_stats + tt = st.total_time() + if tt <= 0: + return None + return f"crypto {st.crypto_time / tt:.0%} — write {st.write_time / tt:.0%}" + elif self.consumer_stats is not None: + cs = self.consumer_stats + tt = cs.total_time() + if tt <= 0: + return None + return f"wait {cs.wait_time / tt:.0%} — write {cs.write_time / tt:.0%}" + return None + + def print_summary(self, verbose: int = 0): + """Print a nice one-liner summary with optional colors.""" + speed_gbs = (self.written / 1_000_000_000) / self.elapsed if self.elapsed > 0 else 0 + size_str = format_size(self.written) + time_str = format_time(self.elapsed) + + # I/O stats for verbose mode + io_stats = self._format_io_stats() if verbose >= 1 else None + stats_fmt = f"\033[0;32m • {io_stats}" if io_stats else "" + status_fmt = " \033[31m(interrupted)\033[0m" if self.interrupted else "" + + # Continue or repeat command + cmd_line = "" + if self.interrupted and self.continue_cmd: + cmd_line = f"\n\033[2mContinue >>>\033[0;34m {self.continue_cmd}\033[0m" + elif not self.interrupted and self.repeat_cmd: + cmd_line = f"\n\033[2mRepeat >>>\033[0;34m {self.repeat_cmd}\033[0m" + + msg = ( + f"\n\033[36m[RandQuik]\033[32m {self.action} \033[1m{size_str}\033[0;32m in " + f"\033[1m{time_str}\033[0;32m @ \033[1;32m{speed_gbs:.2f}GB/s{stats_fmt}\033[0m" + f"{status_fmt}\033[1m{cmd_line}\n" + ) + + if not sys.stderr.isatty(): + msg = re.sub(r"\033\[[0-9;]*m", "", msg) + + sys.stderr.write(msg) + + def print_detailed_stats(self): + """Print detailed worker statistics table (for -vv).""" + if self.worker_stats and self.consumer_stats: + report = format_worker_stats_report(self.worker_stats, self.consumer_stats) + sys.stderr.write(report + "\n") + + +def format_worker_stats_report( + worker_stats: list[WorkerStats], consumer_stats: ConsumerStats +) -> str: + """Format a complete stats report for all workers as a table.""" + if not worker_stats: + return "No worker stats available" + + def ms(val: float) -> str: + return f"{val * 1000:.0f}" + + # Column width for worker data + col_w = 8 + pct_w = 6 # Width for percentage column + + # Total time across all workers for percentage calculation + total_all = sum(s.total_time() for s in worker_stats) + + def pct(val: float) -> str: + return f"{100 * val / total_all:.0f}%" if total_all > 0 else "--" + + # Header row with worker numbers + header = ( + "Worker stats" + + f"{'%':>{pct_w}}" + + "".join(f"{'W' + str(s.worker_id):>{col_w}}" for s in worker_stats) + ) + sep = "-" * len(header) + + # Non-timing rows (counters) + total_blocks = sum(s.blocks_processed for s in worker_stats) + total_cycles = sum(s.wait_cycles for s in worker_stats) + + def cycles_pct() -> str: + return f"{100 * total_cycles / total_blocks:.0f}%" if total_blocks > 0 else "--" + + counter_rows = [ + ("1MiB blocks", "", [str(s.blocks_processed) for s in worker_stats]), + ("wait_cycles", cycles_pct(), [str(s.wait_cycles) for s in worker_stats]), + ] + + # Timing rows with summed values for percentage + timing_rows = [ + ( + "crypto ms", + sum(s.crypto_time for s in worker_stats), + [ms(s.crypto_time) for s in worker_stats], + ), + ( + "lock_acq ms", + sum(s.lock_acquire_time for s in worker_stats), + [ms(s.lock_acquire_time) for s in worker_stats], + ), + ( + "wait_sp ms", + sum(s.lock_wait_space_time for s in worker_stats), + [ms(s.lock_wait_space_time) for s in worker_stats], + ), + ( + "claim ms", + sum(s.lock_claim_time for s in worker_stats), + [ms(s.lock_claim_time) for s in worker_stats], + ), + ( + "notify ms", + sum(s.lock_notify_time for s in worker_stats), + [ms(s.lock_notify_time) for s in worker_stats], + ), + ] + + timing_rows.append(("total ms", total_all, [ms(s.total_time()) for s in worker_stats])) + + lines = [header, sep] + for label, pct_val, values in counter_rows: + row = f"{label:<12}" + f"{pct_val:>{pct_w}}" + "".join(f"{v:>{col_w}}" for v in values) + lines.append(row) + lines.append(sep) + for label, total_val, values in timing_rows: + row = ( + f"{label:<12}" + f"{pct(total_val):>{pct_w}}" + "".join(f"{v:>{col_w}}" for v in values) + ) + lines.append(row) + + # Consumer stats + lines.append(sep) + return "\n".join(lines) diff --git a/randquik/utils.py b/randquik/utils.py index 16ba95c..b380502 100644 --- a/randquik/utils.py +++ b/randquik/utils.py @@ -3,58 +3,15 @@ import pathlib import re import sys -import time __all__ = [ - "format_size", - "format_time", "get_output_size", "get_sector_size", "parse_size", - "print_summary", + "sparse_range", ] -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] = {} @@ -223,15 +180,3 @@ def sparse_range(n: int, max_items: int = 9) -> list[int]: 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 index c5809a3..56f5161 100644 --- a/randquik/workers.py +++ b/randquik/workers.py @@ -1,94 +1,37 @@ """Worker threads and ring buffer management for parallel generation.""" -import logging import os +import sys import threading -from dataclasses import dataclass +import time -from randquik.utils import stopwatch +from randquik.io import open_fd +from randquik.progress import ProgressDisplay +from randquik.stats import ( + ConsumerStats, + RunResult, + SingleThreadedStats, + WorkerStats, + stopwatch, +) __all__ = [ "BLOCK_SIZE", - "FdProducer", - "WorkerStats", + "RunResult", + "run", ] -@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: +class _FdProducer: """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 + - Workers wait on has_space, notify has_data when block is ready + - Consumer waits on has_data, notifies has_space when block is consumed - Crypto runs outside the lock """ @@ -102,7 +45,6 @@ class FdProducer: dry: bool = False, iseek: int = 0, block_size: int = BLOCK_SIZE, - profile: bool = False, ): self.workers = workers self.key = key @@ -112,189 +54,149 @@ 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 - more slots reduce wait time - # With N workers, we want enough buffers so workers rarely wait - self.num_slots = workers * 4 + self.num_slots = workers + 2 # Tested optimal (+1 for I/O and +1 to avoid congestion) self._buf = bytearray(self.num_slots * block_size) - # Single lock with conditions (simpler, faster than semaphore + events) + # Separate conditions for producers and consumer self._lock = threading.Lock() - self._has_data = threading.Condition(self._lock) - self._has_space = threading.Condition(self._lock) + self.has_data = threading.Condition(self._lock) # Consumer waits, workers notify + self.has_space = threading.Condition(self._lock) # Workers wait, consumer notifies + self.lock_blkno = threading.Lock() + self.blkno = self.start_block # next block to generate + self.ready = [False] * self.num_slots # which block number is there ready + self.quit = False - 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.threads: list[threading.Thread] = [] self.written = 0 - self.wait_time = 0.0 - self.write_time = 0.0 + self.consumer_stats = ConsumerStats() # Per-worker stats, collected after threads finish self._worker_stats: list[WorkerStats] = [] self._stats_lock = threading.Lock() def start(self): - """Start worker threads.""" - for i in range(self.workers): - t = threading.Thread(target=self._worker, args=(i,), daemon=True) - self._threads.append(t) + self.threads = [ + threading.Thread(target=self.worker, args=(i,)) for i in range(self.workers) + ] + for t in self.threads: t.start() - 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 - - # 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): + def worker(self, worker_id: int): + assert self.num_slots >= self.workers, "Ring buffer quarantee broken" view = memoryview(self._buf) - stats = WorkerStats(worker_id=worker_id) if self.profile else None + slots = [ + view[i * self.block_size : (i + 1) * self.block_size] for i in range(self.num_slots) + ] + + # Profiling setup + stats = WorkerStats(worker_id=worker_id) + timer = stopwatch() + try: - 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) + slot = -1 # No slot to commit on first iteration + while True: + # Claim next block number + with self.lock_blkno: + blkno = self.blkno + self.blkno += 1 + stats.lock_claim_time += next(timer) + + with self._lock: + stats.lock_acquire_time += next(timer) + # Commit previous block (mark ready + notify consumer) + if slot >= 0: + self.ready[slot] = True + self.has_data.notify() + stats.lock_notify_time += next(timer) + # Wait for the NEXT slot to be free + slot = blkno % self.num_slots + while self.ready[slot] and not self.quit: + stats.wait_cycles += 1 + self.has_space.wait() + stats.lock_wait_space_time += next(timer) + if self.quit: + return + + # Generate block + self.ciph.stream( + self.key, + blkno.to_bytes(self.ciph.NONCEBYTES, "little"), + into=slots[slot], + ) + stats.crypto_time += next(timer) + stats.blocks_processed += 1 + stats.bytes_generated += self.block_size + finally: view.release() - if stats: - with self._stats_lock: - self._worker_stats.append(stats) - with self._lock: - self._quit = True - self._has_data.notify_all() + with self._stats_lock: + self._worker_stats.append(stats) - def run(self, progress_state: dict | None = None): + def consumer(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: - return - self._has_data.wait() - self._ready[slot] = False - self._conpos += 1 - self._has_space.notify() + slots = [ + view[i * self.block_size : (i + 1) * self.block_size] for i in range(self.num_slots) + ] + blkno = self.start_block + slot = blkno % self.num_slots + total = sys.maxsize if self.total_bytes is None else self.total_bytes + # Handle the first block: skip start_offset bytes (note: this is purposefully left out of stats) + with self.has_data: + while not self.ready[slot] and not self.quit: + self.has_data.wait() + if self.quit: + return - 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, - ) + buf = slots[slot][self.start_offset : self.start_offset + total] if not self.dry: - os.write(self.fd, buf[buf_start : buf_start + to_write]) - self.write_time += next(timer) - self.written += to_write + os.write(self.fd, buf) + self.written += len(buf) + # Other blocks + timer = stopwatch() + while self.written < total: + # Take/wait for expected slot + with self._lock: + # Release previous slot and notify workers + self.ready[slot] = False + self.has_space.notify_all() + # Wait for the next buffer to be ready + blkno += 1 + slot = blkno % self.num_slots + while not self.ready[slot] and not self.quit: + self.has_data.wait() + if self.quit: + return + self.consumer_stats.wait_time += next(timer) + buf = slots[slot] + # Last block? Trim to remaining size + if self.written + len(buf) > total: + buf = buf[: total - self.written] + if not self.dry: + os.write(self.fd, buf) + + self.consumer_stats.write_time += next(timer) + self.written += len(buf) if progress_state is not None: progress_state["written"] = self.written + finally: + self.stop() view.release() def stop(self): """Signal workers to stop and wait for them.""" with self._lock: - self._quit = True - self._has_space.notify_all() - for t in self._threads: + self.quit = True + self.has_data.notify_all() + self.has_space.notify_all() + for t in self.threads: t.join() def get_worker_stats(self) -> list[WorkerStats]: @@ -303,21 +205,129 @@ class FdProducer: # 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() + def run(self, progress_state: dict | None = None): + """Run multi-threaded generation.""" + self.start() + try: + self.consumer(progress_state) + finally: + self.stop() - if not stats_list: - return "No worker stats available" - # Per-worker stats - for stats in stats_list: - lines.append(stats.format_report()) +class _SingleThreadedProducer: + """Single-threaded producer for infinite output or workers=0 mode.""" - # 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") + def __init__( + self, + key: bytes, + ciph, + total_bytes: int | None, + fd: int, + dry: bool = False, + block_size: int = BLOCK_SIZE, + ): + self.key = key + self.ciph = ciph + self.total_bytes = total_bytes + self.fd = fd + self.dry = dry + self.block_size = block_size - return "\n".join(lines) + self.written = 0 + self.stats = SingleThreadedStats() + + def run(self, progress_state: dict | None = None): + """Generate and write blocks sequentially.""" + buf = bytearray(self.block_size) + view = memoryview(buf) + nonce = bytearray(self.ciph.NONCEBYTES) + total = sys.maxsize if self.total_bytes is None else self.total_bytes + timer = stopwatch() + + try: + while self.written < total: + size = min(self.block_size, total - self.written) + chunk = view[:size] + self.ciph.stream(self.key, nonce, size, into=chunk) + self.stats.crypto_time += next(timer) + if not self.dry: + os.write(self.fd, chunk) + self.stats.write_time += next(timer) + self.ciph.nonce_increment(nonce) + self.written += size + if progress_state is not None: + progress_state["written"] = self.written + finally: + view.release() + + +def run( + output: str | None, + total_bytes: int | None, + iseek: int, + oseek: int, + key: bytes, + ciph, + workers: int = 1, + dry: bool = False, + quiet: bool = False, + seed_for_display: str | None = None, + action: str = "wrote", + continue_cmd: str | None = None, +) -> RunResult: + """Run random generation with specified number of workers. Returns RunResult. + + Args: + workers: Number of worker threads. 0 for single-threaded mode. + """ + start_time = time.perf_counter() + infinite = total_bytes is None + fd_size = 0 if infinite else total_bytes + + with open_fd(output, fd_size, dry=dry, oseek=oseek) as fd: + if workers == 0: + producer = _SingleThreadedProducer(key, ciph, total_bytes, fd, dry=dry) + else: + producer = _FdProducer(workers, key, ciph, total_bytes, fd, dry=dry, iseek=iseek) + + progress_state = {"written": 0} + progress = ProgressDisplay( + total_bytes, + start_time, + progress_state, + infinite=infinite, + seed=seed_for_display, + ) + if not quiet: + progress.start() + + interrupted = False + try: + producer.run(progress_state) + except (KeyboardInterrupt, BrokenPipeError): + interrupted = True + finally: + progress.stop() + + elapsed = time.perf_counter() - start_time + + # Build result with raw stats + if workers == 0: + return RunResult( + written=producer.written, + elapsed=elapsed, + interrupted=interrupted, + action=action, + singlethreaded_stats=producer.stats, + continue_cmd=continue_cmd, + ) + else: + return RunResult( + written=producer.written, + elapsed=elapsed, + interrupted=interrupted, + action=action, + consumer_stats=producer.consumer_stats, + worker_stats=producer.get_worker_stats(), + continue_cmd=continue_cmd, + )