Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c9cc15ed08 | ||
|
|
e62d26e4f8 | ||
|
|
7b483153af | ||
|
|
8380ab202e | ||
|
|
463a6e939f | ||
|
|
da152f5bd8 | ||
|
|
1ff5037fe2 | ||
|
|
a71bb0f0d4 | ||
|
|
75d42cf0d5 | ||
|
|
f0190b7ef9 | ||
|
|
017095df36 | ||
|
|
52cfcb9872 | ||
|
|
4ac1ddbe2a | ||
|
|
5747e9a183 |
@@ -0,0 +1,7 @@
|
||||
.*
|
||||
*.lock
|
||||
*.egg-info
|
||||
/dist
|
||||
/build
|
||||
__pycache__
|
||||
!.gitignore
|
||||
@@ -1,91 +1,117 @@
|
||||
# The World's Fastest Random Generator
|
||||
# The World's Fastest Random Generator v2
|
||||
|
||||
I was disappointed with the sad state of random number generators. Many languages don't ship anything useful and some are stuck with whatever the OS provides. All existing implementations are slow, typically maxing out at a few hundred megabytes per second.
|
||||

|
||||
*RandQuik running at full speed on a Macbook laptop, writing /dev/null.*
|
||||
|
||||
To give some perspective, my tool reaches 37.8 GB/s, writing /dev/null. Several gigabytes per second also on actual SSDs. And this is using full 20 rounds of shuffling.
|
||||
I was disappointed with the sad state of random number generators. Many languages don't ship anything useful and some are stuck with whatever the OS provides. Most existing implementations are slow, often maxing out at a few hundred megabytes per second, which becomes a real bottleneck in high-throughput systems.
|
||||
|
||||
Secondly, [flaws have been found](https://numpy.org/doc/stable/reference/random/upgrading-pcg64.html) in non-cryptographic algorithms such as the popular Mersenne Twister and PCG64 algorithms.
|
||||
Secondly, [flaws have been found](https://numpy.org/doc/stable/reference/random/upgrading-pcg64.html) in popular non-cryptographic algorithms such as Mersenne Twister and PCG-style generators. These issues range from detectable structure to repeating sequences, making them unsuitable for serious or long-running workloads.
|
||||
|
||||
The cryptographic alternative is simply better, being free of such issues as sequences repeating, and additionally being, well, cryptographically secure. Surprisingly, they appear even to be faster now, so there really ought to be no reason to stick with the old.
|
||||
|
||||
We use the widely used encryption algorithm **ChaCha20** as an extremely fast Cryptographically Secure Pseudorandom Number Generator (CSPRNG).
|
||||
The cryptographic alternative is simply better. Proper CSPRNGs avoid these pitfalls entirely and, on modern hardware, can now be *faster* than legacy non-cryptographic designs. There is little reason left to accept weaker guarantees for worse performance.
|
||||
|
||||
## CLI: randquik
|
||||
This version uses **AEGIS**, a modern authenticated encryption primitive that leverages AES hardware acceleration. AEGIS provides extremely high throughput while retaining strong cryptographic properties, significantly outperforming our legacy implementation and serving as an ideal foundation for a high-performance CSPRNG.
|
||||
|
||||
A simple shell tool that simply produces randomness to a file or pipe. It uses ChaCha20 encryption algorithm to produce a random stream that cannot be predicted unless one knows the key - the seed - being used. Keeping always the same seed may be useful for researchers and such who need repeatable results.
|
||||
## Quick start
|
||||
|
||||
<img src="https://github.com/LeoVasanko/RandQuik/blob/legacy/docs/random.webp?raw=true" width="800" alt="Screenshot">
|
||||
Install [UV](https://docs.astral.sh/uv/getting-started/installation/) and install the CLI tool with it:
|
||||
|
||||
## Installation
|
||||
|
||||
Clone the repository, use Meson build:
|
||||
|
||||
```
|
||||
meson setup build
|
||||
cd build
|
||||
ninja
|
||||
sudo ninja install # optionally
|
||||
randquik > /dev/null
|
||||
```sh
|
||||
uv tool install randquik
|
||||
```
|
||||
|
||||
Alternatively you may compile it by hand with `gcc *.c -o randquik -O3 -pthread`.
|
||||
|
||||
## Python module
|
||||
|
||||
Python module fills any buffers with random data very quickly,
|
||||
|
||||
```python
|
||||
from randquik import Cha, generate
|
||||
import secrets
|
||||
|
||||
key = secrets.token_bytes(32)
|
||||
|
||||
# Allocate bytearray and fill with random
|
||||
data = generate(1_000_000, key)
|
||||
|
||||
# Or into an existing buffer
|
||||
generate_into(data, key)
|
||||
You can try how it performs on your machine and find the optimal parameters:
|
||||
```sh
|
||||
randquik --benchmark
|
||||
```
|
||||
|
||||
Given the same key, the generate functions will on each call produce the same sequence. For incremental updates, create a generator object and extract as many non-identical bytes from it as needed. Re-initializing with the same key of course once again repeats the requence.
|
||||
|
||||
```python
|
||||
rng = Cha(key)
|
||||
|
||||
# Fill some buffer with next bytes iteratively
|
||||
rng(data)
|
||||
rng(data)
|
||||
...
|
||||
Wipe an entire file without altering its size:
|
||||
```sh
|
||||
randquik -o sensitive.dat
|
||||
```
|
||||
|
||||
## Numpy module
|
||||
|
||||
Numpy.Random BitGenerator is also provided for use with Numpy distributions. We do not recommend using it for filling byte buffers, where the Python module and the C code are far faster, but it will still provide better quality random numbers and faster than Numpy's own PCG64.
|
||||
|
||||
```python
|
||||
import numpy as np
|
||||
from nprand import Cha
|
||||
|
||||
gen = np.random.Generator(Cha()) # System random seeding by default
|
||||
gen.normal(size=10)
|
||||
Piping and redirection:
|
||||
```sh
|
||||
randquik --quiet | hexdump -C | head
|
||||
```
|
||||
|
||||
This module needs to be built by hand with Cython for now. Will be packaged properly in randquik module eventually.
|
||||
## Features
|
||||
|
||||
## Tests
|
||||
- Blazing-fast CSPRNG built on AEGIS single or multithreaded
|
||||
- Deterministic seeding: same seed, same byte stream
|
||||
- Seekable random stream and file output
|
||||
- Flexible I/O: piping, files, and `mmap`
|
||||
- Built-in benchmarking and dry-run modes
|
||||
- Full screen console graphics for speed display
|
||||
|
||||
Some pytest tests are provided for verifying ChaCha20 implementation correctness against the Cryptography module. Run by `pytest` in the main folder after installation.
|
||||
Below are the most important features with example commands.
|
||||
|
||||
## Performance
|
||||
### Performance
|
||||
|
||||
ChaCha20 as its name implies uses 20 "rounds" of shuffling for each output block. A lower number of rounds can be used for extra performance, where 8 is the minimum that is considered secure, and 12 provides a balanced option, while 20 has comfortable headroom to stay cryptographically secure.
|
||||
You'll be looking at up to 100 GB/s raw generation speed, making this some orders of magnitude faster than your traditional random number generation. Single-threaded performance still is 10-20 times faster than other options that don't have threading.
|
||||
|
||||
All functions and constructors of this module take `rounds` kwarg for adjusting this. On CLI the equivalent option is `-r`. By default 20 rounds are used.
|
||||
Your output, e.g. writing a file, will always be the bottle neck, not your random generator, but modern SSDs allow up to 10 GB/s write speeds already.
|
||||
|
||||
The CLI uses a configurable number of threads for extremely high performance, while the Python and Numpy modules don't - for now at least.
|
||||
### Size units
|
||||
|
||||
The implementation is optimized for Apple Silicon SIMD (Neon) and x86 CPUs using AVX2 where available, falling back to SSSE3 and ultimately plain C on other platforms. The implementation is loosely based on code from libsodium but runs faster than the library can.
|
||||
Many options such as `--len`, `--seek` accept human-readable units. The table uses conventional, capitalized forms (e.g. `KiB`, `MB`), but you may write them in lower case and with or without the trailing `B` (for example `1m` or `5gi`).
|
||||
|
||||
## Seekability
|
||||
| SI unit | Binary unit | Meaning | Bytes factor |
|
||||
|--------:|------------:|------------------------|-----------------------|
|
||||
| 100 | — | 100 bytes | 1 |
|
||||
| 1kB | 1KiB | Kilobyte / kibibyte | 1_000 / 1_024 |
|
||||
| 1MB | 1MiB | Megabyte / mebibyte | 1_000_000 / 1_048_576 |
|
||||
| 1GB | 1GiB | Gigabyte / gibibyte | 1_000_000_000 / 1_073_741_824 |
|
||||
| 1TB | 1TiB | Terabyte / tebibyte | 1_000_000_000_000 / 1_099_511_627_776 |
|
||||
| 1PB | 1PiB | Petabyte / pebibyte | 1_000_... / 1_125_... |
|
||||
| — | 1sect | Sectors of output device | 512 (typical), 4096 (rarely) |
|
||||
|
||||
It is possible to seek ChaCha to any byte position in the stream without delay. This is implemented in C API only for now, and is not exposed via Numpy, Python or CLI interfaces.
|
||||
### Seeding for repeatable output
|
||||
|
||||
You can provide an explicit seed string, always providing the same output, which can be useful e.g. for memory/disk testing where the data needs to be read back and verified.
|
||||
|
||||
```sh
|
||||
randquik -l 64MiB -s my-seed-string -o chunk.bin
|
||||
```
|
||||
|
||||
If no seed is provided, a secure new random one is created and printed on console (unless hidden by `-q`).
|
||||
|
||||
### Seekable random stream and output file
|
||||
|
||||
It is possible to seek to any byte position in the stream without delay.
|
||||
|
||||
- `--iseek`: seek the input random stream
|
||||
- `--oseek`: seek in the output file
|
||||
- `--seek`: set both input and output to the same position
|
||||
|
||||
Example: resume as if 5 terabytes had already been written, and continue writing to `out.dat`. The seed from the prior invocation should be included:
|
||||
```sh
|
||||
randquik --seek 5T --len 1G -s a5Z8Ew1Hfc2VfEtY -o out.dat
|
||||
```
|
||||
|
||||
Bytes prior to seek position are kept as they were while the file is expanded to fit all the data starting at five terabytes mark (using sparse allocation so it doesn't actually consume 5 terabytes).
|
||||
|
||||
Wipe a specific range of a disk or USB drive (using sector numbers e.g. from gdisk):
|
||||
```sh
|
||||
randquik -oseek 2048sect --len 100MiB -o /dev/sde
|
||||
```
|
||||
|
||||
### Benchmark and dry-run modes
|
||||
|
||||
Benchmark different modes and thread counts. Prints the options that perform the best on your system:
|
||||
```sh
|
||||
randquik --benchmark
|
||||
```
|
||||
|
||||
To do a single run without actually writing anywhere, use `--dry`:
|
||||
```sh
|
||||
randquik --len 50GiB -t8 --dry
|
||||
```
|
||||
|
||||
## Legacy
|
||||
|
||||
The original implementation is preserved in the [legacy](https://github.com/LeoVasanko/RandQuik/tree/legacy) git branch.
|
||||
|
||||
That version was once the fastest CSPRNG available, built around ChaCha20 with SIMD Assembly and C code written by me, making it faster than traditional algorithms without such optimizations and faster than the Linux kernel that also uses ChaCha20 to make random numbers. While historically significant, it has been greatly surpassed by the current AEGIS-based design in performance and features.
|
||||
|
||||
The legacy branch remains available for reference and benchmarking, but version 2 is the recommended implementation.
|
||||
|
||||
@@ -1,55 +0,0 @@
|
||||
#include <openssl/evp.h>
|
||||
#include <stdio.h>
|
||||
#include <string.h>
|
||||
|
||||
int main() {
|
||||
// Key and IV should be appropriately sized for ChaCha20
|
||||
unsigned char key[] = {
|
||||
0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
|
||||
0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
|
||||
0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
|
||||
0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f
|
||||
};
|
||||
unsigned char iv[] = {
|
||||
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
|
||||
0x00, 0x00, 0x00, 0x00
|
||||
};
|
||||
|
||||
// Initialize context
|
||||
EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
|
||||
if (!ctx) {
|
||||
perror("EVP_CIPHER_CTX_new failed");
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Initialize the ChaCha20 cipher
|
||||
if (!EVP_EncryptInit_ex(ctx, EVP_chacha20(), NULL, key, iv)) {
|
||||
perror("EVP_EncryptInit_ex failed");
|
||||
EVP_CIPHER_CTX_free(ctx);
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Buffer for the keystream
|
||||
unsigned char keystream[1000000];
|
||||
memset(keystream, 0, sizeof(keystream));
|
||||
|
||||
for (int i = 0; i < 1000; i++) {
|
||||
// Generate keystream
|
||||
int len;
|
||||
if (!EVP_EncryptUpdate(ctx, keystream, &len, keystream, sizeof keystream)) {
|
||||
perror("EVP_EncryptUpdate failed");
|
||||
EVP_CIPHER_CTX_free(ctx);
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
// Clean up
|
||||
EVP_CIPHER_CTX_free(ctx);
|
||||
|
||||
// Print the generated keystream
|
||||
for (int i = 0; i < 64; i++) {
|
||||
printf("%02x", keystream[i]);
|
||||
}
|
||||
printf("\n");
|
||||
|
||||
return 0;
|
||||
}
|
||||
@@ -1,20 +0,0 @@
|
||||
#include "chacha20.h"
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
int main(void) {
|
||||
uint64_t N = 1000000;
|
||||
uint8_t* buf = malloc(N);
|
||||
const uint8_t key[32] = {0};
|
||||
const uint8_t nonce[16] = {0};
|
||||
for (uint64_t i = 0; i < 1000; ++i) {
|
||||
cha_generate(buf, N, key, nonce);
|
||||
}
|
||||
for (unsigned i = 0; i < 16; ++i)
|
||||
printf("%02X ", buf[i]);
|
||||
for (unsigned i = 0; i < 16; ++i)
|
||||
printf(" %02X", buf[1024 + i]);
|
||||
|
||||
puts("");
|
||||
free(buf);
|
||||
}
|
||||
@@ -1,12 +0,0 @@
|
||||
#include <math.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
int main(void)
|
||||
{
|
||||
printf("RAND_MAX = %d (%.1lf bit)\n", RAND_MAX, log2(RAND_MAX));
|
||||
const uint64_t rounds = 1000000000ull * 8 / (unsigned)log2(RAND_MAX);
|
||||
for (uint64_t i = rounds; i-->0;) {
|
||||
rand();
|
||||
}
|
||||
}
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 66 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 10 KiB |
-15
@@ -1,15 +0,0 @@
|
||||
project('randquik', 'c')
|
||||
executable(
|
||||
'randquik',
|
||||
'src/cli.c',
|
||||
c_args: ['-Wall', '-O3', '-march=native'],
|
||||
install: true,
|
||||
)
|
||||
dependency('threads')
|
||||
|
||||
library(
|
||||
'randquik-chacha20',
|
||||
'src/charandom.c',
|
||||
build_by_default: true,
|
||||
c_args: ['-Wall', '-O3', '-march=native'],
|
||||
)
|
||||
+28
-49
@@ -1,63 +1,42 @@
|
||||
[build-system]
|
||||
requires = ["hatchling", "hatch-vcs", "wheel", "cffi"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "randquik"
|
||||
version = "0.1.0"
|
||||
description = "Extremely fast and cryptographically secure random number generator."
|
||||
dynamic = ["version"]
|
||||
description = "CLI tool for extremely fast random bytes."
|
||||
readme = "README.md"
|
||||
license.text = "Public Domain"
|
||||
authors = [{ name = "Vasanko" }]
|
||||
requires-python = ">=3.13"
|
||||
authors = [
|
||||
{ name = "Leo Vasanko" }
|
||||
]
|
||||
keywords = ["random", "CSPRNG", "AEGIS", "shred", "benchmark"]
|
||||
classifiers = [
|
||||
"Operating System :: POSIX",
|
||||
"Operating System :: Unix",
|
||||
"Development Status :: 5 - Production/Stable",
|
||||
"Environment :: Console",
|
||||
"Intended Audience :: Developers",
|
||||
"Intended Audience :: System Administrators",
|
||||
"Topic :: Security :: Cryptography",
|
||||
"Topic :: Security",
|
||||
"Topic :: Software Development :: Libraries :: Python Modules",
|
||||
"Topic :: Utilities",
|
||||
]
|
||||
dependencies = ["cffi>=1.0.1", "numpy"]
|
||||
requires-python = ">=3.10"
|
||||
keywords = [
|
||||
"random",
|
||||
"generator",
|
||||
"fast",
|
||||
"secure",
|
||||
"cryptographic",
|
||||
"randomness",
|
||||
dependencies = [
|
||||
"aeg>=0.4.3",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
randquik = "randquik.cli:main"
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://git.zi.fi/LeoVasanko/randquik"
|
||||
Repository = "https://github.com/LeoVasanko/randquik"
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = ["pytest", "ruff", "cryptography", "scipy"]
|
||||
[build-system]
|
||||
requires = ["hatchling", "hatch-vcs"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatchling]
|
||||
[tool.hatch.version]
|
||||
source = "vcs"
|
||||
|
||||
[tool.ruff]
|
||||
extend-select = ["I", "W", "UP", "C4", "ISC", "S"]
|
||||
# Worth selecting but still too broken: ASYNC, B, DTZ, FA
|
||||
ignore = [
|
||||
"D100",
|
||||
"D101",
|
||||
"D102",
|
||||
"D103",
|
||||
"E402",
|
||||
"E741",
|
||||
"F811",
|
||||
"F821",
|
||||
# ruff format complains about these:
|
||||
"ISC001",
|
||||
"S101",
|
||||
"S102",
|
||||
"S104",
|
||||
"S311",
|
||||
"S603",
|
||||
"S607",
|
||||
"W191",
|
||||
]
|
||||
show-source = true
|
||||
show-fixes = true
|
||||
line-length = 100
|
||||
target-version = "py313"
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
pythonpath = ["."]
|
||||
[tool.ruff.lint.isort]
|
||||
known-first-party = ["randquik"]
|
||||
|
||||
+24
-2
@@ -1,3 +1,25 @@
|
||||
from randquik.cha import Cha, generate, generate_into
|
||||
"""RandQuik - High-performance cryptographic random data generator.
|
||||
|
||||
__all__ = ["Cha", "generate_into", "generate"]
|
||||
This package provides fast random data generation using AEGIS ciphers,
|
||||
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.stats import format_size, format_time
|
||||
from randquik.utils import parse_size
|
||||
|
||||
try:
|
||||
from randquik._version import __version__
|
||||
except ImportError:
|
||||
__version__ = "0.0.0.dev0"
|
||||
|
||||
__all__ = [
|
||||
"ProgressDisplay",
|
||||
"__version__",
|
||||
"derive_key",
|
||||
"format_size",
|
||||
"format_time",
|
||||
"generate_random_seed",
|
||||
"parse_size",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
"""Entry point for running as `python -m randquik`."""
|
||||
|
||||
from randquik.cli import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,142 @@
|
||||
"""Benchmark functions for measuring performance."""
|
||||
|
||||
import contextlib
|
||||
import os
|
||||
import pathlib
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
|
||||
from randquik.utils import sparse_range
|
||||
|
||||
__all__ = ["bench_mode", "run_benchmark"]
|
||||
|
||||
|
||||
def bench_mode(
|
||||
tcounts: list[int],
|
||||
io_mode: str,
|
||||
length: str,
|
||||
alg: str | None,
|
||||
bench_file: pathlib.Path | None,
|
||||
) -> list[tuple[int, float, list[str]]]:
|
||||
max_repeats = 5
|
||||
max_time = 0.5 # Quit early if more than 500ms has passed
|
||||
results = []
|
||||
|
||||
# Use a file in current folder for file modes
|
||||
if "file" in io_mode:
|
||||
iocmd = ["-o", str(bench_file)]
|
||||
elif "dry" in io_mode:
|
||||
iocmd = ["--dry"]
|
||||
elif "null" in io_mode:
|
||||
iocmd = ["-o", os.devnull]
|
||||
else:
|
||||
raise ValueError(f"Unknown io_mode: {io_mode}")
|
||||
|
||||
# Print iocmd at start of row
|
||||
sys.stdout.write(f"{' '.join(iocmd)[:20]:<20}")
|
||||
sys.stdout.flush()
|
||||
|
||||
for workers in tcounts:
|
||||
speeds = []
|
||||
worker_start = time.perf_counter()
|
||||
for rep in range(max_repeats):
|
||||
if rep > 0 and (time.perf_counter() - worker_start) > max_time:
|
||||
break
|
||||
cmd = [
|
||||
sys.executable,
|
||||
"-m",
|
||||
"randquik",
|
||||
f"-l{length}",
|
||||
f"-t{workers}",
|
||||
*([f"-a{alg}"] if alg else []),
|
||||
*iocmd,
|
||||
]
|
||||
|
||||
try:
|
||||
proc = subprocess.run(
|
||||
cmd,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.PIPE,
|
||||
text=False,
|
||||
)
|
||||
except KeyboardInterrupt:
|
||||
sys.stderr.write(
|
||||
f" Interrupted\n\n>>> {' '.join(cmd)}\n"
|
||||
) # 6 not 8 to account for ^C
|
||||
sys.exit(1)
|
||||
stderr = proc.stderr.decode(errors="ignore")
|
||||
if proc.returncode != 0:
|
||||
sys.stderr.write(f"{'ERROR':>8}\n\n>>> {' '.join(cmd)}\n{stderr}")
|
||||
sys.exit(1)
|
||||
m2 = re.findall(r"([0-9]+\.[0-9]+)\s+GB/s", stderr)
|
||||
if m2:
|
||||
speeds.append(float(m2[-1]))
|
||||
|
||||
if speeds:
|
||||
sorted_speeds = sorted(speeds)
|
||||
median = sorted_speeds[len(speeds) // 2]
|
||||
sys.stdout.write(f"{median:>8.2f}")
|
||||
sys.stdout.flush()
|
||||
results.append((workers, median, iocmd))
|
||||
else:
|
||||
sys.stdout.write(f"{'---':>8}")
|
||||
sys.stdout.flush()
|
||||
|
||||
sys.stdout.write("\n") # newline after row
|
||||
sys.stdout.flush()
|
||||
|
||||
# Cleanup bench file
|
||||
if bench_file:
|
||||
with contextlib.suppress(OSError):
|
||||
bench_file.unlink()
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def run_benchmark(args):
|
||||
"""Run comprehensive benchmark across all I/O modes."""
|
||||
# Check if output file already exists
|
||||
bench_file = pathlib.Path(args.output or "test.dat")
|
||||
if bench_file.exists():
|
||||
if not args.output:
|
||||
raise ValueError(
|
||||
f"File test.dat already exists. Use -o {bench_file} to benchmark over it or choose another name."
|
||||
)
|
||||
bench_file.unlink()
|
||||
|
||||
try:
|
||||
length = "1G" if args.len is None else args.len
|
||||
max_threads = args.threads if args.threads is not None else os.cpu_count()
|
||||
|
||||
all_results = {}
|
||||
tcounts = sparse_range(max_threads)
|
||||
|
||||
# Print header
|
||||
header = f"{'randquik':<20}"
|
||||
for w in tcounts:
|
||||
header += f"{'-t' + str(w):>8}"
|
||||
header += "\n" + "-" * (20 + 8 * len(tcounts)) + "\n"
|
||||
sys.stdout.write(header)
|
||||
|
||||
for io_mode in ["dry", "null", "file"]:
|
||||
results = bench_mode(tcounts, io_mode, length, alg=args.alg, bench_file=bench_file)
|
||||
all_results[io_mode] = results
|
||||
|
||||
sys.stdout.write("-" * (20 + 8 * len(tcounts)) + "\n")
|
||||
|
||||
# Find fastest configuration and RNG speed
|
||||
gen_speed = max(r[1] for res in all_results.values() for r in res)
|
||||
best_speed, best_threads, best_iocmd = max(
|
||||
[(sp, w, iocmd) for w, sp, iocmd in all_results["file"]],
|
||||
)
|
||||
threads = f" -t{best_threads}" if best_threads != 1 else ""
|
||||
sys.stderr.write(
|
||||
f"\n>>> Fastest wrote {best_speed:.2f} GB/s, plain RNG {gen_speed:.0f} GB/s\n"
|
||||
f"randquik {' '.join(best_iocmd)}{threads}\n"
|
||||
)
|
||||
finally:
|
||||
# Cleanup bench file even if interrupted
|
||||
with contextlib.suppress(OSError):
|
||||
bench_file.unlink()
|
||||
-104
@@ -1,104 +0,0 @@
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import cffi
|
||||
|
||||
src = Path(__file__).parent.parent / "src"
|
||||
|
||||
if not src.is_dir():
|
||||
raise RuntimeError("Unable to find RandQuik C sources in {src}")
|
||||
|
||||
ffi = cffi.FFI()
|
||||
ffi.cdef(
|
||||
"""
|
||||
typedef uint64_t (*genfunc)(
|
||||
uint8_t* out, size_t outsize, uint32_t state[16], unsigned rounds
|
||||
);
|
||||
typedef struct cha_ctx {
|
||||
uint32_t input[16];
|
||||
uint8_t unconsumed[512];
|
||||
uint32_t offset, end;
|
||||
unsigned rounds;
|
||||
genfunc gen;
|
||||
} cha_ctx;
|
||||
|
||||
int cha_generate(uint8_t* out, uint64_t outlen, const uint8_t key[32], const uint8_t iv[16], unsigned rounds);
|
||||
|
||||
void cha_init(cha_ctx* ctx, const uint8_t* key, const uint8_t* iv, unsigned rounds);
|
||||
void cha_wipe(cha_ctx* ctx);
|
||||
int cha_update(cha_ctx* ctx, uint8_t* out, uint64_t outlen);
|
||||
"""
|
||||
)
|
||||
libname = "librandquik-chacha20.so"
|
||||
if sys.platform == "darwin":
|
||||
libname = "librandquik-chacha20.dylib"
|
||||
elif sys.platform == "win32":
|
||||
libname = "randquik-chacha20.dll"
|
||||
lib = ffi.dlopen((Path(__file__).parent.parent / f"build/{libname}").as_posix())
|
||||
|
||||
|
||||
def _processKeys(key, iv):
|
||||
key = memoryview(key)
|
||||
iv = memoryview(iv)
|
||||
if key.nbytes != 32:
|
||||
raise ValueError("key must be 32 bytes")
|
||||
if iv.nbytes != 16:
|
||||
# Allow original and IETF nonces with zero counter
|
||||
if iv.nbytes == 8:
|
||||
iv = bytes(8) + iv
|
||||
elif iv.nbytes == 12:
|
||||
iv = bytes(4) + iv
|
||||
else:
|
||||
raise ValueError(
|
||||
"iv lenth must be 8 (ChaCha20 original), 12 (IETF) or 16 (counter in initial 8 bytes)"
|
||||
)
|
||||
return ffi.from_buffer(key), ffi.from_buffer(iv)
|
||||
|
||||
|
||||
def _processBuffer(out):
|
||||
try:
|
||||
outlen = out.nbytes
|
||||
except AttributeError:
|
||||
out = memoryview(out)
|
||||
outlen = out.nbytes
|
||||
if getattr(out, "readonly", None):
|
||||
raise ValueError("The output buffer must be writable, not e.g. `bytes`")
|
||||
return ffi.from_buffer(out), outlen
|
||||
|
||||
|
||||
class Cha:
|
||||
def __init__(self, key: bytes | Any, iv: bytes | Any, *, rounds=20):
|
||||
"""Construct a generator that holds its internal state, moving forward on each call."""
|
||||
key, iv = _processKeys(key, iv)
|
||||
self.ctx = ffi.new("cha_ctx*")
|
||||
lib.cha_init(self.ctx, key, iv, rounds)
|
||||
|
||||
def __del__(self):
|
||||
lib.cha_wipe(self.ctx)
|
||||
|
||||
def __call__(self, out: bytearray | Any):
|
||||
"""Fill the parameter with random bytes"""
|
||||
outbuf, outlen = _processBuffer(out)
|
||||
lib.cha_update(self.ctx, outbuf, outlen)
|
||||
return out
|
||||
|
||||
|
||||
def generate_into(
|
||||
out: bytearray | memoryview | Any,
|
||||
key: bytes | Any,
|
||||
iv: bytes | Any = bytes(16),
|
||||
*,
|
||||
rounds=20,
|
||||
):
|
||||
"""Fill in random bytes into an existing array (buffer interface)"""
|
||||
key, iv = _processKeys(key, iv)
|
||||
outbuf, outlen = _processBuffer(out)
|
||||
lib.cha_generate(outbuf, outlen, key, iv, rounds)
|
||||
return out
|
||||
|
||||
|
||||
def generate(outlen: int, key: bytes | Any, iv: bytes | Any = bytes(16), *, rounds=20):
|
||||
"""Return a bytearray of random bytes"""
|
||||
assert outlen >= 0
|
||||
return generate_into(bytearray(outlen), key, iv, rounds=rounds)
|
||||
+206
@@ -0,0 +1,206 @@
|
||||
"""Command-line interface for RandQuik."""
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import sys
|
||||
|
||||
import aeg
|
||||
|
||||
from randquik.benchmark import run_benchmark
|
||||
from randquik.crypto import derive_key, generate_random_seed
|
||||
from randquik.utils import parse_size
|
||||
from randquik.workers import run
|
||||
|
||||
__all__ = ["main"]
|
||||
|
||||
# Disable GC for performance
|
||||
gc.disable()
|
||||
|
||||
DEFAULT_ALG = "AEGIS-128X2"
|
||||
|
||||
|
||||
def prepare_seed(args):
|
||||
"""Prepare seed and determine if it was generated."""
|
||||
generated_seed = args.seed is None
|
||||
seed = generate_random_seed() if generated_seed else args.seed
|
||||
return seed, generated_seed
|
||||
|
||||
|
||||
def prepare_key(seed, keybytes):
|
||||
"""Derive key from seed."""
|
||||
key = derive_key(seed, keybytes)
|
||||
return key
|
||||
|
||||
|
||||
def parse_seeks(args):
|
||||
"""Parse seek values from arguments."""
|
||||
try:
|
||||
iseek = parse_size(args.iseek, args.output) or 0
|
||||
oseek = parse_size(args.oseek, args.output) or 0
|
||||
if args.seek:
|
||||
iseek = oseek = parse_size(args.seek, args.output) or 0
|
||||
except ValueError as e:
|
||||
raise ValueError(f"Error parsing seek: {e}") from None
|
||||
return iseek, oseek
|
||||
|
||||
|
||||
def _main():
|
||||
"""Internal main function that may raise exceptions."""
|
||||
parser = argparse.ArgumentParser(description="Generate random bytes using AEGIS ciphers")
|
||||
parser.add_argument("-s", "--seed", help="Alphanumeric seed string", type=str)
|
||||
parser.add_argument(
|
||||
"-l",
|
||||
"--len",
|
||||
help="Length to generate (e.g. 1g, 100mi, 1000sect)",
|
||||
type=str,
|
||||
default=None,
|
||||
)
|
||||
parser.add_argument("-o", "--output", help="Output file (default: stdout)", type=str)
|
||||
parser.add_argument(
|
||||
"-t",
|
||||
"--threads",
|
||||
help="Number of worker threads (benchmark: upper limit)",
|
||||
type=int,
|
||||
default=None,
|
||||
)
|
||||
parser.add_argument(
|
||||
"-a",
|
||||
"--alg",
|
||||
help=f"Cipher algorithm (default: {DEFAULT_ALG})",
|
||||
type=str,
|
||||
default=None,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--benchmark",
|
||||
action="store_true",
|
||||
help="Run benchmark (generates 1GB and reports speed)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dry",
|
||||
action="store_true",
|
||||
help="Dry run: open output but skip writes (for benchmarking)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seek",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Seek both input stream and output to position (e.g. 1g, 100mi)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--iseek",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Seek input random stream to position (overrides --seek)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--oseek",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Seek output file to position (overrides --seek)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-q",
|
||||
"--quiet",
|
||||
action="store_true",
|
||||
help="Quiet mode: suppress all output except errors",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-v",
|
||||
"--verbose",
|
||||
action="count",
|
||||
default=0,
|
||||
help="Verbose mode: -v for I/O mode, -vv for worker statistics",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Normalize "-" output to None (stdout)
|
||||
if args.output == "-":
|
||||
args.output = None
|
||||
|
||||
ciph = aeg.cipher(args.alg or DEFAULT_ALG)
|
||||
|
||||
# Validate and process args
|
||||
seed, generated_seed = prepare_seed(args)
|
||||
key = prepare_key(seed, ciph.KEYBYTES)
|
||||
iseek, oseek = parse_seeks(args)
|
||||
total_bytes = parse_size(args.len) # None if not specified, 0 if -l0
|
||||
# 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:
|
||||
raise ValueError("Cannot specify seed in benchmark mode")
|
||||
if iseek or oseek:
|
||||
raise ValueError("Cannot use seek options in benchmark mode")
|
||||
run_benchmark(args)
|
||||
return
|
||||
|
||||
# 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,
|
||||
action=action,
|
||||
)
|
||||
|
||||
# Set repeat command for generated seeds
|
||||
result.repeat_cmd = repeat_cmd
|
||||
|
||||
# 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}"
|
||||
|
||||
# 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()
|
||||
|
||||
if result.interrupted:
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def main():
|
||||
"""Main entry point for the CLI with exception handling."""
|
||||
try:
|
||||
_main()
|
||||
except (KeyboardInterrupt, BrokenPipeError):
|
||||
sys.exit(1)
|
||||
except ValueError as e:
|
||||
sys.stderr.write(f"Error: {e}\n")
|
||||
sys.exit(1)
|
||||
@@ -0,0 +1,22 @@
|
||||
"""Cryptographic functions for key derivation and seed generation."""
|
||||
|
||||
import hashlib
|
||||
import secrets
|
||||
import string
|
||||
|
||||
__all__ = [
|
||||
"derive_key",
|
||||
"generate_random_seed",
|
||||
]
|
||||
|
||||
|
||||
def generate_random_seed() -> str:
|
||||
"""Generate a random alphanumeric seed string."""
|
||||
chars = string.ascii_letters + string.digits
|
||||
return "".join(secrets.choice(chars) for _ in range(16))
|
||||
|
||||
|
||||
def derive_key(seed: str, key_bytes: int) -> bytes:
|
||||
"""Derive a key from a seed string using SHA-512."""
|
||||
assert 16 <= key_bytes <= 64, "Only 128-512 bits supported"
|
||||
return hashlib.sha512(seed.encode()).digest()[:key_bytes]
|
||||
+115
@@ -0,0 +1,115 @@
|
||||
"""File I/O helpers for output handling."""
|
||||
|
||||
import contextlib
|
||||
import errno
|
||||
import os
|
||||
import pathlib
|
||||
import sys
|
||||
from collections.abc import Generator
|
||||
|
||||
__all__ = [
|
||||
"open_fd",
|
||||
"open_memoryview",
|
||||
]
|
||||
|
||||
|
||||
def _open_output(
|
||||
output_path: str,
|
||||
total_bytes: int | None,
|
||||
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)
|
||||
"""
|
||||
created = False
|
||||
path = pathlib.Path(output_path)
|
||||
created = not path.exists()
|
||||
flags = os.O_WRONLY | os.O_CREAT
|
||||
fd = os.open(str(path), flags, 0o644)
|
||||
|
||||
required_size = oseek + (total_bytes if total_bytes is not None else 0)
|
||||
current_size = os.fstat(fd).st_size
|
||||
if required_size > current_size:
|
||||
with contextlib.suppress(OSError):
|
||||
os.ftruncate(fd, required_size)
|
||||
|
||||
# Seek to output position
|
||||
if oseek > 0:
|
||||
try:
|
||||
os.lseek(fd, oseek, os.SEEK_SET)
|
||||
except OSError as e:
|
||||
raise ValueError(
|
||||
f"Cannot oseek in {output_path or 'stdout'}. Use only --iseek or specify a seekable file."
|
||||
) from e
|
||||
|
||||
# macOS: try to bypass unified buffer cache (F_NOCACHE)
|
||||
if sys.platform == "darwin":
|
||||
try:
|
||||
import fcntl
|
||||
|
||||
fcntl.fcntl(fd, fcntl.F_NOCACHE, 1)
|
||||
except (OSError, AttributeError, ImportError):
|
||||
pass
|
||||
|
||||
return fd, created
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def open_fd(
|
||||
output_path: str | None,
|
||||
total_bytes: int | None,
|
||||
dry: bool = False,
|
||||
oseek: int = 0,
|
||||
) -> Generator[int]:
|
||||
"""Context manager for output file descriptor.
|
||||
|
||||
Args:
|
||||
output_path: Path to output file, or None for stdout
|
||||
total_bytes: Total bytes to write
|
||||
dry: If True, skip truncation/preallocation and tty check
|
||||
oseek: Seek position for output
|
||||
|
||||
Yields:
|
||||
Integer file descriptor
|
||||
"""
|
||||
if dry:
|
||||
yield -1
|
||||
return
|
||||
if not output_path:
|
||||
if sys.stdout.isatty():
|
||||
raise ValueError("Refusing to write binary data to terminal. Use -o to specify a file.")
|
||||
yield sys.stdout.fileno()
|
||||
return
|
||||
fd, created = _open_output(output_path, total_bytes, oseek)
|
||||
try:
|
||||
yield fd
|
||||
except OSError as e:
|
||||
if e.errno == errno.ENOSPC:
|
||||
# Clean up file we created on disk full
|
||||
if created:
|
||||
with contextlib.suppress(Exception):
|
||||
os.unlink(output_path)
|
||||
raise ValueError(f"No space left on device: {output_path}") from None
|
||||
raise
|
||||
finally:
|
||||
with contextlib.suppress(Exception):
|
||||
os.close(fd)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def open_memoryview(buf) -> Generator[memoryview]:
|
||||
"""Context manager for memoryview.
|
||||
|
||||
Args:
|
||||
buf: Buffer to create memoryview from
|
||||
|
||||
Yields:
|
||||
memoryview object
|
||||
"""
|
||||
view = memoryview(buf)
|
||||
try:
|
||||
yield view
|
||||
finally:
|
||||
view.release()
|
||||
@@ -0,0 +1,844 @@
|
||||
"""Progress display with speed graph at bottom of terminal."""
|
||||
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
|
||||
from randquik.stats import format_size, format_time
|
||||
|
||||
__all__ = ["ProgressDisplay"]
|
||||
|
||||
# Unicode block characters for graph (8 levels per cell)
|
||||
GRAPH_BLOCKS = " ▁▂▃▄▅▆▇█"
|
||||
|
||||
# Maximum height of progress display in terminal rows
|
||||
MAX_HEIGHT = 10
|
||||
|
||||
|
||||
class ProgressDisplay:
|
||||
"""Progress display with speed graph at bottom of terminal, updated every 100ms.
|
||||
|
||||
Only active when stderr is a tty. Reads progress from a shared state dict
|
||||
with a single 'written' key. All display logic is encapsulated here.
|
||||
|
||||
Uses the bottom portion of the terminal with a scrolling region preserved
|
||||
at the top, allowing normal output to scroll above the progress display.
|
||||
|
||||
The graph fills from left to right as progress advances, doubling as both
|
||||
a progress bar and a speed-over-time visualization.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
total_bytes: int | None,
|
||||
start_time: float,
|
||||
state: dict,
|
||||
infinite: bool | None = None,
|
||||
output_name: str | None = None,
|
||||
oseek: int = 0,
|
||||
):
|
||||
self.total_bytes = total_bytes
|
||||
self.start_time = start_time
|
||||
self.state = state # Must have 'written' key
|
||||
self.infinite = infinite if infinite is not None else total_bytes is None
|
||||
self.output_name = output_name or "<stdout>"
|
||||
self.oseek = oseek
|
||||
self.active = sys.stderr.isatty()
|
||||
self._stop = threading.Event()
|
||||
self._thread = None
|
||||
self._last_written = 0
|
||||
self._last_time = start_time
|
||||
# Speed history for graph - fixed size, filled from left as progress advances
|
||||
self._graph_width = 80 # Will be updated on first render
|
||||
self._speed_history: list[float] = [] # Stores GB/s values, one per column
|
||||
self._max_speed: float = 0.01 # Start with small value to avoid div by zero
|
||||
# For infinite mode: track time of each speed sample
|
||||
self._time_history: list[float] = []
|
||||
# X-axis scale smoothing (hysteresis for estimated total time)
|
||||
self._smoothed_scale_time: float | None = None
|
||||
# Terminal handling
|
||||
self._current_scroll_bottom: int | None = None
|
||||
self._hidden_cursor = False
|
||||
self._first_draw = True
|
||||
|
||||
def start(self):
|
||||
if not self.active:
|
||||
return
|
||||
self._setup_terminal_state()
|
||||
self._thread = threading.Thread(target=self._run, daemon=True)
|
||||
self._thread.start()
|
||||
|
||||
def stop(self):
|
||||
if not self.active or self._thread is None:
|
||||
return
|
||||
self._stop.set()
|
||||
self._thread.join(timeout=0.5)
|
||||
# Final render so the finished state stays on screen
|
||||
if self.active:
|
||||
cols, rows, lines, overlay = self._render_frame()
|
||||
self._draw_frame(cols, rows, lines, overlay)
|
||||
self._restore_terminal_state()
|
||||
|
||||
def _setup_terminal_state(self):
|
||||
"""Prepare terminal: hide cursor and start using bottom reserved block."""
|
||||
sys.stderr.write("\x1b[?25l") # Hide cursor to reduce flicker
|
||||
sys.stderr.flush()
|
||||
self._hidden_cursor = True
|
||||
|
||||
def _get_smoothed_speed(self, window_secs: float = 1.0) -> float:
|
||||
"""Calculate average speed over the last window_secs seconds.
|
||||
|
||||
Returns speed in bytes/sec, averaged from recent samples in _speed_history.
|
||||
Falls back to the most recent sample if not enough history.
|
||||
"""
|
||||
if not self._speed_history or not self._time_history:
|
||||
return 0.0
|
||||
|
||||
current_time = self._time_history[-1]
|
||||
cutoff_time = current_time - window_secs
|
||||
|
||||
# Find samples within the window
|
||||
total_speed = 0.0
|
||||
count = 0
|
||||
for idx in range(len(self._time_history) - 1, -1, -1):
|
||||
if self._time_history[idx] < cutoff_time:
|
||||
break
|
||||
total_speed += self._speed_history[idx]
|
||||
count += 1
|
||||
|
||||
if count == 0:
|
||||
return self._speed_history[-1] * 1_000_000_000
|
||||
|
||||
# Return average in bytes/sec (speed_history stores GB/s)
|
||||
return (total_speed / count) * 1_000_000_000
|
||||
|
||||
def _restore_terminal_state(self):
|
||||
"""Restore terminal scrolling and cursor after progress is done."""
|
||||
# Reset scrolling region to full screen
|
||||
sys.stderr.write("\x1b[r")
|
||||
self._current_scroll_bottom = None
|
||||
|
||||
# Move cursor to a fresh line under the progress block
|
||||
cols, rows = self._get_terminal_size()
|
||||
sys.stderr.write(f"\x1b[{rows};1H\n")
|
||||
|
||||
if self._hidden_cursor:
|
||||
sys.stderr.write("\x1b[?25h")
|
||||
self._hidden_cursor = False
|
||||
sys.stderr.flush()
|
||||
|
||||
def _get_terminal_size(self) -> tuple[int, int]:
|
||||
"""Return (columns, rows)."""
|
||||
try:
|
||||
size = os.get_terminal_size(sys.stderr.fileno())
|
||||
return size.columns, size.lines
|
||||
except (OSError, ValueError):
|
||||
return 80, 24
|
||||
|
||||
def _render_graph_row(
|
||||
self,
|
||||
values: list[float],
|
||||
max_val: float,
|
||||
row: int,
|
||||
total_rows: int,
|
||||
width: int,
|
||||
avg_speed: float = 0,
|
||||
) -> str:
|
||||
"""Render one row of the graph using Unicode blocks.
|
||||
|
||||
Row 0 is top, total_rows-1 is bottom. Each cell can show 8 levels.
|
||||
Values list may be shorter than width (unfilled area shown as dim bar at avg_speed).
|
||||
"""
|
||||
filled_chars = []
|
||||
unfilled_chars = []
|
||||
filled_cols = len(values)
|
||||
# Calculate the row threshold for average speed
|
||||
avg_normalized = (avg_speed / max_val) * total_rows * 8 if max_val > 0 else 0
|
||||
row_bottom = (total_rows - row - 1) * 8
|
||||
row_top = row_bottom + 8
|
||||
|
||||
# Build filled portion
|
||||
for i in range(filled_cols):
|
||||
v = values[i]
|
||||
# Normalize value to 0..total_rows*8 range
|
||||
normalized = (v / max_val) * total_rows * 8 if max_val > 0 else 0
|
||||
if normalized <= row_bottom:
|
||||
filled_chars.append(" ")
|
||||
elif normalized >= row_top:
|
||||
filled_chars.append("█")
|
||||
else:
|
||||
level = math.ceil(normalized - row_bottom)
|
||||
filled_chars.append(GRAPH_BLOCKS[min(level, 8)])
|
||||
|
||||
# Build unfilled portion (dim grey at avg_speed level)
|
||||
unfilled_width = width - filled_cols
|
||||
if unfilled_width > 0:
|
||||
if avg_normalized <= row_bottom:
|
||||
unfilled_char = " "
|
||||
elif avg_normalized >= row_top:
|
||||
unfilled_char = "█"
|
||||
else:
|
||||
level = int(avg_normalized - row_bottom)
|
||||
unfilled_char = GRAPH_BLOCKS[min(level, 8)]
|
||||
unfilled_chars = [unfilled_char] * unfilled_width
|
||||
|
||||
# Combine with color codes only at transitions
|
||||
filled_str = "".join(filled_chars)
|
||||
unfilled_str = "".join(unfilled_chars)
|
||||
if unfilled_str:
|
||||
return f"{filled_str}\x1b[0m\x1b[38;5;235m{unfilled_str}\x1b[0m\x1b[33m"
|
||||
return filled_str
|
||||
|
||||
def _build_header(
|
||||
self,
|
||||
cols: int,
|
||||
written: int,
|
||||
speed: float,
|
||||
elapsed: float,
|
||||
eta: float | None = None,
|
||||
total_bytes: int | None = None,
|
||||
) -> str:
|
||||
"""Build the header line with stats, output name, and position.
|
||||
|
||||
Args:
|
||||
cols: Terminal width
|
||||
written: Bytes written so far
|
||||
speed: Current speed in bytes/sec
|
||||
elapsed: Elapsed time in seconds
|
||||
eta: Estimated time remaining (None for infinite mode)
|
||||
total_bytes: Total bytes to write (None for infinite mode)
|
||||
"""
|
||||
spinner = "\u25d0\u25d3\u25d1\u25d2"[int(elapsed * 4) % 4]
|
||||
written_gb = written / 1_000_000_000
|
||||
speed_gbs = speed / 1_000_000_000
|
||||
|
||||
# Build the fixed stats portion
|
||||
if total_bytes is not None:
|
||||
# Finite mode: show progress and ETA
|
||||
total_gb = total_bytes / 1_000_000_000
|
||||
eta_str = format_time(eta) if eta is not None else "--"
|
||||
stats = (
|
||||
f"\x1b[1;36mRandQuik {spinner}\x1b[0m "
|
||||
f"{written_gb:6.2f}\x1b[2m/\x1b[0m{total_gb:.2f} GB "
|
||||
)
|
||||
stats += (
|
||||
f"\x1b[2m@\x1b[0m {speed_gbs:5.2f} GB/s \x1b[2mest.\x1b[0m {eta_str:<8}"
|
||||
if written < total_bytes
|
||||
else f"\x1b[2m{'done':>27}\x1b[0m"
|
||||
)
|
||||
|
||||
# Visible: "RandQuik X " (14) + "XXXX.XX/XXXX.XX GB " (20) + "@ XX.XX GB/s " (15) + "est. XXXXXXXX" (13) = 62
|
||||
stats_len = 62
|
||||
else:
|
||||
# Infinite mode: show written and elapsed
|
||||
stats = (
|
||||
f" \x1b[1;36mRandQuik {spinner}\x1b[0m "
|
||||
f"{written_gb:6.2f} GB \x1b[2m\u221e\x1b[0m "
|
||||
f"\x1b[2m@\x1b[0m {speed_gbs:5.2f} GB/s "
|
||||
f"\x1b[2m\u2502\x1b[0m {format_time(elapsed):>8}"
|
||||
)
|
||||
# Visible: " RandQuik X " (16) + "XXXX.XX GB \u221e " (14) + "@ XX.XX GB/s " (15) + "\u2502 XXXXXXXX" (11) = 56
|
||||
stats_len = 56
|
||||
|
||||
# Build position suffix if oseek was used
|
||||
if self.oseek > 0:
|
||||
file_pos = self.oseek + written
|
||||
pos_str = format_size(file_pos).replace(" ", "")
|
||||
if total_bytes is not None:
|
||||
pos_suffix = f" \x1b[2m[\x1b[0m{pos_str}\x1b[2m]\x1b[0m"
|
||||
pos_suffix_len = 3 + len(pos_str) # " []" + size
|
||||
else:
|
||||
pos_suffix = f" \x1b[2m@\x1b[0m{pos_str}"
|
||||
pos_suffix_len = 2 + len(pos_str) # " @" + size
|
||||
else:
|
||||
pos_suffix = ""
|
||||
pos_suffix_len = 0
|
||||
|
||||
# Calculate available space for filename
|
||||
# Format: {stats} > {name}{pos_suffix}
|
||||
available = cols - stats_len - 4 - pos_suffix_len # 4 for " > "
|
||||
name = self.output_name
|
||||
if len(name) > available > 3:
|
||||
name = "\u2026" + name[-(available - 1) :]
|
||||
elif available <= 3:
|
||||
name = ""
|
||||
|
||||
if name:
|
||||
return f"{stats} \x1b[2m>\x1b[0m {name}{pos_suffix}"
|
||||
return stats
|
||||
|
||||
def _render_progress_block(
|
||||
self, cols: int, rows: int, max_height: int
|
||||
) -> tuple[list[str], tuple[int, int, str] | None]:
|
||||
"""Render the progress block constrained to max_height lines."""
|
||||
if self.infinite:
|
||||
return self._render_infinite_block(cols, rows, max_height), None
|
||||
return self._render_finite_block(cols, rows, max_height)
|
||||
|
||||
def _render_infinite_block(self, cols: int, rows: int, max_height: int) -> list[str]:
|
||||
"""Render progress block for infinite mode (no known total)."""
|
||||
written = self.state.get("written", 0)
|
||||
now = time.perf_counter()
|
||||
elapsed = now - self.start_time
|
||||
|
||||
# Calculate graph width (leave room for Y-axis labels)
|
||||
graph_width = max(10, cols - 8)
|
||||
self._graph_width = graph_width
|
||||
|
||||
# Calculate speeds
|
||||
overall_speed = written / elapsed if elapsed > 0 else 0
|
||||
dt = now - self._last_time
|
||||
instant_speed = (written - self._last_written) / dt if dt > 0 else 0
|
||||
self._last_written = written
|
||||
self._last_time = now
|
||||
|
||||
# Update speed history
|
||||
speed_gbs = instant_speed / 1_000_000_000
|
||||
self._speed_history.append(speed_gbs)
|
||||
self._time_history.append(elapsed)
|
||||
|
||||
# Update max speed
|
||||
if speed_gbs > self._max_speed:
|
||||
self._max_speed = speed_gbs
|
||||
scale_max = self._nice_scale(self._max_speed)
|
||||
|
||||
# Use smoothed speed for header display
|
||||
display_speed = self._get_smoothed_speed()
|
||||
|
||||
# Build output
|
||||
lines: list[str] = []
|
||||
header = self._build_header(cols, written, display_speed, elapsed)
|
||||
lines.append(header)
|
||||
|
||||
# Calculate graph dimensions based on remaining height
|
||||
# Layout: [header][GB/s label][graph rows][time axis]
|
||||
remaining_after_label = max_height - len(lines) - 2
|
||||
|
||||
if remaining_after_label < 1:
|
||||
# Terminal is too short; fall back to header-only view
|
||||
return lines
|
||||
|
||||
graph_rows = max(1, remaining_after_label)
|
||||
|
||||
# Downsample speed history for display - average samples within each column's time range
|
||||
min_scale_time = 10.0
|
||||
scale_time = max(elapsed, min_scale_time)
|
||||
col_width_time = scale_time / (graph_width - 1) if graph_width > 1 else scale_time
|
||||
|
||||
display_values = []
|
||||
for col in range(graph_width):
|
||||
col_time = col / (graph_width - 1) * scale_time if graph_width > 1 else 0
|
||||
if col_time > elapsed:
|
||||
break
|
||||
# Find all samples within this column's time range
|
||||
col_start = col_time - col_width_time / 2
|
||||
col_end = col_time + col_width_time / 2
|
||||
samples = [
|
||||
self._speed_history[idx]
|
||||
for idx, t in enumerate(self._time_history)
|
||||
if col_start <= t <= col_end
|
||||
]
|
||||
if samples:
|
||||
display_values.append(sum(samples) / len(samples))
|
||||
elif self._speed_history:
|
||||
# Fallback to closest if no samples in range
|
||||
best_idx = 0
|
||||
best_diff = float("inf")
|
||||
for idx, t in enumerate(self._time_history):
|
||||
diff = abs(t - col_time)
|
||||
if diff < best_diff:
|
||||
best_diff = diff
|
||||
best_idx = idx
|
||||
display_values.append(self._speed_history[best_idx])
|
||||
|
||||
avg_speed_gbs = overall_speed / 1_000_000_000
|
||||
|
||||
# Use MB/s scale if max speed < 1 GB/s
|
||||
use_mb = scale_max < 1
|
||||
unit_label = "MB/s" if use_mb else "GB/s"
|
||||
lines.append(f" \x1b[36m{unit_label}\x1b[0m")
|
||||
|
||||
# Compute nice Y-axis tick values and map each to its best row
|
||||
nice_ticks = self._nice_y_ticks(scale_max, graph_rows)
|
||||
row_labels = self._assign_ticks_to_rows(nice_ticks, scale_max, graph_rows, use_mb)
|
||||
|
||||
for row in range(graph_rows):
|
||||
graph_line = self._render_graph_row(
|
||||
display_values,
|
||||
scale_max,
|
||||
row,
|
||||
graph_rows,
|
||||
len(display_values),
|
||||
avg_speed_gbs,
|
||||
)
|
||||
graph_line = graph_line.ljust(graph_width)
|
||||
|
||||
# Y-axis label from pre-computed mapping
|
||||
label = row_labels.get(row, " ")
|
||||
lines.append(f" \x1b[36m{label}\x1b[0m \x1b[33m{graph_line}\x1b[0m")
|
||||
|
||||
# Time axis
|
||||
time_axis = self._build_infinite_time_axis(graph_width, scale_time)
|
||||
lines.append(f" {''.join(time_axis)}")
|
||||
|
||||
return lines
|
||||
|
||||
def _nice_scale(self, max_speed: float) -> float:
|
||||
"""Round up to next nice number for scale."""
|
||||
if max_speed <= 0.01:
|
||||
return 0.01
|
||||
log_val = math.log10(max_speed)
|
||||
power = math.floor(log_val)
|
||||
mantissa = max_speed / (10**power)
|
||||
nice_mantissa = math.ceil(mantissa)
|
||||
if nice_mantissa > 9:
|
||||
nice_mantissa = 1
|
||||
power += 1
|
||||
return nice_mantissa * (10**power)
|
||||
|
||||
def _format_label(self, val: float, use_mb: bool = False) -> str:
|
||||
"""Format Y-axis label.
|
||||
|
||||
Args:
|
||||
val: Value in GB/s (will be converted to MB/s if use_mb is True)
|
||||
use_mb: If True, multiply by 1000 and format as MB/s values
|
||||
"""
|
||||
if use_mb:
|
||||
val = val * 1000 # Convert GB/s to MB/s
|
||||
if val == 0:
|
||||
return "0"
|
||||
elif val >= 1:
|
||||
return f"{val:.0f}"
|
||||
else:
|
||||
return f"{val:.1f}"
|
||||
|
||||
def _nice_y_ticks(self, scale_max: float, graph_rows: int = 10) -> list[float]:
|
||||
"""Return nice Y-axis tick values from 0 to scale_max.
|
||||
|
||||
Chooses a nice interval (1, 2, 5 × 10^N) that gives labels with
|
||||
sufficient spacing (at least 3 rows between labels).
|
||||
"""
|
||||
if scale_max <= 0:
|
||||
return [0]
|
||||
|
||||
# We want at least 5 empty rows between labels for readability
|
||||
min_row_spacing = 5
|
||||
max_ticks = max(2, graph_rows // min_row_spacing)
|
||||
|
||||
# Nice intervals: 0.1, 0.2, 0.5, 1, 2, 5, 10, 20, 50, ...
|
||||
nice_bases = [1, 2, 5]
|
||||
best_interval = scale_max
|
||||
for exp in range(-1, 10):
|
||||
for base in nice_bases:
|
||||
interval = base * (10**exp)
|
||||
num_ticks = scale_max / interval
|
||||
if 2 <= num_ticks <= max_ticks:
|
||||
best_interval = interval
|
||||
break
|
||||
else:
|
||||
continue
|
||||
break
|
||||
|
||||
# Generate ticks from 0 to scale_max at best_interval
|
||||
ticks = []
|
||||
val = 0.0
|
||||
while val <= scale_max + 1e-9:
|
||||
ticks.append(val)
|
||||
val += best_interval
|
||||
return ticks
|
||||
|
||||
def _assign_ticks_to_rows(
|
||||
self, ticks: list[float], scale_max: float, graph_rows: int, use_mb: bool = False
|
||||
) -> dict[int, str]:
|
||||
"""Assign each tick to the row closest to its value.
|
||||
|
||||
Returns a dict mapping row index to formatted label string.
|
||||
Each tick is assigned to exactly one row.
|
||||
|
||||
Args:
|
||||
ticks: List of tick values in GB/s
|
||||
scale_max: Maximum scale value in GB/s
|
||||
graph_rows: Number of rows in the graph
|
||||
use_mb: If True, format labels as MB/s instead of GB/s
|
||||
"""
|
||||
row_labels: dict[int, str] = {}
|
||||
if graph_rows <= 1 or scale_max <= 0:
|
||||
return {0: f"{self._format_label(ticks[0] if ticks else 0, use_mb):>4}"}
|
||||
|
||||
for tick in ticks:
|
||||
# Calculate which row this tick value corresponds to
|
||||
# Row 0 is top (scale_max), row graph_rows-1 is bottom (0)
|
||||
exact_row = (1 - tick / scale_max) * (graph_rows - 1)
|
||||
best_row = round(exact_row)
|
||||
best_row = max(0, min(graph_rows - 1, best_row))
|
||||
|
||||
# Only assign if row is not already taken (first tick wins)
|
||||
if best_row not in row_labels:
|
||||
row_labels[best_row] = f"{self._format_label(tick, use_mb):>4}"
|
||||
|
||||
return row_labels
|
||||
|
||||
def _build_infinite_time_axis(self, graph_width: int, scale_time: float) -> list[str]:
|
||||
"""Build time axis for infinite mode with nice interval labels.
|
||||
|
||||
Shows from 0 to scale_time with nice interval markers.
|
||||
"""
|
||||
time_axis = [" "] * graph_width
|
||||
|
||||
# Nice time intervals
|
||||
nice_intervals = [
|
||||
1,
|
||||
2,
|
||||
5,
|
||||
10,
|
||||
15,
|
||||
30,
|
||||
60,
|
||||
120,
|
||||
300,
|
||||
600,
|
||||
900,
|
||||
1800,
|
||||
3600,
|
||||
7200,
|
||||
18000,
|
||||
36000,
|
||||
]
|
||||
|
||||
def format_time_short(secs):
|
||||
"""Format time for axis label."""
|
||||
if secs == 0:
|
||||
return "0"
|
||||
elif secs < 120:
|
||||
return f"{int(secs)}s"
|
||||
elif secs < 3600:
|
||||
m = int(secs // 60)
|
||||
s = int(secs % 60)
|
||||
if s == 0:
|
||||
return f"{m}m"
|
||||
return f"{m}m{s}s"
|
||||
else:
|
||||
h = int(secs // 3600)
|
||||
m = int((secs % 3600) // 60)
|
||||
if m == 0:
|
||||
return f"{h}h"
|
||||
return f"{h}h{m}m"
|
||||
|
||||
# Find a nice interval that gives us ~4-8 labels
|
||||
interval = nice_intervals[-1]
|
||||
for ni in nice_intervals:
|
||||
if scale_time / ni <= 8:
|
||||
interval = ni
|
||||
break
|
||||
|
||||
# Place labels at nice intervals starting from 0
|
||||
t = 0
|
||||
while t <= scale_time:
|
||||
col = int(t / scale_time * (graph_width - 1)) if scale_time > 0 else 0
|
||||
if 0 <= col < graph_width:
|
||||
label = format_time_short(t)
|
||||
label_start = max(0, col - len(label) // 2)
|
||||
label_end = min(graph_width, label_start + len(label))
|
||||
if all(c == " " for c in time_axis[label_start:label_end]):
|
||||
for i, ch in enumerate(label):
|
||||
if label_start + i < graph_width:
|
||||
time_axis[label_start + i] = ch
|
||||
t += interval
|
||||
|
||||
return time_axis
|
||||
|
||||
def _render_finite_block(
|
||||
self, cols: int, rows: int, max_height: int
|
||||
) -> tuple[list[str], tuple[int, int, str] | None]:
|
||||
"""Render the progress block for finite progress."""
|
||||
written = self.state.get("written", 0)
|
||||
now = time.perf_counter()
|
||||
elapsed = now - self.start_time
|
||||
|
||||
# Calculate graph width (leave room for Y-axis labels)
|
||||
graph_width = max(10, cols - 8)
|
||||
self._graph_width = graph_width
|
||||
|
||||
# Calculate speeds
|
||||
overall_speed = written / elapsed if elapsed > 0 else 0
|
||||
dt = now - self._last_time
|
||||
instant_speed = (written - self._last_written) / dt if dt > 0 else 0
|
||||
self._last_written = written
|
||||
self._last_time = now
|
||||
|
||||
# Collect time-based speed samples (like infinite mode)
|
||||
speed_gbs = instant_speed / 1_000_000_000
|
||||
self._speed_history.append(speed_gbs)
|
||||
self._time_history.append(elapsed)
|
||||
|
||||
# Use smoothed speed for header display and ETA
|
||||
display_speed = self._get_smoothed_speed()
|
||||
|
||||
# ETA and total estimated time
|
||||
remaining = self.total_bytes - written
|
||||
# ETA based on smoothed speed for stability
|
||||
eta = remaining / display_speed if display_speed > 0 else -1
|
||||
# Graph X-axis scaling based on overall average speed for stability
|
||||
avg_eta = remaining / overall_speed if overall_speed > 0 else -1
|
||||
estimated_total_time = elapsed + avg_eta if avg_eta > 0 else elapsed
|
||||
|
||||
# Apply hysteresis to scale_time to prevent jumping
|
||||
# Only update if change is significant (>20%) or if new estimate is larger
|
||||
raw_scale_time = max(estimated_total_time, 1.0)
|
||||
if self._smoothed_scale_time is None:
|
||||
self._smoothed_scale_time = raw_scale_time
|
||||
else:
|
||||
# Always grow immediately, shrink only gradually
|
||||
if raw_scale_time > self._smoothed_scale_time:
|
||||
self._smoothed_scale_time = raw_scale_time
|
||||
else:
|
||||
# Shrink slowly: blend 90% old, 10% new
|
||||
self._smoothed_scale_time = 0.9 * self._smoothed_scale_time + 0.1 * raw_scale_time
|
||||
scale_time = self._smoothed_scale_time
|
||||
|
||||
# Progress percentage (for display)
|
||||
pct = min(100, written * 100 / self.total_bytes) if self.total_bytes > 0 else 0
|
||||
|
||||
# Update max speed
|
||||
if speed_gbs > self._max_speed:
|
||||
self._max_speed = speed_gbs
|
||||
scale_max = self._nice_scale(self._max_speed)
|
||||
|
||||
# Build output
|
||||
lines: list[str] = []
|
||||
header = self._build_header(
|
||||
cols, written, display_speed, elapsed, eta=eta, total_bytes=self.total_bytes
|
||||
)
|
||||
lines.append(header)
|
||||
|
||||
# Calculate graph dimensions based on remaining height
|
||||
# Layout: [header][GB/s label][graph rows][time axis]
|
||||
remaining_after_label = max_height - len(lines) - 2
|
||||
|
||||
if remaining_after_label < 1:
|
||||
# Terminal is too short; fall back to a compact header-only view
|
||||
return lines, None
|
||||
|
||||
graph_rows = max(1, remaining_after_label)
|
||||
|
||||
# Downsample speed history for display - average samples within each column's time range
|
||||
# Graph fills based on elapsed / scale_time (smoothed)
|
||||
col_width_time = scale_time / (graph_width - 1) if graph_width > 1 else scale_time
|
||||
|
||||
# Calculate how many columns should be filled based on elapsed time
|
||||
if scale_time > 0:
|
||||
time_pct = min(100, elapsed * 100 / scale_time)
|
||||
else:
|
||||
time_pct = 100
|
||||
target_cols = min(graph_width, int(graph_width * time_pct / 100) + 1) if time_pct > 0 else 0
|
||||
|
||||
display_values = []
|
||||
for col in range(target_cols):
|
||||
col_time = col / (graph_width - 1) * scale_time if graph_width > 1 else 0
|
||||
if col_time > elapsed:
|
||||
break
|
||||
# Find all samples within this column's time range
|
||||
col_start = col_time - col_width_time / 2
|
||||
col_end = col_time + col_width_time / 2
|
||||
samples = [
|
||||
self._speed_history[idx]
|
||||
for idx, t in enumerate(self._time_history)
|
||||
if col_start <= t <= col_end
|
||||
]
|
||||
if samples:
|
||||
display_values.append(sum(samples) / len(samples))
|
||||
elif self._speed_history:
|
||||
# Fallback to closest if no samples in range
|
||||
best_idx = 0
|
||||
best_diff = float("inf")
|
||||
for idx, t in enumerate(self._time_history):
|
||||
diff = abs(t - col_time)
|
||||
if diff < best_diff:
|
||||
best_diff = diff
|
||||
best_idx = idx
|
||||
display_values.append(self._speed_history[best_idx])
|
||||
|
||||
# Render graph rows with Y-axis
|
||||
avg_speed_gbs = overall_speed / 1_000_000_000
|
||||
|
||||
# Use MB/s scale if max speed < 1 GB/s
|
||||
use_mb = scale_max < 1
|
||||
unit_label = "MB/s" if use_mb else "GB/s"
|
||||
lines.append(f" \x1b[36m{unit_label}\x1b[0m")
|
||||
|
||||
# Compute nice Y-axis tick values and map each to its best row
|
||||
nice_ticks = self._nice_y_ticks(scale_max, graph_rows)
|
||||
row_labels = self._assign_ticks_to_rows(nice_ticks, scale_max, graph_rows, use_mb)
|
||||
|
||||
for row in range(graph_rows):
|
||||
graph_line = self._render_graph_row(
|
||||
display_values,
|
||||
scale_max,
|
||||
row,
|
||||
graph_rows,
|
||||
graph_width,
|
||||
avg_speed_gbs,
|
||||
)
|
||||
graph_line = graph_line.ljust(graph_width)
|
||||
# Y-axis label from pre-computed mapping
|
||||
label = row_labels.get(row, " ")
|
||||
lines.append(f" \x1b[36m{label}\x1b[0m \x1b[33m{graph_line}\x1b[0m")
|
||||
|
||||
# Time labels on X-axis with nice intervals
|
||||
time_axis = self._build_time_axis(graph_width, scale_time)
|
||||
lines.append(f" {''.join(time_axis)}")
|
||||
|
||||
# Position percentage at top of bar at current progress point
|
||||
# Find the height of the bar at current progress (last value in display_values)
|
||||
current_speed_gbs = display_values[-1] if display_values else 0
|
||||
# Normalize to find which row the top of the bar is at
|
||||
# Bar fills from bottom up; row 0 is top, graph_rows-1 is bottom
|
||||
if scale_max > 0 and current_speed_gbs > 0:
|
||||
# How many rows from bottom does the bar fill?
|
||||
bar_height_fraction = current_speed_gbs / scale_max
|
||||
# The top of the bar is at this row (0 = top, graph_rows-1 = bottom)
|
||||
bar_top_row = int(graph_rows * (1 - bar_height_fraction))
|
||||
bar_top_row = max(0, min(graph_rows - 1, bar_top_row))
|
||||
else:
|
||||
bar_top_row = graph_rows - 1 # At bottom if no speed
|
||||
|
||||
# Overlay position inside the progress block
|
||||
progress_col = int(pct / 100 * (graph_width - 1)) + 7
|
||||
pct_label = f"{int(pct)}%"
|
||||
pct_col = max(7, min(cols, progress_col - len(pct_label) // 2))
|
||||
first_graph_row = len(lines) - (graph_rows + 2)
|
||||
pct_row_offset = first_graph_row + bar_top_row
|
||||
pct_position = (pct_row_offset, pct_col, f"\x1b[1;37m{pct_label}\x1b[0m")
|
||||
|
||||
return lines, pct_position
|
||||
|
||||
def _build_time_axis(self, graph_width: int, estimated_total_time: float) -> list[str]:
|
||||
"""Build time axis with nice interval labels."""
|
||||
|
||||
def nice_time_interval(total_secs):
|
||||
"""Return a nice interval for time axis labels."""
|
||||
nice_intervals = [
|
||||
1,
|
||||
2,
|
||||
5,
|
||||
10,
|
||||
15,
|
||||
30,
|
||||
60,
|
||||
120,
|
||||
300,
|
||||
600,
|
||||
900,
|
||||
1800,
|
||||
3600,
|
||||
7200,
|
||||
18000,
|
||||
36000,
|
||||
]
|
||||
for interval in nice_intervals:
|
||||
if total_secs / interval <= 8:
|
||||
return interval
|
||||
return 36000
|
||||
|
||||
def format_time_short(secs):
|
||||
"""Format time for axis label."""
|
||||
if secs == 0:
|
||||
return "0s"
|
||||
elif secs < 120:
|
||||
return f"{int(secs)}s"
|
||||
elif secs < 3600:
|
||||
m = int(secs // 60)
|
||||
s = int(secs % 60)
|
||||
if s == 0:
|
||||
return f"{m}m"
|
||||
return f"{m}m{s}s"
|
||||
else:
|
||||
h = int(secs // 3600)
|
||||
m = int((secs % 3600) // 60)
|
||||
if m == 0:
|
||||
return f"{h}h"
|
||||
return f"{h}h{m}m"
|
||||
|
||||
time_axis = [" "] * graph_width
|
||||
if estimated_total_time > 0:
|
||||
interval = nice_time_interval(estimated_total_time)
|
||||
t = 0
|
||||
while t <= estimated_total_time:
|
||||
col = (
|
||||
int(t / estimated_total_time * (graph_width - 1))
|
||||
if estimated_total_time > 0
|
||||
else 0
|
||||
)
|
||||
if col < graph_width:
|
||||
label = format_time_short(t)
|
||||
start = max(0, col - len(label) // 2)
|
||||
end = min(graph_width, start + len(label))
|
||||
if all(c == " " for c in time_axis[start:end]):
|
||||
for i, ch in enumerate(label):
|
||||
if start + i < graph_width:
|
||||
time_axis[start + i] = ch
|
||||
t += interval
|
||||
else:
|
||||
time_axis = list(f"{'0s':<{graph_width}}")
|
||||
|
||||
return time_axis
|
||||
|
||||
def _render_frame(self) -> tuple[int, int, list[str], tuple[int, int, str] | None]:
|
||||
cols, rows = self._get_terminal_size()
|
||||
max_height = min(MAX_HEIGHT, rows) if rows > 0 else MAX_HEIGHT
|
||||
lines, overlay = self._render_progress_block(cols, rows, max_height)
|
||||
# Guard against empty renders
|
||||
if not lines:
|
||||
lines = [""]
|
||||
return cols, rows, lines, overlay
|
||||
|
||||
def _draw_frame(
|
||||
self, cols: int, rows: int, lines: list[str], overlay: tuple[int, int, str] | None
|
||||
):
|
||||
"""Draw the progress block at the bottom of the terminal."""
|
||||
height = min(len(lines), max(1, rows))
|
||||
progress_top = max(1, rows - height + 1)
|
||||
|
||||
# Build entire frame as a single string
|
||||
buf: list[str] = []
|
||||
|
||||
# On first draw, scroll terminal up to make room for progress block
|
||||
if self._first_draw:
|
||||
self._first_draw = False
|
||||
buf.append("\n" * height)
|
||||
|
||||
# Update scrolling region so other output scrolls above the progress block
|
||||
top = 1
|
||||
bottom = max(1, rows - height)
|
||||
if self._current_scroll_bottom != bottom:
|
||||
buf.append(f"\x1b[{top};{bottom}r")
|
||||
self._current_scroll_bottom = bottom
|
||||
|
||||
# Paint each progress line
|
||||
for idx in range(height):
|
||||
row = progress_top + idx
|
||||
line = lines[idx]
|
||||
buf.append(f"\x1b[{row};1H\x1b[2K{line}")
|
||||
|
||||
# Overlay (e.g., percent marker)
|
||||
if overlay:
|
||||
row_offset, col, text = overlay
|
||||
abs_row = progress_top + row_offset
|
||||
abs_col = max(1, min(cols, col))
|
||||
buf.append(f"\x1b[{abs_row};{abs_col}H{text}")
|
||||
|
||||
# Place cursor back at the bottom of the scrolling region
|
||||
anchor_row = max(1, progress_top - 1)
|
||||
buf.append(f"\x1b[{anchor_row};1H")
|
||||
|
||||
# Single atomic write
|
||||
sys.stderr.write("".join(buf))
|
||||
sys.stderr.flush()
|
||||
|
||||
def _run(self):
|
||||
"""Background thread: update display every 100ms."""
|
||||
while not self._stop.wait(0.1):
|
||||
cols, rows, lines, overlay = self._render_frame()
|
||||
self._draw_frame(cols, rows, lines, overlay)
|
||||
@@ -0,0 +1,269 @@
|
||||
"""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 < 120:
|
||||
return f"{int(seconds)}s"
|
||||
elif seconds < 3600:
|
||||
m = int(seconds // 60)
|
||||
s = int(seconds % 60)
|
||||
if s == 0:
|
||||
return f"{m}m"
|
||||
return f"{m}m{s}s"
|
||||
elif seconds < 172800: # 48 hours
|
||||
h = int(seconds // 3600)
|
||||
m = int((seconds % 3600) // 60)
|
||||
if m == 0:
|
||||
return f"{h}h"
|
||||
return f"{h}h{m}m"
|
||||
else:
|
||||
d = int(seconds // 86400)
|
||||
h = int((seconds % 86400) // 3600)
|
||||
if h == 0:
|
||||
return f"{d}d"
|
||||
return f"{d}d{h}h"
|
||||
|
||||
|
||||
@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"\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} ms"
|
||||
|
||||
# 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",
|
||||
sum(s.crypto_time for s in worker_stats),
|
||||
[ms(s.crypto_time) for s in worker_stats],
|
||||
),
|
||||
(
|
||||
"lock_acq",
|
||||
sum(s.lock_acquire_time for s in worker_stats),
|
||||
[ms(s.lock_acquire_time) for s in worker_stats],
|
||||
),
|
||||
(
|
||||
"wait_sp",
|
||||
sum(s.lock_wait_space_time for s in worker_stats),
|
||||
[ms(s.lock_wait_space_time) for s in worker_stats],
|
||||
),
|
||||
(
|
||||
"claim",
|
||||
sum(s.lock_claim_time for s in worker_stats),
|
||||
[ms(s.lock_claim_time) for s in worker_stats],
|
||||
),
|
||||
(
|
||||
"notify",
|
||||
sum(s.lock_notify_time for s in worker_stats),
|
||||
[ms(s.lock_notify_time) for s in worker_stats],
|
||||
),
|
||||
]
|
||||
|
||||
timing_rows.append(("total", 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)
|
||||
@@ -0,0 +1,182 @@
|
||||
"""Utility functions for formatting and parsing."""
|
||||
|
||||
import pathlib
|
||||
import re
|
||||
import sys
|
||||
|
||||
__all__ = [
|
||||
"get_output_size",
|
||||
"get_sector_size",
|
||||
"parse_size",
|
||||
"sparse_range",
|
||||
]
|
||||
|
||||
|
||||
# Cache for sector size lookup (path -> size)
|
||||
_sector_size_cache: dict[str, int] = {}
|
||||
|
||||
|
||||
def get_sector_size(path: str | pathlib.Path) -> int:
|
||||
"""Get sector size for a block device, or 512 as fallback."""
|
||||
import os
|
||||
import stat
|
||||
|
||||
try:
|
||||
st = pathlib.Path(path).stat()
|
||||
if not stat.S_ISBLK(st.st_mode):
|
||||
return 512
|
||||
except OSError:
|
||||
return 512
|
||||
|
||||
try:
|
||||
fd = os.open(str(path), os.O_RDONLY)
|
||||
try:
|
||||
import fcntl
|
||||
import struct
|
||||
|
||||
if sys.platform == "darwin":
|
||||
# macOS: DKIOCGETBLOCKSIZE = 0x40046418
|
||||
DKIOCGETBLOCKSIZE = 0x40046418
|
||||
buf = fcntl.ioctl(fd, DKIOCGETBLOCKSIZE, b"\x00" * 4)
|
||||
return struct.unpack("I", buf)[0]
|
||||
else:
|
||||
# Linux: BLKSSZGET = 0x1268
|
||||
BLKSSZGET = 0x1268
|
||||
buf = fcntl.ioctl(fd, BLKSSZGET, b"\x00" * 4)
|
||||
return struct.unpack("i", buf)[0]
|
||||
finally:
|
||||
os.close(fd)
|
||||
except (OSError, ImportError, Exception):
|
||||
return 512
|
||||
|
||||
|
||||
def get_output_size(path: str | pathlib.Path | None) -> int | None:
|
||||
"""Get the size of output file or block device.
|
||||
|
||||
Returns None for stdout or non-existent files.
|
||||
Returns the size in bytes for existing files or block devices.
|
||||
"""
|
||||
import os
|
||||
import stat
|
||||
|
||||
if not path:
|
||||
return None # stdout
|
||||
|
||||
try:
|
||||
st = pathlib.Path(path).stat()
|
||||
except OSError:
|
||||
return None # doesn't exist yet
|
||||
|
||||
if stat.S_ISBLK(st.st_mode):
|
||||
# Block device - get size via ioctl
|
||||
try:
|
||||
fd = os.open(str(path), os.O_RDONLY)
|
||||
try:
|
||||
import fcntl
|
||||
import struct
|
||||
|
||||
if sys.platform == "darwin":
|
||||
# macOS: DKIOCGETBLOCKCOUNT and DKIOCGETBLOCKSIZE
|
||||
DKIOCGETBLOCKCOUNT = 0x40086419
|
||||
DKIOCGETBLOCKSIZE = 0x40046418
|
||||
count_buf = fcntl.ioctl(fd, DKIOCGETBLOCKCOUNT, b"\x00" * 8)
|
||||
size_buf = fcntl.ioctl(fd, DKIOCGETBLOCKSIZE, b"\x00" * 4)
|
||||
block_count = struct.unpack("Q", count_buf)[0]
|
||||
block_size = struct.unpack("I", size_buf)[0]
|
||||
return block_count * block_size
|
||||
else:
|
||||
# Linux: BLKGETSIZE64 = 0x80081272
|
||||
BLKGETSIZE64 = 0x80081272
|
||||
buf = fcntl.ioctl(fd, BLKGETSIZE64, b"\x00" * 8)
|
||||
return struct.unpack("Q", buf)[0]
|
||||
finally:
|
||||
os.close(fd)
|
||||
except (OSError, ImportError, Exception):
|
||||
return None
|
||||
elif stat.S_ISREG(st.st_mode):
|
||||
return st.st_size
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def parse_size(length: str | None, output_path: str | pathlib.Path | None = None) -> int | None:
|
||||
"""Parse size string with SI/IEC prefixes.
|
||||
|
||||
Supports:
|
||||
- Plain numbers: 1000, 1_000_000
|
||||
- SI prefixes: k, m, g, t, p (powers of 1000)
|
||||
- IEC prefixes: ki, mi, gi, ti, pi (powers of 1024)
|
||||
- Optional 'b' suffix: kb, kib, mb, mib, etc.
|
||||
- Special unit 'sect' = device sector size (detected, fallback 512)
|
||||
- Case insensitive
|
||||
|
||||
Examples: 1k, 1ki, 1kb, 1kib, 100m, 100mi, 1g, 1gi, 10sect
|
||||
"""
|
||||
if length is None:
|
||||
return None
|
||||
s = length.strip().lower().replace("_", "")
|
||||
|
||||
# Handle sect unit
|
||||
m = re.match(r"^(\d+)\s*sect?s?$", s)
|
||||
if m:
|
||||
if output_path:
|
||||
path_str = str(output_path)
|
||||
if path_str not in _sector_size_cache:
|
||||
_sector_size_cache[path_str] = get_sector_size(output_path)
|
||||
sector_size = _sector_size_cache.get(path_str, 512)
|
||||
else:
|
||||
sector_size = 512
|
||||
return int(m.group(1)) * sector_size
|
||||
|
||||
# SI/IEC prefixes
|
||||
si_prefixes = {"k": 1000, "m": 1000**2, "g": 1000**3, "t": 1000**4, "p": 1000**5}
|
||||
iec_prefixes = {
|
||||
"ki": 1024,
|
||||
"mi": 1024**2,
|
||||
"gi": 1024**3,
|
||||
"ti": 1024**4,
|
||||
"pi": 1024**5,
|
||||
}
|
||||
|
||||
# Try IEC first (ki, mi, etc.) - must check before SI
|
||||
m = re.match(r"^(\d+(?:\.\d+)?)\s*(ki|mi|gi|ti|pi)b?$", s)
|
||||
if m:
|
||||
num, prefix = m.groups()
|
||||
return int(float(num) * iec_prefixes[prefix])
|
||||
|
||||
# Try SI (k, m, g, etc.)
|
||||
m = re.match(r"^(\d+(?:\.\d+)?)\s*([kmgtp])b?$", s)
|
||||
if m:
|
||||
num, prefix = m.groups()
|
||||
return int(float(num) * si_prefixes[prefix])
|
||||
|
||||
# Plain number
|
||||
m = re.match(r"^(\d+)$", s)
|
||||
if m:
|
||||
return int(m.group(1))
|
||||
|
||||
raise ValueError(f"Invalid size format: {length}")
|
||||
|
||||
|
||||
def sparse_range(n: int, max_items: int = 9) -> list[int]:
|
||||
"""Generate a sparse range from 1 to N for benchmarking thread counts."""
|
||||
if n < 1:
|
||||
return [1]
|
||||
if n <= max_items - 1:
|
||||
return list(range(n + 1))
|
||||
|
||||
keep = 3 # dense prefix: 1,2,3
|
||||
out = list(range(keep + 1))
|
||||
|
||||
remaining = max_items - keep
|
||||
step = max(1, n // (remaining - 1))
|
||||
|
||||
for k in range(1, remaining):
|
||||
v = k * step
|
||||
if v > out[-1]:
|
||||
out.append(v)
|
||||
|
||||
if out[-1] != n:
|
||||
out[-1] = n
|
||||
|
||||
return out
|
||||
@@ -0,0 +1,332 @@
|
||||
"""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:
|
||||
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)
|
||||
# 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
|
||||
|
||||
# Commit block (mark ready + notify consumer)
|
||||
with self._lock:
|
||||
self.ready[slot] = True
|
||||
self.has_data.notify()
|
||||
|
||||
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,
|
||||
output_name=output,
|
||||
oseek=oseek,
|
||||
)
|
||||
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,
|
||||
)
|
||||
@@ -1,17 +0,0 @@
|
||||
import os
|
||||
from distutils.core import setup
|
||||
|
||||
import numpy
|
||||
from Cython.Build import cythonize
|
||||
from setuptools import Extension
|
||||
|
||||
os.environ["CFLAGS"] = "-O3 -march=native -Wall -Wextra"
|
||||
extensions = [
|
||||
Extension(
|
||||
"nprand",
|
||||
["src/nprand.pyx"],
|
||||
include_dirs=[numpy.get_include()],
|
||||
define_macros=[("NPY_NO_DEPRECATED_API", "NPY_1_7_API_VERSION")],
|
||||
)
|
||||
]
|
||||
setup(ext_modules=cythonize(extensions))
|
||||
-29
@@ -1,29 +0,0 @@
|
||||
#include <stdint.h>
|
||||
#include <string.h>
|
||||
|
||||
// clang-format off
|
||||
#define QUARTERSTEP(a, b, c, n) a += b; c ^= a; c = (c << n) | (c >> (32 - n))
|
||||
#define QUARTERROUND(a, b, c, d) {\
|
||||
QUARTERSTEP(a, b, d, 16); QUARTERSTEP(c, d, b, 12); \
|
||||
QUARTERSTEP(a, b, d, 8); QUARTERSTEP(c, d, b, 7); }
|
||||
|
||||
static inline uint64_t _cha_block(uint8_t* buf, size_t bufsize, uint32_t state[16], unsigned rounds) {
|
||||
unsigned blocks = bufsize / 64;
|
||||
uint32_t* out = (uint32_t*)buf;
|
||||
uint32_t x[16];
|
||||
for (unsigned b = blocks; b-->0;) {
|
||||
for (unsigned i = 0; i < 16; ++i) x[i] = state[i]; // Faster than memcpy
|
||||
for (unsigned i = rounds / 2; i-->0;) {
|
||||
// Mix columns, then diagonals
|
||||
for (unsigned j = 0; j < 4; ++j) QUARTERROUND(x[j], x[4 + j], x[8 + j], x[12 + j]);
|
||||
for (unsigned j = 0; j < 4; ++j) QUARTERROUND(x[j], x[4 + (j+1)%4], x[8 + (j+2)%4], x[12 + (j+3)%4]);
|
||||
}
|
||||
for (unsigned i = 0; i < 16; ++i) *out++ = x[i] + state[i];
|
||||
++*(uint64_t*)(state + 12); // Increment counter
|
||||
}
|
||||
memset(x, 0, sizeof x);
|
||||
return blocks * CHA_BLOCK_SIZE;
|
||||
}
|
||||
|
||||
#undef QUARTERROUND
|
||||
#undef QUARTERSTEP
|
||||
-106
@@ -1,106 +0,0 @@
|
||||
#include <arm_neon.h>
|
||||
#include <stdint.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
// clang-format off
|
||||
|
||||
#define VEC4_ROT(A, IMM) \
|
||||
vreinterpretq_u32_u8(vorrq_u8(vshlq_n_u32(A, IMM), vshrq_n_u32(A, 32 - IMM)))
|
||||
|
||||
|
||||
/* same, but replace 2 of the shift/shift/or "rotation" by byte shuffles (8 &
|
||||
* 16) (better) */
|
||||
#define VEC4_QUARTERROUND(A, B, C, D) \
|
||||
x[A] = vaddq_u32(x[A], x[B]); \
|
||||
x[D] = vqtbl1q_u8(veorq_u32(x[D], x[A]), rot16); \
|
||||
x[C] = vaddq_u32(x[C], x[D]); \
|
||||
x[B] = VEC4_ROT(veorq_u32(x[B], x[C]), 12); \
|
||||
x[A] = vaddq_u32(x[A], x[B]); \
|
||||
x[D] = vqtbl1q_u8(veorq_u32(x[D], x[A]), rot8); \
|
||||
x[C] = vaddq_u32(x[C], x[D]); \
|
||||
x[B] = VEC4_ROT(veorq_u32(x[B], x[C]), 7)
|
||||
|
||||
#define ONEQUAD(A, B, C, D, OUT) \
|
||||
{ \
|
||||
/* Add original block */ \
|
||||
x[A] = vaddq_u32(x[A], orig[A]); \
|
||||
x[B] = vaddq_u32(x[B], orig[B]); \
|
||||
x[C] = vaddq_u32(x[C], orig[C]); \
|
||||
x[D] = vaddq_u32(x[D], orig[D]); \
|
||||
/* Transpose */ \
|
||||
uint32x4x2_t ab = vtrnq_u32(x[A], x[B]); \
|
||||
uint32x4x2_t cd = vtrnq_u32(x[C], x[D]); \
|
||||
x[A] = vcombine_u32(vget_low_u32(ab.val[0]), vget_low_u32(cd.val[0])); \
|
||||
x[B] = vcombine_u32(vget_low_u32(ab.val[1]), vget_low_u32(cd.val[1])); \
|
||||
x[C] = vcombine_u32(vget_high_u32(ab.val[0]), vget_high_u32(cd.val[0])); \
|
||||
x[D] = vcombine_u32(vget_high_u32(ab.val[1]), vget_high_u32(cd.val[1])); \
|
||||
/* Write out 1/4 of each block */ \
|
||||
vst1q_u32((uint32_t*)(OUT), x[A]); \
|
||||
vst1q_u32((uint32_t*)(OUT + 64), x[B]); \
|
||||
vst1q_u32((uint32_t*)(OUT + 128), x[C]); \
|
||||
vst1q_u32((uint32_t*)(OUT + 192), x[D]); \
|
||||
}
|
||||
|
||||
#define COUNTER_INCREMENT(addv) \
|
||||
{ \
|
||||
orig[12] = vaddq_u32(orig[12], addv); \
|
||||
orig[13] = vaddq_u32(orig[13], vshrq_n_u32(vcltq_u32(orig[12], addv), 31)); \
|
||||
}
|
||||
|
||||
static inline uint64_t
|
||||
_cha_4block(uint8_t* buf, size_t bufsize, uint32_t state[16], unsigned rounds) {
|
||||
/* constant for shuffling bytes (replacing multiple-of-8 rotates) */
|
||||
const uint8x16_t rot16 = {
|
||||
2, 3, 0, 1,
|
||||
6, 7, 4, 5,
|
||||
10, 11, 8, 9,
|
||||
14, 15, 12, 13
|
||||
};
|
||||
const uint8x16_t rot8= {
|
||||
3, 0, 1, 2,
|
||||
7, 4, 5, 6,
|
||||
11, 8, 9, 10,
|
||||
15, 12, 13, 14
|
||||
};
|
||||
// Load state to vectors, duplicate four times, only different counters
|
||||
uint32x4_t orig[16];
|
||||
for (unsigned i = 0; i < 16; ++i) orig[i] = vdupq_n_u32(state[i]);
|
||||
uint32x4_t addv = { 0, 1, 2, 3 };
|
||||
COUNTER_INCREMENT(addv);
|
||||
addv = vdupq_n_u32(4);
|
||||
const unsigned batches = bufsize / 256;
|
||||
for (unsigned b = batches; b-->0;) {
|
||||
uint32x4_t x[16];
|
||||
for (unsigned i = 0; i < 16; ++i) x[i] = orig[i];
|
||||
for (unsigned r = rounds / 2; r-->0;) {
|
||||
// Mix columns
|
||||
VEC4_QUARTERROUND(0, 4, 8, 12);
|
||||
VEC4_QUARTERROUND(1, 5, 9, 13);
|
||||
VEC4_QUARTERROUND(2, 6, 10, 14);
|
||||
VEC4_QUARTERROUND(3, 7, 11, 15);
|
||||
// Mix diagonals
|
||||
VEC4_QUARTERROUND(0, 5, 10, 15);
|
||||
VEC4_QUARTERROUND(1, 6, 11, 12);
|
||||
VEC4_QUARTERROUND(2, 7, 8, 13);
|
||||
VEC4_QUARTERROUND(3, 4, 9, 14);
|
||||
}
|
||||
// Add original block, unpack output
|
||||
ONEQUAD(0, 1, 2, 3, buf);
|
||||
ONEQUAD(4, 5, 6, 7, buf + 16);
|
||||
ONEQUAD(8, 9, 10, 11, buf + 32);
|
||||
ONEQUAD(12, 13, 14, 15, buf + 48);
|
||||
COUNTER_INCREMENT(addv);
|
||||
buf += 256;
|
||||
}
|
||||
// Store counter
|
||||
state[12] = vgetq_lane_u32(orig[12], 0);
|
||||
state[13] = vgetq_lane_u32(orig[13], 0);
|
||||
return batches * 256;
|
||||
}
|
||||
|
||||
#undef COUNTER_INCREMENT
|
||||
#undef ONEQUAD
|
||||
#undef ONEQUAD_TRANSPOSE
|
||||
#undef VEC4_ROT
|
||||
#undef VEC4_QUARTERROUND
|
||||
#undef VEC4_QUARTERROUND_SHUFFLE
|
||||
-104
@@ -1,104 +0,0 @@
|
||||
#if defined(__x86_64__)
|
||||
#include <emmintrin.h> // SSE2
|
||||
#include <tmmintrin.h> // SSSE3
|
||||
#elif defined(__aarch64__)
|
||||
#include "sse2neon.h"
|
||||
#endif
|
||||
#include <stdio.h>
|
||||
// clang-format off
|
||||
|
||||
#define VEC4_ROT(A, IMM) \
|
||||
_mm_or_si128(_mm_slli_epi32(A, IMM), _mm_srli_epi32(A, (32 - IMM)))
|
||||
|
||||
/* same, but replace 2 of the shift/shift/or "rotation" by byte shuffles (8 &
|
||||
* 16) (better) */
|
||||
#define VEC4_QUARTERROUND(A, B, C, D) \
|
||||
x[A] = _mm_add_epi32(x[A], x[B]); \
|
||||
x[D] = _mm_shuffle_epi8(_mm_xor_si128(x[D], x[A]), rot16); \
|
||||
x[C] = _mm_add_epi32(x[C], x[D]); \
|
||||
x[B] = VEC4_ROT(_mm_xor_si128(x[B], x[C]), 12); \
|
||||
x[A] = _mm_add_epi32(x[A], x[B]); \
|
||||
x[D] = _mm_shuffle_epi8(_mm_xor_si128(x[D], x[A]), rot8); \
|
||||
x[C] = _mm_add_epi32(x[C], x[D]); \
|
||||
x[B] = VEC4_ROT(_mm_xor_si128(x[B], x[C]), 7)
|
||||
|
||||
#define ONEQUAD(A, B, C, D, OUT) \
|
||||
{ \
|
||||
/* Add original block */ \
|
||||
x[A] = _mm_add_epi32(x[A], orig[A]); \
|
||||
x[B] = _mm_add_epi32(x[B], orig[B]); \
|
||||
x[C] = _mm_add_epi32(x[C], orig[C]); \
|
||||
x[D] = _mm_add_epi32(x[D], orig[D]); \
|
||||
/* Transpose */ \
|
||||
__m128i abl = _mm_unpacklo_epi32(x[A], x[B]); \
|
||||
__m128i cdl = _mm_unpacklo_epi32(x[C], x[D]); \
|
||||
__m128i abh = _mm_unpackhi_epi32(x[A], x[B]); \
|
||||
__m128i cdh = _mm_unpackhi_epi32(x[C], x[D]); \
|
||||
x[A] = _mm_unpacklo_epi64(abl, cdl); /* a0 b0 c0 d0 */ \
|
||||
x[B] = _mm_unpackhi_epi64(abl, cdl); /* a1 b1 c1 d1 */ \
|
||||
x[C] = _mm_unpacklo_epi64(abh, cdh); /* a2 b2 c2 d2 */ \
|
||||
x[D] = _mm_unpackhi_epi64(abh, cdh); /* a3 b3 c3 d3 */ \
|
||||
/* Write out 1/4 of each block */ \
|
||||
_mm_storeu_si128((__m128i*)(OUT), x[A]); \
|
||||
_mm_storeu_si128((__m128i*)(OUT + 64), x[B]); \
|
||||
_mm_storeu_si128((__m128i*)(OUT + 128), x[C]); \
|
||||
_mm_storeu_si128((__m128i*)(OUT + 192), x[D]); \
|
||||
}
|
||||
|
||||
#define COUNTER_INCREMENT(addv) \
|
||||
{ \
|
||||
__m128i carry = orig[12]; \
|
||||
orig[12] = _mm_add_epi32(orig[12], addv); \
|
||||
carry = _mm_srli_epi32(_mm_and_si128(_mm_xor_si128(orig[12], carry), carry), 31); \
|
||||
orig[13] = _mm_add_epi32(orig[13], carry); \
|
||||
}
|
||||
|
||||
static inline uint64_t
|
||||
_cha_4block(uint8_t* buf, size_t bufsize, uint32_t state[16], unsigned rounds) {
|
||||
/* constant for shuffling bytes (replacing multiple-of-8 rotates) */
|
||||
const __m128i rot16 =
|
||||
_mm_set_epi8(13, 12, 15, 14, 9, 8, 11, 10, 5, 4, 7, 6, 1, 0, 3, 2);
|
||||
const __m128i rot8 =
|
||||
_mm_set_epi8(14, 13, 12, 15, 10, 9, 8, 11, 6, 5, 4, 7, 2, 1, 0, 3);
|
||||
// Load state to vectors, duplicate four times, only different counters
|
||||
__m128i orig[16];
|
||||
for (unsigned i = 0; i < 16; ++i) orig[i] = _mm_set1_epi32(state[i]);
|
||||
__m128i addv = _mm_set_epi32(3, 2, 1, 0);
|
||||
COUNTER_INCREMENT(addv);
|
||||
addv = _mm_set1_epi32(4);
|
||||
const unsigned batches = bufsize / 256;
|
||||
for (unsigned b = batches; b-->0;) {
|
||||
__m128i x[16];
|
||||
for (unsigned i = 0; i < 16; ++i) x[i] = orig[i];
|
||||
for (unsigned r = rounds / 2; r-->0;) {
|
||||
// Mix columns
|
||||
VEC4_QUARTERROUND(0, 4, 8, 12);
|
||||
VEC4_QUARTERROUND(1, 5, 9, 13);
|
||||
VEC4_QUARTERROUND(2, 6, 10, 14);
|
||||
VEC4_QUARTERROUND(3, 7, 11, 15);
|
||||
// Mix diagonals
|
||||
VEC4_QUARTERROUND(0, 5, 10, 15);
|
||||
VEC4_QUARTERROUND(1, 6, 11, 12);
|
||||
VEC4_QUARTERROUND(2, 7, 8, 13);
|
||||
VEC4_QUARTERROUND(3, 4, 9, 14);
|
||||
}
|
||||
// Add original block, unpack output
|
||||
ONEQUAD(0, 1, 2, 3, buf);
|
||||
ONEQUAD(4, 5, 6, 7, buf + 16);
|
||||
ONEQUAD(8, 9, 10, 11, buf + 32);
|
||||
ONEQUAD(12, 13, 14, 15, buf + 48);
|
||||
COUNTER_INCREMENT(addv);
|
||||
buf += 256;
|
||||
}
|
||||
// Store counter
|
||||
state[12] = _mm_cvtsi128_si32(orig[12]);
|
||||
state[13] = _mm_cvtsi128_si32(orig[13]);
|
||||
return batches * 256;
|
||||
}
|
||||
|
||||
#undef COUNTER_INCREMENT
|
||||
#undef ONEQUAD
|
||||
#undef ONEQUAD_TRANSPOSE
|
||||
#undef VEC4_ROT
|
||||
#undef VEC4_QUARTERROUND
|
||||
#undef VEC4_QUARTERROUND_SHUFFLE
|
||||
-130
@@ -1,130 +0,0 @@
|
||||
#include <immintrin.h> // AVX2
|
||||
|
||||
// clang-format off
|
||||
|
||||
#define VEC8_ROT(A, IMM) \
|
||||
_mm256_or_si256(_mm256_slli_epi32(A, IMM), _mm256_srli_epi32(A, (32 - IMM)))
|
||||
|
||||
#define VEC8_LINE1(A, B, C, D) \
|
||||
x[A] = _mm256_add_epi32(x[A], x[B]); \
|
||||
x[D] = _mm256_shuffle_epi8(_mm256_xor_si256(x[D], x[A]), rot16)
|
||||
#define VEC8_LINE2(A, B, C, D) \
|
||||
x[C] = _mm256_add_epi32(x[C], x[D]); \
|
||||
x[B] = VEC8_ROT(_mm256_xor_si256(x[B], x[C]), 12)
|
||||
#define VEC8_LINE3(A, B, C, D) \
|
||||
x[A] = _mm256_add_epi32(x[A], x[B]); \
|
||||
x[D] = _mm256_shuffle_epi8(_mm256_xor_si256(x[D], x[A]), rot8)
|
||||
#define VEC8_LINE4(A, B, C, D) \
|
||||
x[C] = _mm256_add_epi32(x[C], x[D]); \
|
||||
x[B] = VEC8_ROT(_mm256_xor_si256(x[B], x[C]), 7)
|
||||
|
||||
#define VEC8_ROUND( \
|
||||
A1, B1, C1, D1, A2, B2, C2, D2, A3, B3, C3, D3, A4, B4, C4, D4 \
|
||||
) \
|
||||
VEC8_LINE1(A1, B1, C1, D1); \
|
||||
VEC8_LINE1(A2, B2, C2, D2); \
|
||||
VEC8_LINE1(A3, B3, C3, D3); \
|
||||
VEC8_LINE1(A4, B4, C4, D4); \
|
||||
VEC8_LINE2(A1, B1, C1, D1); \
|
||||
VEC8_LINE2(A2, B2, C2, D2); \
|
||||
VEC8_LINE2(A3, B3, C3, D3); \
|
||||
VEC8_LINE2(A4, B4, C4, D4); \
|
||||
VEC8_LINE3(A1, B1, C1, D1); \
|
||||
VEC8_LINE3(A2, B2, C2, D2); \
|
||||
VEC8_LINE3(A3, B3, C3, D3); \
|
||||
VEC8_LINE3(A4, B4, C4, D4); \
|
||||
VEC8_LINE4(A1, B1, C1, D1); \
|
||||
VEC8_LINE4(A2, B2, C2, D2); \
|
||||
VEC8_LINE4(A3, B3, C3, D3); \
|
||||
VEC8_LINE4(A4, B4, C4, D4)
|
||||
|
||||
#define TRANSPOSE(A, B, C, D) \
|
||||
{ \
|
||||
const __m256i t0 = _mm256_unpacklo_epi32(x[A], x[B]), \
|
||||
t1 = _mm256_unpacklo_epi32(x[C], x[D]), \
|
||||
t2 = _mm256_unpackhi_epi32(x[A], x[B]), \
|
||||
t3 = _mm256_unpackhi_epi32(x[C], x[D]); \
|
||||
x[A] = _mm256_unpacklo_epi64(t0, t1); \
|
||||
x[B] = _mm256_unpackhi_epi64(t0, t1); \
|
||||
x[C] = _mm256_unpacklo_epi64(t2, t3); \
|
||||
x[D] = _mm256_unpackhi_epi64(t2, t3); \
|
||||
}
|
||||
|
||||
#define ONEOCTO(A, B, C, D, A2, B2, C2, D2, c) \
|
||||
{ \
|
||||
TRANSPOSE(A, B, C, D); \
|
||||
TRANSPOSE(A2, B2, C2, D2); \
|
||||
_mm256_storeu_si256((__m256i*)(c), _mm256_permute2x128_si256(x[A], x[A2], 0x20)); \
|
||||
_mm256_storeu_si256((__m256i*)(c + 64), _mm256_permute2x128_si256(x[B], x[B2], 0x20)); \
|
||||
_mm256_storeu_si256((__m256i*)(c + 128), _mm256_permute2x128_si256(x[C], x[C2], 0x20)); \
|
||||
_mm256_storeu_si256((__m256i*)(c + 192), _mm256_permute2x128_si256(x[D], x[D2], 0x20)); \
|
||||
_mm256_storeu_si256((__m256i*)(c + 256), _mm256_permute2x128_si256(x[A], x[A2], 0x31)); \
|
||||
_mm256_storeu_si256((__m256i*)(c + 320), _mm256_permute2x128_si256(x[B], x[B2], 0x31)); \
|
||||
_mm256_storeu_si256((__m256i*)(c + 384), _mm256_permute2x128_si256(x[C], x[C2], 0x31)); \
|
||||
_mm256_storeu_si256((__m256i*)(c + 448), _mm256_permute2x128_si256(x[D], x[D2], 0x31)); \
|
||||
}
|
||||
|
||||
#define COUNTER_INCREMENT(addv) \
|
||||
{ \
|
||||
__m256i carry = orig[12]; \
|
||||
orig[12] = _mm256_add_epi32(orig[12], addv); \
|
||||
carry = _mm256_srli_epi32(_mm256_and_si256(_mm256_xor_si256(orig[12], carry), carry), 31); \
|
||||
orig[13] = _mm256_add_epi32(orig[13], carry); \
|
||||
}
|
||||
|
||||
static inline uint64_t
|
||||
_cha_8block(uint8_t* buf, size_t bufsize, uint32_t state[16], unsigned rounds) {
|
||||
unsigned batches = bufsize / 512;
|
||||
/* constant for shuffling bytes (replacing multiple-of-8 rotates) */
|
||||
const __m256i rot16 = _mm256_set_epi8(
|
||||
13, 12, 15, 14,
|
||||
9, 8, 11, 10,
|
||||
5, 4, 7, 6,
|
||||
1, 0, 3, 2,
|
||||
13, 12, 15, 14,
|
||||
9, 8, 11, 10,
|
||||
5, 4, 7, 6,
|
||||
1, 0, 3, 2
|
||||
);
|
||||
const __m256i rot8 = _mm256_set_epi8(
|
||||
14, 13, 12, 15,
|
||||
10, 9, 8, 11,
|
||||
6, 5, 4, 7,
|
||||
2, 1, 0, 3,
|
||||
14, 13, 12, 15,
|
||||
10, 9, 8, 11,
|
||||
6, 5, 4, 7,
|
||||
2, 1, 0, 3
|
||||
);
|
||||
__m256i orig[16];
|
||||
for (int i = 0; i < 16; ++i)
|
||||
orig[i] = _mm256_set1_epi32(state[i]);
|
||||
COUNTER_INCREMENT(_mm256_set_epi32(7, 6, 5, 4, 3, 2, 1, 0));
|
||||
|
||||
for (unsigned b = batches; b-->0;) {
|
||||
__m256i x[16];
|
||||
for (int i = 0; i < 16; ++i) x[i] = orig[i];
|
||||
for (unsigned r = rounds / 2; r-->0;) {
|
||||
VEC8_ROUND(0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15);
|
||||
VEC8_ROUND(0, 5, 10, 15, 1, 6, 11, 12, 2, 7, 8, 13, 3, 4, 9, 14);
|
||||
}
|
||||
for (unsigned i = 0; i < 16; ++i) x[i] = _mm256_add_epi32(x[i], orig[i]);
|
||||
ONEOCTO(0, 1, 2, 3, 4, 5, 6, 7, buf);
|
||||
ONEOCTO(8, 9, 10, 11, 12, 13, 14, 15, buf + 32);
|
||||
COUNTER_INCREMENT(_mm256_set1_epi32(8));
|
||||
buf += 512;
|
||||
}
|
||||
state[12] = _mm256_extract_epi32(orig[12], 0);
|
||||
state[13] = _mm256_extract_epi32(orig[13], 0);
|
||||
return batches * 512;
|
||||
}
|
||||
|
||||
#undef COUNTER_INCREMENT
|
||||
#undef ONEOCTO
|
||||
#undef TRANSPOSE
|
||||
#undef VEC8_ROT
|
||||
#undef VEC8_LINE1
|
||||
#undef VEC8_LINE2
|
||||
#undef VEC8_LINE3
|
||||
#undef VEC8_LINE4
|
||||
#undef VEC8_ROUND
|
||||
@@ -1,18 +0,0 @@
|
||||
#include "charandom.h"
|
||||
|
||||
static uint64_t cha_uint64(void* st) {
|
||||
cha_ctx* ctx = (cha_ctx*)st;
|
||||
if (ctx->offset + sizeof(uint64_t) > ctx->end) {
|
||||
ctx->offset -= ctx->end;
|
||||
ctx->end =
|
||||
ctx->gen(ctx->unconsumed, BATCH_SIZE, ctx->state, ctx->rounds);
|
||||
}
|
||||
uint64_t ret = *(uint64_t*)(ctx->unconsumed + ctx->offset);
|
||||
ctx->offset += sizeof(uint64_t);
|
||||
return ret;
|
||||
}
|
||||
static uint32_t cha_uint32(void* st) { return cha_uint64(st); }
|
||||
static double cha_double(void* st) {
|
||||
// Fast uint64_to_double conversion from numpy/random/_common.pxd
|
||||
return (cha_uint64(st) >> 11) * (1.0 / 9007199254740992.0);
|
||||
}
|
||||
@@ -1,3 +0,0 @@
|
||||
#include "charandom.h"
|
||||
|
||||
// Library build for Python CFFI to use
|
||||
-152
@@ -1,152 +0,0 @@
|
||||
#include <stdbool.h>
|
||||
#include <stdint.h>
|
||||
|
||||
#define CHA_BLOCK_SIZE 64
|
||||
#define BATCH_BLOCKS 8
|
||||
#define BATCH_SIZE (BATCH_BLOCKS * CHA_BLOCK_SIZE)
|
||||
|
||||
#if defined(__x86_64__)
|
||||
#ifdef __GNUC__
|
||||
#pragma GCC target("sse2")
|
||||
#pragma GCC target("ssse3")
|
||||
#pragma GCC target("avx2")
|
||||
#endif
|
||||
#include "cha4ssse3.h"
|
||||
#include "cha8avx2.h"
|
||||
#elif defined(__aarch64__)
|
||||
#include "cha4neon.h"
|
||||
#endif
|
||||
|
||||
#include "cha1c.h"
|
||||
|
||||
#include <assert.h>
|
||||
#include <stdbool.h>
|
||||
#include <stdint.h>
|
||||
#include <stdio.h>
|
||||
#include <string.h>
|
||||
|
||||
#include <pthread.h>
|
||||
#include <time.h>
|
||||
#include <unistd.h>
|
||||
|
||||
typedef uint64_t (*genfunc)(
|
||||
uint8_t* out, size_t outsize, uint32_t state[16], unsigned rounds
|
||||
);
|
||||
typedef struct cha_ctx {
|
||||
uint32_t state[16];
|
||||
uint8_t unconsumed[BATCH_SIZE];
|
||||
uint32_t offset, end;
|
||||
unsigned rounds;
|
||||
genfunc gen;
|
||||
} cha_ctx;
|
||||
|
||||
/// @brief Initialize cha_ctx
|
||||
/// @param ctx holds ChaCha20 state
|
||||
/// @param key 32 byte key
|
||||
/// @param iv 16 bytes, usually the first 4-8 bytes are zeroes, the rest nonce
|
||||
/// @param rounds ChaCha iteration count: 8=fast, 12=balanced, 20=secure
|
||||
void cha_init(
|
||||
cha_ctx* ctx, const uint8_t* key, const uint8_t* iv, unsigned rounds
|
||||
) {
|
||||
ctx->state[0] = 0x61707865;
|
||||
ctx->state[1] = 0x3320646e;
|
||||
ctx->state[2] = 0x79622d32;
|
||||
ctx->state[3] = 0x6b206574;
|
||||
memcpy(ctx->state + 4, key, 32);
|
||||
memcpy(ctx->state + 12, iv, 16);
|
||||
memset(ctx->unconsumed, 0, sizeof ctx->unconsumed);
|
||||
ctx->offset = ctx->end = 0;
|
||||
ctx->rounds = rounds;
|
||||
#if defined(__x86_64__)
|
||||
if (__builtin_cpu_supports("avx2"))
|
||||
ctx->gen = _cha_8block;
|
||||
else if (__builtin_cpu_supports("ssse3"))
|
||||
ctx->gen = _cha_4block;
|
||||
else
|
||||
ctx->gen = _cha_block;
|
||||
#elif defined(__aarch64__)
|
||||
ctx->gen = _cha_4block;
|
||||
#else
|
||||
ctx->gen = _cha_block;
|
||||
#endif
|
||||
}
|
||||
|
||||
/// Dispose of sensitive data within the context
|
||||
void cha_wipe(cha_ctx* ctx) { memset(ctx, 0, sizeof(cha_ctx)); }
|
||||
|
||||
/// @brief Advance or rewind the stream to any arbitrary location
|
||||
/// Keeps byte offset inside batch untouched but clears the buffer
|
||||
/// @param ctx ChaCha context
|
||||
/// @param offset Offset in blocks of 64 bytes (counter change)
|
||||
void cha_seek_blocks(cha_ctx* ctx, int64_t offset) {
|
||||
*(uint64_t*)(ctx->state + 12) += offset - (int64_t)ctx->end / 64;
|
||||
ctx->end = 0;
|
||||
}
|
||||
|
||||
/// @brief Seek a number of bytes forward or backward in stream
|
||||
/// Supports seeking backwards and forwards, even beyond start, but only up to
|
||||
/// 64 bit distance which does not cover the whole counter range.
|
||||
/// @param ctx ChaCha context
|
||||
/// @param offset Positive or negative offset from current byte position
|
||||
void cha_seek(cha_ctx* ctx, int64_t offset) {
|
||||
offset += (int64_t)ctx->offset;
|
||||
ctx->offset = ((offset % 64) + 64) % 64;
|
||||
cha_seek_blocks(ctx, (offset - (int64_t)ctx->offset) / 64);
|
||||
}
|
||||
|
||||
/// @brief Tell current byte position in stream (assuming initial counter 0).
|
||||
/// Result is truncated to int64 range (positive or negative)
|
||||
/// @param ctx ChaCha context
|
||||
int64_t cha_tell(cha_ctx* ctx) {
|
||||
int64_t counter = *(int64_t*)(ctx->state + 12);
|
||||
return counter * CHA_BLOCK_SIZE + ctx->offset - ctx->end;
|
||||
}
|
||||
|
||||
/// @brief Incremental generation, keeps state between calls
|
||||
/// @param ctx ChaCha context
|
||||
/// @param out output buffer
|
||||
/// @param outlen output buffer length
|
||||
void cha_update(cha_ctx* ctx, uint8_t* out, uint64_t outlen) {
|
||||
// The included header will mess with these variables
|
||||
uint8_t* end = out + outlen;
|
||||
if (ctx->offset) {
|
||||
// Need to generate stored buffer?
|
||||
if (ctx->end == 0)
|
||||
ctx->end =
|
||||
ctx->gen(ctx->unconsumed, BATCH_SIZE, ctx->state, ctx->rounds);
|
||||
// Deliver stored bytes first
|
||||
uint64_t N = ctx->end - ctx->offset;
|
||||
if (N > outlen)
|
||||
N = outlen;
|
||||
memcpy(out, ctx->unconsumed + ctx->offset, N);
|
||||
ctx->offset += N;
|
||||
out += N;
|
||||
if (ctx->offset == ctx->end)
|
||||
ctx->offset = ctx->end = 0;
|
||||
if (out == end)
|
||||
return;
|
||||
}
|
||||
out += ctx->gen(out, end - out, ctx->state, ctx->rounds);
|
||||
const uint32_t N = end - out;
|
||||
if (N) {
|
||||
ctx->end =
|
||||
ctx->gen(ctx->unconsumed, BATCH_SIZE, ctx->state, ctx->rounds);
|
||||
memcpy(out, ctx->unconsumed, N);
|
||||
ctx->offset = N;
|
||||
}
|
||||
}
|
||||
|
||||
/// @brief Produce a requested number of random bytes, single shot.
|
||||
/// @param out output buffer
|
||||
/// @param outlen output buffer length
|
||||
/// @param key 32 byte key
|
||||
/// @param iv 16 bytes, where normally initial 4-8 bytes are 0 (counter)
|
||||
void cha_generate(
|
||||
uint8_t* out, uint64_t outlen, const uint8_t key[32], const uint8_t iv[16],
|
||||
unsigned rounds
|
||||
) {
|
||||
cha_ctx ctx;
|
||||
cha_init(&ctx, key, iv, rounds);
|
||||
cha_update(&ctx, out, outlen);
|
||||
cha_wipe(&ctx);
|
||||
}
|
||||
@@ -1,295 +0,0 @@
|
||||
#include <errno.h>
|
||||
#include <inttypes.h>
|
||||
#include <pthread.h>
|
||||
#include <signal.h>
|
||||
#include <stdbool.h>
|
||||
#include <stdint.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
#include <time.h>
|
||||
#include <unistd.h>
|
||||
|
||||
#include "charandom.h"
|
||||
|
||||
static volatile bool quit = false;
|
||||
|
||||
void signal_handler(int sig) {
|
||||
quit = true;
|
||||
signal(SIGINT, SIG_DFL);
|
||||
signal(SIGTERM, SIG_DFL);
|
||||
}
|
||||
|
||||
#define BLOCK_SIZE (1 << 21) // 2 MiB seems optimal for speed
|
||||
|
||||
static const unsigned char default_iv[16] = "\0\0\0\0\0\0\0\0RandQuik";
|
||||
typedef struct thread_args {
|
||||
int index;
|
||||
int done;
|
||||
unsigned char* buf;
|
||||
unsigned char key[32];
|
||||
unsigned workers;
|
||||
unsigned rounds;
|
||||
pthread_mutex_t lock;
|
||||
pthread_cond_t cond;
|
||||
pthread_t thread;
|
||||
} thread_args;
|
||||
|
||||
void* producer_thread(void* a) {
|
||||
thread_args* args = (thread_args*)a;
|
||||
const uint64_t ivstep = args->workers * BATCH_BLOCKS;
|
||||
cha_ctx ctx;
|
||||
cha_init(&ctx, args->key, default_iv, args->rounds);
|
||||
cha_seek_blocks(&ctx, args->index * BLOCK_SIZE / 64);
|
||||
while (!quit) {
|
||||
pthread_mutex_lock(&args->lock);
|
||||
while (args->done) {
|
||||
pthread_cond_wait(&args->cond, &args->lock);
|
||||
}
|
||||
cha_update(&ctx, args->buf, BLOCK_SIZE);
|
||||
cha_seek_blocks(&ctx, ivstep);
|
||||
args->done = 1;
|
||||
pthread_cond_signal(&args->cond);
|
||||
pthread_mutex_unlock(&args->lock);
|
||||
}
|
||||
cha_wipe(&ctx);
|
||||
return NULL;
|
||||
}
|
||||
|
||||
void print_status(
|
||||
uint64_t bytes, uint64_t max_bytes, struct timespec start_time
|
||||
) {
|
||||
struct timespec end_time;
|
||||
clock_gettime(CLOCK_MONOTONIC, &end_time);
|
||||
double t = (end_time.tv_sec - start_time.tv_sec) +
|
||||
1e-9 * (end_time.tv_nsec - start_time.tv_nsec);
|
||||
char buf[64] = {};
|
||||
double speed = bytes / t;
|
||||
char const* unit = "MB";
|
||||
double m = 1e-6;
|
||||
if (speed > 0.5e9) {
|
||||
unit = "GB";
|
||||
m = 1e-9;
|
||||
}
|
||||
if (max_bytes) {
|
||||
snprintf(buf, sizeof buf - 1, " of %'.0lf", m * max_bytes);
|
||||
}
|
||||
|
||||
fprintf(
|
||||
stderr, "\r%5.0lf%s %s written, %.2lf %s/s.\e[K", m * bytes, buf, unit,
|
||||
m * speed, unit
|
||||
);
|
||||
}
|
||||
|
||||
int fast(
|
||||
FILE* f, unsigned workers, uint64_t max_bytes, unsigned char const key[32],
|
||||
unsigned char const iv[16], unsigned rounds
|
||||
) {
|
||||
thread_args args[workers];
|
||||
memset(args, 0, sizeof args);
|
||||
for (int i = 0; i < workers; ++i) {
|
||||
args[i].index = i;
|
||||
args[i].buf = malloc(BLOCK_SIZE);
|
||||
args[i].workers = workers;
|
||||
args[i].rounds = rounds;
|
||||
memcpy(args[i].key, key, 32);
|
||||
pthread_mutex_init(&args[i].lock, NULL);
|
||||
pthread_cond_init(&args[i].cond, NULL);
|
||||
pthread_create(&args[i].thread, NULL, producer_thread, &args[i]);
|
||||
}
|
||||
|
||||
struct timespec start_time;
|
||||
clock_gettime(CLOCK_MONOTONIC, &start_time);
|
||||
|
||||
int i = -1;
|
||||
uint64_t bytes = 0;
|
||||
while (!quit) {
|
||||
i = (i + 1) % workers;
|
||||
pthread_mutex_lock(&args[i].lock);
|
||||
while (!args[i].done) {
|
||||
pthread_cond_wait(&args[i].cond, &args[i].lock);
|
||||
}
|
||||
if (bytes % (1 << 30) == 0 || bytes + BLOCK_SIZE >= max_bytes) {
|
||||
print_status(bytes, max_bytes, start_time);
|
||||
}
|
||||
uint64_t sz = BLOCK_SIZE;
|
||||
if (max_bytes && bytes + sz >= max_bytes) {
|
||||
fprintf(stderr, "\r\e[KMax reached\n");
|
||||
sz = max_bytes - bytes;
|
||||
quit = true;
|
||||
}
|
||||
if (fwrite(args[i].buf, sz, 1, f) != 1) {
|
||||
quit = true;
|
||||
fprintf(stderr, "\r\e[KWrite failed: %s\n", strerror(errno));
|
||||
}
|
||||
bytes += sz;
|
||||
args[i].done = 0;
|
||||
pthread_cond_signal(&args[i].cond);
|
||||
pthread_mutex_unlock(&args[i].lock);
|
||||
}
|
||||
|
||||
print_status(bytes, max_bytes, start_time);
|
||||
for (int i = 0; i < workers; ++i) {
|
||||
args[i].done = 0;
|
||||
pthread_cancel(args[i].thread);
|
||||
pthread_join(args[i].thread, NULL);
|
||||
pthread_mutex_destroy(&args[i].lock);
|
||||
pthread_cond_destroy(&args[i].cond);
|
||||
free(args[i].buf);
|
||||
}
|
||||
fprintf(stderr, "\nRandQuik wrote %" PRIu64 " bytes!\n\n", bytes);
|
||||
return 0;
|
||||
}
|
||||
|
||||
bool parse_hex(char* str, unsigned char* buf, size_t len) {
|
||||
for (size_t i = 0; i < len; ++i) {
|
||||
int sz = 0;
|
||||
if (sscanf(str, "%2hhx%n", buf + i, &sz) != 1) {
|
||||
if (*str) {
|
||||
fprintf(stderr, "Unable to read seed at `%s`\n\n", str);
|
||||
return false;
|
||||
}
|
||||
return true; // Shorter than key length is OK
|
||||
}
|
||||
str += sz;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
void print_hex(unsigned char* buf, size_t len) {
|
||||
for (size_t i = 0; i < len; ++i) {
|
||||
fprintf(stderr, "%02hhx", buf[i]);
|
||||
}
|
||||
}
|
||||
|
||||
void help(char** argv) {
|
||||
fprintf(
|
||||
stderr,
|
||||
"Usage: %s [-t #threads] [-s hexseed] [-b #bytes] [-r #rounds] [-o "
|
||||
"outputfile]\n\n",
|
||||
argv[0]
|
||||
);
|
||||
}
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
unsigned char key[32] = {};
|
||||
unsigned char iv[16] = {};
|
||||
unsigned int workers = 8;
|
||||
unsigned int rounds = 20;
|
||||
char* output = NULL;
|
||||
uint64_t max_bytes = 0;
|
||||
bool seeded = false;
|
||||
for (char opt; (opt = getopt(argc, argv, "bostr")) != -1;) {
|
||||
if (opt == 't') {
|
||||
if (optind >= argc || sscanf(argv[optind++], "%u", &workers) != 1) {
|
||||
fprintf(
|
||||
stderr, "Expected the number of worker threads after -t\n"
|
||||
);
|
||||
return 1;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if (opt == 'r') {
|
||||
if (optind >= argc || sscanf(argv[optind++], "%u", &rounds) != 1) {
|
||||
fprintf(
|
||||
stderr,
|
||||
"Expected the number ChaCha rounds (8, 12 or 20) after -r\n"
|
||||
);
|
||||
return 1;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if (opt == 's') {
|
||||
if (optind >= argc || !parse_hex(argv[optind++], key, 32)) {
|
||||
fprintf(stderr, "Expected a hex seed string after -s\n");
|
||||
return 1;
|
||||
}
|
||||
seeded = true;
|
||||
continue;
|
||||
}
|
||||
if (opt == 'o') {
|
||||
if (optind >= argc) {
|
||||
fprintf(stderr, "Expected output filename after -s\n");
|
||||
return 1;
|
||||
}
|
||||
if (strcmp(argv[optind], "-") != 0) {
|
||||
output = argv[optind++];
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if (opt == 'b') {
|
||||
char unit[16] = {};
|
||||
if (optind >= argc || sscanf(argv[optind++], "%" SCNu64 "%15s", &max_bytes, unit) < 1) {
|
||||
fprintf(
|
||||
stderr,
|
||||
"Expected a maximum number of bytes to read after -b\n"
|
||||
);
|
||||
return 1;
|
||||
}
|
||||
if (strcasecmp(unit, "k") == 0 || strcasecmp(unit, "kb") == 0)
|
||||
max_bytes *= 1000ull;
|
||||
else if (strcasecmp(unit, "m") == 0 || strcasecmp(unit, "mb") == 0)
|
||||
max_bytes *= 1000000ull;
|
||||
else if (strcasecmp(unit, "g") == 0 || strcasecmp(unit, "gb") == 0)
|
||||
max_bytes *= 1000000000ull;
|
||||
else if (strcasecmp(unit, "t") == 0 || strcasecmp(unit, "tb") == 0)
|
||||
max_bytes *= 1000000000000ull;
|
||||
else if (strcasecmp(unit, "ki") == 0 || strcasecmp(unit, "kib") == 0)
|
||||
max_bytes <<= 10;
|
||||
else if (strcasecmp(unit, "mi") == 0 || strcasecmp(unit, "mib") == 0)
|
||||
max_bytes <<= 20;
|
||||
else if (strcasecmp(unit, "gi") == 0 || strcasecmp(unit, "gib") == 0)
|
||||
max_bytes <<= 30;
|
||||
else if (strcasecmp(unit, "ti") == 0 || strcasecmp(unit, "tib") == 0)
|
||||
max_bytes <<= 40;
|
||||
continue;
|
||||
}
|
||||
help(argv);
|
||||
return 1;
|
||||
}
|
||||
FILE* f = stdout;
|
||||
if (output) {
|
||||
f = fopen(output, "wb");
|
||||
if (!f) {
|
||||
fprintf(stderr, "Failed to open %s for writing.\n", output);
|
||||
return 1;
|
||||
}
|
||||
} else if (isatty(1)) {
|
||||
fprintf(
|
||||
stderr,
|
||||
"Won't print random on console. Pipe me to another program or "
|
||||
"file instead.\n\n"
|
||||
);
|
||||
help(argv);
|
||||
return 1;
|
||||
}
|
||||
if (!seeded) {
|
||||
FILE* urand = fopen("/dev/urandom", "rb");
|
||||
if (!urand || fread(key, 32, 1, urand) != 1) {
|
||||
fprintf(
|
||||
stderr, "Failed to seed from /dev/urandom. Use -s hexstring for "
|
||||
"manual seeding.\n"
|
||||
);
|
||||
fclose(urand);
|
||||
return 1;
|
||||
}
|
||||
fclose(urand);
|
||||
fprintf(
|
||||
stderr,
|
||||
"Random seed generated. This sequence may be repeated by:\n%s ",
|
||||
argv[0]
|
||||
);
|
||||
if (rounds != 20)
|
||||
fprintf(stderr, "-r %u -s ", rounds);
|
||||
else
|
||||
fprintf(stderr, "-s ");
|
||||
|
||||
print_hex(key, 32);
|
||||
fprintf(stderr, "\n\n");
|
||||
}
|
||||
signal(SIGINT, signal_handler);
|
||||
signal(SIGTERM, signal_handler);
|
||||
int ret = fast(f, workers, max_bytes, key, iv, rounds);
|
||||
fclose(f);
|
||||
return ret;
|
||||
}
|
||||
@@ -1,51 +0,0 @@
|
||||
#cython: language_level=3
|
||||
|
||||
from libc.stdint cimport int64_t, uint32_t, uint8_t, uint64_t
|
||||
from cpython.pycapsule cimport PyCapsule_IsValid, PyCapsule_GetPointer
|
||||
import numpy as np
|
||||
cimport numpy as np
|
||||
cimport cython
|
||||
import secrets
|
||||
|
||||
from numpy.random cimport BitGenerator
|
||||
|
||||
np.import_array()
|
||||
|
||||
cdef extern from "chanumpy.h":
|
||||
struct cha_ctx:
|
||||
uint32_t state[16]
|
||||
uint8_t unconsumed[512]
|
||||
uint32_t offset, end;
|
||||
unsigned rounds;
|
||||
|
||||
void cha_init(cha_ctx* ctx, const uint8_t* key, const uint8_t* iv, unsigned rounds) nogil
|
||||
void cha_seek(cha_ctx* ctx, int64_t offset)
|
||||
int64_t cha_tell(cha_ctx* ctx)
|
||||
|
||||
uint64_t cha_uint64(void *state) nogil
|
||||
uint32_t cha_uint32(void *state) nogil
|
||||
double cha_double(void *state) nogil
|
||||
|
||||
|
||||
cdef class Cha(BitGenerator):
|
||||
cdef cha_ctx rng_state
|
||||
|
||||
def __init__(self, seed=None, *, rounds=20):
|
||||
BitGenerator.__init__(self, seed)
|
||||
self._bitgen.state = <void *>&self.rng_state
|
||||
self._bitgen.next_uint64 = &cha_uint64
|
||||
self._bitgen.next_uint32 = &cha_uint32
|
||||
self._bitgen.next_double = &cha_double
|
||||
self._bitgen.next_raw = &cha_uint64
|
||||
# Generated state is ChaCha20 key
|
||||
key = self._seed_seq.generate_state(4, np.uint64)
|
||||
cha_init(&self.rng_state, <uint8_t *>np.PyArray_DATA(key), bytes(16) + b"NumpRand", rounds)
|
||||
|
||||
def advance(self, delta):
|
||||
cha_seek(&self.rng_state, delta << 3)
|
||||
|
||||
def tell(self):
|
||||
return cha_tell(&self.rng_state) >> 3;
|
||||
|
||||
def state(self):
|
||||
return self.rng_state.state[12], self.rng_state.state[13]
|
||||
@@ -1,75 +0,0 @@
|
||||
from secrets import randbelow, token_bytes
|
||||
|
||||
import pytest
|
||||
from cryptography.hazmat.primitives.ciphers import Cipher
|
||||
from cryptography.hazmat.primitives.ciphers.algorithms import ChaCha20
|
||||
|
||||
from randquik import cha
|
||||
|
||||
|
||||
def test_cipherstreams_fullblocks():
|
||||
"""Requests in multiple of ChaCha20 block size 64 bytes"""
|
||||
key = token_bytes(32)
|
||||
iv = token_bytes(16)
|
||||
c0 = Cipher(ChaCha20(key, iv), None, None).encryptor()
|
||||
c1 = cha.Cha(key, iv)
|
||||
|
||||
for i in range(2048):
|
||||
N = 64 * (1 + randbelow(2048))
|
||||
ct0 = c0.update(bytes(N))
|
||||
ct1 = c1(bytearray(N))
|
||||
assert len(ct0) == len(ct1)
|
||||
assert ct0.hex() == ct1.hex(), f"{i=} {N=}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"counter",
|
||||
[
|
||||
b"\x00\x00\x00\x00\x00\x00\x00\x00",
|
||||
b"\xFF\xFF\xFF\xFF\x00\x00\x00\x00",
|
||||
b"\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF",
|
||||
b"\xFF\xFF\xFF\x7F\xFF\xFF\xFF\xFF",
|
||||
b"\x00\x00\x00\x80\xFF\xFF\xFF\xFF",
|
||||
],
|
||||
)
|
||||
def test_counter_wrap(counter):
|
||||
"""Tests carry handling of counter increments"""
|
||||
key = bytes(32)
|
||||
iv = counter + b"--------"
|
||||
c0 = Cipher(ChaCha20(key, iv), None, None).encryptor()
|
||||
c1 = cha.Cha(key, iv)
|
||||
N = 128
|
||||
ct0 = c0.update(bytes(N))
|
||||
ct1 = c1(bytearray(N))
|
||||
assert len(ct0) == len(ct1)
|
||||
assert ct0.hex() == ct1.hex()
|
||||
|
||||
|
||||
def test_cipherstreams_partial_updates():
|
||||
"""Odd-sized requests that retain leftover buffers"""
|
||||
key = token_bytes(32)
|
||||
iv = token_bytes(16)
|
||||
c0 = Cipher(ChaCha20(key, iv), None, None).encryptor()
|
||||
c1 = cha.Cha(key, iv)
|
||||
|
||||
for i in range(2048):
|
||||
N = 1 + randbelow(2048)
|
||||
ct0 = c0.update(bytes(N))
|
||||
ct1 = c1(bytearray(N))
|
||||
assert len(ct0) == len(ct1)
|
||||
assert ct0.hex() == ct1.hex(), f"{i=} {N=}"
|
||||
|
||||
|
||||
def test_cipherstreams_32leftover():
|
||||
"""Test the special case where the second response comes entirely from the leftover buffer"""
|
||||
key = token_bytes(32)
|
||||
nonce = token_bytes(12) # IETF nonce size
|
||||
iv = bytes(4) + nonce
|
||||
c0 = Cipher(ChaCha20(key, iv), None, None).encryptor()
|
||||
c1 = cha.Cha(key, nonce)
|
||||
|
||||
for N in [512 - 32, 32]:
|
||||
ct0 = c0.update(bytes(N))
|
||||
ct1 = c1(bytearray(N))
|
||||
assert len(ct0) == len(ct1)
|
||||
assert ct0.hex() == ct1.hex()
|
||||
Reference in New Issue
Block a user