Add the ability to reseed rng in uuid_utils (#86)

Author: ewilliamson-patreon

README.md

@@ -74,6 +74,12 @@ in comparison with the `uuid_utils` default behaviour, but is still faster than
 UUID('ffe95fcc-b818-4aca-a350-e0a35b9de6ec')
 ```
 
+## Interaction with Fork
+
+The underlying rng of this library is not guaranteed to be reset when you fork this process. This would mean that `uuid_utils` calls post fork could result in the same value across processes.
+
+If you plan to use this library alongside forking you will want to explicitly redeed the rng post fork. Either by calling `reseed_rng` manually or registering it in `os.register_at_fork`.
+
 ## Benchmarks
 
 |        Benchmark | Min     | Max     | Mean    | Min (+)         | Max (+)         | Mean (+)        |

python/uuid_utils/__init__.py

@@ -12,6 +12,7 @@
     UUID,
     __version__,
     getnode,
+    reseed_rng,
     uuid1,
     uuid3,
     uuid4,
@@ -34,6 +35,7 @@
     "SafeUUID",
     "__version__",
     "getnode",
+    "reseed_rng",
     "uuid1",
     "uuid3",
     "uuid4",

python/uuid_utils/__init__.pyi

@@ -124,6 +124,17 @@ class UUID:
     def __ge__(self, other: UUID) -> bool: ...
 
 def getnode() -> int: ...
+def reseed_rng() -> None:
+    """
+    Reseeds the underlying rng.
+    This is useful in cases where you fork, as without reseeding the
+    generated uuids may be identical. This can be called manually in the child process,
+    or automatically run after fork with:
+
+    os.register_at_fork(after_in_child=uuid_utils.reseed_rng)
+    """
+    ...
+
 def uuid1(node: int | None = None, clock_seq: int | None = None) -> UUID:
     """Generate a UUID from a host ID, sequence number, and the current time.
     If 'node' is not given, getnode() is used to obtain the hardware

src/lib.rs

@@ -2,7 +2,7 @@ use ahash::AHasher;
 use mac_address::get_mac_address;
 use pyo3::{
     prelude::*, ffi,
-    exceptions::{PyTypeError, PyValueError},
+    exceptions::{PyOSError, PyTypeError, PyValueError},
     pyclass::CompareOp,
     types::{PyBytes, PyDict},
     IntoPyObjectExt
@@ -448,6 +448,13 @@ fn getnode() -> PyResult<u64> {
     Ok(_getnode())
 }
 
+#[pyfunction]
+fn reseed_rng() -> PyResult<()> {
+    rand::rng()
+        .reseed()
+        .map_err(|err| PyOSError::new_err(err.to_string()))
+}
+
 #[pymodule]
 fn _uuid_utils(m: &Bound<'_, PyModule>) -> PyResult<()> {
     let safe_uuid_unknown = Python::attach(|py| {
@@ -472,6 +479,7 @@ fn _uuid_utils(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(uuid7, m)?)?;
     m.add_function(wrap_pyfunction!(uuid8, m)?)?;
     m.add_function(wrap_pyfunction!(getnode, m)?)?;
+    m.add_function(wrap_pyfunction!(reseed_rng, m)?)?;
     m.add("NAMESPACE_DNS", UUID::NAMESPACE_DNS)?;
     m.add("NAMESPACE_URL", UUID::NAMESPACE_URL)?;
     m.add("NAMESPACE_OID", UUID::NAMESPACE_OID)?;

tests/test_uuid.py

@@ -1,4 +1,5 @@
 import copy
+import os
 import pickle
 import sys
 from datetime import datetime
@@ -208,3 +209,29 @@ def test_is_safe() -> None:
 @pytest.mark.xfail(sys.platform == "linux", reason="Might fail in Github Actions")
 def test_getnode() -> None:
     assert uuid_utils.getnode() == getnode()
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="Does not run on Windows")
+def test_reseed_rng_with_fork() -> None:
+    read_end, write_end = os.pipe()
+    # forcibly generate uuid before fork to have rng state
+    uuid_utils.uuid4()
+
+    pid = os.fork()
+    if pid == 0:
+        os.close(read_end)
+        # explicity reseed in the child
+        uuid_utils.reseed_rng()
+        next_uuid_child = str(uuid_utils.uuid4())
+        with os.fdopen(write_end, "w") as write_pipe:
+            write_pipe.write(next_uuid_child)
+        os._exit(0)
+
+    os.close(write_end)
+    next_parent_uuid = uuid_utils.uuid4()
+    os.waitpid(pid, 0)
+    with os.fdopen(read_end) as read_pipe:
+        uuid_from_pipe = uuid_utils.UUID(read_pipe.read())
+
+    # the uuids should be different because we reseeded
+    assert next_parent_uuid != uuid_from_pipe