299 lines
10 KiB
Python
299 lines
10 KiB
Python
"""Worker threads and ring buffer management for parallel generation."""
|
|
|
|
import logging
|
|
import os
|
|
import sys
|
|
import threading
|
|
|
|
from randquik.io import madvise_file
|
|
from randquik.utils import stopwatch
|
|
|
|
__all__ = [
|
|
"BLOCK_SIZE",
|
|
"FdProducer",
|
|
"MmapProducer",
|
|
]
|
|
|
|
BLOCK_SIZE = 1 << 20
|
|
|
|
|
|
class FdProducer:
|
|
"""Multi-threaded producer with ring buffer for sequential file output."""
|
|
|
|
def __init__(
|
|
self,
|
|
workers: int,
|
|
key: bytes,
|
|
ciph,
|
|
total_bytes: int | None,
|
|
fd: int,
|
|
dry: bool = False,
|
|
iseek: int = 0,
|
|
block_size: int = BLOCK_SIZE,
|
|
):
|
|
self.workers = workers
|
|
self.key = key
|
|
self.ciph = ciph
|
|
self.total_bytes = total_bytes
|
|
self.fd = fd
|
|
self.dry = dry
|
|
self.iseek = iseek
|
|
self.block_size = block_size
|
|
|
|
# iseek handling: which block to start at, and offset within first block
|
|
self.start_block = iseek // block_size
|
|
self.start_offset = iseek % block_size
|
|
|
|
# Ring buffer state
|
|
self.num_slots = workers + 1
|
|
self._buf = bytearray(self.num_slots * block_size)
|
|
self._lock = threading.Lock()
|
|
self.needdata = threading.Condition(self._lock)
|
|
self.needspace = threading.Condition(self._lock)
|
|
self._ready = [False] * self.num_slots
|
|
self._genpos = 0
|
|
self._conpos = 0
|
|
self._quit = threading.Event()
|
|
|
|
self._threads: list[threading.Thread] = []
|
|
self.written = 0
|
|
self.wait_time = 0.0
|
|
self.write_time = 0.0
|
|
|
|
def start(self):
|
|
"""Start worker threads."""
|
|
for _ in range(self.workers):
|
|
t = threading.Thread(target=self._worker, daemon=True)
|
|
self._threads.append(t)
|
|
t.start()
|
|
|
|
def _worker_round(self, view) -> bool:
|
|
"""Process one block. Returns True if work was done, False if should quit."""
|
|
with self._lock:
|
|
while self._genpos - self._conpos >= self.num_slots:
|
|
if self._quit.is_set():
|
|
return False
|
|
self.needspace.wait()
|
|
if self._quit.is_set():
|
|
return False
|
|
block_num = self._genpos
|
|
self._genpos += 1
|
|
slot = block_num % self.num_slots
|
|
buf = view[slot * self.block_size : (slot + 1) * self.block_size]
|
|
# Use start_block + block_num as actual nonce to handle iseek
|
|
actual_block = self.start_block + block_num
|
|
self.ciph.stream(self.key, actual_block.to_bytes(self.ciph.NONCEBYTES, "little"), into=buf)
|
|
with self._lock:
|
|
self._ready[slot] = True
|
|
self.needdata.notify_all()
|
|
return True
|
|
|
|
def _worker(self):
|
|
# Thread-local memoryview of the shared buffer
|
|
view = memoryview(self._buf)
|
|
try:
|
|
while self._worker_round(view):
|
|
pass
|
|
except BaseException as e:
|
|
logging.exception("Worker thread exception: %s", e)
|
|
finally:
|
|
view.release()
|
|
self._quit.set()
|
|
with self._lock:
|
|
self.needdata.notify_all()
|
|
|
|
def run(self, progress_state: dict | None = None):
|
|
"""Consume blocks and write to fd. Call start() first."""
|
|
view = memoryview(self._buf)
|
|
timer = stopwatch()
|
|
is_first_block = True
|
|
try:
|
|
while self.total_bytes is None or self.written < self.total_bytes:
|
|
with self._lock:
|
|
slot = self._conpos % self.num_slots
|
|
while not self._ready[slot]:
|
|
if self._quit.is_set():
|
|
return
|
|
self.needdata.wait()
|
|
self._ready[slot] = False
|
|
|
|
self.wait_time += next(timer)
|
|
buf = view[slot * self.block_size : (slot + 1) * self.block_size]
|
|
# Handle first block: skip start_offset bytes
|
|
if is_first_block and self.start_offset > 0:
|
|
buf_start = self.start_offset
|
|
is_first_block = False
|
|
else:
|
|
buf_start = 0
|
|
|
|
to_write = min(
|
|
self.block_size - buf_start,
|
|
self.total_bytes - self.written
|
|
if self.total_bytes is not None
|
|
else self.block_size - buf_start,
|
|
)
|
|
if not self.dry:
|
|
os.write(self.fd, buf[buf_start : buf_start + to_write])
|
|
self.write_time += next(timer)
|
|
self.written += to_write
|
|
if progress_state is not None:
|
|
progress_state["written"] = self.written
|
|
|
|
with self._lock:
|
|
self._conpos += 1
|
|
self.needspace.notify_all()
|
|
finally:
|
|
view.release()
|
|
|
|
def stop(self):
|
|
"""Signal workers to stop and wait for them."""
|
|
with self._lock:
|
|
self._quit.set()
|
|
self.needspace.notify_all()
|
|
for t in self._threads:
|
|
t.join()
|
|
|
|
|
|
class MmapProducer:
|
|
"""Multi-threaded producer that writes directly into mmap."""
|
|
|
|
def __init__(
|
|
self,
|
|
workers: int,
|
|
key: bytes,
|
|
ciph,
|
|
total_bytes: int | None,
|
|
mm_raw=None,
|
|
use_madvise: bool = False,
|
|
oseek: int = 0,
|
|
iseek: int = 0,
|
|
block_size: int = BLOCK_SIZE,
|
|
dry: bool = False,
|
|
):
|
|
self.workers = workers
|
|
self.key = key
|
|
self.ciph = ciph
|
|
self.total_bytes = total_bytes
|
|
self.use_madvise = use_madvise
|
|
self.oseek = oseek
|
|
self.iseek = iseek
|
|
self.block_size = block_size
|
|
self.dry = dry
|
|
self.written = 0
|
|
|
|
# For dry runs, create an anonymous mmap
|
|
import mmap as mmap_module
|
|
|
|
if dry:
|
|
length = oseek + (total_bytes if total_bytes is not None else 0)
|
|
self.mm_raw = mmap_module.mmap(-1, length)
|
|
self._owns_mmap = True
|
|
else:
|
|
self.mm_raw = mm_raw
|
|
self._owns_mmap = False
|
|
|
|
# iseek handling: which block to start at, and offset within first block
|
|
self.start_block = iseek // block_size
|
|
self.start_offset = iseek % block_size
|
|
|
|
if total_bytes is None:
|
|
self.num_blocks = sys.maxsize // block_size
|
|
else:
|
|
self.num_blocks = (total_bytes + block_size - 1) // block_size
|
|
# If first block is partial, we need one more block to cover total_bytes
|
|
if self.start_offset > 0 and total_bytes > 0:
|
|
self.num_blocks = (self.start_offset + total_bytes + block_size - 1) // block_size
|
|
self._lock = threading.Lock()
|
|
self._next_block = 0
|
|
self._quit = threading.Event()
|
|
self._threads: list[threading.Thread] = []
|
|
self.progress_state: dict | None = None
|
|
|
|
def start(self):
|
|
"""Start worker threads."""
|
|
for _ in range(self.workers):
|
|
t = threading.Thread(target=self._worker, daemon=True)
|
|
self._threads.append(t)
|
|
t.start()
|
|
|
|
def _worker(self):
|
|
# Create thread-local memoryview from mmap
|
|
view = memoryview(self.mm_raw)
|
|
# Temporary buffer for first/last partial blocks
|
|
tmp_buf = bytearray(self.block_size)
|
|
try:
|
|
while True:
|
|
with self._lock:
|
|
if self._next_block >= self.num_blocks:
|
|
return
|
|
block_num = self._next_block
|
|
self._next_block += 1
|
|
|
|
# Calculate actual nonce (accounting for iseek)
|
|
actual_block = self.start_block + block_num
|
|
nonce = actual_block.to_bytes(self.ciph.NONCEBYTES, "little")
|
|
|
|
# Calculate byte range within output
|
|
# block_num=0 corresponds to output byte 0
|
|
# But if start_offset > 0, first block is partial
|
|
if block_num == 0 and self.start_offset > 0:
|
|
# First partial block: generate full block, copy from start_offset
|
|
self.ciph.stream(self.key, nonce, self.block_size, into=tmp_buf)
|
|
if self.total_bytes is None:
|
|
copy_len = self.block_size - self.start_offset
|
|
else:
|
|
copy_len = min(self.block_size - self.start_offset, self.total_bytes)
|
|
out_start = self.oseek
|
|
out_end = out_start + copy_len
|
|
view[out_start:out_end] = tmp_buf[
|
|
self.start_offset : self.start_offset + copy_len
|
|
]
|
|
written = copy_len
|
|
else:
|
|
# Full block or last partial block
|
|
# Output position: account for first block being partial
|
|
if self.start_offset > 0:
|
|
out_start = (
|
|
self.oseek
|
|
+ (self.block_size - self.start_offset)
|
|
+ (block_num - 1) * self.block_size
|
|
)
|
|
else:
|
|
out_start = self.oseek + block_num * self.block_size
|
|
if self.total_bytes is None:
|
|
out_end = out_start + self.block_size
|
|
else:
|
|
out_end = min(out_start + self.block_size, self.oseek + self.total_bytes)
|
|
size = out_end - out_start
|
|
self.ciph.stream(self.key, nonce, size, into=view[out_start:out_end])
|
|
written = size
|
|
|
|
if self.use_madvise and self.mm_raw:
|
|
madvise_file(self.mm_raw, out_start, written)
|
|
with self._lock:
|
|
self.written += written
|
|
if self.progress_state is not None:
|
|
self.progress_state["written"] = self.written
|
|
except BaseException as e:
|
|
logging.exception("Worker thread exception: %s", e)
|
|
finally:
|
|
view.release()
|
|
self._quit.set()
|
|
self._quit.wait()
|
|
|
|
def join(self):
|
|
"""Wait for all worker threads to finish."""
|
|
for t in self._threads:
|
|
t.join()
|
|
|
|
def stop(self):
|
|
"""Signal workers to stop."""
|
|
self._quit.set()
|
|
|
|
def cleanup(self):
|
|
"""Release references to mmap resources."""
|
|
if self._owns_mmap and self.mm_raw is not None:
|
|
self.mm_raw.close()
|
|
self.mm_raw = None
|
|
self._threads.clear()
|