Fix undeterministic output bug, clean up workers code, make prettier status messages and refactor cli/workers/stats division.
This commit is contained in:
@@ -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.crypto import derive_key, generate_random_seed
|
||||||
from randquik.progress import ProgressDisplay
|
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:
|
try:
|
||||||
from randquik._version import __version__
|
from randquik._version import __version__
|
||||||
|
|||||||
+58
-99
@@ -2,32 +2,20 @@
|
|||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
import gc
|
import gc
|
||||||
import os
|
|
||||||
import sys
|
import sys
|
||||||
import time
|
|
||||||
|
|
||||||
import aeg
|
import aeg
|
||||||
|
|
||||||
from randquik.benchmark import run_benchmark
|
from randquik.benchmark import run_benchmark
|
||||||
from randquik.crypto import derive_key, generate_random_seed
|
from randquik.crypto import derive_key, generate_random_seed
|
||||||
from randquik.io import open_fd
|
from randquik.utils import parse_size
|
||||||
from randquik.progress import ProgressDisplay
|
from randquik.workers import run
|
||||||
from randquik.utils import (
|
|
||||||
parse_size,
|
|
||||||
print_summary,
|
|
||||||
)
|
|
||||||
from randquik.workers import (
|
|
||||||
BLOCK_SIZE,
|
|
||||||
FdProducer,
|
|
||||||
)
|
|
||||||
|
|
||||||
__all__ = ["main"]
|
__all__ = ["main"]
|
||||||
|
|
||||||
# Disable GC for performance
|
# Disable GC for performance
|
||||||
gc.disable()
|
gc.disable()
|
||||||
|
|
||||||
ciph: aeg.Cipher = None # type: ignore (set in main after parsing args)
|
|
||||||
|
|
||||||
DEFAULT_ALG = "AEGIS-128X2"
|
DEFAULT_ALG = "AEGIS-128X2"
|
||||||
|
|
||||||
|
|
||||||
@@ -56,56 +44,6 @@ def parse_seeks(args):
|
|||||||
return iseek, oseek
|
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():
|
def _main():
|
||||||
"""Internal main function that may raise exceptions."""
|
"""Internal main function that may raise exceptions."""
|
||||||
parser = argparse.ArgumentParser(description="Generate random bytes using AEGIS ciphers")
|
parser = argparse.ArgumentParser(description="Generate random bytes using AEGIS ciphers")
|
||||||
@@ -169,8 +107,9 @@ def _main():
|
|||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"-v",
|
"-v",
|
||||||
"--verbose",
|
"--verbose",
|
||||||
action="store_true",
|
action="count",
|
||||||
help="Verbose mode: show I/O mode and timing statistics",
|
default=0,
|
||||||
|
help="Verbose mode: -v for I/O mode, -vv for worker statistics",
|
||||||
)
|
)
|
||||||
|
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
@@ -179,7 +118,6 @@ def _main():
|
|||||||
if args.output == "-":
|
if args.output == "-":
|
||||||
args.output = None
|
args.output = None
|
||||||
|
|
||||||
global ciph
|
|
||||||
ciph = aeg.cipher(args.alg or DEFAULT_ALG)
|
ciph = aeg.cipher(args.alg or DEFAULT_ALG)
|
||||||
|
|
||||||
# Validate and process args
|
# Validate and process args
|
||||||
@@ -187,10 +125,8 @@ def _main():
|
|||||||
key = prepare_key(seed, ciph.KEYBYTES)
|
key = prepare_key(seed, ciph.KEYBYTES)
|
||||||
iseek, oseek = parse_seeks(args)
|
iseek, oseek = parse_seeks(args)
|
||||||
total_bytes = parse_size(args.len) # None if not specified, 0 if -l0
|
total_bytes = parse_size(args.len) # None if not specified, 0 if -l0
|
||||||
# Seed hint for generated seeds
|
# Always track the seed for commands, but only show repeat for generated seeds
|
||||||
seed_for_display = seed if generated_seed else None
|
seed_for_display = seed
|
||||||
|
|
||||||
start_time = time.perf_counter()
|
|
||||||
|
|
||||||
if args.benchmark:
|
if args.benchmark:
|
||||||
if args.seed is not None:
|
if args.seed is not None:
|
||||||
@@ -200,41 +136,64 @@ def _main():
|
|||||||
run_benchmark(args)
|
run_benchmark(args)
|
||||||
return
|
return
|
||||||
|
|
||||||
# Single-threaded mode (workers == 0)
|
# Build continue command for interruption
|
||||||
if args.threads == 0:
|
action = "generated" if args.dry else "wrote"
|
||||||
return singlethreaded(args, total_bytes, oseek, start_time, key, seed_for_display)
|
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
|
workers = args.threads if args.threads is not None else 1
|
||||||
|
result = run(
|
||||||
# File mode with ring buffers
|
output=args.output,
|
||||||
with open_fd(args.output, total_bytes, dry=args.dry, oseek=oseek) as fd:
|
total_bytes=total_bytes,
|
||||||
producer = FdProducer(
|
iseek=iseek,
|
||||||
workers, key, ciph, total_bytes, fd, dry=args.dry, iseek=iseek, profile=args.verbose
|
oseek=oseek,
|
||||||
|
key=key,
|
||||||
|
ciph=ciph,
|
||||||
|
workers=workers,
|
||||||
|
dry=args.dry,
|
||||||
|
quiet=args.quiet,
|
||||||
|
seed_for_display=seed_for_display,
|
||||||
|
action=action,
|
||||||
)
|
)
|
||||||
|
|
||||||
if args.verbose:
|
# Set repeat command for generated seeds
|
||||||
dry_str = " (dry run)" if args.dry else ""
|
result.repeat_cmd = repeat_cmd
|
||||||
print(
|
|
||||||
f"I/O mode: {workers} workers, {producer.num_slots} buffers, sequential writes{dry_str}",
|
# Update continue command with actual written bytes
|
||||||
file=sys.stderr,
|
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}
|
# Print summary
|
||||||
progress = ProgressDisplay(total_bytes, start_time, progress_state, seed=seed_for_display)
|
show_summary = not args.quiet or args.verbose >= 1 or result.interrupted
|
||||||
if not args.quiet:
|
if show_summary and (total_bytes is not None or result.interrupted):
|
||||||
progress.start()
|
result.print_summary(verbose=args.verbose)
|
||||||
|
if args.verbose >= 2:
|
||||||
|
result.print_detailed_stats()
|
||||||
|
|
||||||
try:
|
if result.interrupted:
|
||||||
producer.start()
|
sys.exit(1)
|
||||||
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)
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
|
|||||||
@@ -6,7 +6,7 @@ import sys
|
|||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
|
|
||||||
from randquik.utils import format_time
|
from randquik.stats import format_time
|
||||||
|
|
||||||
__all__ = ["ProgressDisplay"]
|
__all__ = ["ProgressDisplay"]
|
||||||
|
|
||||||
|
|||||||
@@ -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)
|
||||||
+1
-56
@@ -3,58 +3,15 @@
|
|||||||
import pathlib
|
import pathlib
|
||||||
import re
|
import re
|
||||||
import sys
|
import sys
|
||||||
import time
|
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"format_size",
|
|
||||||
"format_time",
|
|
||||||
"get_output_size",
|
"get_output_size",
|
||||||
"get_sector_size",
|
"get_sector_size",
|
||||||
"parse_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)
|
# Cache for sector size lookup (path -> size)
|
||||||
_sector_size_cache: dict[str, int] = {}
|
_sector_size_cache: dict[str, int] = {}
|
||||||
|
|
||||||
@@ -223,15 +180,3 @@ def sparse_range(n: int, max_items: int = 9) -> list[int]:
|
|||||||
out[-1] = n
|
out[-1] = n
|
||||||
|
|
||||||
return out
|
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()
|
|
||||||
|
|||||||
+232
-222
@@ -1,94 +1,37 @@
|
|||||||
"""Worker threads and ring buffer management for parallel generation."""
|
"""Worker threads and ring buffer management for parallel generation."""
|
||||||
|
|
||||||
import logging
|
|
||||||
import os
|
import os
|
||||||
|
import sys
|
||||||
import threading
|
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__ = [
|
__all__ = [
|
||||||
"BLOCK_SIZE",
|
"BLOCK_SIZE",
|
||||||
"FdProducer",
|
"RunResult",
|
||||||
"WorkerStats",
|
"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
|
BLOCK_SIZE = 1 << 20
|
||||||
|
|
||||||
|
|
||||||
class FdProducer:
|
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:
|
Uses efficient synchronization:
|
||||||
- Single lock with two conditions (has_data, has_space)
|
- 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
|
- Crypto runs outside the lock
|
||||||
"""
|
"""
|
||||||
|
|
||||||
@@ -102,7 +45,6 @@ class FdProducer:
|
|||||||
dry: bool = False,
|
dry: bool = False,
|
||||||
iseek: int = 0,
|
iseek: int = 0,
|
||||||
block_size: int = BLOCK_SIZE,
|
block_size: int = BLOCK_SIZE,
|
||||||
profile: bool = False,
|
|
||||||
):
|
):
|
||||||
self.workers = workers
|
self.workers = workers
|
||||||
self.key = key
|
self.key = key
|
||||||
@@ -112,189 +54,149 @@ class FdProducer:
|
|||||||
self.dry = dry
|
self.dry = dry
|
||||||
self.iseek = iseek
|
self.iseek = iseek
|
||||||
self.block_size = block_size
|
self.block_size = block_size
|
||||||
self.profile = profile
|
|
||||||
|
|
||||||
# iseek handling: which block to start at, and offset within first block
|
# iseek handling: which block to start at, and offset within first block
|
||||||
self.start_block = iseek // block_size
|
self.start_block = iseek // block_size
|
||||||
self.start_offset = iseek % block_size
|
self.start_offset = iseek % block_size
|
||||||
|
|
||||||
# Ring buffer state - more slots reduce wait time
|
self.num_slots = workers + 2 # Tested optimal (+1 for I/O and +1 to avoid congestion)
|
||||||
# With N workers, we want enough buffers so workers rarely wait
|
|
||||||
self.num_slots = workers * 4
|
|
||||||
self._buf = bytearray(self.num_slots * block_size)
|
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._lock = threading.Lock()
|
||||||
self._has_data = threading.Condition(self._lock)
|
self.has_data = threading.Condition(self._lock) # Consumer waits, workers notify
|
||||||
self._has_space = threading.Condition(self._lock)
|
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.threads: list[threading.Thread] = []
|
||||||
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.written = 0
|
||||||
self.wait_time = 0.0
|
self.consumer_stats = ConsumerStats()
|
||||||
self.write_time = 0.0
|
|
||||||
|
|
||||||
# Per-worker stats, collected after threads finish
|
# Per-worker stats, collected after threads finish
|
||||||
self._worker_stats: list[WorkerStats] = []
|
self._worker_stats: list[WorkerStats] = []
|
||||||
self._stats_lock = threading.Lock()
|
self._stats_lock = threading.Lock()
|
||||||
|
|
||||||
def start(self):
|
def start(self):
|
||||||
"""Start worker threads."""
|
self.threads = [
|
||||||
for i in range(self.workers):
|
threading.Thread(target=self.worker, args=(i,)) for i in range(self.workers)
|
||||||
t = threading.Thread(target=self._worker, args=(i,), daemon=True)
|
]
|
||||||
self._threads.append(t)
|
for t in self.threads:
|
||||||
t.start()
|
t.start()
|
||||||
|
|
||||||
def _worker_fast(self, view):
|
def worker(self, worker_id: int):
|
||||||
"""Fast worker loop without profiling."""
|
assert self.num_slots >= self.workers, "Ring buffer quarantee broken"
|
||||||
while True:
|
view = memoryview(self._buf)
|
||||||
with self._lock:
|
slots = [
|
||||||
# Wait for a slot
|
view[i * self.block_size : (i + 1) * self.block_size] for i in range(self.num_slots)
|
||||||
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
|
# Profiling setup
|
||||||
slot = block_num % self.num_slots
|
stats = WorkerStats(worker_id=worker_id)
|
||||||
buf = view[slot * self.block_size : (slot + 1) * self.block_size]
|
timer = stopwatch()
|
||||||
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:
|
try:
|
||||||
# Measure time waiting for space
|
slot = -1 # No slot to commit on first iteration
|
||||||
while self._genpos - self._conpos >= self.num_slots:
|
while True:
|
||||||
if self._quit:
|
# Claim next block number
|
||||||
self._lock.release()
|
with self.lock_blkno:
|
||||||
return
|
blkno = self.blkno
|
||||||
next(timer)
|
self.blkno += 1
|
||||||
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)
|
stats.lock_claim_time += next(timer)
|
||||||
finally:
|
|
||||||
self._lock.release()
|
|
||||||
|
|
||||||
slot = block_num % self.num_slots
|
with self._lock:
|
||||||
buf = view[slot * self.block_size : (slot + 1) * self.block_size]
|
stats.lock_acquire_time += next(timer)
|
||||||
actual_block = self.start_block + block_num
|
# 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
|
||||||
|
|
||||||
# Measure crypto
|
# Generate block
|
||||||
next(timer)
|
|
||||||
self.ciph.stream(
|
self.ciph.stream(
|
||||||
self.key, actual_block.to_bytes(self.ciph.NONCEBYTES, "little"), into=buf
|
self.key,
|
||||||
|
blkno.to_bytes(self.ciph.NONCEBYTES, "little"),
|
||||||
|
into=slots[slot],
|
||||||
)
|
)
|
||||||
stats.crypto_time += next(timer)
|
stats.crypto_time += next(timer)
|
||||||
stats.blocks_processed += 1
|
stats.blocks_processed += 1
|
||||||
stats.bytes_generated += self.block_size
|
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:
|
|
||||||
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:
|
finally:
|
||||||
view.release()
|
view.release()
|
||||||
if stats:
|
|
||||||
with self._stats_lock:
|
with self._stats_lock:
|
||||||
self._worker_stats.append(stats)
|
self._worker_stats.append(stats)
|
||||||
with self._lock:
|
|
||||||
self._quit = True
|
|
||||||
self._has_data.notify_all()
|
|
||||||
|
|
||||||
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."""
|
"""Consume blocks and write to fd. Call start() first."""
|
||||||
view = memoryview(self._buf)
|
view = memoryview(self._buf)
|
||||||
timer = stopwatch()
|
|
||||||
is_first_block = True
|
|
||||||
try:
|
try:
|
||||||
while self.total_bytes is None or self.written < self.total_bytes:
|
slots = [
|
||||||
with self._lock:
|
view[i * self.block_size : (i + 1) * self.block_size] for i in range(self.num_slots)
|
||||||
slot = self._conpos % self.num_slots
|
]
|
||||||
while not self._ready[slot]:
|
blkno = self.start_block
|
||||||
if self._quit:
|
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
|
return
|
||||||
self._has_data.wait()
|
|
||||||
self._ready[slot] = False
|
|
||||||
self._conpos += 1
|
|
||||||
self._has_space.notify()
|
|
||||||
|
|
||||||
self.wait_time += next(timer)
|
buf = slots[slot][self.start_offset : self.start_offset + total]
|
||||||
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:
|
if not self.dry:
|
||||||
os.write(self.fd, buf[buf_start : buf_start + to_write])
|
os.write(self.fd, buf)
|
||||||
self.write_time += next(timer)
|
self.written += len(buf)
|
||||||
self.written += to_write
|
# 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:
|
if progress_state is not None:
|
||||||
progress_state["written"] = self.written
|
progress_state["written"] = self.written
|
||||||
|
|
||||||
finally:
|
finally:
|
||||||
|
self.stop()
|
||||||
view.release()
|
view.release()
|
||||||
|
|
||||||
def stop(self):
|
def stop(self):
|
||||||
"""Signal workers to stop and wait for them."""
|
"""Signal workers to stop and wait for them."""
|
||||||
with self._lock:
|
with self._lock:
|
||||||
self._quit = True
|
self.quit = True
|
||||||
self._has_space.notify_all()
|
self.has_data.notify_all()
|
||||||
for t in self._threads:
|
self.has_space.notify_all()
|
||||||
|
for t in self.threads:
|
||||||
t.join()
|
t.join()
|
||||||
|
|
||||||
def get_worker_stats(self) -> list[WorkerStats]:
|
def get_worker_stats(self) -> list[WorkerStats]:
|
||||||
@@ -303,21 +205,129 @@ class FdProducer:
|
|||||||
# Sort by worker_id for consistent output
|
# Sort by worker_id for consistent output
|
||||||
return sorted(self._worker_stats, key=lambda s: s.worker_id)
|
return sorted(self._worker_stats, key=lambda s: s.worker_id)
|
||||||
|
|
||||||
def format_stats_report(self) -> str:
|
def run(self, progress_state: dict | None = None):
|
||||||
"""Format a complete stats report for all workers."""
|
"""Run multi-threaded generation."""
|
||||||
lines = []
|
self.start()
|
||||||
stats_list = self.get_worker_stats()
|
try:
|
||||||
|
self.consumer(progress_state)
|
||||||
|
finally:
|
||||||
|
self.stop()
|
||||||
|
|
||||||
if not stats_list:
|
|
||||||
return "No worker stats available"
|
|
||||||
|
|
||||||
# Per-worker stats
|
class _SingleThreadedProducer:
|
||||||
for stats in stats_list:
|
"""Single-threaded producer for infinite output or workers=0 mode."""
|
||||||
lines.append(stats.format_report())
|
|
||||||
|
|
||||||
# Consumer stats
|
def __init__(
|
||||||
lines.append("Consumer:")
|
self,
|
||||||
lines.append(f" wait time: {self.wait_time * 1000:.1f}ms")
|
key: bytes,
|
||||||
lines.append(f" write time: {self.write_time * 1000:.1f}ms")
|
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,
|
||||||
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user