Skip to main content

reifydb_runtime/context/rng/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4use std::sync::{Arc, Mutex};
5
6use getrandom::fill as getrandom_fill;
7use rand::{Rng as RandRng, RngExt, SeedableRng, rngs::StdRng};
8
9/// A random number generator that can be either OS-backed or seeded/deterministic.
10#[derive(Clone, Default)]
11pub enum Rng {
12	/// Uses OS entropy via `getrandom` (non-deterministic).
13	#[default]
14	Os,
15	/// Uses a seeded PRNG for deterministic output (e.g. testing).
16	Seeded(SeededRng),
17}
18
19impl Rng {
20	/// Create a deterministic RNG from the given seed.
21	pub fn seeded(seed: u64) -> Self {
22		Rng::Seeded(SeededRng::new(seed))
23	}
24
25	/// Generate 16 random bytes (suitable for UUID v4).
26	pub fn bytes_16(&self) -> [u8; 16] {
27		match self {
28			Rng::Os => {
29				let mut buf = [0u8; 16];
30				getrandom_fill(&mut buf).expect("getrandom failed");
31				buf
32			}
33			Rng::Seeded(seeded) => {
34				let mut buf = [0u8; 16];
35				let mut rng = seeded.inner.lock().unwrap();
36				rng.fill_bytes(&mut buf);
37				buf
38			}
39		}
40	}
41
42	/// Generate 32 random bytes (suitable for token generation).
43	pub fn bytes_32(&self) -> [u8; 32] {
44		match self {
45			Rng::Os => {
46				let mut buf = [0u8; 32];
47				getrandom_fill(&mut buf).expect("getrandom failed");
48				buf
49			}
50			Rng::Seeded(seeded) => {
51				let mut buf = [0u8; 32];
52				let mut rng = seeded.inner.lock().unwrap();
53				rng.fill_bytes(&mut buf);
54				buf
55			}
56		}
57	}
58
59	/// Generate 10 random bytes (suitable for UUID v7 random portion).
60	pub fn bytes_10(&self) -> [u8; 10] {
61		match self {
62			Rng::Os => {
63				let mut buf = [0u8; 10];
64				getrandom_fill(&mut buf).expect("getrandom failed");
65				buf
66			}
67			Rng::Seeded(seeded) => {
68				let mut buf = [0u8; 10];
69				let mut rng = seeded.inner.lock().unwrap();
70				rng.fill_bytes(&mut buf);
71				buf
72			}
73		}
74	}
75
76	/// Generate 10 random bytes from the infrastructure RNG stream.
77	///
78	/// Uses a separate RNG stream so that infrastructure operations (like
79	/// transaction ID generation) do not perturb the primary RNG state.
80	/// This ensures deterministic test output regardless of how many internal
81	/// transactions each test runner creates.
82	pub fn infra_bytes_10(&self) -> [u8; 10] {
83		match self {
84			Rng::Os => {
85				let mut buf = [0u8; 10];
86				getrandom_fill(&mut buf).expect("getrandom failed");
87				buf
88			}
89			Rng::Seeded(seeded) => {
90				let mut buf = [0u8; 10];
91				let mut rng = seeded.infra.lock().unwrap();
92				rng.fill_bytes(&mut buf);
93				buf
94			}
95		}
96	}
97
98	/// Generate 32 random bytes from the infrastructure RNG stream.
99	///
100	/// Uses a separate RNG stream for infrastructure operations (like
101	/// session token generation) that should not affect deterministic
102	/// test output.
103	pub fn infra_bytes_32(&self) -> [u8; 32] {
104		match self {
105			Rng::Os => {
106				let mut buf = [0u8; 32];
107				getrandom_fill(&mut buf).expect("getrandom failed");
108				buf
109			}
110			Rng::Seeded(seeded) => {
111				let mut buf = [0u8; 32];
112				let mut rng = seeded.infra.lock().unwrap();
113				rng.fill_bytes(&mut buf);
114				buf
115			}
116		}
117	}
118
119	pub fn infra_u64_inclusive(&self, max_inclusive: u64) -> u64 {
120		if max_inclusive == 0 {
121			return 0;
122		}
123		match self {
124			Rng::Os => {
125				let mut buf = [0u8; 8];
126				getrandom_fill(&mut buf).expect("getrandom failed");
127				let raw = u64::from_le_bytes(buf);
128				if max_inclusive == u64::MAX {
129					raw
130				} else {
131					raw % (max_inclusive + 1)
132				}
133			}
134			Rng::Seeded(seeded) => {
135				let mut rng = seeded.infra.lock().unwrap();
136				rng.random_range(0..=max_inclusive)
137			}
138		}
139	}
140}
141
142/// A seeded, deterministic RNG backed by `StdRng` wrapped in `Arc<Mutex<..>>`.
143#[derive(Clone)]
144pub struct SeededRng {
145	inner: Arc<Mutex<StdRng>>,
146	/// Separate RNG stream for infrastructure use (e.g. transaction IDs).
147	infra: Arc<Mutex<StdRng>>,
148}
149
150impl SeededRng {
151	/// Create a new seeded RNG from the given seed.
152	pub fn new(seed: u64) -> Self {
153		Self {
154			inner: Arc::new(Mutex::new(StdRng::seed_from_u64(seed))),
155			infra: Arc::new(Mutex::new(StdRng::seed_from_u64(seed ^ 0x5A5A5A5A5A5A5A5A))),
156		}
157	}
158}