332 lines
11 KiB
Python
332 lines
11 KiB
Python
"""Worker threads and ring buffer management for parallel generation."""
|
|
|
|
import os
|
|
import sys
|
|
import threading
|
|
import time
|
|
|
|
from randquik.io import open_fd
|
|
from randquik.progress import ProgressDisplay
|
|
from randquik.stats import (
|
|
ConsumerStats,
|
|
RunResult,
|
|
SingleThreadedStats,
|
|
WorkerStats,
|
|
stopwatch,
|
|
)
|
|
|
|
__all__ = [
|
|
"BLOCK_SIZE",
|
|
"RunResult",
|
|
"run",
|
|
]
|
|
|
|
|
|
BLOCK_SIZE = 1 << 20
|
|
|
|
|
|
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, notify has_data when block is ready
|
|
- Consumer waits on has_data, notifies has_space when block is consumed
|
|
- Crypto runs outside the lock
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
workers: int,
|
|
key: bytes,
|
|
ciph,
|
|
total_bytes: int | None,
|
|
fd: int,
|
|
dry: bool = False,
|
|
iseek: int = 0,
|
|
block_size: int = BLOCK_SIZE,
|
|
):
|
|
self.workers = workers
|
|
self.key = key
|
|
self.ciph = ciph
|
|
self.total_bytes = total_bytes
|
|
self.fd = fd
|
|
self.dry = dry
|
|
self.iseek = iseek
|
|
self.block_size = block_size
|
|
|
|
# iseek handling: which block to start at, and offset within first block
|
|
self.start_block = iseek // block_size
|
|
self.start_offset = iseek % block_size
|
|
|
|
self.num_slots = workers + 2 # Tested optimal (+1 for I/O and +1 to avoid congestion)
|
|
self._buf = bytearray(self.num_slots * block_size)
|
|
|
|
# Separate conditions for producers and consumer
|
|
self._lock = threading.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.threads: list[threading.Thread] = []
|
|
self.written = 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):
|
|
self.threads = [
|
|
threading.Thread(target=self.worker, args=(i,)) for i in range(self.workers)
|
|
]
|
|
for t in self.threads:
|
|
t.start()
|
|
|
|
def worker(self, worker_id: int):
|
|
assert self.num_slots >= self.workers, "Ring buffer quarantee broken"
|
|
view = memoryview(self._buf)
|
|
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:
|
|
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()
|
|
with self._stats_lock:
|
|
self._worker_stats.append(stats)
|
|
|
|
def consumer(self, progress_state: dict | None = None):
|
|
"""Consume blocks and write to fd. Call start() first."""
|
|
view = memoryview(self._buf)
|
|
try:
|
|
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
|
|
|
|
buf = slots[slot][self.start_offset : self.start_offset + total]
|
|
if not self.dry:
|
|
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_data.notify_all()
|
|
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 run(self, progress_state: dict | None = None):
|
|
"""Run multi-threaded generation."""
|
|
self.start()
|
|
try:
|
|
self.consumer(progress_state)
|
|
finally:
|
|
self.stop()
|
|
|
|
|
|
class _SingleThreadedProducer:
|
|
"""Single-threaded producer for infinite output or workers=0 mode."""
|
|
|
|
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
|
|
|
|
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,
|
|
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,
|
|
)
|
|
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,
|
|
)
|